source/src/bgc_retrieval/campaign_analysis.py

"""Reproducible, explicitly post-hoc analysis of completed campaign artifacts."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
from scipy import stats

from .artifacts import sha256_file, write_json_immutable
from .statistics import holm_adjust, paired_family_test


INTERNAL_METRICS = ["recall@50", "mrr", "map", "ndcg@50"]


def _aligned_edges(
    first: pd.DataFrame, second: pd.DataFrame
) -> tuple[pd.DataFrame, pd.DataFrame]:
    required = {"record_a", "record_b", "score", "structural_similarity"}
    for name, frame in (("first", first), ("second", second)):
        if missing := required.difference(frame.columns):
            raise ValueError(f"{name} edge table is missing columns: {sorted(missing)}")
    keys = ["record_a", "record_b"]
    left = first.sort_values(keys).reset_index(drop=True)
    right = second.sort_values(keys).reset_index(drop=True)
    if not left[keys].equals(right[keys]):
        raise ValueError("Paired methods must contain identical edges")
    if not np.allclose(left["structural_similarity"], right["structural_similarity"]):
        raise ValueError("Paired methods disagree on structural truth")
    return left, right


def paired_correlation_difference_bootstrap(
    first: pd.DataFrame,
    second: pd.DataFrame,
    samples: int = 10000,
    confidence: float = 0.95,
    seed: int = 0,
) -> dict[str, float | int]:
    """Paired two-endpoint BGC cluster bootstrap for a Spearman difference.

    Ranks are fixed on the complete paired edge set. Each dyad contributes half
    its sufficient-statistic weight to each endpoint block. Both methods use the
    same resampled endpoint multiplicities in every replicate.
    """
    if samples < 1:
        raise ValueError("Bootstrap samples must be positive")
    if not 0.0 < confidence < 1.0:
        raise ValueError("Confidence must be between zero and one")
    left, right = _aligned_edges(first, second)
    if len(left) < 3:
        raise ValueError("At least three paired edges are required")

    score_first = stats.rankdata(left["score"].to_numpy(float), method="average")
    score_second = stats.rankdata(right["score"].to_numpy(float), method="average")
    truth = stats.rankdata(left["structural_similarity"].to_numpy(float), method="average")
    identifiers = pd.Index(
        sorted(set(left["record_a"].astype(str)) | set(left["record_b"].astype(str)))
    )
    index_left = identifiers.get_indexer(left["record_a"].astype(str))
    index_right = identifiers.get_indexer(left["record_b"].astype(str))

    def block_sum(values: np.ndarray) -> np.ndarray:
        return 0.5 * (
            np.bincount(index_left, weights=values, minlength=len(identifiers))
            + np.bincount(index_right, weights=values, minlength=len(identifiers))
        )

    weights = 0.5 * (
        np.bincount(index_left, minlength=len(identifiers))
        + np.bincount(index_right, minlength=len(identifiers))
    )
    blocks = np.column_stack(
        [
            weights,
            block_sum(score_first),
            block_sum(score_second),
            block_sum(truth),
            block_sum(score_first**2),
            block_sum(score_second**2),
            block_sum(truth**2),
            block_sum(score_first * truth),
            block_sum(score_second * truth),
        ]
    )

    random_state = np.random.default_rng(seed)
    differences: list[float] = []
    for _ in range(samples):
        selected = random_state.integers(0, len(blocks), size=len(blocks))
        weight, sum_a, sum_b, sum_y, sum_aa, sum_bb, sum_yy, sum_ay, sum_by = (
            blocks[selected].sum(axis=0)
        )
        variance_y = sum_yy - sum_y * sum_y / weight

        def correlation(sum_x: float, sum_xx: float, sum_xy: float) -> float:
            covariance = sum_xy - sum_x * sum_y / weight
            variance_x = sum_xx - sum_x * sum_x / weight
            denominator = np.sqrt(max(variance_x, 0.0) * max(variance_y, 0.0))
            return covariance / denominator if denominator > 0.0 else float("nan")

        difference = correlation(sum_a, sum_aa, sum_ay) - correlation(
            sum_b, sum_bb, sum_by
        )
        if np.isfinite(difference):
            differences.append(float(difference))
    if not differences:
        raise ValueError("No finite paired bootstrap differences were produced")

    values = np.asarray(differences)
    tail = (1.0 - confidence) / 2.0
    lower, upper = np.quantile(values, [tail, 1.0 - tail])
    first_rho = float(stats.spearmanr(left["score"], left["structural_similarity"]).statistic)
    second_rho = float(stats.spearmanr(right["score"], right["structural_similarity"]).statistic)
    sign_probability = 2.0 * min(
        (np.count_nonzero(values <= 0.0) + 1) / (len(values) + 1),
        (np.count_nonzero(values >= 0.0) + 1) / (len(values) + 1),
    )
    return {
        "pairs": len(left),
        "bgcs": len(identifiers),
        "spearman_first": first_rho,
        "spearman_second": second_rho,
        "delta_spearman": first_rho - second_rho,
        "ci_lower": float(lower),
        "ci_upper": float(upper),
        "bootstrap_two_sided_sign_probability": float(min(sign_probability, 1.0)),
        "bootstrap_samples": len(values),
    }


def internal_paired_comparisons(main: pd.DataFrame, no_phase1: pd.DataFrame) -> pd.DataFrame:
    comparisons: list[dict[str, Any]] = []
    ensemble_methods = sorted(
        method
        for method in main["method"].astype(str).unique()
        if method.startswith("ensemble_validation_alpha_")
    )
    if len(ensemble_methods) != 1:
        raise ValueError(
            "Expected exactly one validation-selected ensemble method; "
            f"found {ensemble_methods}"
        )
    ensemble_method = ensemble_methods[0]
    families = [
        ("setnet", "raw_esm_mean", "main_vs_raw", main),
        ("setnet", "pfam_jaccard_max", "main_vs_pfam", main),
        (
            ensemble_method,
            "pfam_jaccard_max",
            "ensemble_vs_pfam",
            main,
        ),
    ]
    phase_ablation = pd.concat(
        [
            main[main["method"] == "setnet"].assign(method="setnet_phase1"),
            no_phase1[no_phase1["method"] == "setnet"].assign(
                method="setnet_no_phase1"
            ),
        ],
        ignore_index=True,
    )
    families.append(
        ("setnet_phase1", "setnet_no_phase1", "phase1_ablation", phase_ablation)
    )
    for method, baseline, family, frame in families:
        rows = [paired_family_test(frame, method, baseline, metric) for metric in INTERNAL_METRICS]
        adjusted = holm_adjust(row["p_value"] for row in rows)
        for row, corrected in zip(rows, adjusted):
            row.update(family=family, p_value_holm=corrected, analysis_status="post_hoc")
            comparisons.append(row)
    return pd.DataFrame(comparisons)


def _methods(path: Path) -> dict[str, pd.DataFrame]:
    frame = pd.read_csv(path)
    return {
        str(method): rows.drop(columns="method").reset_index(drop=True)
        for method, rows in frame.groupby("method", sort=False)
    }


def _training_summary(path: Path, metric: str, maximize: bool) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as handle:
        history = json.load(handle)
    best = (max if maximize else min)(history, key=lambda row: row[metric])
    return {"epochs": len(history), "best": best, "first": history[0], "last": history[-1]}


def analyze_campaign(
    artifact_root: str | Path,
    campaign_tag: str,
    output_dir: str | Path,
    bootstrap_samples: int = 10000,
    confidence: float = 0.95,
    seed: int = 20260810,
) -> Path:
    root = Path(artifact_root)
    output = Path(output_dir)
    output.mkdir(parents=True, exist_ok=False)
    main_internal = root / f"{campaign_tag}-main-evaluation/group_results.csv"
    no_internal = root / f"{campaign_tag}-no-phase1-evaluation/group_results.csv"
    main_external = root / f"{campaign_tag}-main-external/external_pair_scores.csv"
    no_external = root / f"{campaign_tag}-no-phase1-external/external_pair_scores.csv"
    required = [main_internal, no_internal, main_external, no_external]
    if missing := [str(path) for path in required if not path.is_file()]:
        raise FileNotFoundError(f"Campaign artifacts are missing: {missing}")

    internal = internal_paired_comparisons(
        pd.read_csv(main_internal), pd.read_csv(no_internal)
    )
    internal.to_csv(output / "internal_paired_comparisons.csv", index=False)
    main_methods = _methods(main_external)
    no_methods = _methods(no_external)
    external_rows: list[dict[str, Any]] = []
    comparisons = [
        ("main_setnet_vs_raw", main_methods["setnet"], main_methods["raw_esm_mean"]),
        (
            "main_setnet_vs_raw_bigscape_edges",
            main_methods["setnet_on_bigscape_edges"],
            main_methods["raw_esm_mean_on_bigscape_edges"],
        ),
        ("phase1_ablation", main_methods["setnet"], no_methods["setnet"]),
        (
            "phase1_ablation_bigscape_edges",
            main_methods["setnet_on_bigscape_edges"],
            no_methods["setnet_on_bigscape_edges"],
        ),
        (
            "bigscape_vs_main_setnet",
            main_methods["bigscape"],
            main_methods["setnet_on_bigscape_edges"],
        ),
    ]
    for name, first, second in comparisons:
        for subset_name, subset in (
            ("all", first),
            ("cross_genus", first[first["cross_genus"]]),
        ):
            keys = set(zip(subset["record_a"], subset["record_b"]))
            paired_second = second[
                [pair in keys for pair in zip(second["record_a"], second["record_b"])]
            ]
            result = paired_correlation_difference_bootstrap(
                subset,
                paired_second,
                samples=bootstrap_samples,
                confidence=confidence,
                seed=seed,
            )
            external_rows.append(
                {
                    "comparison": name,
                    "subset": subset_name,
                    **result,
                    "analysis_status": "post_hoc",
                }
            )
    pd.DataFrame(external_rows).to_csv(
        output / "external_paired_comparisons.csv", index=False
    )

    exact_rows: list[dict[str, Any]] = []
    for run_name in (f"{campaign_tag}-main-external", f"{campaign_tag}-no-phase1-external"):
        for method in ("setnet", "raw_esm_mean"):
            path = root / run_name / f"{method}_exact_product_retrieval.csv"
            frame = pd.read_csv(path)
            exact_rows.append(
                {
                    "run": run_name,
                    "method": method,
                    "references": len(frame),
                    "cross_genus_references": int(
                        (frame["cross_genus_positive_count"] > 0).sum()
                    ),
                    **{metric: float(frame[metric].mean()) for metric in INTERNAL_METRICS},
                    "precision@50": float(frame["precision@50"].mean()),
                }
            )
    pd.DataFrame(exact_rows).to_csv(output / "exact_product_summary.csv", index=False)

    training = {
        "phase1": _training_summary(
            root / f"{campaign_tag}-phase1/phase1_history.json",
            "validation_loss",
            maximize=False,
        ),
        "phase2_main": _training_summary(
            root / f"{campaign_tag}-main/phase2_history.json",
            "validation_recall@50",
            maximize=True,
        ),
        "phase2_no_phase1": _training_summary(
            root / f"{campaign_tag}-no-phase1/phase2_history.json",
            "validation_recall@50",
            maximize=True,
        ),
    }
    write_json_immutable(output / "training_summary.json", training)
    metadata = {
        "schema_version": 1,
        "campaign_tag": campaign_tag,
        "analysis_status": "post_hoc_exploratory",
        "bootstrap_samples": bootstrap_samples,
        "confidence": confidence,
        "seed": seed,
        "pair_bootstrap": "paired_two_endpoint_bgc_cluster_fixed_ranks",
        "input_sha256": {str(path.relative_to(root)): sha256_file(path) for path in required},
    }
    write_json_immutable(output / "analysis_metadata.json", metadata)
    return output