"""
Core machinery for the genetic-code robustness analysis.

Everything here is code-intrinsic: standard code table, amino-acid distance
metrics, single-nucleotide neighbour graph, expected mutational cost Phi,
and the two null ensembles.

No external data required. Validation target: reproduce the Haig & Hurst (1991)
result for polar requirement under block-preserving randomisation.
"""

import itertools
import numpy as np

BASES = "UCAG"

# ---------------------------------------------------------------- code table
# Standard genetic code (NCBI translation table 1), RNA alphabet.
_AAS = (
    "FFLLSSSSYY**CC*W"
    "LLLLPPPPHHQQRRRR"
    "IIIMTTTTNNKKSSRR"
    "VVVVAAAADDEEGGGG"
)
# NCBI table-1 order is  UUU UUC UUA UUG UCU ...  i.e. base1 slowest.
# Rebuild explicitly to avoid transcription errors.
STANDARD_CODE = {}
_ncbi = (
    "FFLLSSSSYY**CC*WLLLLPPPPHHQQRRRRIIIMTTTTNNKKSSRRVVVVAAAADDEEGGGG"
)
_i = 0
for b1 in BASES:
    for b2 in BASES:
        for b3 in BASES:
            STANDARD_CODE[b1 + b2 + b3] = _ncbi[_i]
            _i += 1

CODONS = sorted(STANDARD_CODE)
SENSE = [c for c in CODONS if STANDARD_CODE[c] != "*"]
STOPS = [c for c in CODONS if STANDARD_CODE[c] == "*"]

AA20 = sorted(set(STANDARD_CODE.values()) - {"*"})


# ------------------------------------------------------------------- metrics
# Woese polar requirement, as used by Haig & Hurst (1991).
POLAR_REQUIREMENT = {
    "A": 7.0, "R": 9.1, "N": 10.0, "D": 13.0, "C": 4.8,
    "Q": 8.6, "E": 12.5, "G": 7.9, "H": 8.4, "I": 4.9,
    "L": 4.9, "K": 10.1, "M": 5.3, "F": 5.0, "P": 6.6,
    "S": 7.5, "T": 6.6, "W": 5.2, "Y": 5.4, "V": 5.6,
}

KYTE_DOOLITTLE = {
    "A": 1.8, "R": -4.5, "N": -3.5, "D": -3.5, "C": 2.5,
    "Q": -3.5, "E": -3.5, "G": -0.4, "H": -3.2, "I": 4.5,
    "L": 3.8, "K": -3.9, "M": 1.9, "F": 2.8, "P": -1.6,
    "S": -0.8, "T": -0.7, "W": -0.9, "Y": -1.3, "V": 4.2,
}

# Grantham (1974): composition c, polarity p, volume v.
_GRANTHAM_CPV = {
    "S": (1.42, 9.2, 32), "R": (0.65, 10.5, 124), "L": (0.00, 4.9, 111),
    "P": (0.39, 8.0, 32.5), "T": (0.71, 8.6, 61), "A": (0.00, 8.1, 31),
    "V": (0.00, 5.9, 84), "G": (0.74, 9.0, 3), "I": (0.00, 5.2, 111),
    "F": (0.00, 5.2, 132), "Y": (0.20, 6.2, 136), "C": (2.75, 5.5, 55),
    "H": (0.58, 10.4, 96), "Q": (0.89, 10.5, 85), "N": (1.33, 11.6, 56),
    "K": (0.33, 11.3, 119), "D": (1.38, 13.0, 54), "E": (0.92, 12.3, 83),
    "M": (0.00, 5.7, 105), "W": (0.13, 5.4, 170),
}
_G_ALPHA, _G_BETA, _G_GAMMA = 1.833, 0.1018, 0.000399


def grantham_matrix():
    """Grantham distance, scaled so the mean over all 190 pairs is 100."""
    d = {}
    raw = []
    for a, b in itertools.combinations(AA20, 2):
        ca, pa, va = _GRANTHAM_CPV[a]
        cb, pb, vb = _GRANTHAM_CPV[b]
        val = np.sqrt(
            _G_ALPHA * (ca - cb) ** 2
            + _G_BETA * (pa - pb) ** 2
            + _G_GAMMA * (va - vb) ** 2
        )
        d[(a, b)] = val
        raw.append(val)
    rho = 100.0 / np.mean(raw)
    out = {}
    for (a, b), v in d.items():
        out[(a, b)] = out[(b, a)] = rho * v
    for a in AA20:
        out[(a, a)] = 0.0
    return out


