"""Allow locus-independent candidate choices alongside linked training templates.

The final hidden state is an unlinked component. At each locus it draws an allele
from the supplied candidate distribution independently of previous loci, even
when this hidden state persists. Its mass and candidate distributions must be
fixed or selected within training data, never using held-out truth.
"""
import math
from linkage_hmm import posterior

def distribution(values):
    result=[float(x) for x in values]
    if not result or any(not math.isfinite(x) or x<0 for x in result):
        raise ValueError('Invalid candidate distribution')
    if not math.isclose(math.fsum(result),1,rel_tol=1e-12,abs_tol=1e-12):
        raise ValueError('Candidate probabilities must sum to one')
    return result

def infer(likelihoods, templates, template_weights, candidate_weights,
          unlinked_mass, switch_probabilities):
    """Return allele-pair marginals, including candidates absent from templates.

    likelihoods[locus][allele1][allele2] is a common, nonnegative diploid
    likelihood shared by every ablation. templates[haplotype][locus] is an allele
    index. rho=1 is the independent-locus control with identical allele priors.
    The output is uncalibrated if the input scores are uncalibrated.
    """
    if not math.isfinite(unlinked_mass) or not 0<unlinked_mass<1:
        raise ValueError('Unlinked mass must lie strictly between zero and one')
    weights=distribution(template_weights)
    if any(x<=0 for x in weights):raise ValueError('Template weights must be positive')
    h=len(templates);loci=len(likelihoods)
    if h!=len(weights) or not loci or len(candidate_weights)!=loci:
        raise ValueError('Input dimension mismatch')
    if any(len(t)!=loci for t in templates):raise ValueError('Incomplete template')
    priors=[distribution(x) for x in candidate_weights];scaled=[];offset=0
    for k,matrix in enumerate(likelihoods):
        n=len(priors[k])
        if len(matrix)!=n or any(len(row)!=n for row in matrix):raise ValueError('Likelihood dimension mismatch')
        m=[[float(x) for x in row] for row in matrix]
        if any(not math.isfinite(x) or x<0 for row in m for x in row):raise ValueError('Invalid likelihood')
        if any(not isinstance(t[k],int) or not 0<=t[k]<n for t in templates):raise ValueError('Invalid template allele')
        peak=max(x for row in m for x in row)
        if peak<=0:raise ValueError('Unsupported locus')
        offset+=math.log(peak);scaled.append([[x/peak for x in row] for row in m])
    emissions=[]
    for k,m in enumerate(scaled):
        q=priors[k];n=len(q);a=[t[k] for t in templates]
        row=[math.fsum(m[i][j]*q[j] for j in range(n)) for i in range(n)]
        col=[math.fsum(q[i]*m[i][j] for i in range(n)) for j in range(n)]
        both=math.fsum(q[i]*row[i] for i in range(n))
        e=[[m[a[i]][a[j]] for j in range(h)]+[row[a[i]]] for i in range(h)]
        e.append([col[a[j]] for j in range(h)]+[both]);emissions.append(e)
    hidden=posterior(emissions,[(1-unlinked_mass)*w for w in weights]+[unlinked_mass],switch_probabilities)
    output=[]
    for k,(m,e,p) in enumerate(zip(scaled,emissions,hidden['ordered_pair_marginals'])):
        q=priors[k];n=len(q);a=[t[k] for t in templates];result=[[0.0]*n for _ in range(n)]
        for i in range(h):
            for j in range(h):result[a[i]][a[j]]+=p[i][j]
        for i in range(h):
            if e[i][h]:
                for b in range(n):result[a[i]][b]+=p[i][h]*q[b]*m[a[i]][b]/e[i][h]
            if e[h][i]:
                for b in range(n):result[b][a[i]]+=p[h][i]*q[b]*m[b][a[i]]/e[h][i]
        if e[h][h]:
            for a0 in range(n):
                for b0 in range(n):result[a0][b0]+=p[h][h]*q[a0]*q[b0]*m[a0][b0]/e[h][h]
        if not math.isclose(math.fsum(x for row in result for x in row),1,abs_tol=1e-10):raise ArithmeticError('Posterior mass lost')
        output.append(result)
    return {'ordered_allele_pair_marginals':output,'ordered_template_pair_marginals':hidden['ordered_pair_marginals'],
            'log_evidence':hidden['log_evidence']+offset,
            'unlinked_mass':unlinked_mass,'absent_from_all_candidates_recoverable':False}