"""
Mechanistic decomposition of the code's advantage (protocol v2, section 6).

Splits Delta-Phi = mean_null Phi - Phi_SGC into additive components, so we can
answer: WHICH parts of the mutation spectrum make the SGC's advantage rise
or fall?

Decomposition axes:
  - codon position (1 / 2 / 3)
  - transition vs transversion
  - the 12 directional substitution types
  - functional class (synonymous / missense / nonsense)

No external data. Run on synthetic spectra to generate pre-registered
quantitative predictions before the real spectra are loaded.
"""

import numpy as np
from codecore import (
    BASES, STANDARD_CODE, SENSE, NEIGHBOURS, AA20,
    POLAR_REQUIREMENT, KYTE_DOOLITTLE, grantham_matrix,
    scalar_metric_matrix, random_code_block_preserving,
)

PURINES = set("AG")


def is_transition(a, b):
    return (a in PURINES) == (b in PURINES)


def substitution_events(code):
    """Enumerate every (codon, position, from_base, to_base, class, cost_key).

    Yields the structural facts that do not depend on the spectrum, so the
    ensemble loop can reuse them.
    """
    for c in SENSE:
        aa_from = code[c]
        for c2, pos in NEIGHBOURS[c]:
            aa_to = code[c2]
            b_from, b_to = c[pos], c2[pos]
            if aa_to == "*":
                klass = "nonsense"
            elif aa_to == aa_from:
                klass = "synonymous"
            else:
                klass = "missense"
            yield c, pos, b_from, b_to, klass, aa_from, aa_to


def decompose(code, dist, spectrum, pi=None):
    """Weighted cost split by position, ts/tv, substitution type, class.

    Returns (total_cost, weight, dict_of_component_dicts).
    Nonsense excluded from the cost (reported separately as a weight share).
    """
    by_pos = {0: [0.0, 0.0], 1: [0.0, 0.0], 2: [0.0, 0.0]}
    by_tstv = {"ts": [0.0, 0.0], "tv": [0.0, 0.0]}
    by_sub = {}
    by_class = {"synonymous": 0.0, "missense": 0.0, "nonsense": 0.0}
    num = den = 0.0

    for c, pos, bf, bt, klass, aaf, aat in substitution_events(code):
        w_c = 1.0 if pi is None else pi.get(c, 0.0)
        rate = spectrum.get((bf, bt), 0.0)
        w = w_c * rate
        if w == 0.0:
            continue
        by_class[klass] += w
        if klass == "nonsense":
            continue
        cost = dist[(aaf, aat)]
        num += w * cost
        den += w
        by_pos[pos][0] += w * cost
        by_pos[pos][1] += w
        k = "ts" if is_transition(bf, bt) else "tv"
        by_tstv[k][0] += w * cost
        by_tstv[k][1] += w
        sk = f"{bf}>{bt}"
        s = by_sub.setdefault(sk, [0.0, 0.0])
        s[0] += w * cost
        s[1] += w

    return num / den, den, {
        "position": by_pos,
        "tstv": by_tstv,
        "substitution": by_sub,
        "class": by_class,
    }


def advantage_decomposition(dist, spectrum, n_null=3000, seed=17, pi=None):
    """Delta-Phi decomposed into per-component contributions.

    For each component k:  delta_k = (mean_null numerator_k - SGC numerator_k)
                                     / total_weight
    Components sum to Delta-Phi exactly, because the weights are identical
    across codes (the ensemble only permutes amino acids, never rates).
    """
    rng = np.random.default_rng(seed)
    sgc_total, weight, sgc = decompose(STANDARD_CODE, dist, spectrum, pi)

    acc = {"position": {p: 0.0 for p in (0, 1, 2)},
           "tstv": {"ts": 0.0, "tv": 0.0},
           "substitution": {k: 0.0 for k in sgc["substitution"]}}
    totals = []
    for _ in range(n_null):
        rc = random_code_block_preserving(rng)
        tot, _w, comp = decompose(rc, dist, spectrum, pi)
        totals.append(tot)
        for p in acc["position"]:
            acc["position"][p] += comp["position"][p][0]
        for k in acc["tstv"]:
            acc["tstv"][k] += comp["tstv"][k][0]
        for k in acc["substitution"]:
            acc["substitution"][k] += comp["substitution"].get(k, [0.0, 0.0])[0]

    totals = np.array(totals)
    out = {"Phi_SGC": sgc_total,
           "Phi_null_mean": totals.mean(),
           "Phi_null_sd": totals.std(ddof=1),
           "weight": weight,
           "class_weights": sgc["class"]}
    out["delta_total"] = totals.mean() - sgc_total

    for axis in ("position", "tstv", "substitution"):
        d = {}
        for k, s in acc[axis].items():
            null_num = s / n_null
            sgc_num = sgc[axis][k][0] if axis != "substitution" \
                else sgc["substitution"][k][0]
            d[k] = (null_num - sgc_num) / weight
        out[axis] = d
    return out
