"""Tests for the bunbun Python interface.

Run after `maturin develop`:
    pytest tests/test_bunbun.py -v
"""

from pathlib import Path

import numpy as np
import polars as pl
import pytest


@pytest.fixture
def plink_fileset(tmp_path: Path):
    """Create a minimal PLINK 1.x fileset for testing."""
    # ── .fam ──────────────────────────────────────────────────────
    fam_lines = [
        "FAM1\tSAM1\t0\t0\t1\t-9\n",
        "FAM2\tSAM2\t0\t0\t2\t1.5\n",
        "FAM3\tSAM3\t0\t0\t0\t-9\n",
        "FAM4\tSAM4\t0\t0\t1\t2.0\n",
    ]
    fam_path = tmp_path / "test.fam"
    fam_path.write_text("".join(fam_lines))

    # ── .bim ──────────────────────────────────────────────────────
    bim_lines = [
        "1\trs1\t0.0\t100\tA\tG\n",
        "1\trs2\t0.0\t200\tC\tT\n",
        "2\trs3\t0.5\t300\tG\tA\n",
    ]
    bim_path = tmp_path / "test.bim"
    bim_path.write_text("".join(bim_lines))

    # ── .bed ──────────────────────────────────────────────────────
    # SNP-major format, 4 samples → 1 byte per variant
    # Genotype encoding (2 bits): 00=hom-ref, 01=missing, 10=het, 11=hom-alt
    #
    # Variant 0: SAM1=hom-ref(00), SAM2=het(10), SAM3=hom-alt(11), SAM4=missing(01)
    # byte = 01_11_10_00 = 0b01111000 = 0x78
    #
    # Variant 1: all hom-ref
    # byte = 00_00_00_00 = 0x00
    #
    # Variant 2: SAM1=het(10), SAM2=het(10), SAM3=hom-ref(00), SAM4=hom-alt(11)
    # byte = 11_00_10_10 = 0b11001010 = 0xCA

    magic = bytes([0x6C, 0x1B, 0x01])  # PLINK magic + SNP-major mode
    bed_data = magic + bytes([0x78, 0x00, 0xCA])
    bed_path = tmp_path / "test.bed"
    bed_path.write_bytes(bed_data)

    return tmp_path / "test"


@pytest.fixture
def vcf_file(tmp_path: Path):
    """Create a minimal VCF file."""
    vcf_content = """##fileformat=VCFv4.2
##FORMAT=<ID=GT,Number=1,Type=String>
#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tSAM1\tSAM2\tSAM3
1\t100\trs1\tA\tG\t30\tPASS\t.\tGT\t0/0\t0/1\t1/1
1\t200\trs2\tC\tT\t25\tPASS\t.\tGT\t0/1\t0/0\t./.
2\t300\t.\tG\tA,C\t40\tPASS\t.\tGT\t0/0\t1/1\t0/1
"""
    vcf_path = tmp_path / "test.vcf"
    vcf_path.write_text(vcf_content)
    return vcf_path


class TestPlinkReader:
    def test_basic_read(self, plink_fileset):
        import bunbun

        data = bunbun.read(str(plink_fileset) + ".bed")

        assert data.n_samples == 4
        assert data.n_variants == 3
        assert data.shape == (4, 3)
        assert data.source_format == "plink"

    def test_samples_dataframe(self, plink_fileset):
        import bunbun

        data = bunbun.read_plink(str(plink_fileset) + ".bed")
        samples = data.samples

        assert isinstance(samples, pl.DataFrame)
        assert "sample_id" in samples.columns
        assert "sex" in samples.columns
        assert samples["sample_id"].to_list() == ["SAM1", "SAM2", "SAM3", "SAM4"]
        assert samples["sex"].to_list() == ["male", "female", "unknown", "male"]

    def test_variants_dataframe(self, plink_fileset):
        import bunbun

        data = bunbun.read_plink(str(plink_fileset) + ".bed")
        variants = data.variants

        assert isinstance(variants, pl.DataFrame)
        assert "chrom" in variants.columns
        assert "pos" in variants.columns
        assert "ref" in variants.columns
        assert variants["id"].to_list() == ["rs1", "rs2", "rs3"]
        assert variants["chrom"].to_list() == ["1", "1", "2"]

    def test_dosage_matrix(self, plink_fileset):
        import bunbun

        data = bunbun.read_plink(str(plink_fileset) + ".bed")
        dosages = data.dosages

        assert isinstance(dosages, np.ndarray)
        assert dosages.dtype == np.float32
        assert dosages.shape == (4, 3)

        # Variant 0: SAM1=0, SAM2=1, SAM3=2, SAM4=NaN
        assert dosages[0, 0] == 0.0
        assert dosages[1, 0] == 1.0
        assert dosages[2, 0] == 2.0
        assert np.isnan(dosages[3, 0])

        # Variant 1: all hom-ref
        np.testing.assert_array_equal(dosages[:, 1], [0.0, 0.0, 0.0, 0.0])

    def test_allele_frequencies(self, plink_fileset):
        import bunbun

        data = bunbun.read_plink(str(plink_fileset) + ".bed")
        freqs = data.allele_frequencies()

        assert isinstance(freqs, np.ndarray)
        assert len(freqs) == 3
        # Variant 0: (0+1+2) / (2*3) = 0.5 (SAM4 is missing)
        assert abs(freqs[0] - 0.5) < 0.01

    def test_mean_impute(self, plink_fileset):
        import bunbun

        data = bunbun.read_plink(str(plink_fileset) + ".bed")
        assert np.any(np.isnan(data.dosages))

        data.mean_impute()
        assert not np.any(np.isnan(data.dosages))

    def test_repr(self, plink_fileset):
        import bunbun

        data = bunbun.read_plink(str(plink_fileset) + ".bed")
        r = repr(data)
        assert "GenotypeData" in r
        assert "n_samples=4" in r
        assert "n_variants=3" in r


class TestVcfReader:
    def test_basic_read(self, vcf_file):
        import bunbun

        data = bunbun.read_vcf(str(vcf_file))
        assert data.n_samples == 3
        assert data.n_variants == 3
        assert data.source_format == "vcf"

    def test_genotypes(self, vcf_file):
        import bunbun

        data = bunbun.read_vcf(str(vcf_file))
        d = data.dosages

        # Variant 0: SAM1=0/0→0, SAM2=0/1→1, SAM3=1/1→2
        assert d[0, 0] == 0.0
        assert d[1, 0] == 1.0
        assert d[2, 0] == 2.0

        # Variant 1: SAM3=./.→NaN
        assert np.isnan(d[2, 1])

    def test_multiallelic_id(self, vcf_file):
        import bunbun

        data = bunbun.read_vcf(str(vcf_file))
        ids = data.variants["id"].to_list()
        # Third variant has no rsID → should be chrom:pos:ref:alt
        assert ids[2] == "2:300:G:A,C"


class TestAutoDetect:
    def test_plink_detection(self, plink_fileset):
        import bunbun

        # Should work with any of the three extensions
        data = bunbun.read(str(plink_fileset) + ".bed")
        assert data.source_format == "plink"

    def test_vcf_detection(self, vcf_file):
        import bunbun

        data = bunbun.read(str(vcf_file))
        assert data.source_format == "vcf"


class TestSupportedFormats:
    def test_returns_dict(self):
        import bunbun

        fmts = bunbun.supported_formats()
        assert isinstance(fmts, dict)
        assert "plink" in fmts
        assert "vcf" in fmts
        assert ".bed" in fmts["plink"]
