"""LaTeX tables generated from the estimates (no manual numeric editing)."""
import os, sys
sys.path.insert(0, os.path.dirname(__file__))
from common import *
TAB = os.path.join(REPO, "results", "tables"); os.makedirs(TAB, exist_ok=True)

def fmt(x, nd=3):
    if x is None or (isinstance(x, float) and np.isnan(x)): return "--"
    return f"{x:.{nd}f}" if isinstance(x, (float, np.floating)) else str(x)

def table_arms(split):
    at = pd.read_csv(os.path.join(EST, f"{split}_arm_table.csv"))
    lines = [r"\begin{tabular}{lcccccccc}", r"\toprule", r"Arm & Runs & Tasks & $\hat p_a$ & Correct (any) & Viol.\ rate & Repro. & Quality & Cost/run (USD) \\", r"\midrule"]
    for _, r in at.iterrows():
        lines.append(f"{r['arm']} & {int(r['n_runs'])} & {int(r['n_tasks'])} & {fmt(r['Y_task_weighted'])} & {fmt(r['correct_any_task_weighted'])} & {fmt(r['violation_rate'])} & {fmt(r['reproducible_rate'])} & {fmt(r['quality_mean'])} & {fmt(r['cost_mean_usd'],2)} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    open(os.path.join(TAB, f"tab_arms_{split}.tex"), "w").write("\n".join(lines) + "\n")

def table_contrasts(split):
    ct = pd.read_csv(os.path.join(EST, f"{split}_contrasts.csv"))
    lines = [r"\begin{tabular}{llcccccc}", r"\toprule", r"Endpoint & Contrast & Role & $J$ & $\hat\Delta$ & 95\% CI & Perm.\ $p$ & Holm $p$ \\", r"\midrule"]
    names = {"Y": "Primary $Y$", "Y_any": "Correct (any)", "viol": "Violation", "quality_score": "Quality"}
    for _, r in ct.iterrows():
        lines.append(f"{names.get(r['endpoint'], r['endpoint'])} & {r['contrast']} & {r['role']} & {int(r['n_task_clusters'])} & {fmt(r['estimate'])} & [{fmt(r['ci_low'])}, {fmt(r['ci_high'])}] & {fmt(r['perm_p'])} & {fmt(r['holm_p']) if not pd.isna(r['holm_p']) else '--'} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    open(os.path.join(TAB, f"tab_contrasts_{split}.tex"), "w").write("\n".join(lines) + "\n")

def table_cost(split):
    at = pd.read_csv(os.path.join(EST, f"{split}_arm_table.csv"))
    lines = [r"\begin{tabular}{lccccccccc}", r"\toprule", r"Arm & Mean & Median & P90 & Total & Reviewer & Wall med.\ (s) & Wall P90 & Timeouts & Cost/success \\", r"\midrule"]
    for _, r in at.iterrows():
        lines.append(f"{r['arm']} & {fmt(r['cost_mean_usd'],2)} & {fmt(r['cost_median_usd'],2)} & {fmt(r['cost_p90_usd'],2)} & {fmt(r['cost_total_usd'],2)} & {fmt(r['reviewer_usd_total'],2)} & {fmt(r['wall_median_s'],0)} & {fmt(r['wall_p90_s'],0)} & {int(r['n_timeout'])} & {fmt(r['cost_per_success_usd'],2) if not pd.isna(r['cost_per_success_usd']) else 'undefined'} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    open(os.path.join(TAB, f"tab_cost_{split}.tex"), "w").write("\n".join(lines) + "\n")

def table_family(split):
    fam = pd.read_csv(os.path.join(EST, f"{split}_family_table.csv")); fams = sorted(fam["family"].unique())
    lines = [r"\begin{tabular}{l" + "c" * (2 * len(ARMS) + 1) + "}", r"\toprule", "Family & Tasks & " + " & ".join(f"$\\hat p_{a}$ & viol.\\ {a}" for a in ARMS) + r" \\", r"\midrule"]
    for f in fams:
        d = fam[fam["family"] == f].set_index("treatment")
        lines.append(f"{f} & {int(d['n_tasks'].iloc[0])} & " + " & ".join(f"{fmt(d.loc[a, 'Y'],2) if a in d.index else '--'} & {fmt(d.loc[a, 'viol'],2) if a in d.index else '--'}" for a in ARMS) + r" \\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    open(os.path.join(TAB, f"tab_family_{split}.tex"), "w").write("\n".join(lines) + "\n")

def table_accounting():
    df = pd.read_csv(os.path.join(DATA, "runs.csv"))
    lines = [r"\begin{tabular}{lcccccccc}", r"\toprule", r"Split & Assigned & Launched & Completed & Timeout & Budget & Infra.\ error & Evaluated & Successes \\", r"\midrule"]
    for s in ["development", "pilot", "confirmatory", "transfer", "recovery", "ablation"]:
        d = df[df["split"] == s]
        if d.empty: continue
        tr = d["terminal_reason"].fillna("")
        lines.append(f"{s} & {len(d)} & {int((d['status'] != 'not_started').sum())} & {int(tr.str.startswith('executor_finished').sum() + tr.eq('max_turns_reached').sum())} & {int(tr.eq('wall_clock_limit').sum())} & {int(tr.str.contains('budget').sum())} & {int(tr.str.startswith('infrastructure').sum())} & {int((d['status'] == 'evaluated').sum())} & {int((d['scientific_success'] == 1).sum())} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    open(os.path.join(TAB, "tab_accounting.tex"), "w").write("\n".join(lines) + "\n")

def table_robustness(split):
    rb = pd.read_csv(os.path.join(EST, f"{split}_robustness.csv"))
    lines = [r"\begin{tabular}{lcccc}", r"\toprule", r"Analysis (D$-$A) & $J$ & $\hat\Delta$ & 95\% CI & Perm.\ $p$ \\", r"\midrule"]
    for _, r in rb.iterrows():
        lines.append(f"{r['analysis'].replace('_', ' ')} & {int(r['n_tasks'])} & {fmt(r['estimate'])} & [{fmt(r['ci_low'])}, {fmt(r['ci_high'])}] & {fmt(r['perm_p'])} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    open(os.path.join(TAB, "tab_robustness.tex"), "w").write("\n".join(lines) + "\n")

def table_order_node(split):
    on = pd.read_csv(os.path.join(EST, f"{split}_order_node.csv"))
    lines = [r"\begin{tabular}{llcccc}", r"\toprule", r"Factor & Level & A & B & C & D \\", r"\midrule"]
    for fac in ["order", "node"]:
        d = on[on.factor == fac]
        for lev in sorted(d["half"].unique()):
            dd = d[d["half"] == lev].set_index("treatment")
            lines.append(f"{fac} & {lev} & " + " & ".join(fmt(dd.loc[a, 'Y'], 2) if a in dd.index else "--" for a in ARMS) + r" \\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    open(os.path.join(TAB, "tab_order_node.tex"), "w").write("\n".join(lines) + "\n")

def table_claims(split):
    c = pd.read_csv(os.path.join(EST, f"{split}_claims_by_arm.csv"))
    lines = [r"\begin{tabular}{lcccccc}", r"\toprule", r"Arm & Runs & Claims/run & Supported/run & Unsupported ``supported''/run & Runs with $\geq$1 & Not assessable/run \\", r"\midrule"]
    for _, r in c.iterrows():
        lines.append(f"{r['treatment']} & {int(r['runs'])} & {fmt(r['claims_per_run'],2)} & {fmt(r['supported_per_run'],2)} & {fmt(r['unsupported_per_run'],3)} & {int(r['runs_with_unsupported'])} & {fmt(r['not_assessable_per_run'],2)} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    open(os.path.join(TAB, f"tab_claims_{split}.tex"), "w").write("\n".join(lines) + "\n")

def table_postfix(split=None):
    """Post-fix verification runs (amendment A-006): smoke task × 4 arms and T1-forecast arm C with the corrected launcher; never pooled."""
    r = pd.read_csv(os.path.join(DATA, "runs.csv")); r = r[r["notes"].fillna("").str.contains("seed 102|seed 402")]
    lines = [r"\begin{tabular}{llcccccc}", r"\toprule", r"Run & Task & Arm & Status & Success & Viol. & Sessions & USD \\", r"\midrule"]
    for _, x in r.sort_values("run_id").iterrows():
        lines.append(f"{x.run_id} & {x.task_id} & {x.treatment} & {str(x.terminal_reason or x.status).replace('_', ' ')[:20]} & {'yes' if x.scientific_success == 1 else ('no' if x.scientific_success == 0 else '--')} & {int(x.n_violations) if pd.notna(x.n_violations) else '--'} & {int(x.n_sessions) if pd.notna(x.n_sessions) else '--'} & {x.spent_usd:.2f} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    open(os.path.join(TAB, "tab_postfix.tex"), "w").write("\n".join(lines) + "\n")

if __name__ == "__main__":
    split = sys.argv[1] if len(sys.argv) > 1 else "confirmatory"
    for f in (table_arms, table_contrasts, table_cost, table_family, table_robustness, table_order_node, table_claims, table_postfix):
        try: f(split); print("ok", f.__name__)
        except Exception as e: print("skip", f.__name__, e)
    table_accounting()