"""Post-hoc nested cross-fitted calibration and transport failure map."""

from __future__ import annotations

import csv
import json
import math
import random
import sys
from collections import defaultdict
from pathlib import Path
from typing import Iterable, Sequence

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.patches import Rectangle


HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parents[1] / "common"))

from plot_style import COLORS, configure, label_panels, save_figure

RESULTS = HERE / "results"
FIGURES = HERE / "figures"
METHODS = ("trace", "semantic")
BOOTSTRAP_REPLICATES = 10_000
BOOTSTRAP_SEED = 20260804


def read_predictions(path: Path) -> list[dict[str, object]]:
    rows: list[dict[str, object]] = []
    with path.open(encoding="utf-8", newline="") as handle:
        for source in csv.DictReader(handle):
            if source["cohort"] != "official" or source["method"] not in METHODS:
                continue
            rows.append(
                {
                    "row_id": source["row_id"],
                    "group_id": source["group_id"],
                    "model_id": source["model_id"],
                    "domain": source["domain"],
                    "task_fold": int(source["task_fold"]),
                    "success": int(source["success"]),
                    "method": source["method"],
                    "probability": float(source["probability"]),
                }
            )
    return rows


def read_frozen_primary_contrast(path: Path) -> tuple[float, float, float]:
    with path.open(encoding="utf-8", newline="") as handle:
        matches = [
            row
            for row in csv.DictReader(handle)
            if row["cohort"] == "official"
            and row["contrast"] == "semantic_minus_trace"
            and row["metric"] == "brier_difference"
        ]
    if len(matches) != 1:
        raise ValueError("expected one frozen official semantic-minus-trace Brier contrast")
    row = matches[0]
    return float(row["difference"]), float(row["ci_lower"]), float(row["ci_upper"])


def _logit(probability: float) -> float:
    clipped = min(max(probability, 1e-6), 1.0 - 1e-6)
    return math.log(clipped / (1.0 - clipped))


def _sigmoid(value: np.ndarray) -> np.ndarray:
    output = np.empty_like(value, dtype=float)
    positive = value >= 0
    output[positive] = 1.0 / (1.0 + np.exp(-value[positive]))
    exp_value = np.exp(value[~positive])
    output[~positive] = exp_value / (1.0 + exp_value)
    return output


def fit_logit_calibrator(
    probabilities: Sequence[float], labels: Sequence[int], ridge: float = 1e-4
) -> tuple[float, float]:
    if len(probabilities) != len(labels) or len(probabilities) < 10:
        raise ValueError("calibration requires aligned samples and at least ten rows")
    x = np.column_stack(
        [np.ones(len(probabilities)), np.array([_logit(value) for value in probabilities])]
    )
    y = np.asarray(labels, dtype=float)
    beta = np.array([0.0, 1.0], dtype=float)
    penalty = np.diag([0.0, ridge])
    for _ in range(100):
        fitted = _sigmoid(x @ beta)
        weights = np.clip(fitted * (1.0 - fitted), 1e-8, None)
        gradient = x.T @ (y - fitted) - penalty @ beta
        information = x.T @ (weights[:, None] * x) + penalty
        step = np.linalg.solve(information, gradient)
        beta += step
        if float(np.max(np.abs(step))) < 1e-10:
            break
    return float(beta[0]), float(beta[1])


def apply_calibrator(probability: float, intercept: float, slope: float) -> float:
    value = intercept + slope * _logit(probability)
    if value >= 0:
        return 1.0 / (1.0 + math.exp(-value))
    exp_value = math.exp(value)
    return exp_value / (1.0 + exp_value)


def nested_crossfit(rows: Sequence[dict[str, object]]) -> tuple[list[dict[str, object]], list[dict[str, object]]]:
    output: list[dict[str, object]] = []
    parameters: list[dict[str, object]] = []
    models = sorted({str(row["model_id"]) for row in rows})
    folds = sorted({int(row["task_fold"]) for row in rows})
    for method in METHODS:
        method_rows = [row for row in rows if row["method"] == method]
        for held_model in models:
            for held_fold in folds:
                training = [
                    row
                    for row in method_rows
                    if row["model_id"] != held_model and int(row["task_fold"]) != held_fold
                ]
                testing = [
                    row
                    for row in method_rows
                    if row["model_id"] == held_model and int(row["task_fold"]) == held_fold
                ]
                model_overlap = sum(row["model_id"] == held_model for row in training)
                task_fold_overlap = sum(int(row["task_fold"]) == held_fold for row in training)
                if model_overlap or task_fold_overlap:
                    raise AssertionError("nested calibration training overlaps a held axis")
                intercept, slope = fit_logit_calibrator(
                    [float(row["probability"]) for row in training],
                    [int(row["success"]) for row in training],
                )
                parameters.append(
                    {
                        "method": method,
                        "held_model": held_model,
                        "held_task_fold": held_fold,
                        "training_rows": len(training),
                        "test_rows": len(testing),
                        "training_model_overlap": model_overlap,
                        "training_task_fold_overlap": task_fold_overlap,
                        "intercept": f"{intercept:.12f}",
                        "slope": f"{slope:.12f}",
                    }
                )
                for row in testing:
                    output.append(
                        {
                            **row,
                            "recalibrated_probability": apply_calibrator(
                                float(row["probability"]), intercept, slope
                            ),
                        }
                    )
    output.sort(key=lambda row: (str(row["row_id"]), str(row["method"])))
    return output, parameters


