"""
Tests for the deterministic half of the pipeline.

The model outputs cannot be asserted on, but the glue can -- and the glue is
where a silent error does the most damage. An off-by-one in the TIS coordinate
shifts every reading frame and still produces a plausible-looking protein, so
these cases exist mainly to pin that down.
"""

import orf
import pipeline


class TestTranslate:
    def test_basic(self):
        assert orf.translate("ATGGCCTAA") == "MA"

    def test_stops_at_stop(self):
        assert orf.translate("ATGGCCTAAATGGGG") == "MA"

    def test_can_read_through(self):
        assert orf.translate("ATGTAAATG", stop_at_stop=False) == "M*M"

    def test_rna_input(self):
        assert orf.translate("AUGGCCUAA") == "MA"

    def test_n_becomes_x(self):
        assert orf.translate("ATGNNNGCC") == "MXA"

    def test_trailing_partial_codon_ignored(self):
        assert orf.translate("ATGGCCA") == "MA"


class TestExtractORF:
    def test_finds_in_frame_stop(self):
        o = orf.extract_orf("AAATGGCCTAAGGG", 2)
        assert o.start_codon == "ATG"
        assert o.protein == "MA"
        assert o.stop_pos == 8
        assert o.has_stop

    def test_skips_out_of_frame_stop(self):
        # ATG CTA AGC TAA -- there is a TAA at index 4 which is out of frame
        # and must be ignored in favour of the in-frame one at index 9.
        o = orf.extract_orf("ATGCTAAGCTAA", 0)
        assert o.has_stop
        assert o.stop_pos == 9
        assert o.protein == "MLS"

    def test_runs_off_end(self):
        o = orf.extract_orf("ATGGCCGCC", 0)
        assert not o.has_stop
        assert o.stop_pos is None
        assert o.protein == "MAA"

    def test_microprotein_threshold(self):
        short = orf.extract_orf("ATG" + "GCC" * 50 + "TAA", 0)
        assert short.aa_len == 51
        assert short.is_microprotein
        long = orf.extract_orf("ATG" + "GCC" * 120 + "TAA", 0)
        assert not long.is_microprotein


class TestCoordinateConversion:
    """TIS Transformer reports TIS_pos 1-based; everything here is 0-based."""

    def test_one_based_to_zero_based(self):
        assert orf.tis_pos_to_zero_based(1) == 0
        assert orf.tis_pos_to_zero_based(60) == 59

    def test_round_trip_against_real_values(self):
        # Observed: INS TIS_pos 60 is the annotated CDS start at 0-based 59.
        assert orf.tis_pos_to_zero_based(60) == 59   # INS
        assert orf.tis_pos_to_zero_based(379) == 378  # TYK2
        assert orf.tis_pos_to_zero_based(85) == 84    # ACTB


class TestClassifyORF:
    # A transcript with CDS at [100, 400).
    CDS_START, CDS_END = 100, 400

    def _at(self, pos, stop=None):
        o = orf.extract_orf("ATG" + "GCC" * 10 + "TAA", 0)
        o.tis_pos = pos
        o.stop_pos = stop if stop is not None else pos + 30
        return o

    def test_canonical(self):
        assert orf.classify_orf(self._at(100), self.CDS_START, self.CDS_END) == "CDS"

    def test_uorf_entirely_upstream(self):
        assert orf.classify_orf(self._at(10, 40), self.CDS_START, self.CDS_END) == "uORF"

    def test_n_extension_is_in_frame_upstream(self):
        # 100 - 91 = 9, divisible by 3, and it overlaps the CDS.
        assert orf.classify_orf(self._at(91, 200), self.CDS_START, self.CDS_END) == "N-ext"

    def test_uoorf_is_out_of_frame_upstream_overlap(self):
        assert orf.classify_orf(self._at(92, 200), self.CDS_START, self.CDS_END) == "uoORF"

    def test_n_truncation_in_frame_inside_cds(self):
        assert orf.classify_orf(self._at(163), self.CDS_START, self.CDS_END) == "N-trunc"

    def test_internal_orf_out_of_frame(self):
        assert orf.classify_orf(self._at(164), self.CDS_START, self.CDS_END) == "intORF"

    def test_downstream_orf(self):
        assert orf.classify_orf(self._at(450), self.CDS_START, self.CDS_END) == "dORF"

    def test_unannotated_is_ncorf(self):
        assert orf.classify_orf(self._at(50), None, None) == "ncORF"


