"""Sequence-first diploid fitting of commonly scored fragment evidence.

This is an experimental candidate-limited likelihood model. Score gaps are not
calibrated posterior probabilities. Native calls are a fallback, not fixed truth.
"""
import collections
import copy
import math
import re
import itertools
import numpy as np
from joint_model import fit
from sequence_identity import group_signature


def allele_fields(label, fields):
    label = label.removeprefix('HLA-')
    match = re.fullmatch(r'([^*]+)\*(\d+(?::\d+)*)([A-Z]?)', label)
    if not match or len(match[2].split(':')) < fields:
        return None
    return match[1] + '*' + ':'.join(match[2].split(':')[:fields])


def emission_rows(fragments, metadata, context=False, include_graph=True):
    """One row per original fragment; never sum native and graph as extra reads."""
    identifiers, rows = [], []
    for fragment in fragments:
        values = {}
        for name, placements in fragment['candidates'].items():
            if name not in metadata:
                raise ValueError('Unknown candidate')
            if not include_graph and metadata[name]['source'] == 'graph':
                continue
            eligible = [p['score'] for p in placements if p['gene_overlap'] and (context or p['internal'])]
            if eligible:
                # Duplicate retrieval routes are the same observation. Maximum
                # alignment is an explicit approximation, not independent support.
                values[name] = max(eligible)
        if values:
            maximum = max(values.values())
            rows.append({p: math.exp(max(-100, score - maximum)) for p, score in values.items()})
            identifiers.append(fragment['fragment'])
    if len(set(identifiers)) != len(identifiers):
        raise ValueError('Repeated fragment')
    return identifiers, rows


def gene_evidence_rows(fragments, identifiers, minimum_gene_bases=30):
    indexed = {}
    wanted = set(identifiers)
    for fragment in fragments:
        if fragment['fragment'] not in wanted:
            continue
        if fragment['fragment'] in indexed:
            raise ValueError('Repeated gene-evidence fragment')
        values = {}
        for path, placements in fragment['candidates'].items():
            if any('gene_penalty' not in p or 'gene_bases' not in p for p in placements):
                raise ValueError('Context calls require gene-discrimination evidence; old pilot evidence is insufficient')
            scores = [p['gene_penalty'] for p in placements if p['gene_bases'] >= minimum_gene_bases]
            if scores:
                values[path] = max(scores)
        indexed[fragment['fragment']] = values
    if set(indexed) != wanted:
        raise ValueError('Missing gene-evidence fragments')
    return [indexed[name] for name in identifiers]


def hard_locus_rows(rows, metadata, minimum_log_gap=10):
    """Ablation: discard ambiguous loci instead of jointly assigning fragments."""
    result = []
    for row in rows:
        best = {}
        for path, value in row.items():
            gene = metadata[path]['gene']
            best[gene] = max(best.get(gene, 0), value)
        ordered = sorted(best, key=lambda g: (-best[g], g))
        if not ordered or (len(ordered) > 1 and best[ordered[1]] >= best[ordered[0]] * math.exp(-minimum_log_gap)):
            result.append({})
        else:
            result.append({p: value for p, value in row.items() if metadata[p]['gene'] == ordered[0]})
    return result