def mean(values: Sequence[float]) -> float:
    return sum(values) / len(values)


def brier(rows: Sequence[dict[str, object]], probability_key: str) -> float:
    return mean(
        [
            (float(row[probability_key]) - int(row["success"])) ** 2
            for row in rows
        ]
    )


def calibration_gap(rows: Sequence[dict[str, object]], probability_key: str) -> float:
    return mean([float(row[probability_key]) for row in rows]) - mean(
        [float(row["success"]) for row in rows]
    )


def paired_difference(rows: Sequence[dict[str, object]], probability_key: str) -> float:
    by_key = {(str(row["row_id"]), str(row["method"])): row for row in rows}
    row_ids = sorted({str(row["row_id"]) for row in rows})
    differences = []
    for row_id in row_ids:
        semantic = by_key[(row_id, "semantic")]
        trace = by_key[(row_id, "trace")]
        label = int(trace["success"])
        differences.append(
            (float(semantic[probability_key]) - label) ** 2
            - (float(trace[probability_key]) - label) ** 2
        )
    return mean(differences)


def task_bootstrap_difference(
    rows: Sequence[dict[str, object]], probability_key: str, replicates: int, seed: int
) -> tuple[float, float, float]:
    by_row = {(str(row["row_id"]), str(row["method"])): row for row in rows}
    by_group: dict[str, list[float]] = defaultdict(list)
    for row_id in sorted({str(row["row_id"]) for row in rows}):
        trace = by_row[(row_id, "trace")]
        semantic = by_row[(row_id, "semantic")]
        label = int(trace["success"])
        difference = (
            (float(semantic[probability_key]) - label) ** 2
            - (float(trace[probability_key]) - label) ** 2
        )
        by_group[str(trace["group_id"])].append(difference)
    groups = sorted(by_group)
    rng = random.Random(seed)
    draws = []
    for _ in range(replicates):
        sampled = [groups[rng.randrange(len(groups))] for _ in groups]
        values = [value for group in sampled for value in by_group[group]]
        draws.append(mean(values))
    draws.sort()

    def percentile(probability: float) -> float:
        position = (len(draws) - 1) * probability
        lower = math.floor(position)
        upper = math.ceil(position)
        if lower == upper:
            return draws[lower]
        weight = position - lower
        return draws[lower] * (1.0 - weight) + draws[upper] * weight

    return paired_difference(rows, probability_key), percentile(0.025), percentile(0.975)


def reliability_bins(
    rows: Sequence[dict[str, object]], probability_key: str, bins: int = 10
) -> list[dict[str, object]]:
    ordered = sorted(rows, key=lambda row: (float(row[probability_key]), str(row["row_id"])))
    output = []
    for bin_index in range(bins):
        start = len(ordered) * bin_index // bins
        end = len(ordered) * (bin_index + 1) // bins
        selected = ordered[start:end]
        output.append(
            {
                "bin": bin_index + 1,
                "rows": len(selected),
                "mean_probability": mean([float(row[probability_key]) for row in selected]),
                "success_rate": mean([float(row["success"]) for row in selected]),
            }
        )
    return output


