"""
Context-integrated mechanistic decomposition (protocol v4, sections 5.1 and 6).

Replaces the context-free decomposition, whose magnitudes were inflated ~2x
because the CpG boost was applied to every C>T rather than only to C>T in
genuine CpG context.

Key structural fact that makes the decomposition exactly additive:
stop codons are FIXED in every randomized code (v4 stop rule 4), so the
sense/stop structure of the mutation graph -- and therefore every weight,
including the per-codon-position renormalisation -- is identical across the
whole ensemble. Only the amino-acid labels change. So

    Delta-Phi = sum_k [ (null-mean numerator_k) - (SGC numerator_k) ] / W

splits cleanly by any partition k of the edges.
"""

import numpy as np
from codecore import BASES, SENSE, STANDARD_CODE
from context5 import EXPOSURE, canonical

PURINES = set("AG")


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


def build_edges(spectrum, pi=None, exposure=None):
    """Pre-compute every weighted edge once. Weights are code-independent.

    Returns (edges, total_weight, nonsense_share) where each edge is
    (codon_from, codon_to, weight, position, base_from, base_to).
    Sense->stop edges are excluded from both numerator and denominator and the
    remainder renormalised per (codon, position) -- v4 stop rule 2.
    """
    exposure = exposure or EXPOSURE
    edges = []
    total = 0.0
    nonsense = 0.0

    for c in SENSE:
        w_c = 1.0 if pi is None else pi.get(c, 0.0)
        if w_c == 0.0:
            continue
        for j in (0, 1, 2):
            local, local_ns = [], 0.0
            for ctx, pq in exposure[(c, j)].items():
                for b2 in BASES:
                    if b2 == c[j]:
                        continue
                    ck, bk = canonical(ctx, b2)
                    rate = spectrum.get((ck, bk), 0.0)
                    if rate == 0.0:
                        continue
                    c2 = c[:j] + b2 + c[j + 1:]
                    w = w_c * pq * rate
                    if STANDARD_CODE[c2] == "*":     # stops are code-invariant
                        local_ns += w
                        continue
                    local.append((c, c2, w, j, c[j], b2))
            nonsense += local_ns
            # NO per-block rescaling: nonsense is excluded from numerator and
            # denominator, which is a single GLOBAL normalisation by the total
            # sense weight. Rescaling each (codon, position) block to equal mass
            # would force every position to contribute equally and destroy the
            # real relative weighting of positions -- a different estimand.
            for e in local:
                edges.append(e)
                total += e[2]

    return edges, total, nonsense / (total + nonsense)


def numerators(code, dist, edges):
    """Cost numerators split by position, ts/tv and directional substitution."""
    pos = {0: 0.0, 1: 0.0, 2: 0.0}
    tstv = {"ts": 0.0, "tv": 0.0}
    sub = {}
    tot = 0.0
    for (c, c2, w, j, bf, bt) in edges:
        cost = dist[(code[c], code[c2])]
        wc = w * cost
        tot += wc
        pos[j] += wc
        tstv["ts" if is_transition(bf, bt) else "tv"] += wc
        sub[f"{bf}>{bt}"] = sub.get(f"{bf}>{bt}", 0.0) + wc
    return tot, pos, tstv, sub


def decompose_advantage(dist, spectrum, null_codes, pi=None):
    """Delta-Phi split into additive components under context integration."""
    edges, W, nonsense_share = build_edges(spectrum, pi=pi)
    sgc_tot, sgc_pos, sgc_ts, sgc_sub = numerators(STANDARD_CODE, dist, edges)

    acc_tot = 0.0
    acc_pos = {0: 0.0, 1: 0.0, 2: 0.0}
    acc_ts = {"ts": 0.0, "tv": 0.0}
    acc_sub = {k: 0.0 for k in sgc_sub}
    n = len(null_codes)
    for rc_code in null_codes:
        t, p, ts, sb = numerators(rc_code, dist, edges)
        acc_tot += t
        for k in acc_pos:
            acc_pos[k] += p[k]
        for k in acc_ts:
            acc_ts[k] += ts[k]
        for k in acc_sub:
            acc_sub[k] += sb.get(k, 0.0)

    out = {
        "Phi_SGC": sgc_tot / W,
        "Phi_null": (acc_tot / n) / W,
        "nonsense_share": nonsense_share,
    }
    out["delta"] = out["Phi_null"] - out["Phi_SGC"]
    out["R"] = out["delta"] / out["Phi_null"]
    out["position"] = {k: (acc_pos[k] / n - sgc_pos[k]) / W for k in acc_pos}
    out["tstv"] = {k: (acc_ts[k] / n - sgc_ts[k]) / W for k in acc_ts}
    out["substitution"] = {k: (acc_sub[k] / n - sgc_sub[k]) / W for k in acc_sub}
    return out
