import argparse
import csv
import json
from pathlib import Path
from typing import Any, Dict, List, Optional


def _get_by_path(d: Dict[str, Any], path: str) -> Any:
    cur: Any = d
    for p in str(path).split("."):
        if not isinstance(cur, dict) or p not in cur:
            return None
        cur = cur[p]
    return cur


def _as_str(x: Any) -> Optional[str]:
    if x is None:
        return None
    return str(x)


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--out_root", required=True)
    p.add_argument("--out_csv", required=True)
    args = p.parse_args()

    out_root = Path(args.out_root).resolve()
    rows: List[Dict[str, Any]] = []

    for run_dir in sorted([p for p in out_root.iterdir() if p.is_dir()]):
        summary_path = run_dir / "summary.json"
        if not summary_path.exists():
            continue

        data = json.loads(summary_path.read_text())
        cfg = data.get("config", {})
        summ = data.get("summary", {})

        row: Dict[str, Any] = {}
        row["run_dir"] = run_dir.name
        row["seed"] = _get_by_path(cfg, "seed")
        row["graph_type"] = _get_by_path(cfg, "graph.type")
        row["graph_n"] = _get_by_path(cfg, "graph.n")

        row["graph_sizes"] = _get_by_path(cfg, "graph.sizes")
        row["graph_shuffle_nodes"] = _get_by_path(cfg, "graph.shuffle_nodes")
        row["graph_p_in"] = _get_by_path(cfg, "graph.p_in")
        row["graph_p_out"] = _get_by_path(cfg, "graph.p_out")
        row["graph_sbm_avg_degree"] = _get_by_path(cfg, "graph.sbm_avg_degree")
        row["graph_sbm_mu"] = _get_by_path(cfg, "graph.sbm_mu")

        row["control_steps"] = _get_by_path(cfg, "run.control_steps")
        row["window"] = _get_by_path(cfg, "run.window")

        row["q_edge_sample_prob"] = _get_by_path(cfg, "observation.edge_sample_prob")
        row["noise_type"] = _get_by_path(cfg, "observation.count_noise.type")
        row["noise_sigma"] = _get_by_path(cfg, "observation.count_noise.sigma")

        row["mu_min"] = _get_by_path(cfg, "ethics.mu_min")
        row["eta_max"] = _get_by_path(cfg, "ethics.eta_max")
        row["v_max"] = _get_by_path(cfg, "ethics.v_max")
        row["mu_disp_max"] = _get_by_path(cfg, "ethics.mu_disp_max")
        row["mu_disp_mode"] = _get_by_path(cfg, "ethics.mu_disp_mode")

        row["meso_k_mean"] = _get_by_path(summ, "meso.k.mean")
        row["meso_k_min"] = _get_by_path(summ, "meso.k.min")
        row["meso_k_max"] = _get_by_path(summ, "meso.k.max")
        row["meso_k_std"] = _get_by_path(summ, "meso.k.std")
        row["meso_mu_disp_nonzero_fraction"] = _get_by_path(summ, "meso.mu_disp.nonzero_fraction")
        row["meso_mu_disp_violation_fraction"] = _get_by_path(summ, "meso.mu_disp.violation_fraction")

        row["top_act_enabled"] = _get_by_path(cfg, "topology_actuator.enabled")
        row["top_act_type"] = _get_by_path(cfg, "topology_actuator.type")
        row["top_act_cooldown_steps"] = _get_by_path(cfg, "topology_actuator.cooldown_steps")
        row["top_act_budget_max"] = _get_by_path(cfg, "topology_actuator.budget_max")

        row["shock_count"] = len(_get_by_path(cfg, "stress.shocks") or [])
        row["rewire_count"] = len(_get_by_path(cfg, "stress.rewires") or [])

        row["graph_m"] = _get_by_path(summ, "graph.m")
        row["graph_avg_degree"] = _get_by_path(summ, "graph.avg_degree")

        row["time_in_e"] = _get_by_path(summ, "time_in_e")
        row["det_brier"] = _get_by_path(summ, "detector.brier")
        row["det_log_loss"] = _get_by_path(summ, "detector.log_loss")
        row["det_brier_mu"] = _get_by_path(summ, "detector.brier_mu")
        row["det_brier_eta"] = _get_by_path(summ, "detector.brier_eta")
        row["det_brier_v"] = _get_by_path(summ, "detector.brier_v")
        row["det_brier_disp"] = _get_by_path(summ, "detector.brier_disp")

        row["u_mean"] = _get_by_path(summ, "intervention.mean")
        row["u_sum"] = _get_by_path(summ, "intervention.sum")
        row["u_du_abs_sum"] = _get_by_path(summ, "intervention.du_abs_sum")

        shock_recs = _get_by_path(summ, "recovery.shocks") or []
        rewire_recs = _get_by_path(summ, "recovery.rewires") or []
        row["shock_recovery_first"] = _get_by_path(shock_recs[0], "recovery") if shock_recs else None
        row["rewire_recovery_first"] = _get_by_path(rewire_recs[0], "recovery") if rewire_recs else None

        top_s = _get_by_path(summ, "topology_actuator")
        row["top_events_count"] = _get_by_path(top_s, "events.count") if isinstance(top_s, dict) else None
        row["top_rewired_total"] = _get_by_path(top_s, "events.rewired_total") if isinstance(top_s, dict) else None
        row["top_budget_total"] = _get_by_path(top_s, "events.budget_total") if isinstance(top_s, dict) else None
        row["top_fraction_steps_active"] = _get_by_path(top_s, "fraction_steps_active") if isinstance(top_s, dict) else None

        stuck_s = _get_by_path(summ, "diagnostics.stuck")
        row["stuck_fraction"] = _get_by_path(stuck_s, "fraction") if isinstance(stuck_s, dict) else None
        row["u_saturation_fraction"] = _get_by_path(summ, "diagnostics.u_saturation.fraction")

        rows.append(row)

    out_csv = Path(args.out_csv).resolve()
    out_csv.parent.mkdir(parents=True, exist_ok=True)

    fieldnames = sorted({k for r in rows for k in r.keys()})
    with out_csv.open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames)
        w.writeheader()
        for r in rows:
            w.writerow({k: _as_str(r.get(k)) for k in fieldnames})

    print(f"WROTE {len(rows)} rows to {out_csv}")


if __name__ == "__main__":
    main()