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

if [[ $# -lt 1 || $# -gt 2 ]]; then
  echo "Usage: $0 aha|l2a_style [GPU_ID]" >&2
  exit 2
fi

ARM="$1"
GPU="${2:-0}"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
REPO="$(cd "$RECIPE/.." && pwd)"
VANILLA_DIR="${VANILLA_DIR:-$REPO}"
DATA_DIR="${DATA_DIR:-$RECIPE/data/am_distilled_long_mix}"
OUTPUT_ROOT="${OUTPUT_ROOT:-$REPO/outputs}"

case "$ARM" in
  aha) GRANULARITY=token_kv_head ;;
  l2a_style) GRANULARITY=token ;;
  *) echo "Unknown arm: $ARM (expected aha or l2a_style)" >&2; exit 2 ;;
esac

ARM_OUT="$OUTPUT_ROOT/$ARM"
HOTSTART="$ARM_OUT/hotstart"
STAGE1="$ARM_OUT/stage1"
STAGE2="$ARM_OUT/stage2"
LOG_DIR="$ARM_OUT/logs"
mkdir -p "$LOG_DIR"

export CUDA_VISIBLE_DEVICES="$GPU"
export PYTHONUNBUFFERED=1
export TOKENIZERS_PARALLELISM=false
export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"

if [[ ! -f "$VANILLA_DIR/model.safetensors" ]]; then
  echo "Missing tuned-vanilla checkpoint: $VANILLA_DIR/model.safetensors" >&2
  exit 2
fi

if [[ ! -f "$HOTSTART/config.json" ]]; then
  python "$SCRIPT_DIR/export_tuned_vanilla_aha_hotstart.py" \
    --vanilla-path "$VANILLA_DIR" \
    --output-path "$HOTSTART" \
    --window-size 128 \
    --local-kind sink_recent \
    --gate-init-full-prob 0.90 \
    --router-granularity "$GRANULARITY" \
    2>&1 | tee "$LOG_DIR/hotstart.log"
fi

if [[ ! -f "$STAGE1/checkpoint-300/config.json" ]]; then
  AHA_TRAIN_GATE_HARD_THRESHOLD=0.50 \
  python "$RECIPE/dynamic_duo_train.py" \
    --aha_checkpoint "$HOTSTART" \
    --model_path "$VANILLA_DIR" \
    --output_dir "$STAGE1" \
    --data_source am_distilled \
    --am_dataset_path "$DATA_DIR" \
    --am_dataset_split train \
    --am_label_mode full \
    --max_length 8192 \
    --num_steps 300 \
    --warmup_ratio 0.10 \
    --lr 3e-5 \
    --reg_weight 0.1 \
    --ce_weight 0.0 \
    --batch_size 1 \
    --grad_accum 1 \
    --save_steps 100 \
    --log_steps 10 \
    --seed 42 \
    --dtype bfloat16 \
    --attn_impl sdpa \
    --aha_local_kind sink_recent \
    --router_granularity "$GRANULARITY" \
    2>&1 | tee "$LOG_DIR/stage1.log"
fi

if [[ ! -f "$STAGE2/checkpoint-75/config.json" ]]; then
  MODEL_PATH="$VANILLA_DIR" \
  AHA_CHECKPOINT_PATH="$STAGE1/checkpoint-300" \
  DATASET_PATH="$DATA_DIR" \
  DATASET_SPLIT=train \
  OUTPUT_DIR="$STAGE2" \
  AHA_MODE=dynamic \
  AHA_ROUTER_GRANULARITY="$GRANULARITY" \
  AHA_LOCAL_KIND=sink_recent \
  AHA_CE_WEIGHT=1.0 \
  AHA_DISTILL_WEIGHT=0.5 \
  AHA_REG_WEIGHT=0.01 \
  AHA_TRAIN_GATE_HARD_THRESHOLD=0.58 \
  GROUPED_LR=1 \
  GATE_ONLY=0 \
  LEARNING_RATE=3e-6 \
  GATE_LEARNING_RATE=3e-6 \
  BACKBONE_LEARNING_RATE=3e-7 \
  LR_SCHEDULER_TYPE=constant_with_warmup \
  WARMUP_RATIO=0.10 \
  WEIGHT_DECAY=0.0 \
  MAX_SEQ_LENGTH=8192 \
  PER_DEVICE_TRAIN_BATCH_SIZE=1 \
  GRADIENT_ACCUMULATION_STEPS=1 \
  MAX_STEPS=75 \
  LOGGING_STEPS=5 \
  SAVE_STEPS=25 \
  SAVE_TOTAL_LIMIT=3 \
  FREEZE_EMBEDDINGS_LM_HEAD=1 \
  REPORT_TO=none \
  SEED=47 \
  python "$RECIPE/sft.py" \
    2>&1 | tee "$LOG_DIR/stage2.log"
fi

python - "$STAGE2/checkpoint-25" "$GRANULARITY" <<'PY'
import json
import sys
from pathlib import Path

checkpoint = Path(sys.argv[1])
expected = sys.argv[2]
config = json.loads((checkpoint / "config.json").read_text())
actual = config.get("aha_router_granularity", "token_kv_head")
if actual != expected:
    raise RuntimeError(f"checkpoint granularity {actual!r} != {expected!r}")
print(f"selected_checkpoint={checkpoint} router_granularity={actual}")
PY