def scalar_metric_matrix(values, squared=True):
    """Distance matrix from a 1-D amino-acid scale.

    squared=True gives (x_i - x_j)^2, which is the Haig & Hurst / Freeland &
    Hurst convention for polar requirement (mean *square* change).
    """
    out = {}
    for a in AA20:
        for b in AA20:
            diff = values[a] - values[b]
            out[(a, b)] = diff ** 2 if squared else abs(diff)
    return out


# ------------------------------------------------------- neighbour structure
def single_nt_neighbours(codon):
    """The 9 codons one substitution away, with the position changed."""
    out = []
    for pos in range(3):
        for b in BASES:
            if b != codon[pos]:
                out.append((codon[:pos] + b + codon[pos + 1:], pos))
    return out


NEIGHBOURS = {c: single_nt_neighbours(c) for c in CODONS}


# ------------------------------------------------------------------ blocks
def code_blocks(code=None):
    """Map amino acid -> tuple of codons encoding it (stops excluded)."""
    code = code or STANDARD_CODE
    blocks = {}
    for c, aa in code.items():
        if aa == "*":
            continue
        blocks.setdefault(aa, []).append(c)
    return {aa: tuple(sorted(v)) for aa, v in blocks.items()}


STANDARD_BLOCKS = code_blocks()


# ---------------------------------------------------------------------- Phi
def phi(code, dist, spectrum=None, pi=None, include_nonsense=False):
    """Expected cost of a random single-nucleotide substitution.

    code      : dict codon -> aa
    dist      : dict (aa, aa) -> distance
    spectrum  : dict (from_base, to_base) -> relative rate, or None for uniform
    pi        : dict codon -> weight, or None for uniform over sense codons
    include_nonsense : if False, sense->stop substitutions are excluded from
                       both numerator and denominator and reported separately.
    """
    num = 0.0
    den = 0.0
    nonsense_w = 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
        aa_from = code[c]
        for c2, pos in NEIGHBOURS[c]:
            rate = 1.0 if spectrum is None else spectrum.get((c[pos], c2[pos]), 0.0)
            if rate == 0.0:
                continue
            w = w_c * rate
            aa_to = code[c2]
            if aa_to == "*":
                nonsense_w += w
                if not include_nonsense:
                    continue
                raise ValueError("include_nonsense=True needs an explicit stop cost")
            num += w * dist[(aa_from, aa_to)]
            den += w
    return num / den, nonsense_w / (den + nonsense_w)


# ------------------------------------------------------------ null ensembles
def random_code_block_preserving(rng, code=None):
    """Permute the 20 amino acids among the 20 standard blocks. Stops fixed."""
    blocks = code_blocks(code)
    aas = list(blocks)
    perm = rng.permutation(len(aas))
    new = {c: "*" for c in STOPS}
    for i, aa in enumerate(aas):
        target = aas[perm[i]]
        for c in blocks[aa]:
            new[c] = target
    return new


def random_code_degeneracy_only(rng):
    """Preserve only the multiset of family sizes; scatter codons freely."""
    sizes = [len(v) for v in STANDARD_BLOCKS.values()]
    aas = list(STANDARD_BLOCKS)
    pool = list(SENSE)
    rng.shuffle(pool)
    new = {c: "*" for c in STOPS}
    k = 0
    order = rng.permutation(len(aas))
    for idx in order:
        aa = aas[idx]
        n = sizes[idx]
        for c in pool[k:k + n]:
            new[c] = aa
        k += n
    return new


def null_distribution(dist, n, generator, spectrum=None, pi=None, seed=0):
    rng = np.random.default_rng(seed)
    out = np.empty(n)
    for i in range(n):
        rc = generator(rng)
        out[i], _ = phi(rc, dist, spectrum=spectrum, pi=pi)
    return out


def z_score(dist, null, spectrum=None, pi=None):
    obs, _ = phi(STANDARD_CODE, dist, spectrum=spectrum, pi=pi)
    return (null.mean() - obs) / null.std(ddof=1), obs


if __name__ == "__main__":
    # Pipeline validation against the published benchmark. Reports the standard
    # code's cost under uniform mutation and how many randomized codes beat it.
    import numpy as _np
    _pr = scalar_metric_matrix(POLAR_REQUIREMENT, True)
    _obs, _ = phi(STANDARD_CODE, _pr)
    _null = null_distribution(_pr, 10000, random_code_block_preserving,
                              seed=20260727)
    _beat = int((_null < _obs).sum())
    print(f"Phi_SGC (uniform mutation, squared polar requirement) = {_obs:.4f}")
    print(f"randomized codes with lower cost: {_beat} of 10,000")
    print("expected: 5.194 and 3 (Haig and Hurst 1991 reported ~2 in 10,000)")
    _ok = abs(_obs - 5.1936) < 5e-4 and _beat == 3
    print("BENCHMARK:", "PASS" if _ok else "FAIL")
    raise SystemExit(0 if _ok else 1)
