#!/usr/bin/env bash
set -euo pipefail

usage() {
  cat <<'EOF'
Usage:
  run_qscaild_with_macer.sh -w WORKDIR [options]

Required:
  -w, --workdir DIR          qSCAILD working directory (contains POSCAR/SPOSCAR/INCAR/POTCAR/parameters)

Options:
  -r, --repo DIR             macer repository root (default: parent of this script)
      --ff NAME              macer force field (default: chgnet)
      --device DEV           macer device (default: cpu)
      --max-cycles N         safety cap for submit/fit cycles (default: 200)
      --no-local-phonopy     do not prepend REPO to PYTHONPATH
  -h, --help                 show this help

What this script does:
  1) Run qscaild submit once to generate config-* and to_calc
  2) For each path in to_calc, run:
       macer relax -p POSCAR --isif 0 ... --outcar OUTCAR --vasprun vasprun.xml
  3) Run qscaild submit again to fit/update force constants
  4) Repeat 2-3 until file "finished" appears

Outputs in WORKDIR:
  - finished
  - FORCE_CONSTANTS_CURRENT
  - qscaild-submit-prepare-*.log
  - qscaild-submit-fit-*.log
  - macer-relax-*.log
EOF
}

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
WORKDIR=""
FF="chgnet"
DEVICE="cpu"
MAX_CYCLES=200
USE_LOCAL_PHONOPY=1

while [[ $# -gt 0 ]]; do
  case "$1" in
    -w|--workdir)
      WORKDIR="$2"
      shift 2
      ;;
    -r|--repo)
      REPO_DIR="$2"
      shift 2
      ;;
    --ff)
      FF="$2"
      shift 2
      ;;
    --device)
      DEVICE="$2"
      shift 2
      ;;
    --max-cycles)
      MAX_CYCLES="$2"
      shift 2
      ;;
    --no-local-phonopy)
      USE_LOCAL_PHONOPY=0
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Unknown option: $1" >&2
      usage
      exit 2
      ;;
  esac
done

if [[ -z "$WORKDIR" ]]; then
  echo "ERROR: --workdir is required." >&2
  usage
  exit 2
fi

WORKDIR="$(cd "$WORKDIR" && pwd)"
if [[ ! -d "$WORKDIR" ]]; then
  echo "ERROR: WORKDIR does not exist: $WORKDIR" >&2
  exit 2
fi

for req in POSCAR SPOSCAR INCAR POTCAR parameters; do
  if [[ ! -f "$WORKDIR/$req" ]]; then
    echo "ERROR: missing required file: $WORKDIR/$req" >&2
    exit 2
  fi
done

if [[ ! -f "$REPO_DIR/qscaild-master/submit_qscaild.py" ]]; then
  echo "ERROR: submit_qscaild.py not found under repo: $REPO_DIR" >&2
  exit 2
fi

echo "REPO_DIR=$REPO_DIR"
echo "WORKDIR=$WORKDIR"
echo "FF=$FF DEVICE=$DEVICE MAX_CYCLES=$MAX_CYCLES"

if [[ "$USE_LOCAL_PHONOPY" -eq 1 ]]; then
  export PYTHONPATH="$REPO_DIR${PYTHONPATH:+:$PYTHONPATH}"
  echo "PYTHONPATH prepended with repo for local phonopy compatibility."
fi

submit_qscaild() {
  local log_file="$1"
  (
    cd "$WORKDIR"
    python "$REPO_DIR/qscaild-master/submit_qscaild.py"
  ) >"$WORKDIR/$log_file" 2>&1
}

run_macer_on_to_calc() {
  local tag="$1"
  local log_file="$WORKDIR/macer-relax-$tag.log"
  : >"$log_file"

  if [[ ! -f "$WORKDIR/to_calc" ]]; then
    echo "ERROR: to_calc not found in $WORKDIR" >&2
    return 1
  fi

  while IFS= read -r calc_dir; do
    [[ -z "$calc_dir" ]] && continue
    if [[ ! -d "$calc_dir" ]]; then
      echo "ERROR: calc dir not found: $calc_dir" | tee -a "$log_file" >&2
      return 1
    fi
    if [[ ! -f "$calc_dir/POSCAR" ]]; then
      echo "ERROR: POSCAR missing in $calc_dir" | tee -a "$log_file" >&2
      return 1
    fi

    echo "[macer] $calc_dir" | tee -a "$log_file"
    (
      cd "$calc_dir"
      macer relax -p POSCAR --isif 0 \
        --ff "$FF" \
        --device "$DEVICE" \
        --output-dir . \
        --outcar OUTCAR \
        --contcar CONTCAR \
        --vasprun vasprun.xml \
        --no-pdf \
        --quiet
    ) >>"$log_file" 2>&1
  done < "$WORKDIR/to_calc"
}

prepare_idx=1
fit_idx=1
cycle=1

if [[ ! -f "$WORKDIR/finished" ]]; then
  echo "==> Initial qscaild submit"
  submit_qscaild "qscaild-submit-prepare-${prepare_idx}.log"
  prepare_idx=$((prepare_idx + 1))
fi

while [[ ! -f "$WORKDIR/finished" ]]; do
  if [[ "$cycle" -gt "$MAX_CYCLES" ]]; then
    echo "ERROR: reached MAX_CYCLES=$MAX_CYCLES without finished." >&2
    exit 1
  fi

  echo "==> Cycle $cycle: run macer on to_calc"
  run_macer_on_to_calc "$cycle"

  echo "==> Cycle $cycle: qscaild fit/update submit"
  submit_qscaild "qscaild-submit-fit-${fit_idx}.log"
  fit_idx=$((fit_idx + 1))
  cycle=$((cycle + 1))
done

echo "==> Done"
echo "finished:"
cat "$WORKDIR/finished"
if [[ -f "$WORKDIR/FORCE_CONSTANTS_CURRENT" ]]; then
  echo "FORCE_CONSTANTS_CURRENT: $WORKDIR/FORCE_CONSTANTS_CURRENT"
fi

