"""Recompute the reviewer-requested paired bootstrap table from frozen predictions.

The identity-grouped task resamples molecular-connectivity groups.  The
scaffold-aware task resamples the connected scaffold components used to form
the held-out partition.  No model is refitted and no prediction is changed.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score


TASKS = {
    "canonical_grouped": ("Identity-grouped", "structure_group"),
    "scaffold_aware": ("Scaffold-aware", "scaffold_component_group"),
}
REFERENCES = {
    "FPNN": "prediction_fpnn",
    "GAT": "prediction_gat",
    "GCN": "prediction_gcn",
}
CANDIDATE = "prediction_stack_all_plus_descriptors"
METRICS = ("mae", "rmse", "r2")


def metric(name: str, y_true: np.ndarray, y_pred: np.ndarray) -> float:
    if name == "mae":
        return float(mean_absolute_error(y_true, y_pred))
    if name == "rmse":
        return float(np.sqrt(mean_squared_error(y_true, y_pred)))
    if name == "r2":
        return float(r2_score(y_true, y_pred))
    raise ValueError(name)


def paired_bootstrap(
    y_true: np.ndarray,
    candidate: np.ndarray,
    reference: np.ndarray,
    groups: np.ndarray,
    *,
    n_resamples: int,
    seed: int,
) -> dict[str, dict[str, float | int]]:
    unique_groups = np.unique(groups.astype(str))
    positions = {group: np.flatnonzero(groups == group) for group in unique_groups}
    rng = np.random.default_rng(seed)
    samples = {name: [] for name in METRICS}

    for _ in range(n_resamples):
        sampled_groups = rng.choice(unique_groups, size=len(unique_groups), replace=True)
        sampled_positions = np.concatenate([positions[group] for group in sampled_groups])
        for name in METRICS:
            delta = metric(name, y_true[sampled_positions], candidate[sampled_positions])
            delta -= metric(name, y_true[sampled_positions], reference[sampled_positions])
            if np.isfinite(delta):
                samples[name].append(delta)

    result: dict[str, dict[str, float | int]] = {}
    for name in METRICS:
        point = metric(name, y_true, candidate) - metric(name, y_true, reference)
        low, high = np.quantile(np.asarray(samples[name]), [0.025, 0.975])
        result[name] = {
            "difference_point": float(point),
            "ci_low": float(low),
            "ci_high": float(high),
            "n_resampling_units": int(len(unique_groups)),
            "n_valid_resamples": int(len(samples[name])),
        }
    return result


def fmt(value: float, low: float, high: float, digits: int) -> str:
    return f"{value:.{digits}f} [{low:.{digits}f}, {high:.{digits}f}]"


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--artifacts", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--table", type=Path, required=True)
    parser.add_argument("--n-resamples", type=int, default=2000)
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)

    nested: dict[str, dict[str, object]] = {}
    rows: list[dict[str, object]] = []
    for task_dir, (task_label, group_column) in TASKS.items():
        nested[task_dir] = {}
        for repeat, seed in enumerate((123456, 123457, 123458), start=1):
            path = args.artifacts / task_dir / f"seed_{seed}" / "neural_stack" / "test_predictions.csv"
            frame = pd.read_csv(path)
            y_true = frame["RT"].to_numpy(dtype=float)
            candidate = frame[CANDIDATE].to_numpy(dtype=float)
            groups = frame[group_column].astype(str).to_numpy()
            repeat_result: dict[str, object] = {
                "seed": seed,
                "resampling_unit": group_column,
                "n_rows": int(len(frame)),
                "n_resampling_units": int(pd.Series(groups).nunique()),
                "comparisons": {},
            }
            for reference_label, reference_column in REFERENCES.items():
                result = paired_bootstrap(
                    y_true,
                    candidate,
                    frame[reference_column].to_numpy(dtype=float),
                    groups,
                    n_resamples=args.n_resamples,
                    seed=seed,
                )
                repeat_result["comparisons"][reference_label] = result
                row: dict[str, object] = {
                    "task": task_label,
                    "resampling_unit": group_column,
                    "reference": reference_label,
                    "repeat": repeat,
                    "seed": seed,
                    "n_rows": len(frame),
                    "n_resampling_units": pd.Series(groups).nunique(),
                }
                for name in METRICS:
                    for key, value in result[name].items():
                        row[f"{name}_{key}"] = value
                rows.append(row)
            nested[task_dir][f"seed_{seed}"] = repeat_result

    (args.output_dir / "paired_bootstrap_corrected.json").write_text(
        json.dumps(nested, indent=2), encoding="utf-8"
    )
    frame = pd.DataFrame(rows)
    frame.to_csv(args.output_dir / "paired_bootstrap_corrected.csv", index=False)

    latex = [
        r"\begin{table*}[htbp]",
        r"\centering",
        r"\small",
        r"\caption{Paired cluster-bootstrap differences for the full stack relative to each neural base learner. Identity-grouped repeats resample molecular-connectivity groups; scaffold-aware repeats resample scaffold components. Each interval uses 2,000 resamples. Negative $\Delta$MAE and $\Delta$RMSE and positive $\Delta R^2$ favor the full stack.}",
        r"\label{tab:paired-bootstrap}",
        r"\begin{tabular}{lllccc}",
        r"\toprule",
        r"Task & Reference & Repeat & $\Delta$MAE (95\% interval), min & $\Delta$RMSE (95\% interval), min & $\Delta R^2$ (95\% interval) \\",
        r"\midrule",
    ]
    for task_label in ("Identity-grouped", "Scaffold-aware"):
        task_rows = frame.loc[frame["task"] == task_label]
        for reference_label in REFERENCES:
            for _, row in task_rows.loc[task_rows["reference"] == reference_label].iterrows():
                mae = fmt(row.mae_difference_point, row.mae_ci_low, row.mae_ci_high, 3)
                rmse = fmt(row.rmse_difference_point, row.rmse_ci_low, row.rmse_ci_high, 3)
                r2 = fmt(row.r2_difference_point, row.r2_ci_low, row.r2_ci_high, 3)
                latex.append(
                    f"{task_label} & {reference_label} & {int(row['repeat'])} & {mae} & {rmse} & {r2} \\\\"
                )
    latex.extend([r"\bottomrule", r"\end{tabular}", r"\end{table*}", ""])
    args.table.write_text("\n".join(latex), encoding="utf-8")


if __name__ == "__main__":
    main()