"""The controlled predictive test — does a pre-merge functional-conflict measure predict merge damage?

(A) The framework-motivated predictor: confidence-weighted functional conflict (bilateral confident
disagreement, measured before merging — a proposed proxy for merge-relevant interactions, not a
measured epistasis) against the merge penalty (oracle parent potential − merged achieved; ordering is
sensitive to the outcome reference — see stats_llm_epistasis.py).

(B) The geometry predictor on the same outcome: weight divergence (LoRA-delta L2) — the
matched-divergence contrast: the duration axis spans large weight divergence at ~zero penalty, while
the conflict axis generates penalty at modest divergence. Distance is not what breaks merging.

(C) The league table: |Spearman rho| against the merge penalty for every pre-merge predictor,
including the internal ablation (raw disagreement, which the theory predicts must mislead because it
counts harmless complementation as conflict).

Usage:  python figures/plot_llm_epistasis.py
"""

from __future__ import annotations

import sys
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import spearmanr

sys.path.insert(0, str(Path(__file__).parent))
from _figlib import load_bundle, savefig  # noqa: E402

PREDICTORS = [("epi_conf", "conf-weighted\nfunctional conflict"),
              ("dis_raw", "raw\ndisagreement"),
              ("grad_cos", "gradient\nalignment"),
              ("delta_cos", "delta\ncosine"),
              ("delta_l2", "delta\nL2"),
              ("cross_perf", "cross-family\naccuracy")]
OUTCOME = "merge_penalty"


def _scatter(ax, df, xcol, xlabel, title):
    for mode, color, marker in (("conflict", "#d62728", "o"), ("duration", "#2c7fb8", "s"),
                                ("compat", "#41ab5d", "^")):
        sub = df[df["mode"] == mode]
        ax.scatter(sub[xcol], sub[OUTCOME], c=color, marker=marker, s=42, alpha=0.75,
                   label=f"{mode} axis")
    rho, p = spearmanr(df[xcol], df[OUTCOME])
    ax.set(xlabel=xlabel, ylabel="merge penalty (parent potential − merged)",
           title=f"{title}\nSpearman ρ = {rho:.2f} (p = {p:.1g}, n = {len(df)})")
    ax.axhline(0, color="#999", lw=0.6)
    ax.legend(frameon=False, fontsize=8)


def main() -> None:
    import pandas as pd
    df, _ = load_bundle("results/llm_epistasis")
    try:                                             # the overlap-without-conflict control axis
        compat, _ = load_bundle("results/llm_epistasis_compat")
        df = pd.concat([df, compat], ignore_index=True)
    except Exception:
        pass

    fig, axes = plt.subplots(1, 3, figsize=(16.5, 4.9))
    _scatter(axes[0], df, "epi_conf", "confidence-weighted functional conflict (pre-merge)",
             "(A) the framework-motivated predictor")
    _scatter(axes[1], df, "delta_cos", "LoRA-delta cosine similarity (pre-merge)",
             "(B) the geometry predictor — does it detect\nincompatibility, or just task overlap?")

    ax = axes[2]
    rhos, labels = [], []
    for col, label in PREDICTORS:
        rho, _ = spearmanr(df[col], df[OUTCOME])
        rhos.append(abs(rho)); labels.append(label)
    colors = ["#d62728" if c == "epi_conf" else ("#fc9272" if c == "dis_raw" else "#9ecae1")
              for c, _ in PREDICTORS]
    x = np.arange(len(rhos))
    ax.bar(x, rhos, 0.62, color=colors)
    ax.set_xticks(x); ax.set_xticklabels(labels, fontsize=7)
    ax.set(ylabel="|Spearman ρ| vs merge penalty", ylim=(0, 1),
           title="(C) predictor league table (pre-merge only)")

    fig.suptitle("A controlled predictive test: across this task grid, pre-merge functional disagreement "
                 "predicted merge penalties; the selected weight-geometry baselines did not "
                 "(three axes decorrelated by construction; 13 conditions x 3 seeds)", y=1.03, fontsize=11.5)
    fig.tight_layout()
    savefig(fig, "results/llm_epistasis", "llm_epistasis")


if __name__ == "__main__":
    main()