#!/bin/bash
# ============================================================
# make_vesta_arrow.sh
# Create VESTA with displacement arrows AND (optionally) magCIF.
#  - Computes displacements between two POSCARs
#  - Writes VECTR/VECTT into a reference .vesta
#  - Optionally converts that .vesta to .mcif (magnetic CIF)
#  - Threshold applies to *unscaled* displacement; --scale only scales arrows
# ============================================================

set -euo pipefail

# ===== Defaults =====
ARROW_R=0
ARROW_G=0
ARROW_B=255
SCALE=1.0
VESTA_OUT="reference-arrow.vesta"
WRITE_MCIF=1
MCIF_OUT=""           # default derived from VESTA_OUT
FORCE=0

usage() {
  cat <<EOF
Usage:
  $0 -i <initial_POSCAR> -f <final_POSCAR> -t <threshold> -v <reference.vesta>
     [--color R G B] [--scale FACTOR] [--mcif-out PATH] [--no-mcif] [--force]

Example:
  $0 -i ./perfect/CONTCAR-finish -f ./CONTCAR-finish -t 0.03 -v reference.vesta \\
     --color 0 255 0 --scale 2.0 --mcif-out arrows.mcif

Defaults:
  Arrow color (R,G,B): 0, 0, 255
  Scale factor       : 1.0
  VESTA output       : ${VESTA_OUT}
  MCIF output        : <VESTA_OUT basename>.mcif (unless --no-mcif)
EOF
  exit 1
}

# ----- parse args -----
perfect=""
distorted=""
thresh=""
reference=""
positional=()

while [[ $# -gt 0 ]]; do
  case "$1" in
    -i|--initial)    perfect="$2"; shift 2 ;;
    -f|--final)      distorted="$2"; shift 2 ;;
    -t|--threshold)  thresh="$2"; shift 2 ;;
    -v|--vesta)      reference="$2"; shift 2 ;;
    --color)         ARROW_R="$2"; ARROW_G="$3"; ARROW_B="$4"; shift 4 ;;
    --scale)         SCALE="$2"; shift 2 ;;
    --mcif-out)      MCIF_OUT="$2"; shift 2 ;;
    --no-mcif)       WRITE_MCIF=0; shift ;;
    --force)         FORCE=1; shift ;;
    -h|--help)       usage ;;
    -*)              echo "❌ Unknown option: $1"; usage ;;
    *)               positional+=("$1"); shift ;;
  esac
done
# ignore stray positionals

# ----- required -----
[[ -n "$perfect" && -n "$distorted" && -n "$thresh" && -n "$reference" ]] || { echo "❌ Missing required arguments."; usage; }

# ----- file checks -----
[[ -f "$perfect"   ]] || { echo "❌ Not found: $perfect"; exit 1; }
[[ -f "$distorted" ]] || { echo "❌ Not found: $distorted"; exit 1; }
[[ -f "$reference" ]] || { echo "❌ Not found: $reference"; exit 1; }

# ----- atom count match -----
atoms1=$(awk 'NR==6{getline; split($0,a," "); s=0; for(i in a)s+=a[i]; print s}' "$perfect")
atoms2=$(awk 'NR==6{getline; split($0,a," "); s=0; for(i in a)s+=a[i]; print s}' "$distorted")
if [[ "$atoms1" -ne "$atoms2" ]]; then
  echo "❌ Error: Number of atoms differs between files!"
  echo "   $perfect  : $atoms1 atoms"
  echo "   $distorted : $atoms2 atoms"
  exit 1
fi

# ----- decide outputs & safety -----
# VESTA_OUT stays default unless the user wants a specific path: we keep default name
VESTA_OUT="${VESTA_OUT}"

# prevent overwriting reference .vesta if user points VESTA_OUT to same path (we keep fixed name, but check anyway)
if [[ "$(realpath -m "$reference")" = "$(realpath -m "$VESTA_OUT")" ]]; then
  base="${reference%.*}"
  VESTA_OUT="${base}.arrow.vesta"
  echo "🛑 VESTA output equals input reference. Writing to: $VESTA_OUT"
fi