def gene_consistency(matrix, groups, metadata, fitted, core_rows):
    """Flanks may assign a locus, but gene sequence must discriminate its allele pair."""
    if len(core_rows) != len(matrix):
        raise ValueError('Gene/context fragment mismatch')
    components = {}
    for gene, result in fitted['genes'].items():
        a, b = result['pairs'][0]
        components[gene] = result['gene_fraction'] * .5 * (matrix[:, a] + matrix[:, b])
    total = 1e-6 + sum(components.values(), np.zeros(len(matrix)))
    result = {}
    for gene, fit_result in fitted['genes'].items():
        indices = [i for i, g in enumerate(groups) if g['gene'] == gene]
        local = {i: j for j, i in enumerate(indices)}
        penalties = np.full((len(matrix), len(indices)), -np.inf)
        for i, row in enumerate(core_rows):
            for j, column in enumerate(indices):
                available = [row[p] for p in groups[column]['paths'] if p in row]
                if available:
                    penalties[i, j] = max(available)
        valid = np.isfinite(penalties).any(axis=1)
        core = np.zeros_like(penalties)
        core[valid] = np.exp(np.maximum(-100, penalties[valid] - penalties[valid].max(axis=1)[:, None]))
        # Missing placements are zero, not exp(-100) support.
        core[~np.isfinite(penalties)] = 0
        weights = components[gene] / total
        def signature(pair):
            return tuple(sorted(group_signature(groups[i], metadata, decode_group(groups[i], metadata, 4)) for i in pair))
        def score(pair):
            a, b = [local[i] for i in pair]
            return float((weights * np.log(.01 + .99 * .5 * (core[:, a] + core[:, b]))).sum())
        pair = fit_result['pairs'][0]
        proposed_signature = signature(pair)
        alternatives = [score(other) for other in itertools.combinations_with_replacement(indices, 2)
                        if signature(other) != proposed_signature]
        proposed = score(pair)
        result[gene] = dict(weighted_gene_fragments=float(weights[valid].sum()),
                            score_gap=proposed - max(alternatives) if alternatives else None,
                            rule='Context-weighted within-gene error evidence; flank matches supply no allele-discrimination bonus')
    return result


def prepare(rows, metadata, native, maximum_groups=32, minimum_mass=3):
    """Collapse identical evidence within loci before bounded pair enumeration.

    All native-called alleles are protected. Ties at the pruning boundary expand
    the candidate set. Report pruning explicitly; this is not exhaustive all-IPD
    inference, even though candidate retrieval retains all native assignments.
    """
    if maximum_groups < 2 or minimum_mass < 0:
        raise ValueError('Invalid candidate bounds')
    observed = sorted({p for row in rows for p in row})
    if not observed:
        return np.empty((len(rows), 0)), [], {}
    positions = {p: i for i, p in enumerate(observed)}
    matrix = np.zeros((len(rows), len(observed)), dtype=float)
    for i, row in enumerate(rows):
        for p, value in row.items():
            matrix[i, positions[p]] = value
    # Hash exact arrays, then verify equality on collision. Full sequence aliases
    # and indistinguishable contexts therefore cannot vote multiple times.
    import hashlib
    groups, buckets = [], {}
    for j, path in enumerate(observed):
        values = matrix[:, j]
        key = (metadata[path]['gene'], hashlib.sha256(values.tobytes()).hexdigest())
        if key in buckets:
            group = groups[buckets[key]]
            if not np.array_equal(values, matrix[:, group['column']]):
                raise ValueError('Emission digest collision')
            group['paths'].append(path)
        else:
            buckets[key] = len(groups)
            groups.append(dict(gene=metadata[path]['gene'], paths=[path], column=j,
                               mass=float(values.sum())))
    protected = set()
    for call in native.values():
        for slot in call.get('resolutions', {}).get('4', []):
            protected.update(slot.get('alleles', []))
    by_gene = collections.defaultdict(list)
    for group in groups:
        group['protected'] = any(a in protected for p in group['paths'] for a in metadata[p]['labels'])
        by_gene[group['gene']].append(group)
    selected, audit = [], {}
    for gene, values in sorted(by_gene.items()):
        ranked = sorted(values, key=lambda g: (-g['mass'], g['paths']))
        cutoff = ranked[min(maximum_groups, len(ranked)) - 1]['mass']
        kept = [g for g in ranked if g['protected'] or g['mass'] >= max(minimum_mass, cutoff)]
        # Low-evidence loci remain explicit, but cannot produce a confident call.
        if not kept:
            kept = [g for g in ranked if g['mass'] == ranked[0]['mass']]
        selected.extend(kept)
        audit[gene] = dict(observed_paths=sum(len(g['paths']) for g in values),
                           evidence_groups=len(values), retained_groups=len(kept),
                           retained_paths=sum(len(g['paths']) for g in kept), cutoff=cutoff)
    return matrix[:, [g['column'] for g in selected]], selected, audit


