"""Diploid copying-model forward/backward kernel for fixed per-gene likelihoods.

States are ordered pairs of training haplotypes. Each chromosome independently
retains its template with probability 1-rho, or draws from fixed training weights
with probability rho. Input emissions must share identical read evidence between
arms. This kernel does not construct read likelihoods, tune rho, or assign types.
"""
import math

def _matrix(values, n):
    if len(values) != n or any(len(row) != n for row in values):
        raise ValueError('Expected square haplotype-pair matrix')
    result = [[float(v) for v in row] for row in values]
    if any(not math.isfinite(v) or v < 0 for row in result for v in row):
        raise ValueError('Likelihoods must be finite and nonnegative')
    return result

def _normalize(a):
    z = math.fsum(v for row in a for v in row)
    if not math.isfinite(z) or z <= 0:
        raise ValueError('No supported state or invalid likelihood scale')
    return [[v/z for v in row] for row in a], z

def forward_transition(a, weights, rho):
    """T^T A T in O(H^2), for T_ij=(1-rho)delta_ij+rho*w_j."""
    n=len(weights);keep=1-rho
    rows=[math.fsum(row) for row in a]
    cols=[math.fsum(a[i][j] for i in range(n)) for j in range(n)]
    total=math.fsum(rows)
    return [[keep*keep*a[i][j]+keep*rho*(weights[i]*cols[j]+rows[i]*weights[j])+rho*rho*weights[i]*weights[j]*total for j in range(n)] for i in range(n)]

def backward_transition(a, weights, rho):
    """T A T^T, the adjoint transition for backward messages."""
    n=len(weights);keep=1-rho
    rows=[math.fsum(a[i][j]*weights[j] for j in range(n)) for i in range(n)]
    cols=[math.fsum(weights[i]*a[i][j] for i in range(n)) for j in range(n)]
    total=math.fsum(weights[i]*rows[i] for i in range(n))
    return [[keep*keep*a[i][j]+keep*rho*(cols[j]+rows[i])+rho*rho*total for j in range(n)] for i in range(n)]

def posterior(emissions, weights, switch_probabilities):
    """Return ordered template-pair marginals and evidence log-likelihood.

    rho=1 gives independent genes with the same marginal training weights.
    rho=0 restricts copying to one pair throughout. Positive rho permits template
    changes; it cannot invent a gene sequence absent from all templates. Missing
    observations use an all-ones matrix; all-zero observations are rejected.
    """
    weights=[float(x) for x in weights];n=len(weights)
    if not n or any(not math.isfinite(x) or x<=0 for x in weights):
        raise ValueError('Training weights must be positive and finite')
    if not math.isclose(math.fsum(weights),1,rel_tol=1e-12,abs_tol=1e-12):
        raise ValueError('Training weights must sum to one')
    if not emissions or len(switch_probabilities)!=len(emissions)-1:
        raise ValueError('Require one transition probability between loci')
    if any(not math.isfinite(r) or not 0<=r<=1 for r in switch_probabilities):
        raise ValueError('Invalid switch probability')
    e=[_matrix(m,n) for m in emissions]
    # Remove a per-locus scale to avoid underflow while retaining its contribution
    # to total evidence. This does not change posterior probabilities.
    logs=[]
    for k,m in enumerate(e):
        scale=max(v for row in m for v in row)
        if scale<=0:raise ValueError('No supported state at locus')
        logs.append(math.log(scale));e[k]=[[v/scale for v in row] for row in m]
    prior=[[weights[i]*weights[j] for j in range(n)] for i in range(n)]
    forward=[];scales=[]
    for k,m in enumerate(e):
        if k:prior=forward_transition(forward[-1],weights,switch_probabilities[k-1])
        current,z=_normalize([[prior[i][j]*m[i][j] for j in range(n)] for i in range(n)])
        forward.append(current);scales.append(z)
    marginals=[None]*len(e);back=[[1.0]*n for _ in range(n)]
    for k in range(len(e)-1,-1,-1):
        marginals[k],_= _normalize([[forward[k][i][j]*back[i][j] for j in range(n)] for i in range(n)])
        if k:
            weighted=[[e[k][i][j]*back[i][j] for j in range(n)] for i in range(n)]
            back=backward_transition(weighted,weights,switch_probabilities[k-1])
            peak=max(v for row in back for v in row)
            if peak<=0:raise ValueError('No supported backward state')
            back=[[v/peak for v in row] for row in back]
    return {'ordered_pair_marginals':marginals,'log_evidence':math.fsum(logs)+math.fsum(math.log(z) for z in scales)}