class TestTilePeptides:
    def test_lengths(self):
        peps = orf.tile_peptides("ACDEFGHIKLMNP", lengths=(9,))
        assert all(len(p) == 9 for p in peps)
        assert "ACDEFGHIK" in peps

    def test_too_short_yields_nothing(self):
        assert orf.tile_peptides("ACDEF", lengths=(9,)) == []

    def test_initiator_met_removal(self):
        # Without Met removal the 8mer starting at position 1 is unreachable.
        peps = orf.tile_peptides("MACDEFGHI", lengths=(8,), drop_initiator_met=True)
        assert "ACDEFGHI" in peps
        assert "MACDEFGH" in peps

    def test_drops_peptides_with_x(self):
        peps = orf.tile_peptides("ACDXFGHIKLMN", lengths=(9,))
        assert all("X" not in p for p in peps)


class TestHeaderAnnotation:
    def test_parses_real_header(self):
        header = (
            "INS_ENST00000381330.5 Homo sapiens insulin (INS), MANE Select, "
            "cDNA len=465nt 5UTR=59nt CDS=333nt 3UTR=73nt | Ensembl GRCh38"
        )
        start, end = orf.parse_cds_annotation(header)
        assert start == 59
        assert end == 59 + 333

    def test_missing_tags_returns_none(self):
        assert orf.parse_cds_annotation("just a name") == (None, None)


class TestFastaParsing:
    def test_name_is_first_token(self):
        recs, headers = pipeline.parse_fasta_with_headers(
            ">INS_ENST1 some description here 5UTR=10nt CDS=30nt\nATGGCC\n"
        )
        assert recs[0][0] == "INS_ENST1"
        assert "5UTR=10nt" in headers["INS_ENST1"]

    def test_bare_sequence_gets_a_header(self):
        recs, _ = pipeline.parse_fasta_with_headers("ATGGCCTAA")
        assert recs == [("query", "ATGGCCTAA")]

    def test_rna_normalised(self):
        recs, _ = pipeline.parse_fasta_with_headers(">x\nAUGGCC")
        assert recs[0][1] == "ATGGCC"

    def test_rejects_protein_sequence(self):
        try:
            pipeline.parse_fasta_with_headers(">x\nMAEEGFKL")
            assert False, "should have rejected amino acids"
        except ValueError as e:
            assert "unexpected characters" in str(e)


class TestNovelty:
    def test_self_peptide_flagged_not_novel(self):
        import pandas as pd

        canonical = {"tx": "MAEEGFKLDYVPT"}
        df = pd.DataFrame({"peptide": ["EEGFKLDY", "QQQWWWYY"]})
        out = pipeline.flag_peptide_novelty(df, canonical)
        assert out.loc[0, "novel"] is False or not out.loc[0, "novel"]
        assert out.loc[1, "novel"]

    def test_no_false_join_across_proteins(self):
        """A peptide spanning two concatenated proteins must not count as self."""
        import pandas as pd

        canonical = {"a": "MAEEG", "b": "FKLDY"}
        df = pd.DataFrame({"peptide": ["AEEGFKLD"]})
        out = pipeline.flag_peptide_novelty(df, canonical)
        assert out.loc[0, "novel"]


class TestCutoffSplit:
    """Stage 1 runs at the floor, so the cutoff is applied here instead."""

    @staticmethod
    def _table():
        import pandas as pd

        return pd.DataFrame({
            pipeline.TIS_SCORE_COL: [0.9, 0.03, 0.0037, 0.0008],
            "orf_type": ["CDS", "uORF", "N-trunc", "dORF"],
        })

    def test_cutoff_is_inclusive(self):
        above, below = pipeline.split_at_cutoff(self._table(), 0.03)
        assert list(above["orf_type"]) == ["CDS", "uORF"]

    def test_rejected_sites_come_back_ranked(self):
        _, below = pipeline.split_at_cutoff(self._table(), 0.03)
        assert list(below["orf_type"]) == ["N-trunc", "dORF"]

    def test_empty_input_gives_two_empty_frames(self):
        import pandas as pd

        above, below = pipeline.split_at_cutoff(pd.DataFrame(), 0.03)
        assert above.empty and below.empty