def write_csv(path: Path, rows: Iterable[dict[str, object]], fields: Sequence[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(fields), lineterminator="\n")
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    frozen = read_predictions(RESULTS / "crossfit_predictions.csv")
    recalibrated, parameters = nested_crossfit(frozen)
    expected = len({str(row["row_id"]) for row in frozen}) * len(METHODS)
    if len(recalibrated) != expected:
        raise AssertionError("every trajectory must receive one nested cross-fitted prediction per method")
    write_csv(
        RESULTS / "posthoc_recalibrated_predictions.csv",
        (
            {
                **row,
                "probability": f"{float(row['probability']):.12f}",
                "recalibrated_probability": f"{float(row['recalibrated_probability']):.12f}",
            }
            for row in recalibrated
        ),
        (
            "row_id",
            "group_id",
            "model_id",
            "domain",
            "task_fold",
            "success",
            "method",
            "probability",
            "recalibrated_probability",
        ),
    )
    write_csv(
        RESULTS / "posthoc_calibration_parameters.csv",
        parameters,
        (
            "method",
            "held_model",
            "held_task_fold",
            "training_rows",
            "test_rows",
            "training_model_overlap",
            "training_task_fold_overlap",
            "intercept",
            "slope",
        ),
    )

    cell_rows: list[dict[str, object]] = []
    models = sorted({str(row["model_id"]) for row in recalibrated})
    domains = sorted({str(row["domain"]) for row in recalibrated})
    for model in models:
        for domain in domains:
            cell = [
                row
                for row in recalibrated
                if row["model_id"] == model and row["domain"] == domain
            ]
            for probability_key, stage in (
                ("probability", "frozen"),
                ("recalibrated_probability", "nested_recalibrated"),
            ):
                trace_rows = [row for row in cell if row["method"] == "trace"]
                semantic_rows = [row for row in cell if row["method"] == "semantic"]
                difference, lower, upper = task_bootstrap_difference(
                    cell,
                    probability_key,
                    replicates=2_000,
                    seed=BOOTSTRAP_SEED + 100 * models.index(model) + 10 * domains.index(domain) + (stage == "nested_recalibrated"),
                )
                cell_rows.append(
                    {
                        "model_id": model,
                        "domain": domain,
                        "stage": stage,
                        "trajectories": len(trace_rows),
                        "task_clusters": len({str(row["group_id"]) for row in trace_rows}),
                        "success_rate": f"{mean([float(row['success']) for row in trace_rows]):.9f}",
                        "trace_brier": f"{brier(trace_rows, probability_key):.9f}",
                        "semantic_brier": f"{brier(semantic_rows, probability_key):.9f}",
                        "semantic_minus_trace_brier": f"{difference:.9f}",
                        "contrast_ci_lower": f"{lower:.9f}",
                        "contrast_ci_upper": f"{upper:.9f}",
                        "trace_calibration_gap": f"{calibration_gap(trace_rows, probability_key):.9f}",
                        "semantic_calibration_gap": f"{calibration_gap(semantic_rows, probability_key):.9f}",
                    }
                )
    write_csv(
        RESULTS / "posthoc_transport_cells.csv",
        cell_rows,
        (
            "model_id",
            "domain",
            "stage",
            "trajectories",
            "task_clusters",
            "success_rate",
            "trace_brier",
            "semantic_brier",
            "semantic_minus_trace_brier",
            "contrast_ci_lower",
            "contrast_ci_upper",
            "trace_calibration_gap",
            "semantic_calibration_gap",
        ),
    )

    frozen_difference, frozen_lower, frozen_upper = read_frozen_primary_contrast(
        RESULTS / "paired_contrasts.csv"
    )
    global_results = {}
    for probability_key, stage in (
        ("probability", "frozen"),
        ("recalibrated_probability", "nested_recalibrated"),
    ):
        if stage == "frozen":
            difference = paired_difference(recalibrated, probability_key)
            if not math.isclose(difference, frozen_difference, rel_tol=0.0, abs_tol=1e-12):
                raise AssertionError("frozen predictions do not reproduce the primary contrast")
            lower, upper = frozen_lower, frozen_upper
            interval_source = "frozen primary paired_contrasts.csv; 10,000 task-cluster draws; seed 20260803"
        else:
            difference, lower, upper = task_bootstrap_difference(
                recalibrated,
                probability_key,
                BOOTSTRAP_REPLICATES,
                BOOTSTRAP_SEED + 1,
            )
            interval_source = "post-hoc nested calibration; 10,000 task-cluster draws; seed 20260805"
        global_results[stage] = {
            "trace_brier": brier([row for row in recalibrated if row["method"] == "trace"], probability_key),
            "semantic_brier": brier([row for row in recalibrated if row["method"] == "semantic"], probability_key),
            "semantic_minus_trace_brier": difference,
            "task_bootstrap_95_interval": [lower, upper],
            "interval_source": interval_source,
        }
    calibrated_cells = [row for row in cell_rows if row["stage"] == "nested_recalibrated"]
    summary = {
        "analysis_role": "post-hoc secondary nested calibration and exhaustive model-by-domain failure map",
        "crossfit_guard": "calibrator training excludes the held agent model and held task fold",
        "counts": {
            "trajectories": len({str(row["row_id"]) for row in recalibrated}),
            "model_domain_cells": len(calibrated_cells),
            "calibrator_fits": len(parameters),
        },
        "global": global_results,
        "calibrated_cells_semantic_better": sum(
            float(row["semantic_minus_trace_brier"]) < 0 for row in calibrated_cells
        ),
        "calibrated_cells_interval_excludes_zero": sum(
            float(row["contrast_ci_lower"]) > 0 or float(row["contrast_ci_upper"]) < 0
            for row in calibrated_cells
        ),
        "limitations": [
            "The calibrator and cell map were selected after the frozen primary analysis.",
            "Cells are the exhaustive four-model by three-domain grid, not a searched subgroup list.",
            "Nested recalibration uses released outcomes from other task folds and other agent models; it is not external validation.",
        ],
    }
    (RESULTS / "posthoc_transport_summary.json").write_text(
        json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8", newline="\n"
    )

    configure()
    fig, axes = plt.subplots(
        1,
        3,
        figsize=(7.09, 2.65),
        constrained_layout=False,
        gridspec_kw={"width_ratios": [0.82, 1, 1]},
    )
    fig.subplots_adjust(left=0.075, right=0.91, bottom=0.22, top=0.84, wspace=0.34)
    stages = ["Frozen", "Nested\ncalibrated"]
    method_styles = (
        ("trace", COLORS["blue"], "o", "-", "Trace"),
        ("semantic", COLORS["vermillion"], "s", "--", "Trace + semantic"),
    )
    for method, color, marker, linestyle, label in method_styles:
        values = [global_results["frozen"][f"{method}_brier"], global_results["nested_recalibrated"][f"{method}_brier"]]
        axes[0].plot(
            [0, 1],
            values,
            color=color,
            marker=marker,
            linestyle=linestyle,
            markerfacecolor="white" if method == "semantic" else color,
            markeredgecolor=color,
            label=label,
        )
    axes[0].set_xticks([0, 1], stages)
    axes[0].set_xlim(-0.2, 1.2)
    axes[0].set_ylim(0.128, 0.142)
    axes[0].set_ylabel("Brier score")
    axes[0].set_title("Global calibration", loc="left")
    axes[0].legend(loc="upper right")

    display_models = ["claude-sonnet-4-5", "glm-5", "gpt-5-2-high", "qwen3-5-397b"]
    display_domains = ["airline", "retail", "telecom"]
    bound = max(abs(float(row["semantic_minus_trace_brier"])) for row in cell_rows)
    cmap = LinearSegmentedColormap.from_list(
        "contrast_diverging", [COLORS["blue"], "white", COLORS["vermillion"]]
    )
    images = []
    for axis, stage, title in zip(
        axes[1:],
        ("frozen", "nested_recalibrated"),
        ("Frozen cell contrast", "Nested-calibrated cell contrast"),
    ):
        matrix = np.array(
            [
                [
                    float(
                        next(
                            row["semantic_minus_trace_brier"]
                            for row in cell_rows
                            if row["stage"] == stage and row["model_id"] == model and row["domain"] == domain
                        )
                    )
                    for domain in display_domains
                ]
                for model in display_models
            ]
        )
        image = axis.imshow(matrix, cmap=cmap, vmin=-bound, vmax=bound, aspect="auto")
        images.append(image)
        for row_index in range(matrix.shape[0]):
            for column_index in range(matrix.shape[1]):
                value = matrix[row_index, column_index]
                text_color = "white" if abs(value) > 0.55 * bound else COLORS["ink"]
                axis.text(
                    column_index,
                    row_index,
                    f"{value:+.3f}",
                    ha="center",
                    va="center",
                    fontsize=5.8,
                    color=text_color,
                )
                cell = next(
                    row for row in cell_rows
                    if row["stage"] == stage and row["model_id"] == display_models[row_index] and row["domain"] == display_domains[column_index]
                )
                if float(cell["contrast_ci_lower"]) > 0 or float(cell["contrast_ci_upper"]) < 0:
                    axis.add_patch(
                        Rectangle(
                            (column_index - 0.47, row_index - 0.47),
                            0.94,
                            0.94,
                            fill=False,
                            edgecolor=COLORS["ink"],
                            linewidth=0.7,
                        )
                    )
        axis.set_xticks(range(len(display_domains)), display_domains)
        axis.set_yticks(range(len(display_models)), ["Claude", "GLM", "GPT", "Qwen"])
        axis.set_title(title, loc="left")
        axis.set_xlabel("Held-out domain")
        axis.tick_params(length=0)
        axis.spines[:].set_visible(False)
    colorbar = fig.colorbar(images[-1], ax=axes[1:], shrink=0.82, pad=0.025)
    colorbar.set_label("Δ Brier (semantic − trace)")
    colorbar.outline.set_edgecolor(COLORS["gray"])
    colorbar.outline.set_linewidth(0.6)
    label_panels(axes)
    save_figure(fig, FIGURES, "posthoc_transport_calibration")
    print(json.dumps(summary, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()