# refuse to overwrite existing outputs unless --force
if [[ -e "$VESTA_OUT" && "$FORCE" -ne 1 ]]; then
  echo "❗ Output already exists: $VESTA_OUT"
  echo "   Re-run with --force or remove the file."
  exit 1
fi

# mcif name
if [[ "$WRITE_MCIF" -eq 1 ]]; then
  if [[ -z "$MCIF_OUT" ]]; then
    base="${VESTA_OUT##*/}"; stem="${base%.*}"
    MCIF_OUT="$(dirname "$VESTA_OUT")/${stem}.mcif"
  fi
  # avoid writing mcif into the same path as any input accidentally
  if [[ "$(realpath -m "$MCIF_OUT")" = "$(realpath -m "$reference")" ]]; then
    MCIF_OUT="${MCIF_OUT%.mcif}.converted.mcif"
    echo "🛑 MCIF output equals input reference. Writing to: $MCIF_OUT"
  fi
  if [[ -e "$MCIF_OUT" && "$FORCE" -ne 1 ]]; then
    echo "❗ Output already exists: $MCIF_OUT"
    echo "   Re-run with --force or choose another --mcif-out."
    exit 1
  fi
fi

# ----- helpers -----
tmp_c1="$(mktemp -t mvesta_coords1.XXXXXX)"
tmp_c2="$(mktemp -t mvesta_coords2.XXXXXX)"
tmp_diff="$(mktemp -t mvesta_diff.XXXXXX)"
tmp_vect="$(mktemp -t mvesta_vect.XXXXXX)"
cleanup() { rm -f "$tmp_c1" "$tmp_c2" "$tmp_diff" "$tmp_vect"; }
trap cleanup EXIT

# ----- diffcon (in-bash) -----
diffcon() {
  local pos1="$1" pos2="$2"

  parse_poscar() {
    local file="$1"
    awk '
      NR==2 {scale=$1}
      NR>=3 && NR<=5 {for(i=1;i<=3;i++) lattice[NR-2,i]=$i}
      NR==6 {getline; split($0,nums," "); total=0; for(i in nums) total+=nums[i]}
      NR>8 && NF>=3 { x[++n]=$1; y[n]=$2; z[n]=$3 }
      END { for(i=1;i<=n;i++) printf("%f %f %f\n",x[i],y[i],z[i]) }
    ' "$file"
  }

  coords1="$(parse_poscar "$pos1")"
  coords2="$(parse_poscar "$pos2")"

  read -r a11 a12 a13 <<< "$(awk 'NR==3{print $0}' "$pos1")"
  read -r a21 a22 a23 <<< "$(awk 'NR==4{print $0}' "$pos1")"
  read -r a31 a32 a33 <<< "$(awk 'NR==5{print $0}' "$pos1")"
  scale=$(awk 'NR==2{print $1}' "$pos1")

  echo "$coords1" > "$tmp_c1"
  echo "$coords2" > "$tmp_c2"

  paste "$tmp_c1" "$tmp_c2" | awk -v a11="$a11" -v a12="$a12" -v a13="$a13" \
      -v a21="$a21" -v a22="$a22" -v a23="$a23" \
      -v a31="$a31" -v a32="$a32" -v a33="$a33" \
      -v scale="$scale" -v s="$SCALE" '
    {
      dx = $4 - $1; dy = $5 - $2; dz = $6 - $3
      if (dx > 0.5) dx -= 1; if (dx < -0.5) dx += 1
      if (dy > 0.5) dy -= 1; if (dy < -0.5) dy += 1
      if (dz > 0.5) dz -= 1; if (dz < -0.5) dz += 1

      # base displacement
      cx0 = scale * (a11*dx + a21*dy + a31*dz)
      cy0 = scale * (a12*dx + a22*dy + a32*dz)
      cz0 = scale * (a13*dx + a23*dy + a33*dz)
      norm = sqrt(cx0*cx0 + cy0*cy0 + cz0*cz0)

      # scaled for visualization
      cx = s * cx0; cy = s * cy0; cz = s * cz0

      printf("%12.6f %12.6f %12.6f : %10.6f  %4d\n", cx, cy, cz, norm, NR)
    }'
}

echo "➡️  Calculating displacements..."
diffcon "$perfect" "$distorted" > "$tmp_diff"

