"""
3-mer context machinery -- Cohort B (CORAL), protocol v5 section 5.3.1.

Same strand canonicalisation as the 5-mer version (central base C or T;
reverse complement for A/G), but on trinucleotide windows:

    j=0 : prev[2]  c[0]   c[1]
    j=1 : c[0]     c[1]   c[2]      (entirely inside the codon)
    j=2 : c[1]     c[2]   next[0]

Position 2 needs no flanking codon at all.

CORAL uses the DNA alphabet and the category naming L[REF>ALT]R with REF in
{C, T} -- which is already the canonical form. Internally this module keeps the
RNA alphabet of codecore (U instead of T) and converts at the boundary.
"""

from collections import defaultdict
from codecore import BASES, SENSE, STANDARD_CODE

COMP = {"A": "U", "U": "A", "C": "G", "G": "C"}
PYRIMIDINES = {"C", "U"}


def rc3(seq):
    return "".join(COMP[b] for b in reversed(seq))


def canonical3(context, to_base):
    """Canonical (3-mer, to_base) with the central base a pyrimidine."""
    if context[1] in PYRIMIDINES:
        return context, to_base
    return rc3(context), COMP[to_base]


def to_coral_key(context, to_base):
    """Canonical RNA form -> CORAL column name, e.g. 'A[C>T]G'."""
    ctx = context.replace("U", "T")
    alt = to_base.replace("U", "T")
    return f"{ctx[0]}[{ctx[1]}>{alt}]{ctx[2]}"


def from_coral_key(key):
    """'A[C>T]G' -> (RNA 3-mer, RNA alt base)."""
    left, rest = key.split("[")
    sub, right = rest.split("]")
    ref, alt = sub.split(">")
    ctx = (left + ref + right).replace("T", "U")
    return ctx, alt.replace("T", "U")


def _flank_marginals():
    """Uniform over the 61 sense codons: marginal of position 2 and position 0."""
    n = len(SENSE)
    m2 = defaultdict(float)   # last base of the preceding codon
    m0 = defaultdict(float)   # first base of the following codon
    for c in SENSE:
        m2[c[2]] += 1.0 / n
        m0[c[0]] += 1.0 / n
    return dict(m2), dict(m0)


_M2, _M0 = _flank_marginals()


def context_exposure3(codon, j):
    """q0(x | c, j) for trinucleotide windows. Exact, code-intrinsic."""
    out = defaultdict(float)
    if j == 0:
        for a, p in _M2.items():
            out[a + codon[0] + codon[1]] += p
    elif j == 1:
        out[codon] = 1.0
    elif j == 2:
        for e, p in _M0.items():
            out[codon[1] + codon[2] + e] += p
    else:
        raise ValueError("j must be 0, 1 or 2")
    return dict(out)


EXPOSURE3 = {(c, j): context_exposure3(c, j) for c in SENSE for j in (0, 1, 2)}


def build_edges3(spectrum, pi=None, exposure=None):
    """Weighted sense->sense edges under 3-mer integration.

    spectrum : dict {(canonical RNA 3-mer, alt base): rate}
    Global normalisation only (protocol v5 stop rule): sense->stop excluded from
    numerator and denominator, no per-block rescaling.
    Returns (edges, total_weight, nonsense_share); weights are code-independent
    because stop codons are fixed across the ensemble.
    """
    exposure = exposure or EXPOSURE3
    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):
            for ctx, pq in exposure[(c, j)].items():
                for b2 in BASES:
                    if b2 == c[j]:
                        continue
                    ck, bk = canonical3(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] == "*":
                        nonsense += w
                        continue
                    edges.append((c, c2, w, j, c[j], b2))
                    total += w
    return edges, total, nonsense / (total + nonsense)
