"""Tests for gbid vectorized functions:
    encode_variants_vectorized, encode_positions_vectorized,
    encode_spdi_variants_vectorized.
"""

import os

import numpy as np
import pandas as pd
import pytest

from genebe.gbid import (
    InvalidPositionException,
    encode_positions_vectorized,
    encode_spdi_variants_vectorized,
    encode_spdi_variant_gbid,
    encode_variants_vectorized,
    encode_vcf_position_gbid,
    encode_vcf_variant_gbid,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_df(rows):
    """Build a minimal DataFrame with chr/pos/ref/alt columns."""
    return pd.DataFrame(rows, columns=["chr", "pos", "ref", "alt"])


def _scalar_results(df):
    """Reference results computed via the scalar function."""
    return [
        encode_vcf_variant_gbid(row["chr"], row["pos"], row["ref"], row["alt"])
        for _, row in df.iterrows()
    ]


# ---------------------------------------------------------------------------
# 1. Basic correctness – vectorized == scalar for every row
# ---------------------------------------------------------------------------


class TestCorrectnessVsScalar:
    """The vectorized function must produce identical results to the scalar one."""

    def test_snvs(self):
        df = _make_df(
            [
                ("1", 12, "A", "C"),
                ("1", 100, "G", "T"),
                ("2", 50000, "C", "A"),
                ("X", 1000000, "T", "G"),
                ("22", 17000000, "A", "T"),
            ]
        )
        np.testing.assert_array_equal(
            encode_variants_vectorized(df), _scalar_results(df)
        )

    def test_pure_insertions(self):
        df = _make_df(
            [
                ("1", 100, "", "A"),
                ("1", 200, "", "ACGT"),
                ("2", 300, "", "TTTT"),
            ]
        )
        np.testing.assert_array_equal(
            encode_variants_vectorized(df), _scalar_results(df)
        )

    def test_pure_deletions(self):
        df = _make_df(
            [
                ("1", 100, "A", ""),
                ("1", 200, "ACGT", ""),
                ("2", 300, "GGG", ""),
            ]
        )
        np.testing.assert_array_equal(
            encode_variants_vectorized(df), _scalar_results(df)
        )

    def test_mixed_variants(self):
        df = _make_df(
            [
                ("1", 16044378, "C", "CACACACACAT"),   # known insertion
                ("1", 248956422, "A", "T"),             # near end of chr
                ("M", 100, "G", "A"),                   # mitochondrial
                ("Y", 200, "T", "C"),
                ("1", 500, "GTA", "C"),                 # complex
                ("1", 600, "C", "GTT"),                 # complex
            ]
        )
        np.testing.assert_array_equal(
            encode_variants_vectorized(df), _scalar_results(df)
        )

    def test_shared_prefix_stripping(self):
        """Leading common bases must be stripped correctly."""
        df = _make_df(
            [
                ("1", 100, "ACCC", "ACGT"),   # strip 'A' prefix → pos 101, CCC vs CGT
                ("1", 200, "GA", "G"),         # strip 'G' prefix → pos 201, A vs ""
                ("1", 300, "TTAG", "TTAC"),    # strip 'TTA' → pos 303, G vs C
            ]
        )
        np.testing.assert_array_equal(
            encode_variants_vectorized(df), _scalar_results(df)
        )

    def test_chr_prefix_stripping(self):
        """Chromosomes written as 'chr1', 'CHR1', 'Chr1' should all normalise to '1'."""
        df_no_prefix = _make_df([("1", 12345, "A", "G")])
        df_chr       = _make_df([("chr1", 12345, "A", "G")])
        df_CHR       = _make_df([("CHR1", 12345, "A", "G")])
        df_Chr       = _make_df([("Chr1", 12345, "A", "G")])

        expected = encode_variants_vectorized(df_no_prefix)

        np.testing.assert_array_equal(encode_variants_vectorized(df_chr), expected)
        np.testing.assert_array_equal(encode_variants_vectorized(df_CHR), expected)
        np.testing.assert_array_equal(encode_variants_vectorized(df_Chr), expected)


# ---------------------------------------------------------------------------
# 2. Hash-path variants (large ins/del, N in alt)
# ---------------------------------------------------------------------------


class TestHashPath:
    """Rows that exceed size limits or contain N must use the change-hash path."""

    def test_large_insertion_triggers_hash(self):
        """Insertion longer than MAX_INS_LENGTH=9 should use hash and match scalar."""
        df = _make_df([("1", 16044378, "C", "CACACACACAT")])  # 10-char alt after strip
        np.testing.assert_array_equal(
            encode_variants_vectorized(df), _scalar_results(df)
        )

    def test_large_deletion_triggers_hash(self):
        """Deletion of > 127 bases should use hash path."""
        df = _make_df([("1", 1000, "A" * 150, "")])
        np.testing.assert_array_equal(
            encode_variants_vectorized(df), _scalar_results(df)
        )

    def test_n_in_alt_triggers_hash(self):
        """Alt sequence containing N should use hash path."""
        df = _make_df([("1", 500, "A", "N"), ("1", 600, "C", "ANTG")])
        np.testing.assert_array_equal(
            encode_variants_vectorized(df), _scalar_results(df)
        )

    def test_mixed_direct_and_hash_in_same_batch(self):
        """Batch with both direct-path and hash-path rows."""
        df = _make_df(
            [
                ("1", 100, "A", "G"),              # direct
                ("1", 200, "A" * 150, ""),          # hash (large del)
                ("1", 300, "C", "T"),              # direct
                ("1", 400, "", "ATGCATGCATGC"),    # hash (large ins)
            ]
        )
        np.testing.assert_array_equal(
            encode_variants_vectorized(df), _scalar_results(df)
        )


# ---------------------------------------------------------------------------
# 3. Return type and shape
# ---------------------------------------------------------------------------


class TestReturnType:
    def test_returns_numpy_array(self):
        df = _make_df([("1", 100, "A", "G")])
        result = encode_variants_vectorized(df)
        assert isinstance(result, np.ndarray)

    def test_dtype_is_int64(self):
        df = _make_df([("1", 100, "A", "G")])
        result = encode_variants_vectorized(df)
        assert result.dtype == np.int64

    def test_length_matches_input(self):
        rows = [("1", 100 + i, "A", "G") for i in range(50)]
        df = _make_df(rows)
        result = encode_variants_vectorized(df)
        assert len(result) == 50

    def test_single_row(self):
        df = _make_df([("1", 500, "C", "T")])
        result = encode_variants_vectorized(df)
        assert len(result) == 1
        assert result[0] == encode_vcf_variant_gbid("1", 500, "C", "T")


# ---------------------------------------------------------------------------
# 4. Cross-validation against the reference CSV
# ---------------------------------------------------------------------------


CSV_PATH = os.path.join(os.path.dirname(__file__), "variant-ids-hg38.csv")


@pytest.mark.skipif(not os.path.exists(CSV_PATH), reason="CSV fixture not found")
class TestAgainstReferenceCSV:
    """Validate against the pre-computed reference file."""

    def test_all_rows_match_expected_id(self):
        csv = pd.read_csv(CSV_PATH)
        # Only keep rows where both ref and alt are non-empty strings (NaN = empty allele in CSV)
        csv["ref"] = csv["ref"].fillna("")
        csv["alt"] = csv["alt"].fillna("")
        # Build DataFrame in the shape encode_variants_vectorized expects
        df = pd.DataFrame(
            {
                "chr": csv["chr"].astype(str),
                "pos": csv["pos1based"].astype(int),
                "ref": csv["ref"].astype(str),
                "alt": csv["alt"].astype(str),
            }
        )
        result = encode_variants_vectorized(df)
        expected = csv["id"].to_numpy(dtype=np.int64)
        np.testing.assert_array_equal(result, expected)

    def test_scalar_and_vectorized_agree_on_full_csv(self):
        """Scalar and vectorized must produce identical outputs for every CSV row."""
        csv = pd.read_csv(CSV_PATH)
        csv["ref"] = csv["ref"].fillna("")
        csv["alt"] = csv["alt"].fillna("")
        df = pd.DataFrame(
            {
                "chr": csv["chr"].astype(str),
                "pos": csv["pos1based"].astype(int),
                "ref": csv["ref"].astype(str),
                "alt": csv["alt"].astype(str),
            }
        )
        vectorized = encode_variants_vectorized(df)
        scalar = np.array(_scalar_results(df), dtype=np.int64)
        np.testing.assert_array_equal(vectorized, scalar)


# ---------------------------------------------------------------------------
# 5. Error handling
# ---------------------------------------------------------------------------


class TestErrorHandling:
    def test_invalid_chromosome_raises(self):
        df = _make_df([("25", 100, "A", "G")])  # chr 25 does not exist
        with pytest.raises(InvalidPositionException):
            encode_variants_vectorized(df)

    def test_invalid_bases_in_alt_raises(self):
        """Alt with non-ACGNT character should raise InvalidPositionException."""
        df = _make_df([("1", 100, "A", "X")])
        with pytest.raises(InvalidPositionException):
            encode_variants_vectorized(df)


# ===========================================================================
# encode_positions_vectorized
# ===========================================================================

POS_CSV_PATH = os.path.join(os.path.dirname(__file__), "position-ids-hg38.csv")


def _make_pos_df(rows):
    """Build a minimal DataFrame with chr/pos columns (1-based)."""
    return pd.DataFrame(rows, columns=["chr", "pos"])


class TestEncodePositionsVectorized:
    """encode_positions_vectorized must match encode_vcf_position_gbid row by row."""

    def test_basic_snvs(self):
        df = _make_pos_df(
            [
                ("1", 1),
                ("1", 248956422),
                ("2", 1),
                ("X", 1000000),
                ("Y", 50000),
                ("M", 100),
            ]
        )
        expected = [encode_vcf_position_gbid(r["chr"], r["pos"]) for _, r in df.iterrows()]
        np.testing.assert_array_equal(
            encode_positions_vectorized(df), np.array(expected, dtype=np.int64)
        )

    def test_chr_prefix_normalisation(self):
        """chr1 / CHR1 / Chr1 must all yield the same ID as '1'."""
        base = _make_pos_df([("1", 12345)])
        expected = encode_positions_vectorized(base)

        for prefix in ("chr1", "CHR1", "Chr1"):
            df = _make_pos_df([(prefix, 12345)])
            np.testing.assert_array_equal(encode_positions_vectorized(df), expected)

    def test_returns_numpy_int64(self):
        df = _make_pos_df([("1", 100)])
        result = encode_positions_vectorized(df)
        assert isinstance(result, np.ndarray)
        assert result.dtype == np.int64

    def test_length_matches_input(self):
        rows = [("1", 1 + i) for i in range(100)]
        df = _make_pos_df(rows)
        assert len(encode_positions_vectorized(df)) == 100

    def test_unsupported_chromosome_sentinel(self):
        """Chromosome '25' has no known offset → sentinel CHR_NOT_SUPPORTED (-1)."""
        df = _make_pos_df([("25", 100)])
        result = encode_positions_vectorized(df)
        assert result[0] == -1

    def test_out_of_range_position_sentinel(self):
        """Position beyond chromosome length → ERROR_WRONG_CHR_POSITION (-2)."""
        df = _make_pos_df([("1", 999_999_999)])   # chr1 is only ~249 Mbp
        result = encode_positions_vectorized(df)
        assert result[0] == -2

    @pytest.mark.skipif(not os.path.exists(POS_CSV_PATH), reason="CSV fixture not found")
    def test_all_csv_rows_match_expected_id(self):
        csv = pd.read_csv(POS_CSV_PATH)
        df = pd.DataFrame({"chr": csv["chr"].astype(str), "pos": csv["pos1based"].astype(int)})
        result = encode_positions_vectorized(df)
        expected = csv["id"].to_numpy(dtype=np.int64)
        np.testing.assert_array_equal(result, expected)

    @pytest.mark.skipif(not os.path.exists(POS_CSV_PATH), reason="CSV fixture not found")
    def test_scalar_and_vectorized_agree_on_full_csv(self):
        csv = pd.read_csv(POS_CSV_PATH)
        df = pd.DataFrame({"chr": csv["chr"].astype(str), "pos": csv["pos1based"].astype(int)})
        vectorized = encode_positions_vectorized(df)
        scalar = np.array(
            [encode_vcf_position_gbid(r["chr"], r["pos"]) for _, r in df.iterrows()],
            dtype=np.int64,
        )
        np.testing.assert_array_equal(vectorized, scalar)


# ===========================================================================
# encode_spdi_variants_vectorized
# ===========================================================================

VARIANT_CSV_PATH = os.path.join(os.path.dirname(__file__), "variant-ids-hg38.csv")


def _vcf_to_spdi_rows(chr_, pos1, ref, alt):
    """Convert a VCF record to SPDI (0-based pos, del_length, ins) for testing."""
    pos0 = pos1 - 1
    ref_up = ref.upper()
    alt_up = alt.upper()
    # Strip shared leading bases (mirrors VariantIdEncoder.encode1based)
    while ref_up and alt_up and ref_up[0] == alt_up[0]:
        pos0 += 1
        ref_up = ref_up[1:]
        alt_up = alt_up[1:]
    return chr_, pos0, len(ref_up), alt_up


def _make_spdi_df(rows):
    """Build a DataFrame with chr/pos/del_length/ins columns."""
    return pd.DataFrame(rows, columns=["chr", "pos", "del_length", "ins"])


class TestEncodeSpdiVariantsVectorized:
    """encode_spdi_variants_vectorized must match encode_spdi_variant_gbid row by row."""

    def _scalar(self, df):
        return [
            encode_spdi_variant_gbid(r["chr"], r["pos"], r["del_length"], r["ins"])
            for _, r in df.iterrows()
        ]

    def test_snvs(self):
        # SNVs in SPDI: ref length == 1, pos is 0-based
        rows = [_vcf_to_spdi_rows("1", pos, "A", alt) for pos, alt in [(12, "C"), (100, "G"), (200, "T")]]
        df = _make_spdi_df(rows)
        np.testing.assert_array_equal(encode_spdi_variants_vectorized(df), self._scalar(df))

    def test_pure_deletions(self):
        rows = [_vcf_to_spdi_rows("1", 100, "AAA", "")]
        df = _make_spdi_df(rows)
        np.testing.assert_array_equal(encode_spdi_variants_vectorized(df), self._scalar(df))

    def test_pure_insertions(self):
        rows = [_vcf_to_spdi_rows("1", 200, "", "ACGT")]
        df = _make_spdi_df(rows)
        np.testing.assert_array_equal(encode_spdi_variants_vectorized(df), self._scalar(df))

    def test_large_insertion_hash_path(self):
        """Insertion >MAX_INS_LENGTH (9) must take the change-hash path."""
        rows = [_vcf_to_spdi_rows("1", 16044378, "C", "CACACACACAT")]
        df = _make_spdi_df(rows)
        np.testing.assert_array_equal(encode_spdi_variants_vectorized(df), self._scalar(df))

    def test_large_deletion_hash_path(self):
        """Deletion >127 bases must take the change-hash path."""
        rows = [_vcf_to_spdi_rows("1", 1000, "A" * 150, "")]
        df = _make_spdi_df(rows)
        np.testing.assert_array_equal(encode_spdi_variants_vectorized(df), self._scalar(df))

    def test_n_in_ins_hash_path(self):
        rows = [_vcf_to_spdi_rows("1", 500, "A", "N"), _vcf_to_spdi_rows("1", 600, "C", "ANTG")]
        df = _make_spdi_df(rows)
        np.testing.assert_array_equal(encode_spdi_variants_vectorized(df), self._scalar(df))

    def test_mixed_direct_and_hash(self):
        rows = [
            _vcf_to_spdi_rows("1", 100, "A", "G"),
            _vcf_to_spdi_rows("1", 200, "A" * 150, ""),
            _vcf_to_spdi_rows("1", 300, "C", "T"),
            _vcf_to_spdi_rows("1", 400, "", "ATGCATGCATGC"),
        ]
        df = _make_spdi_df(rows)
        np.testing.assert_array_equal(encode_spdi_variants_vectorized(df), self._scalar(df))

    def test_chr_prefix_normalisation(self):
        base = _make_spdi_df([("1", 11, 1, "C")])
        expected = encode_spdi_variants_vectorized(base)
        for prefix in ("chr1", "CHR1", "Chr1"):
            df = _make_spdi_df([(prefix, 11, 1, "C")])
            np.testing.assert_array_equal(encode_spdi_variants_vectorized(df), expected)

    def test_lowercase_ins_is_normalised(self):
        """Lowercase ins must produce the same result as uppercase."""
        df_upper = _make_spdi_df([("1", 11, 1, "ACGT")])
        df_lower = _make_spdi_df([("1", 11, 1, "acgt")])
        np.testing.assert_array_equal(
            encode_spdi_variants_vectorized(df_lower),
            encode_spdi_variants_vectorized(df_upper),
        )

    def test_returns_numpy_int64(self):
        df = _make_spdi_df([("1", 11, 1, "C")])
        result = encode_spdi_variants_vectorized(df)
        assert isinstance(result, np.ndarray)
        assert result.dtype == np.int64

    def test_length_matches_input(self):
        rows = [("1", 100 + i, 1, "G") for i in range(50)]
        df = _make_spdi_df(rows)
        assert len(encode_spdi_variants_vectorized(df)) == 50

    def test_invalid_chromosome_raises(self):
        df = _make_spdi_df([("25", 100, 1, "G")])
        with pytest.raises(InvalidPositionException):
            encode_spdi_variants_vectorized(df)

    def test_invalid_bases_raises(self):
        df = _make_spdi_df([("1", 100, 1, "X")])
        with pytest.raises(InvalidPositionException):
            encode_spdi_variants_vectorized(df)

    def test_agrees_with_vcf_vectorized_on_snvs(self):
        """Same SNVs encoded via VCF path and SPDI path must produce identical IDs."""
        vcf_rows = [
            ("1", 12, "A", "C"),
            ("2", 50000, "G", "T"),
            ("X", 100000, "C", "A"),
        ]
        vcf_df = _make_df(vcf_rows)
        spdi_rows = [_vcf_to_spdi_rows(*r) for r in vcf_rows]
        spdi_df = _make_spdi_df(spdi_rows)
        np.testing.assert_array_equal(
            encode_variants_vectorized(vcf_df),
            encode_spdi_variants_vectorized(spdi_df),
        )

    @pytest.mark.skipif(not os.path.exists(VARIANT_CSV_PATH), reason="CSV fixture not found")
    def test_full_csv_matches_vcf_vectorized(self):
        """SPDI-variant encoding of the full reference CSV must match VCF-variant encoding."""
        csv = pd.read_csv(VARIANT_CSV_PATH)
        csv["ref"] = csv["ref"].fillna("")
        csv["alt"] = csv["alt"].fillna("")
        vcf_df = pd.DataFrame(
            {
                "chr": csv["chr"].astype(str),
                "pos": csv["pos1based"].astype(int),
                "ref": csv["ref"].astype(str),
                "alt": csv["alt"].astype(str),
            }
        )
        spdi_rows = [
            _vcf_to_spdi_rows(r["chr"], r["pos"], r["ref"], r["alt"])
            for _, r in vcf_df.iterrows()
        ]
        spdi_df = _make_spdi_df(spdi_rows)

        np.testing.assert_array_equal(
            encode_spdi_variants_vectorized(spdi_df),
            encode_variants_vectorized(vcf_df),
        )
