"""
Band-combination sweep: train + evaluate each combo, save checkpoints, collect
results. Runs on the patches in --data-dir (default all_data/patches). Patch size
is auto-detected by train.py / evaluate.py, so the SAME runner works for the
64 px (30 m) set now and the 128 px (10 m) set later — just point --data-dir at
the new patches.

    python run_experiments.py --gpu 1
    python run_experiments.py --gpu 1 --only f17_proven,f17_standard,f17_optical
    python run_experiments.py --gpu 1 --data-dir all_data/patches --force

Each combo writes  all_data/checkpoint_<name>/  and  all_data/result_<name>/.
Already-finished combos are skipped (resume-friendly) unless --force.
At the end it prints + saves the comparison table (all_data/runs_comparison.csv).
"""

import os
import sys
import argparse
import subprocess

HERE = os.path.dirname(os.path.abspath(__file__))
RUN_GPU = os.path.join(HERE, "run_gpu.sh")

# name -> band list (None = all 35 bands). Edit / extend freely.
COMBOS = {
    # --- three 17-feature variants (the focus); first is the current best 17 ---
    "f17_proven":     ["TPI", "ROUGHNESS", "LST_winter", "NDWI", "LST_summer", "REL_RELIEF",
                       "ASPECT", "LST_anomaly", "ELEV_COP", "B2", "HILLSHADE", "B8A", "NDMI", "B12", "B3", "B4", "B8"],
    "f17_standard":   ["TPI", "SLOPE", "LST_winter", "NDVI", "LST_summer", "REL_RELIEF",
                       "ASPECT", "LST_anomaly", "ELEV_COP", "B2", "HILLSHADE", "B8A", "NDMI", "B11", "B3", "B4", "B8"],
    "f17_optical":    ["B2", "B3", "B4", "B8", "B8A", "B11", "B12", "NDVI", "NDWI", "NDMI", "NDSI", "BSI",
                       "ELEV_COP", "SLOPE", "TPI", "LST_winter", "LST_summer"],
    # --- extra combos to round out the sweep ---
    "f17_terrain":    ["TPI", "ROUGHNESS", "TRI", "SLOPE", "ASPECT", "HILLSHADE", "CURVATURE", "REL_RELIEF",
                       "ELEV_COP", "LST_winter", "LST_summer", "LST_anomaly", "B2", "B4", "B8", "NDWI", "NDVI"],
    "f17_no_thermal": ["TPI", "ROUGHNESS", "SLOPE", "ASPECT", "REL_RELIEF", "HILLSHADE", "CURVATURE",
                       "ELEV_COP", "B2", "B3", "B4", "B8", "B8A", "B11", "B12", "NDWI", "NDVI"],
    "f16_balanced":   ["TPI", "ROUGHNESS", "REL_RELIEF", "ELEV_COP", "SLOPE", "ASPECT", "LST_winter", "LST_summer",
                       "NDWI", "NDVI", "NDMI", "B2", "B3", "B4", "B8", "B12"],
    "f20_rich":       ["TPI", "ROUGHNESS", "LST_winter", "NDWI", "LST_summer", "REL_RELIEF", "ASPECT", "LST_anomaly",
                       "ELEV_COP", "B2", "HILLSHADE", "B8A", "NDMI", "B12", "B3", "B4", "B8", "NDVI", "SLOPE", "B11"],
    "f14_lean":       ["TPI", "ROUGHNESS", "LST_winter", "NDWI", "LST_summer", "REL_RELIEF", "ASPECT",
                       "ELEV_COP", "B2", "B8A", "NDMI", "B12", "B4", "B8"],
    "f12_core":       ["TPI", "ROUGHNESS", "LST_winter", "LST_summer", "REL_RELIEF", "ASPECT",
                       "ELEV_COP", "B2", "B8", "NDWI", "NDMI", "B12"],
    # --- reference baseline (auto-skips if already trained) ---
    "all":            None,
}


def dirs_for(name):
    ck = "all_data/checkpoint" if name == "all" else f"all_data/checkpoint_{name}"
    rs = "all_data/result" if name == "all" else f"all_data/result_{name}"
    return os.path.join(HERE, ck), os.path.join(HERE, rs)


def run_step(script, env, log_path):
    os.makedirs(os.path.dirname(log_path), exist_ok=True)
    with open(log_path, "w") as lf:
        r = subprocess.run([RUN_GPU, script], cwd=HERE, env=env,
                           stdout=lf, stderr=subprocess.STDOUT)
    return r.returncode


def run_one(name, bands, gpu, data_dir, force):
    ckpt_dir, result_dir = dirs_for(name)
    weights = os.path.join(ckpt_dir, "best_model.weights.h5")
    metrics = os.path.join(result_dir, "metrics.json")

    env = os.environ.copy()
    env["CUDA_VISIBLE_DEVICES"] = str(gpu)
    env["DATA_DIR"] = data_dir
    env["RUN_NAME"] = name
    env.pop("BANDS", None)
    if bands is None:
        env.pop("BAND_LIST", None)
    else:
        env["BAND_LIST"] = ",".join(bands)

    nb = "all" if bands is None else len(bands)
    if force or not os.path.exists(weights):
        print(f"\n[TRAIN] {name}  ({nb} bands)  GPU {gpu}")
        if run_step("train.py", env, os.path.join(ckpt_dir, "train_log.txt")) != 0:
            print(f"  TRAIN FAILED -> {ckpt_dir}/train_log.txt")
            return False
    else:
        print(f"\n[skip train] {name}  (weights exist)")

    if force or not os.path.exists(metrics):
        print(f"[EVAL ] {name}")
        if run_step("evaluate.py", env, os.path.join(result_dir, "eval_log.txt")) != 0:
            print(f"  EVAL FAILED -> {result_dir}/eval_log.txt")
            return False
    else:
        print(f"[skip eval ] {name}  (metrics exist)")
    return True


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--gpu", default="1", help="CUDA device index (default 1)")
    ap.add_argument("--data-dir", default="all_data/patches",
                    help="patches root (default all_data/patches)")
    ap.add_argument("--only", default=None,
                    help="comma list of combo names to run (default: all in COMBOS)")
    ap.add_argument("--force", action="store_true",
                    help="re-run even if checkpoint/metrics already exist")
    args = ap.parse_args()

    names = list(COMBOS)
    if args.only:
        wanted = [n.strip() for n in args.only.split(",")]
        bad = [n for n in wanted if n not in COMBOS]
        if bad:
            sys.exit(f"Unknown combo(s): {bad}\nAvailable: {list(COMBOS)}")
        names = wanted

    print(f"Running {len(names)} combo(s) on GPU {args.gpu}, data={args.data_dir}:")
    print("  " + ", ".join(names))
    ok, failed = [], []
    for n in names:
        (ok if run_one(n, COMBOS[n], args.gpu, args.data_dir, args.force) else failed).append(n)

    print(f"\nDone. ok={ok}  failed={failed}")
    print("\n" + "=" * 60 + "\n  RESULTS\n" + "=" * 60)
    subprocess.run([sys.executable, os.path.join(HERE, "results_table.py")], cwd=HERE)


if __name__ == "__main__":
    main()