def decode_group(group, metadata, fields):
    labels, unresolved = set(), False
    for path in group['paths']:
        source = metadata[path]['labels']
        if fields == 2 and not source:
            source = metadata[path].get('cds_labels', [])
        if not source:
            unresolved = True
        for label in source:
            value = allele_fields(label, fields)
            if value is None:
                unresolved = True
            else:
                labels.add(value)
    return dict(raw='CONTEXT:' + ','.join(group['paths']), alleles=sorted(labels), unresolved=unresolved)


def infer(rows, metadata, native, sequence_first=True, minimum_gap=10, minimum_support=5,
          maximum_groups=32, core_rows=None):
    matrix, groups, audit = prepare(rows, metadata, native, maximum_groups)
    fitted = fit(matrix, [g['gene'] for g in groups], maximum_iterations=100)
    direct = gene_consistency(matrix, groups, metadata, fitted, core_rows) if core_rows is not None else {}
    calls = copy.deepcopy(native)
    decisions = {}
    for gene, result in fitted['genes'].items():
        decision = dict(reason='fallback', pairs=result['pairs'], score_gap=result['score_gap'])
        decisions[gene] = decision
        if not fitted['converged']:
            decision['reason'] = 'joint_fit_not_converged'
            continue
        if len(result['pairs']) != 1:
            decision['reason'] = 'tied_sequence_pairs'
            continue
        if result['score_gap'] is None or result['score_gap'] < minimum_gap:
            decision['reason'] = 'insufficient_sequence_pair_gap'
            continue
        a, b = result['pairs'][0]
        # Require distinct evidence for each copy in a heterozygote. Shared
        # placements alone cannot create a confident second allele.
        if a == b:
            support = [float(matrix[:, a].sum())] * 2
        else:
            support = [float(np.maximum(matrix[:, a] - matrix[:, b], 0).sum()),
                       float(np.maximum(matrix[:, b] - matrix[:, a], 0).sum())]
        decision['distinct_support'] = support
        if min(support) < minimum_support:
            decision['reason'] = 'insufficient_distinct_copy_support'
            continue
        if core_rows is not None:
            decision['gene_discrimination'] = direct[gene]
            if direct[gene]['score_gap'] is None or direct[gene]['score_gap'] < minimum_gap:
                decision['reason'] = 'insufficient_within_gene_discrimination'
                continue
        resolutions = {str(f): [decode_group(groups[i], metadata, f) for i in (a, b)] for f in (2, 4)}
        if not sequence_first:
            old = native.get(gene, {}).get('resolutions', {})
            if any(s.get('unresolved', True) for s in old.get('4', [{}])) or len(old.get('4', [])) != 2:
                decision['reason'] = 'native_four_field_unresolved_control'
                continue
            old_pair = sorted(tuple(s.get('alleles', [])) for s in old.get('2', []))
            new_pair = sorted(tuple(s['alleles']) for s in resolutions['2'])
            if old_pair != new_pair or any(s['unresolved'] for s in resolutions['2']):
                decision['reason'] = 'coarse_anchor_control'
                continue
        calls[gene] = dict(copy_count=1 if a == b else 2, resolutions=resolutions)
        decision['reason'] = 'sequence_pair_selected'
        decision['selected_paths'] = [groups[i]['paths'] for i in (a, b)]
    return calls, dict(decisions=decisions, candidate_audit=audit, fit=fitted,
                       limitations='Candidate-limited coordinate ascent; score gaps are uncalibrated; tied groups preserve label uncertainty.')