experiments/source_selection_payoff.py

"""Experiment 1: End-to-end source-selection payoff.

For each held-out target dataset, pick a source under three strategies:
  (a) LWM-CDE distance     : argmin_s D_CDE[s, t]
  (b) Frozen-LWM distance  : argmin_s D_LWM[s, t]
  (c) Random               : random s != t
  (d) Oracle (P-row)       : argmin_s P[s, t]   <-- upper bound

Achieved transfer is P[s_picked, t] (lower = better, since P = 1 - F1).
Report mean +/- std across the held-out targets, plus the gap to oracle and
the gap to random expectation.

Usage:
  python experiments/source_selection_payoff.py \
      --task los-nlos --n_train 400 --n_test 400 --epochs 12

Output: results JSON next to this file plus a summary table to stdout.
"""
from __future__ import annotations

import argparse, contextlib, io, json, os, sys, time
from datetime import datetime

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))   # the repository root

import numpy as np
import torch

from dataset_distancing import (
    DatasetCollection,
    DistanceConfig,
    TransferConfig,
    compute_distances,
    compute_performance_matrix,
)
from dataset_distancing.embeddings import register_lwm_path, load_runner


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--lwm_path",   type=str, default="third_party/lwm-1.1")
    p.add_argument("--cache_dir",  type=str, default="cache/exp1")
    p.add_argument("--gpu",        type=int, default=1)
    p.add_argument("--seed",       type=int, default=42)
    p.add_argument("--n_train",    type=int, default=400)
    p.add_argument("--n_val",      type=int, default=400)
    p.add_argument("--n_test",     type=int, default=400)
    p.add_argument("--task",       type=str, default="los-nlos",
                   choices=["los-nlos", "beam-pred"])
    p.add_argument("--n_ant",      type=int, default=32)
    p.add_argument("--n_sub",      type=int, default=32)
    p.add_argument("--n_beams",    type=int, default=32)
    p.add_argument("--n_cities",   type=int, default=8,
                   help="number of distinct cities (each contributes n_bs base-stations)")
    p.add_argument("--n_bs",       type=int, default=3,
                   help="number of base stations per city (1..3); multi-BS adds within-city variety")
    p.add_argument("--epochs",     type=int, default=12)
    p.add_argument("--lr",         type=float, default=1e-3)
    p.add_argument("--batch_size", type=int, default=128)
    p.add_argument("--cde_epochs", type=int, default=20)
    p.add_argument("--n_random_trials", type=int, default=20,
                   help="random source picks per target (averaged)")
    p.add_argument("--out_json",   type=str, default=None)
    return p.parse_args()


def _select_pool_scenarios(task: str, n_cities: int, n_bs: int):
    """Return a deliberately *diverse* pool of (city, BS) DeepMIMO datasets.

    Mirrors the original paper's recipe of mixing high-LoS US cities with
    low-LoS international cities so the transfer matrix has enough variance to
    distinguish good from bad source choices.
    """
    import deepmimo as dm
    all_scs = dm.search()
    sub6 = [s for s in all_scs if "_3p5" in s and s.startswith("city_")]
    n_low  = n_cities // 2
    n_high = n_cities - n_low
    # First-indexed scenarios = US cities, mostly high LoS
    # Later-indexed scenarios = international cities (Singapore, Dubai, …) low LoS
    cities = sub6[:n_high] + sub6[max(len(sub6)-n_low, n_high):][:n_low]
    pool = []
    for s in cities:
        for tx in range(1, n_bs + 1):
            pool.append({"name": s, "tx_set": tx})
    return pool


