#!/bin/bash
set -euo pipefail
DEMO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CONFIG=""
MODE="run"
while [[ $# -gt 0 ]]; do
    case "$1" in
        --config) CONFIG="$2"; shift 2 ;;
        --config=*) CONFIG="${1#*=}"; shift ;;
        --submit) MODE="submit"; shift ;;
        -h|--help) echo "Usage: bash run.sh --config configs/<name>.yaml [--submit]"; exit 0 ;;
        *) echo "Unknown argument: $1" >&2; exit 2 ;;
    esac
done
[[ -n "$CONFIG" ]] || { echo "Please specify --config configs/<name>.yaml" >&2; exit 2; }
[[ "$CONFIG" = /* ]] || CONFIG="$DEMO_DIR/$CONFIG"
[[ -f "$CONFIG" ]] || { echo "Config not found: $CONFIG" >&2; exit 2; }
if [[ -z "${CONDA_PREFIX:-}" || -z "${ONESCIENCE_MODELS_DIR:-}" ]]; then
    export MATCHEM_CONDA_NAME="${MATCHEM_CONDA_NAME:-onescience-mattergen-source}"
    source "$DEMO_DIR/../../matchem_env.sh"
fi
CMD=$(python3 - "$CONFIG" "$DEMO_DIR/.." <<'PY'
import os, shlex, sys, yaml
import json
cfg = yaml.safe_load(open(sys.argv[1]))
root = os.path.abspath(sys.argv[2])
task = cfg.get("task", "train")
script = {"train": "train.py", "finetune": "finetune.py", "generate": "generate.py"}.get(task)
if not script: raise SystemExit(f"Unsupported task: {task}")
args = []
for key, value in cfg.get("args", {}).items():
    if isinstance(value, bool):
        if value: args.append(f"--{key.replace('_', '-')}")
    elif value is not None:
        value = json.dumps(value) if isinstance(value, (dict, list)) else str(value)
        value = os.path.expandvars(value)
        args.extend([f"--{key.replace('_', '-')}", value] if task == "generate" else [f"{key}={value}"])
print(" ".join(map(shlex.quote, [sys.executable, os.path.join(root, script), *args])))
PY
)
echo "Running MatterGen: $CMD"
if [[ "$MODE" == "submit" ]]; then
    TASK=$(python3 - "$CONFIG" <<'PY'
import sys, yaml
print(yaml.safe_load(open(sys.argv[1])).get("task", "train"))
PY
)
    [[ "$TASK" == "train" ]] || { echo "--submit currently supports train configs only" >&2; exit 2; }
    DEVICES=$(python3 - "$CONFIG" <<'PY'
import sys, yaml
print(yaml.safe_load(open(sys.argv[1])).get("args", {}).get("trainer.devices", 8))
PY
)
    sbatch \
        --gres="dcu:$DEVICES" \
        --export="ALL,CONFIG_PATH=$CONFIG,SCRIPT_DIR=$DEMO_DIR/.." \
        "$DEMO_DIR/../submit_train.sh"
    exit 0
fi
cd "$DEMO_DIR/.."
eval "$CMD"