#!/usr/bin/env python3
import argparse, json, os, platform, socket, subprocess, time
from pathlib import Path
import numpy as np
import sys
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from spectra_rsi import SyntheticWorld, SpectraConfig, SpectraRSILoop
from spectra_rsi.metrics import support_f1, normalized_delta_error, regression_recall

def gpu_info():
    info = {
        "cuda_visible_devices": os.getenv("CUDA_VISIBLE_DEVICES"),
        "hostname": socket.gethostname(),
        "platform": platform.platform(),
    }

    try:
        import torch
        info.update(
            torch_version=torch.__version__,
            cuda_available=torch.cuda.is_available(),
            gpu_count=torch.cuda.device_count(),
        )
        if torch.cuda.is_available():
            info["gpus"] = [
                torch.cuda.get_device_name(i)
                for i in range(torch.cuda.device_count())
            ]
    except Exception as e:
        info["torch_probe_error"] = str(e)

    # Lightweight fallback for NVIDIA systems where PyTorch is not installed.
    if not info.get("gpus"):
        try:
            out = subprocess.check_output(
                [
                    "nvidia-smi",
                    "--query-gpu=name,driver_version",
                    "--format=csv,noheader",
                ],
                text=True,
                stderr=subprocess.DEVNULL,
                timeout=5,
            )
            rows = [line.strip() for line in out.splitlines() if line.strip()]
            if rows:
                names = []
                drivers = []
                for row in rows:
                    parts = [x.strip() for x in row.split(",", 1)]
                    names.append(parts[0])
                    if len(parts) > 1:
                        drivers.append(parts[1])

                info["gpu_count"] = len(names)
                info["gpus"] = names
                info["nvidia_driver"] = drivers[0] if drivers else None
                info["gpu_probe"] = "nvidia-smi"
        except Exception as e:
            info["nvidia_smi_probe_error"] = str(e)

    return info

def main():
    p=argparse.ArgumentParser(description="SPECTRA-RSI reproducible scaling benchmark")
    p.add_argument("--n-slices",type=int,default=400); p.add_argument("--n-experts",type=int,default=16)
    p.add_argument("--rank",type=int,default=2); p.add_argument("--m-coarse",type=int,default=80); p.add_argument("--m-focused",type=int,default=120)
    p.add_argument("--items-per-row",type=int,default=600); p.add_argument("--bootstrap-reps",type=int,default=30)
    p.add_argument("--lambda-l1",type=float,default=2e-3); p.add_argument("--lambda-group",type=float,default=8e-3)
    p.add_argument("--seed",type=int,default=7); p.add_argument("--candidate-seed",type=int,default=99)
    p.add_argument("--candidate",choices=["single_gain","single_regression","canceling_mixture","broad_noncompressible","off_dictionary"],default="single_gain")
    p.add_argument("--scale",type=float,default=0.4); p.add_argument("--output",required=True)
    a=p.parse_args(); Path(a.output).parent.mkdir(parents=True,exist_ok=True)
    cfg=SpectraConfig(n_slices=a.n_slices,n_experts=a.n_experts,rank_per_expert=a.rank,m_coarse=a.m_coarse,m_focused=a.m_focused,items_per_row=a.items_per_row,bootstrap_reps=a.bootstrap_reps,
                      lambda_l1=a.lambda_l1,lambda_group=a.lambda_group,
                      seed=a.seed,audit_dir=str(Path(a.output).parent/"audit_logs"))
    world=SyntheticWorld(cfg.n_slices,cfg.n_experts,cfg.rank_per_expert,seed=cfg.seed,offband_leak=0.01)
    rng=np.random.default_rng(a.candidate_seed)
    ex={"single_gain":[max(0,a.n_experts//3)],"single_regression":[max(0,2*a.n_experts//3)],"canceling_mixture":[max(0,a.n_experts//4),max(0,3*a.n_experts//4)]}.get(a.candidate)
    kw={"scale":a.scale,"rng":rng};
    if ex is not None: kw["experts"]=ex
    cand=world.make_candidate(a.candidate,**kw)
    t=time.perf_counter(); rep=SpectraRSILoop(world,cfg).run_iteration(cand); wall=time.perf_counter()-t
    row={"benchmark":"spectra-rsi-scaling-v1","candidate":a.candidate,"config":vars(a),"system":gpu_info(),"wall_seconds":wall,"dense_fallback":bool(rep.dense_fallback),"pilot_residual":float(rep.pilot_residual),"items_probe":int(rep.items_probe),"items_sense":int(rep.items_sense),"items_anchor":int(rep.items_anchor)}
    if not rep.dense_fallback:
        truth=world.true_delta(cand); row.update(gate_decision=rep.gate_decision,support_f1=float(support_f1(rep.recovered_experts,cand.true_support)),normalized_delta_error=float(normalized_delta_error(rep.delta_hat,truth)),regression_recall=float(regression_recall(rep.delta_hat,truth,cfg.tau_margin)),true_experts=sorted(map(int,cand.true_support)),
                   nominated_experts=sorted(map(int,rep.nominated_experts)),
                   recovered_experts=sorted(map(int,rep.recovered_experts)),
                   bootstrap_frequencies=[float(x) for x in rep.bootstrap_frequencies])
    Path(a.output).write_text(json.dumps(row,indent=2)+"\n"); print(json.dumps(row,indent=2))
if __name__=="__main__": main()