"""
track1_reference.py

Local Track 1 reference module for renunney.

This is the historically faithful Nunney-style simulation kernel used by the
local analysis, threshold, and orchestration layers.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import numpy as np


@dataclass(frozen=True, init=False)
class Track1Parameters:
    """Reference parameters for the Track 1 baseline."""

    K: int
    N0: int
    n: int
    u: float
    R: float
    T: int
    epochs: int = 8
    p: float = 0.5
    a_max: Optional[int] = None

    def __init__(
        self,
        K: int,
        N0: int,
        n: int,
        u: float | None = None,
        R: float = 10.0,
        T: int = 300,
        epochs: int = 8,
        p: float = 0.5,
        a_max: Optional[int] = None,
        M: float | None = None,
    ) -> None:
        if u is None:
            if M is None:
                raise ValueError("Track1Parameters requires u, or convenience M to derive u.")
            u = float(M / (2.0 * K))
        object.__setattr__(self, "K", int(K))
        object.__setattr__(self, "N0", int(N0))
        object.__setattr__(self, "n", int(n))
        object.__setattr__(self, "u", float(u))
        object.__setattr__(self, "R", float(R))
        object.__setattr__(self, "T", int(T))
        object.__setattr__(self, "epochs", int(epochs))
        object.__setattr__(self, "p", float(p))
        object.__setattr__(self, "a_max", a_max)

    @property
    def M(self) -> float:
        return float(2.0 * self.K * self.u)

    def resolved_a_max(self) -> int:
        if self.a_max is not None:
            return self.a_max
        return self.epochs

    def intrinsic_growth_rate(self) -> float:
        return float(np.log(self.R / 2.0))

    def total_generations(self) -> int:
        return self.epochs * self.T + (self.T // 2)


@dataclass
class PopulationState:
    """Diploid population state for one generation."""

    genomes: np.ndarray
    sexes: np.ndarray

    @property
    def size(self) -> int:
        return int(self.genomes.shape[0])


@dataclass(frozen=True)
class GenerationSummary:
    """Per-generation diagnostics for Track 1."""

    t: int
    N: int
    female_fraction: float
    male_count: int
    female_count: int
    fecundity: float
    mean_fitness: float
    mean_expected_female_productivity: float
    target_value: float
    mean_allele_value: float
    mean_genotype_value: float
    mean_tracking_gap: float
    paper_M: float
    expected_mutations_current_N: float
    realized_mutation_count: int
    realized_mutation_rate_per_allele: float
    birth_count: int
    surviving_offspring_count: int
    ne_approx: float
    extinct: bool


def initialize_population(params: Track1Parameters, rng: np.random.Generator) -> PopulationState:
    genomes = np.zeros((params.N0, 2, params.n), dtype=np.int16)
    sexes = rng.binomial(1, 1.0 - params.p, size=params.N0).astype(np.int8)
    return PopulationState(genomes=genomes, sexes=sexes)


def female_fecundity(r: float, N: int, K: int) -> float:
    return float(2.0 * np.exp(r * (1.0 - (N / K) ** (1.0 / r))))


def genotype_fitness(genomes: np.ndarray, r: float, T: int, t: int) -> np.ndarray:
    target = t / T
    locus_means = 0.5 * (genomes[:, 0, :] + genomes[:, 1, :])
    squared_distance = np.square(locus_means - target, dtype=np.float64)
    return np.exp(-(r / genomes.shape[2]) * np.sum(squared_distance, axis=1))


def expected_female_productivity(fecundity: float, fitness: np.ndarray) -> np.ndarray:
    return fecundity * fitness


def paper_mutation_supply_M(K: int, u: float) -> float:
    return float(2.0 * K * u)


def expected_mutations_for_population(N: int, n: int, u: float) -> float:
    return float(2.0 * N * n * u)


def allele_tracking_metrics(genomes: np.ndarray, T: int, t: int) -> tuple[float, float, float, float]:
    target_value = float(t / T)
    if genomes.size == 0:
        return target_value, 0.0, 0.0, -target_value
    mean_allele_value = float(np.mean(genomes))
    genotype_values = 0.5 * (genomes[:, 0, :] + genomes[:, 1, :])
    mean_genotype_value = float(np.mean(genotype_values))
    mean_tracking_gap = mean_genotype_value - target_value
    return target_value, mean_allele_value, mean_genotype_value, mean_tracking_gap


def approximate_ne(N: int, female_fraction: float) -> float:
    male_fraction = 1.0 - female_fraction
    denom = female_fraction + male_fraction
    if N <= 0 or female_fraction <= 0.0 or male_fraction <= 0.0 or denom == 0.0:
        return 0.0
    return float((2.0 * N * female_fraction * male_fraction) / denom)


def female_fraction(sexes: np.ndarray) -> float:
    if sexes.size == 0:
        return 0.0
    return float(np.sum(sexes == 0) / sexes.size)


def is_extinct(state: PopulationState) -> bool:
    if state.size == 0:
        return True
    female_count = int(np.sum(state.sexes == 0))
    male_count = int(np.sum(state.sexes == 1))
    return female_count == 0 or male_count == 0


def choose_gamete(parent: np.ndarray, rng: np.random.Generator) -> np.ndarray:
    n = parent.shape[1]
    picks = rng.integers(0, 2, size=n)
    gamete = np.empty(n, dtype=parent.dtype)
    gamete[:] = parent[picks, np.arange(n)]
    return gamete


def choose_mutant_allele(current: int, a_max: int, rng: np.random.Generator) -> int:
    if a_max <= 0:
        return current
    draw = int(rng.integers(0, a_max))
    return draw if draw < current else draw + 1


def mutate_zygote(zygote: np.ndarray, params: Track1Parameters, rng: np.random.Generator) -> np.ndarray:
    a_max = params.resolved_a_max()
    out = zygote.copy()
    for chrom in range(out.shape[0]):
        for locus in range(out.shape[1]):
            if rng.random() <= params.u:
                out[chrom, locus] = choose_mutant_allele(int(out[chrom, locus]), a_max, rng)
    return out


def produce_offspring(
    mother: np.ndarray,
    father: np.ndarray,
    params: Track1Parameters,
    rng: np.random.Generator,
    locus_indices: np.ndarray,
) -> tuple[np.ndarray, int]:
    n = params.n
    offspring = np.empty((2, n), dtype=mother.dtype)
    maternal_picks = rng.integers(0, 2, size=n)
    paternal_picks = rng.integers(0, 2, size=n)
    offspring[0, :] = mother[maternal_picks, locus_indices]
    offspring[1, :] = father[paternal_picks, locus_indices]
    a_max = params.resolved_a_max()
    mutation_count = 0
    for chrom in range(2):
        for locus in range(n):
            if rng.random() <= params.u:
                offspring[chrom, locus] = choose_mutant_allele(int(offspring[chrom, locus]), a_max, rng)
                mutation_count += 1
    return offspring, mutation_count


def realize_birth_counts(fecundity: float, sexes: np.ndarray, rng: np.random.Generator) -> np.ndarray:
    female_mask = sexes == 0
    counts = np.zeros(sexes.shape[0], dtype=np.int32)
    counts[female_mask] = rng.poisson(fecundity, size=int(np.sum(female_mask)))
    return counts


def summarize_generation(state: PopulationState, params: Track1Parameters, t: int) -> GenerationSummary:
    if state.size == 0:
        return GenerationSummary(
            t=t,
            N=0,
            female_fraction=0.0,
            male_count=0,
            female_count=0,
            fecundity=0.0,
            mean_fitness=0.0,
            mean_expected_female_productivity=0.0,
            target_value=float(t / params.T),
            mean_allele_value=0.0,
            mean_genotype_value=0.0,
            mean_tracking_gap=float(-(t / params.T)),
            paper_M=params.M,
            expected_mutations_current_N=0.0,
            realized_mutation_count=0,
            realized_mutation_rate_per_allele=0.0,
            birth_count=0,
            surviving_offspring_count=0,
            ne_approx=0.0,
            extinct=True,
        )
    fit, fec, exp_fp, ff, female_count, male_count = generation_metrics(state, params, t)
    mean_expected_fp = float(np.mean(exp_fp[state.sexes == 0])) if female_count > 0 else 0.0
    target_value, mean_allele_value, mean_genotype_value, mean_tracking_gap = allele_tracking_metrics(
        state.genomes,
        T=params.T,
        t=t,
    )
    return GenerationSummary(
        t=t,
        N=state.size,
        female_fraction=ff,
        male_count=male_count,
        female_count=female_count,
        fecundity=fec,
        mean_fitness=float(np.mean(fit)),
        mean_expected_female_productivity=mean_expected_fp,
        target_value=target_value,
        mean_allele_value=mean_allele_value,
        mean_genotype_value=mean_genotype_value,
        mean_tracking_gap=mean_tracking_gap,
        paper_M=params.M,
        expected_mutations_current_N=expected_mutations_for_population(state.size, params.n, params.u),
        realized_mutation_count=0,
        realized_mutation_rate_per_allele=0.0,
        birth_count=0,
        surviving_offspring_count=0,
        ne_approx=approximate_ne(state.size, ff),
        extinct=is_extinct(state),
    )


def generation_metrics(
    state: PopulationState,
    params: Track1Parameters,
    t: int,
) -> tuple[np.ndarray, float, np.ndarray, float, int, int]:
    r = params.intrinsic_growth_rate()
    fit = genotype_fitness(state.genomes, r=r, T=params.T, t=t)
    fec = female_fecundity(r=r, N=state.size, K=params.K)
    exp_fp = expected_female_productivity(fec, fit)
    female_count = int(np.sum(state.sexes == 0))
    male_count = int(state.size - female_count)
    ff = float(female_count / state.size) if state.size > 0 else 0.0
    return fit, fec, exp_fp, ff, female_count, male_count


def simulate_one_generation(
    state: PopulationState,
    params: Track1Parameters,
    t: int,
    rng: np.random.Generator,
) -> tuple[PopulationState, GenerationSummary]:
    if state.size == 0:
        summary = summarize_generation(state, params, t)
        return state, summary

    fit, fec, exp_fp, ff, female_count, male_count = generation_metrics(state, params, t)
    extinct = female_count == 0 or male_count == 0
    target_value, mean_allele_value, mean_genotype_value, mean_tracking_gap = allele_tracking_metrics(
        state.genomes,
        T=params.T,
        t=t,
    )
    expected_mutations_current_N = expected_mutations_for_population(state.size, params.n, params.u)
    summary = GenerationSummary(
        t=t,
        N=state.size,
        female_fraction=ff,
        male_count=male_count,
        female_count=female_count,
        fecundity=fec,
        mean_fitness=float(np.mean(fit)),
        mean_expected_female_productivity=float(np.mean(exp_fp[state.sexes == 0])) if female_count > 0 else 0.0,
        target_value=target_value,
        mean_allele_value=mean_allele_value,
        mean_genotype_value=mean_genotype_value,
        mean_tracking_gap=mean_tracking_gap,
        paper_M=paper_mutation_supply_M(params.K, params.u),
        expected_mutations_current_N=expected_mutations_current_N,
        realized_mutation_count=0,
        realized_mutation_rate_per_allele=0.0,
        birth_count=0,
        surviving_offspring_count=0,
        ne_approx=approximate_ne(state.size, ff),
        extinct=extinct,
    )
    if extinct:
        return state, summary

    birth_counts = realize_birth_counts(fec, state.sexes, rng)

    female_indices = np.flatnonzero(state.sexes == 0)
    male_indices = np.flatnonzero(state.sexes == 1)
    total_births = int(np.sum(birth_counts[female_indices]))

    new_genomes = np.zeros((total_births, 2, params.n), dtype=np.int16)
    new_sexes = np.zeros(total_births, dtype=np.int8)
    locus_indices = np.arange(params.n)
    r = params.intrinsic_growth_rate()
    offspring_t = t + 1

    survivor_cursor = 0
    realized_mutation_count = 0
    for mother_index in female_indices:
        count = int(birth_counts[mother_index])
        if count == 0:
            continue
        father_index = int(male_indices[rng.integers(0, male_indices.size)])
        for _ in range(count):
            offspring, offspring_mutations = produce_offspring(
                state.genomes[mother_index],
                state.genomes[father_index],
                params,
                rng,
                locus_indices,
            )
            realized_mutation_count += offspring_mutations
            offspring_fitness = float(
                genotype_fitness(offspring[np.newaxis, :, :], r=r, T=params.T, t=offspring_t)[0]
            )
            if rng.random() <= offspring_fitness:
                new_genomes[survivor_cursor] = offspring
                new_sexes[survivor_cursor] = int(rng.binomial(1, 1.0 - params.p))
                survivor_cursor += 1

    allele_exposures = 2 * total_births * params.n
    next_state = PopulationState(genomes=new_genomes[:survivor_cursor], sexes=new_sexes[:survivor_cursor])
    summary = GenerationSummary(
        t=summary.t,
        N=summary.N,
        female_fraction=summary.female_fraction,
        male_count=summary.male_count,
        female_count=summary.female_count,
        fecundity=summary.fecundity,
        mean_fitness=summary.mean_fitness,
        mean_expected_female_productivity=summary.mean_expected_female_productivity,
        target_value=summary.target_value,
        mean_allele_value=summary.mean_allele_value,
        mean_genotype_value=summary.mean_genotype_value,
        mean_tracking_gap=summary.mean_tracking_gap,
        paper_M=summary.paper_M,
        expected_mutations_current_N=summary.expected_mutations_current_N,
        realized_mutation_count=realized_mutation_count,
        realized_mutation_rate_per_allele=0.0
        if allele_exposures == 0
        else float(realized_mutation_count / allele_exposures),
        birth_count=total_births,
        surviving_offspring_count=survivor_cursor,
        ne_approx=summary.ne_approx,
        extinct=summary.extinct,
    )
    return next_state, summary


def simulate_run(
    params: Track1Parameters,
    seed: Optional[int] = None,
) -> list[GenerationSummary]:
    rng = np.random.default_rng(seed)
    state = initialize_population(params, rng)
    t = -(params.T // 2)
    summaries: list[GenerationSummary] = []

    for _ in range(params.total_generations()):
        state, summary = simulate_one_generation(state, params, t, rng)
        summaries.append(summary)
        if is_extinct(state):
            terminal_t = t if summary.extinct else t + 1
            terminal_summary = summarize_generation(state, params, terminal_t)
            if summaries[-1] != terminal_summary:
                summaries.append(terminal_summary)
            break
        t += 1

    return summaries