echo "➡️  Generating VECTR/VECTT block..."
awk -v t="$thresh" -v R="$ARROW_R" -v G="$ARROW_G" -v B="$ARROW_B" '
/:/ {
  split($0, arr, ":")
  split(arr[1], v)
  split(arr[2], w)
  x=v[1]; y=v[2]; z=v[3]
  norm=w[1]
  i++
  # threshold on *unscaled* norm: if smaller, zero the arrow vector
  if (norm < t) {x=y=z=0}
  # VECTR block (per atom index i)
  printf("%5d %12.6f %12.6f %12.6f  0\n%5d   0    0    0    0\n 0 0 0 0 0\n", i, x, y, z, i)
  n=i
}
END {
  print " 0 0 0 0 0"
  print "VECTT"
  for (i=1; i<=n; i++)
    printf("%5d  0.500 %3d %3d %3d 1\n", i, R, G, B)
  print " 0 0 0 0 0"
}' "$tmp_diff" > "$tmp_vect"

echo "➡️  Updating VECTR/VECTT section in $reference → $VESTA_OUT"
awk '
  FNR==1 && FNR!=NR {in_v=0}
  FILENAME==ARGV[1] {block[NR]=$0; next}
  /^VECTR/ {
    print "VECTR"
    for (i=1;i<=length(block);i++) print block[i]
    in_v=1
    next
  }
  in_v && /^SPLAN/ {in_v=0}
  !in_v {print $0}
' "$tmp_vect" "$reference" > "$VESTA_OUT"

echo "✅ VESTA done: $VESTA_OUT"

# -------------------- MCIF writer (from VESTA_OUT) --------------------
if [[ "$WRITE_MCIF" -eq 1 ]]; then
  echo "➡️  Writing magCIF from $VESTA_OUT → $MCIF_OUT"

  atoms_tsv="$(mktemp -t vesta_atoms.XXXXXX)"
  vectors_tsv="$(mktemp -t vesta_vectors.XXXXXX)"
  joined_tsv="$(mktemp -t vesta_joined.XXXXXX)"
  trap 'rm -f "$atoms_tsv" "$vectors_tsv" "$joined_tsv"; cleanup' EXIT

  # CELLP
  read -r a b c alpha beta gamma <<EOF
