#!/bin/bash
# Sweep all v10_r2 ckpts (except 956 which is already done) to find the true
# best step. Each ckpt: same 8-rank protocol as run_eval_v10_r2_ckpt956.sh.
# Sequential to avoid GPU contention; ~10 min per ckpt × 8 ckpts = ~1.5 h.
set -e

cd /mnt/local-fast/zhangt/forensics_grpo

export PATH="/mnt/local-fast/zhangt/torch_env/bin:$PATH"
export LD_LIBRARY_PATH="/opt/conda/lib:${LD_LIBRARY_PATH}"
export PYTHONPATH=".:$PYTHONPATH"

STEPS=(240 270 480 510 720 750 780 930)

for S in "${STEPS[@]}"; do
    MODEL=outputs_forensics/v10_r2/checkpoint-${S}
    OUT=eval_v10_r2_ckpt${S}
    if [ -f "$OUT/grounding_metrics.txt" ]; then
        echo "=== ckpt-${S} already done, skipping ==="
        continue
    fi
    mkdir -p "$OUT/logs"
    echo
    echo "================================================================"
    echo "ckpt-${S}: evaluating $MODEL -> $OUT"
    echo "================================================================"
    for R in 0 1 2 3 4 5 6 7; do
      CUDA_VISIBLE_DEVICES=$R python evaluate_forensics.py \
          --model_path "$MODEL" \
          --rank $R --world_size 8 --device 0 \
          --out_dir "$OUT" \
          --cot false --max_new_tokens 64 --temperature 0.0 \
          > "$OUT/logs/rank_${R}.log" 2>&1 &
    done
    wait
    python evaluate_grounding_metrics.py --out_dir "$OUT" | tee "$OUT/grounding_metrics.txt"
done

echo
echo "================================================================"
echo "SUMMARY across all v10_r2 ckpts (incl. 956)"
echo "================================================================"
python3 <<'PY'
import json, os
from collections import Counter
STEPS = [240, 270, 480, 510, 720, 750, 780, 930, 956]
print(f"{'step':>6} {'mIoU':>7} {'F1strict':>9} {'F1@0.5':>7} {'F1@0.7':>7} {'F1@0.85':>8} {'F1@0.95':>8} | {'1-seg':>7} {'multi':>7} {'collapse%':>10} {'K2_match%':>10}")
for S in STEPS:
    d = f"eval_v10_r2_ckpt{S}"
    metrics_path = f"{d}/grounding_metrics.txt"
    if not os.path.exists(metrics_path):
        print(f"{S:>6}  (no eval)")
        continue
    overall = {}
    with open(metrics_path) as fh:
        for line in fh:
            if "mIoU" in line and "=" in line and "%" in line:
                try:
                    overall["mIoU"] = float(line.split("=")[1].strip().rstrip("%"))
                except: pass
            if "F1@0.3 / F1@0.5 / F1@0.7" in line:
                parts = line.split("=")[1].split("/")
                overall["F1@0.5"] = float(parts[1].strip().rstrip("%"))
                overall["F1@0.7"] = float(parts[2].strip().rstrip("%"))
            if "F1@0.85 / F1@0.95" in line:
                parts = line.split("=")[1].split("/")
                overall["F1@0.85"] = float(parts[0].strip().rstrip("%"))
                overall["F1@0.95"] = float(parts[1].strip().rstrip("%"))
            if "mean F1@strict" in line:
                overall["F1strict"] = float(line.split("=")[1].strip().rstrip("%"))
    # K-decision diagnostics from jsonl
    recs = []
    for r in range(8):
        p = f"{d}/rank_{r}.jsonl"
        if os.path.exists(p):
            with open(p) as fh:
                recs.extend(json.loads(l) for l in fh)
    if recs:
        single = [x for x in recs if x["n_gt"] == 1]
        multi  = [x for x in recs if x["n_gt"] >  1]
        k2     = [x for x in recs if x["n_gt"] == 2]
        single_iou = sum(x["hungarian_iou"] for x in single)/len(single)*100 if single else 0
        multi_iou  = sum(x["hungarian_iou"] for x in multi )/len(multi )*100 if multi  else 0
        collapse_pct = sum(1 for x in multi if x["n_pred"] <= 1)/len(multi)*100 if multi else 0
        k2_match     = sum(1 for x in k2 if x["n_pred"] == 2)/len(k2)*100 if k2 else 0
    else:
        single_iou = multi_iou = collapse_pct = k2_match = 0
    print(f"{S:>6} {overall.get('mIoU',0):>7.2f} {overall.get('F1strict',0):>9.2f} {overall.get('F1@0.5',0):>7.2f} {overall.get('F1@0.7',0):>7.2f} {overall.get('F1@0.85',0):>8.2f} {overall.get('F1@0.95',0):>8.2f} | {single_iou:>7.2f} {multi_iou:>7.2f} {collapse_pct:>10.2f} {k2_match:>10.2f}")
PY