def main():
    args = parse_args()
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    register_lwm_path(args.lwm_path)

    out_dir = os.path.dirname(os.path.abspath(__file__))
    out_json = args.out_json or os.path.join(
        out_dir, f"exp1_{args.task}_c{args.n_cities}b{args.n_bs}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    )

    # ── Pool ──────────────────────────────────────────────────────────────
    scenarios = _select_pool_scenarios(args.task, args.n_cities, args.n_bs)
    print(f"[Exp1] Pool: {len(scenarios)} scenarios for task={args.task} ({args.n_cities} cities × {args.n_bs} BS)")
    for s in scenarios:
        print("  -", s["name"])

    full = DatasetCollection.from_scenarios(
        scenarios,
        task=args.task,
        n_ant=args.n_ant,
        n_sub=args.n_sub,
        n_beams=args.n_beams,
        cache_dir=args.cache_dir,
    )

    # ── Split ─────────────────────────────────────────────────────────────
    train_ds, val_ds, test_ds = [], [], []
    for ds in full:
        tr, va, te = ds.split(
            train=args.n_train, val=args.n_val, test=args.n_test, seed=args.seed,
        )
        train_ds.append(tr); val_ds.append(va); test_ds.append(te)
    train_col = DatasetCollection(datasets=train_ds)
    val_col   = DatasetCollection(datasets=val_ds)
    test_col  = DatasetCollection(datasets=test_ds)

    # ── Runner ────────────────────────────────────────────────────────────
    with contextlib.redirect_stdout(io.StringIO()):
        runner = load_runner(args.lwm_path, weights_file=None)

    # ── P matrix (full pool, symmetric=False to keep the directed truth) ──
    print("[Exp1] Computing transfer matrix P …")
    t0 = time.time()
    P = compute_performance_matrix(
        train_col,
        cfg=TransferConfig(
            task=args.task, epochs=args.epochs, lr=args.lr,
            batch_size=args.batch_size, n_beams=args.n_beams,
            seed=args.seed, symmetrize=False, verbose=False,
        ),
        val_collection=val_col, test_collection=test_col,
    )
    P_np = P.detach().cpu().numpy()
    print(f"[Exp1] P done in {time.time()-t0:.1f}s, shape={P_np.shape}")

    # ── Frozen-LWM distances ──────────────────────────────────────────────
    print("[Exp1] Computing frozen-LWM distances …")
    D_lwm = compute_distances(train_col, runner, cfg=DistanceConfig(
        distance="swd", label_aware=True, n_per_dataset=args.n_train,
        sw_projections=64, asymmetric=False,
    )).detach().cpu().numpy()

    # ── CDE fine-tune on the same pool, then distances ────────────────────
    print("[Exp1] Fine-tuning LWM-CDE on the pool …")
    from dataset_distancing.cde import CDETrainer
    from dataset_distancing import CDEConfig
    cde_cfg = CDEConfig(
        epochs=args.cde_epochs, lr=5e-4,
        n_per_dataset=64, n_per_dataset_eval=128, corr_n_per_dataset=64,
        asymmetric=False,
        distance=DistanceConfig(sw_projections=32, label_aware=True),
    )
    trainer = CDETrainer(runner, cde_cfg)
    trainer.fit(train_col, P, eval_collection=train_col, perf_eval=P, verbose=False)

    print("[Exp1] Computing CDE-refined distances …")
    D_cde = compute_distances(train_col, runner, cfg=DistanceConfig(
        distance="swd", label_aware=True, n_per_dataset=args.n_train,
        sw_projections=64, asymmetric=False,
    )).detach().cpu().numpy()

    # ── Selection payoff ──────────────────────────────────────────────────
    N = P_np.shape[0]
    rng = np.random.default_rng(args.seed)

    def _pick_argmin(M, t):
        col = M[:, t].copy()
        col[t] = np.inf  # forbid self-selection
        return int(np.argmin(col))

    def _pick_random(t):
        choices = [s for s in range(N) if s != t]
        return int(rng.choice(choices))

    def _expected_random(t):
        choices = [s for s in range(N) if s != t]
        return float(np.mean(P_np[choices, t]))

    rows = []
    for t in range(N):
        s_oracle = _pick_argmin(P_np,  t)
        s_lwm    = _pick_argmin(D_lwm, t)
        s_cde    = _pick_argmin(D_cde, t)
        # Random: sample n_random_trials picks and average achieved P
        rand_pcs = [P_np[_pick_random(t), t] for _ in range(args.n_random_trials)]
        rows.append({
            "target":       t,
            "target_name":  test_col[t].name,
            "P_oracle":     float(P_np[s_oracle, t]),
            "P_lwm":        float(P_np[s_lwm,    t]),
            "P_cde":        float(P_np[s_cde,    t]),
            "P_rand_mean":  float(np.mean(rand_pcs)),
            "P_rand_std":   float(np.std(rand_pcs)),
            "P_rand_expected": _expected_random(t),
            "s_oracle":     s_oracle, "s_lwm": s_lwm, "s_cde": s_cde,
        })

    summary = {
        "task":       args.task,
        "n_pool":     N,
        "scenarios":  [s["name"] for s in scenarios],
        "rows":       rows,
        "agg": {
            "oracle_P_mean":  float(np.mean([r["P_oracle"] for r in rows])),
            "oracle_F1_mean": float(1.0 - np.mean([r["P_oracle"] for r in rows])),
            "lwm_P_mean":     float(np.mean([r["P_lwm"]    for r in rows])),
            "lwm_F1_mean":    float(1.0 - np.mean([r["P_lwm"]    for r in rows])),
            "cde_P_mean":     float(np.mean([r["P_cde"]    for r in rows])),
            "cde_F1_mean":    float(1.0 - np.mean([r["P_cde"]    for r in rows])),
            "rand_P_mean":    float(np.mean([r["P_rand_expected"] for r in rows])),
            "rand_F1_mean":   float(1.0 - np.mean([r["P_rand_expected"] for r in rows])),
        },
    }

    print("\n──────── Source-Selection Payoff ────────")
    print(f"task = {args.task}, pool = {N}")
    a = summary["agg"]
    print(f"  Random      : F1 = {a['rand_F1_mean']:.4f}   (P = {a['rand_P_mean']:.4f})")
    print(f"  Frozen LWM  : F1 = {a['lwm_F1_mean']:.4f}    (P = {a['lwm_P_mean']:.4f})")
    print(f"  LWM-CDE     : F1 = {a['cde_F1_mean']:.4f}    (P = {a['cde_P_mean']:.4f})")
    print(f"  Oracle (P)  : F1 = {a['oracle_F1_mean']:.4f} (P = {a['oracle_P_mean']:.4f})")
    if a["rand_F1_mean"] != a["oracle_F1_mean"]:
        gap = a["oracle_F1_mean"] - a["rand_F1_mean"]
        cde_recovered = (a["cde_F1_mean"] - a["rand_F1_mean"]) / gap if gap > 0 else float("nan")
        lwm_recovered = (a["lwm_F1_mean"] - a["rand_F1_mean"]) / gap if gap > 0 else float("nan")
        print(f"  Random→Oracle F1 gap: {gap:+.4f}")
        print(f"  Frozen LWM recovers {100*lwm_recovered:.1f}% of the gap")
        print(f"  LWM-CDE   recovers {100*cde_recovered:.1f}% of the gap")

    with open(out_json, "w") as f:
        json.dump(summary, f, indent=2)
    print(f"\n[Exp1] saved → {out_json}")


if __name__ == "__main__":
    main()