import sys
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
SRC_DIR = ROOT / "src"
if str(SRC_DIR) not in sys.path:
    sys.path.insert(0, str(SRC_DIR))

import renunney.track1_reference as ref


def test_zero_mutation_preserves_zero_alleles_in_one_generation():
    params = ref.Track1Parameters(K=5000, N0=2, n=2, u=0.0, R=10.0, T=20)
    state = ref.PopulationState(
        genomes=np.zeros((2, 2, 2), dtype=np.int16),
        sexes=np.array([0, 1], dtype=np.int8),
    )
    next_state, summary = ref.simulate_one_generation(state, params, t=0, rng=np.random.default_rng(1))
    assert summary.mean_allele_value == 0.0
    assert int(next_state.genomes.sum()) == 0


def test_zero_offspring_fitness_prevents_survival(monkeypatch):
    params = ref.Track1Parameters(K=5000, N0=2, n=1, u=0.0, R=10.0, T=20)
    state = ref.PopulationState(
        genomes=np.zeros((2, 2, 1), dtype=np.int16),
        sexes=np.array([0, 1], dtype=np.int8),
    )

    def fake_fitness(genomes, r, T, t):
        if genomes.shape[0] == 1:
            return np.zeros(1, dtype=float)
        return np.ones(genomes.shape[0], dtype=float)

    monkeypatch.setattr(ref, "genotype_fitness", fake_fitness)
    monkeypatch.setattr(ref, "female_fecundity", lambda r, N, K: 4.0)
    monkeypatch.setattr(ref, "realize_birth_counts", lambda fecundity, sexes, rng: np.array([3, 0], dtype=np.int32))

    next_state, _ = ref.simulate_one_generation(state, params, t=0, rng=np.random.default_rng(2))
    assert next_state.size == 0


def test_unit_offspring_fitness_keeps_all_births(monkeypatch):
    params = ref.Track1Parameters(K=5000, N0=2, n=1, u=0.0, R=10.0, T=20)
    state = ref.PopulationState(
        genomes=np.zeros((2, 2, 1), dtype=np.int16),
        sexes=np.array([0, 1], dtype=np.int8),
    )

    monkeypatch.setattr(ref, "genotype_fitness", lambda genomes, r, T, t: np.ones(genomes.shape[0], dtype=float))
    monkeypatch.setattr(ref, "female_fecundity", lambda r, N, K: 4.0)
    monkeypatch.setattr(ref, "realize_birth_counts", lambda fecundity, sexes, rng: np.array([3, 0], dtype=np.int32))

    next_state, summary = ref.simulate_one_generation(state, params, t=0, rng=np.random.default_rng(2))
    assert summary.female_count == 1
    assert next_state.size == 3


def test_one_father_is_used_per_female_reproductive_event(monkeypatch):
    params = ref.Track1Parameters(K=5000, N0=3, n=2, u=0.0, R=10.0, T=20)
    state = ref.PopulationState(
        genomes=np.array(
            [
                [[0, 0], [0, 0]],
                [[1, 1], [1, 1]],
                [[2, 2], [2, 2]],
            ],
            dtype=np.int16,
        ),
        sexes=np.array([0, 1, 1], dtype=np.int8),
    )

    monkeypatch.setattr(ref, "genotype_fitness", lambda genomes, r, T, t: np.ones(genomes.shape[0], dtype=float))
    monkeypatch.setattr(ref, "female_fecundity", lambda r, N, K: 4.0)
    monkeypatch.setattr(ref, "realize_birth_counts", lambda fecundity, sexes, rng: np.array([3, 0, 0], dtype=np.int32))

    next_state, _ = ref.simulate_one_generation(state, params, t=0, rng=np.random.default_rng(3))
    paternal_alleles = next_state.genomes[:, 1, :]
    assert next_state.size == 3
    assert np.unique(paternal_alleles).size == 1


def test_absence_of_males_or_females_counts_as_extinction():
    female_only = ref.PopulationState(
        genomes=np.zeros((2, 2, 1), dtype=np.int16),
        sexes=np.array([0, 0], dtype=np.int8),
    )
    male_only = ref.PopulationState(
        genomes=np.zeros((2, 2, 1), dtype=np.int16),
        sexes=np.array([1, 1], dtype=np.int8),
    )
    assert ref.is_extinct(female_only) is True
    assert ref.is_extinct(male_only) is True