$(awk '
  /^CELLP[[:space:]]*$/ { if (getline L) { split(L,t); printf("%s %s %s %s %s %s\n",t[1],t[2],t[3],t[4],t[5],t[6]); exit } }
  /^CELLP[[:space:]]+/   { printf("%s %s %s %s %s %s\n",$2,$3,$4,$5,$6,$7); exit }
' "$VESTA_OUT")
EOF
  : "${a:=10.0}" ; : "${b:=10.0}" ; : "${c:=10.0}"
  : "${alpha:=90}" ; : "${beta:=90}" ; : "${gamma:=90}"

  # GROUP
  spg=$(awk '
    /^GROUP[[:space:]]*$/ { if (getline L) { split(L,t); for (i=1;i<=length(L);i++) if (t[i] ~ /^[PIFCR]$/) { printf("%s 1\n", t[i]); exit } print "P 1"; exit } }
  ' "$VESTA_OUT")
  : "${spg:=P 1}"

  # STRUC atoms (1st line per atom)
  awk '
    BEGIN { in_s=0 }
    /^STRUC/ { in_s=1; next }
    in_s {
      if ($0 ~ /^[[:space:]]*0([[:space:]]+0){2,}/) { in_s=0; next }
      if ($1 ~ /^[0-9]+$/ && $2 ~ /^[A-Za-z][A-Za-z0-9]*$/ && NF>=8) {
        idx=$1; sym=$2; lab=$3; occ=$4; fx=$5; fy=$6; fz=$7
        gsub(/[^A-Za-z0-9_]/,"",lab)
        if (lab=="") lab = sym idx
        printf("%d\t%s\t%s\t%g\t%f\t%f\t%f\n", idx, sym, lab, occ, fx, fy, fz)
      }
    }
  ' "$VESTA_OUT" | sort -n -k1,1 > "$atoms_tsv"

  natoms=$(wc -l < "$atoms_tsv" | tr -d ' ')
  [[ "$natoms" -gt 0 ]] || { echo "❌ No atoms parsed from STRUC in $VESTA_OUT"; exit 1; }

  # VECTR (support A/B)
  awk -v NATOMS="$natoms" '
    BEGIN { in_v=0 }
    /^VECTR[[:space:]]*$/ { in_v=1; next }
    in_v && /^VECTT[[:space:]]*$/ { in_v=0; next }
    in_v {
      if ($0 ~ /^[[:space:]]*0([[:space:]]+0){4,}[[:space:]]*$/) next
      n=split($0,t)
      # B: id atom dx dy dz 0
      if (n>=6 && t[1] ~ /^[0-9]+$/ && t[2] ~ /^[0-9]+$/) {
        atom=t[2]; x=t[3]; y=t[4]; z=t[5]
        if (atom>=1 && atom<=NATOMS && !(atom in seen)) { print atom "\t" x "\t" y "\t" z; seen[atom]=1 }
        next
      }
      # A: atom dx dy dz 0
      if (n>=5 && t[1] ~ /^[0-9]+$/) {
        atom=t[1]; x=t[2]; y=t[3]; z=t[4]
        if (atom>=1 && atom<=NATOMS && !(atom in seen)) { print atom "\t" x "\t" y "\t" z; seen[atom]=1 }
        next
      }
    }
  ' "$VESTA_OUT" | sort -n -k1,1 > "$vectors_tsv"

  nvec=$(wc -l < "$vectors_tsv" | tr -d ' ')
  if [[ "$nvec" -ne "$natoms" ]]; then
    echo "⚠️  Note: VECTR entries ($nvec) != atoms ($natoms). Missing indices get 0,0,0."
  fi

  # join
  awk -F'\t' 'FNR==NR{vec[$1]=$2"\t"$3"\t"$4; next}
  {
    idx=$1; sym=$2; lab=$3; occ=$4; fx=$5; fy=$6; fz=$7
    if (idx in vec) print idx, sym, lab, occ, fx, fy, fz, vec[idx];
    else            print idx, sym, lab, occ, fx, fy, fz, "0\t0\t0";
  }' OFS='\t' "$vectors_tsv" "$atoms_tsv" > "$joined_tsv"

  # write mcif
  cat > "$MCIF_OUT" <<EOF
#======================================================================
# CRYSTAL DATA (generated from VESTA)
#----------------------------------------------------------------------
data_VESTA_phase_1

_chemical_name_common                  'unknown system'
_cell_length_a                         $a
_cell_length_b                         $b
_cell_length_c                         $c
_cell_angle_alpha                      $alpha
_cell_angle_beta                       $beta
_cell_angle_gamma                      $gamma
_space_group_name_H-M_alt              '$spg'
_space_group_IT_number                 1

loop_
_space_group_symop_operation_xyz
   'x, y, z'

loop_
   _atom_site_label
   _atom_site_occupancy
   _atom_site_fract_x
   _atom_site_fract_y
   _atom_site_fract_z
   _atom_site_adp_type
   _atom_site_U_iso_or_equiv
   _atom_site_type_symbol
EOF

  awk -F'\t' '
  {
    idx=$1; sym=$2; lab=$3; occ=$4; fx=$5; fy=$6; fz=$7
    label = (lab!="" ? lab : sym idx)
    printf("   %-10s %6.3f %12.6f %12.6f %12.6f    Uiso  ? %s\n",
           label, occ, fx, fy, fz, sym)
  }' "$joined_tsv" >> "$MCIF_OUT"

  cat >> "$MCIF_OUT" <<'EOF'

loop_
   _atom_site_moment_label
   _atom_site_moment_crystalaxis_x
   _atom_site_moment_crystalaxis_y
   _atom_site_moment_crystalaxis_z
EOF

  awk -F'\t' '
  {
    idx=$1; sym=$2; lab=$3; mx=$8; my=$9; mz=$10
    label = (lab!="" ? lab : sym idx)
    printf("   %-10s %12.6f %12.6f %12.6f\n", label, mx, my, mz)
  }' "$joined_tsv" >> "$MCIF_OUT"

  echo "✅ MCIF done: $MCIF_OUT"
fi

echo "All done."

