"""
In-silico code optimisation under fixed transition/transversion bias.

Question, fixed before running: if a code is optimised for a spectrum with a
given kappa, do the optima resemble the STANDARD code more as kappa rises?

Search space  : permutations of the 20 amino acids over the 20 standard blocks
                (the same space as the frozen null ensemble)
Cost          : Phi under 3-mer context integration, fixed code-intrinsic
                exposure, squared Woese polar requirement (the locked metric)
Similarity    : fraction of the 20 blocks assigned the same amino acid as the
                standard code. Chance level = 1/20 = 5%.
Prediction    : similarity increases monotonically with kappa.

Key speedup: the cost only depends on the assignment through block pairs, so
weights are pre-aggregated into a 20x20 block-pair matrix and each candidate
code costs one 20x20 contraction instead of a sweep over 1576 edges.
"""

import itertools
import numpy as np

from codecore import SENSE, STANDARD_CODE, AA20, STANDARD_BLOCKS
from context3 import EXPOSURE3, canonical3

BLOCK_OF = {}
for bi, aa in enumerate(sorted(STANDARD_BLOCKS)):
    for c in STANDARD_BLOCKS[aa]:
        BLOCK_OF[c] = bi
BLOCK_AA = sorted(STANDARD_BLOCKS)                      # block index -> standard aa
STANDARD_ASSIGN = np.arange(20)                          # identity = standard code


def kappa_spectrum(kappa):
    """Canonical 3-mer spectrum with transition/transversion ratio kappa.

    No CpG term: this isolates the transition-bias axis, which is the one the
    data identified as the driver.
    """
    pur = {"A", "G"}
    spec = {}
    for ctx in map("".join, itertools.product("ACGU", repeat=3)):
        if ctx[1] not in {"C", "U"}:
            continue
        for alt in "ACGU":
            if alt == ctx[1]:
                continue
            spec[(ctx, alt)] = kappa if (ctx[1] in pur) == (alt in pur) else 1.0
    return spec


def block_pair_weights(spectrum):
    """Aggregate sense->sense edge weights into a 20x20 block-pair matrix."""
    B = np.zeros((20, 20))
    total = 0.0
    for c in SENSE:
        bi = BLOCK_OF[c]
        for j in (0, 1, 2):
            for ctx, pq in EXPOSURE3[(c, j)].items():
                for b2 in "ACGU":
                    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:]
                    if STANDARD_CODE[c2] == "*":
                        continue
                    w = pq * rate
                    B[bi, BLOCK_OF[c2]] += w
                    total += w
    return B / total


def cost(assign, B, D):
    """Phi for one assignment: block -> amino-acid index."""
    return float((B * D[np.ix_(assign, assign)]).sum())


def anneal(B, D, rng, iters=200000, t0=None):
    """Simulated annealing over permutations, by pairwise block swaps."""
    a = rng.permutation(20)
    c = cost(a, B, D)
    best_a, best_c = a.copy(), c
    if t0 is None:
        t0 = 0.05 * abs(c) + 1e-12
    for k in range(iters):
        T = t0 * (1.0 - k / iters) + 1e-15
        i, j = rng.integers(0, 20, 2)
        if i == j:
            continue
        a[i], a[j] = a[j], a[i]
        c2 = cost(a, B, D)
        if c2 <= c or rng.random() < np.exp(-(c2 - c) / T):
            c = c2
            if c < best_c:
                best_c, best_a = c, a.copy()
        else:
            a[i], a[j] = a[j], a[i]
    return best_a, best_c


def similarity(assign):
    """Fraction of blocks assigned the same amino acid as the standard code."""
    return float(np.mean(assign == STANDARD_ASSIGN))
