#!/usr/bin/env python3
# ---------------------------------------------------------------
# Fixed-orthogonal-mask scheme across user counts U = 1..6, for
# overlaying the independent reference on Fig. 2 and Fig. 3 of the
# WCL revision.
#
# Reuses the EXACT training/eval path of eval_task_oriented.py
# (fixed QR masks, transceiver trained with MSE+CosSim, seed 0,
# matched 100-epoch x 200-step budget) so the U=4 numbers reproduce
# Table II. Evaluates the full SNR sweep per U so the same file
# serves Fig 3(a) (CosSim vs SNR, U=4), Fig 2(b) (U=4, 20 dB) and
# Fig 3(b) (throughput vs U at 10 dB).
#
# Output: ../results_sweeps/task_oriented/fixed_orth_byU.csv
# ---------------------------------------------------------------
import os, csv, argparse
from types import SimpleNamespace
import torch

from eval_task_oriented import (load_emb, train_transceiver, sweep_metrics,
                                build_fixed_orthogonal_masks)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--users-list", type=int, nargs="+",
                    default=[1, 2, 3, 4, 5, 6])
    ap.add_argument("--epochs", type=int, default=100)
    ap.add_argument("--steps-per-epoch", type=int, default=200)
    ap.add_argument("--eval-trials", type=int, default=200)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--embed-file", type=str, default="../bert_agnews_8000.pt")
    ap.add_argument("--out", type=str,
                    default="../results_sweeps/task_oriented/fixed_orth_byU.csv")
    a = ap.parse_args()

    device = torch.device("mps" if torch.backends.mps.is_available()
                          else ("cuda" if torch.cuda.is_available() else "cpu"))
    print(f"[INFO] device={device}")
    d_bert, mux = 768, 4
    d_s = d_bert * mux
    emb = load_emb(a.embed_file, d_bert, 8000, device)

    rows_all = []
    for U in a.users_list:
        print(f"\n=== Fixed-Orth scheme, U={U} ===", flush=True)
        args = SimpleNamespace(
            users=U, users_max=8, mux_factor=mux, d_bert=d_bert,
            hidden=256, rank=64, channel="rayleigh",
            train_snr=[0, 5, 10, 15, 20, 25], lr=1e-3, ce_tau=16.0,
            epochs=a.epochs, steps_per_epoch=a.steps_per_epoch,
            eval_trials=a.eval_trials, seed=a.seed)
        Mfix = build_fixed_orthogonal_masks(U, d_s, seed=0)
        trx, mfn = train_transceiver(emb, args, device, "fixed",
                                     fixed_masks=Mfix)
        rows = sweep_metrics(trx, emb, mfn, device, U, a.eval_trials)
        for r in rows:
            rows_all.append((U, *r))

    os.makedirs(os.path.dirname(a.out), exist_ok=True)
    with open(a.out, "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["users", "snr_db", "cos_sim", "orthogonality", "top1_acc"])
        for r in rows_all:
            w.writerow(r)
    print(f"\n[DONE] wrote {a.out}", flush=True)


if __name__ == "__main__":
    main()