import unittest
from context_inference import emission_rows, prepare, infer, decode_group, hard_locus_rows


class InferenceTests(unittest.TestCase):
    def test_joint_ablation_discards_cross_locus_ambiguity(self):
        meta = {'a': {'gene': 'A'}, 'e': {'gene': 'E'}}
        self.assertEqual(hard_locus_rows([{'a': 1, 'e': 1}, {'a': 1}], meta), [{}, {'a': 1}])

    def test_flank_only_preference_cannot_change_gene_call(self):
        meta = {p: dict(gene='A', labels=[label]) for p, label in [('a', 'A*01:01:01:01'), ('b', 'A*02:01:01:01')]}
        rows = [{'a': 1, 'b': .2}] * 30
        # Perfect conserved gene sequence for both candidates; the full score
        # advantage comes from context alone and must not name a different allele.
        core = [{'a': 0, 'b': 0}] * 30
        calls, report = infer(rows, meta, {}, core_rows=core)
        self.assertEqual(calls, {})
        self.assertEqual(report['decisions']['A']['reason'], 'insufficient_within_gene_discrimination')
        calls, _ = infer(rows, meta, {}, core_rows=[{'a': 0, 'b': -10}] * 30)
        self.assertEqual(calls['A']['resolutions']['4'][0]['alleles'], ['A*01:01:01:01'])

    def test_native_and_graph_routes_are_one_fragment(self):
        meta = {'n': {'source': 'ipd'}, 'g': {'source': 'graph'}}
        record = dict(fragment='r', candidates={p: [dict(score=20, internal=True, gene_overlap=True)] for p in meta})
        names, rows = emission_rows([record], meta)
        self.assertEqual(names, ['r'])
        self.assertEqual(rows, [{'n': 1, 'g': 1}])

    def test_identical_evidence_does_not_multiply_candidates(self):
        meta = {p: dict(gene='A', labels=['A*01:01:01:01']) for p in ['n', 'g']}
        matrix, groups, audit = prepare([{'n': 1, 'g': 1}] * 10, meta, {})
        self.assertEqual(matrix.shape, (10, 1))
        self.assertEqual(groups[0]['paths'], ['g', 'n'])

    def test_unknown_genomic_label_is_not_fabricated(self):
        meta = {'g': dict(labels=[], cds_labels=['A*01:01:01:01'])}
        group = dict(paths=['g'])
        self.assertTrue(decode_group(group, meta, 4)['unresolved'])
        self.assertEqual(decode_group(group, meta, 2)['alleles'], ['A*01:01'])

    def test_resolves_without_native_four_field_pair(self):
        meta = {p: dict(gene='A', labels=[label]) for p, label in [('a', 'A*01:01:01:01'), ('b', 'A*02:01:01:01')]}
        rows = [{'a': 1}] * 30 + [{'b': 1}] * 30
        calls, report = infer(rows, meta, {}, sequence_first=True)
        self.assertEqual(calls['A']['resolutions']['4'][0]['alleles'], ['A*01:01:01:01'])
        self.assertEqual(calls['A']['resolutions']['4'][1]['alleles'], ['A*02:01:01:01'])
        anchored, _ = infer(rows, meta, {}, sequence_first=False)
        self.assertEqual(anchored, {})


if __name__ == '__main__':
    unittest.main()