"""Simulation benchmark track."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd

from babappa.config import BABAPPAConfig
from babappa.io.fasta import CodonAlignment, read_codon_alignment
from babappa.io.newick import read_newick
from babappa.null.model import make_default_null_model
from babappa.null.simulator import simulate_alignment
from babappa.workflow import fit_null_bundle, score_genes


def _mutate_alignment(aln: CodonAlignment, seed: int, frac_sites: float, strong: bool) -> CodonAlignment:
    rng = np.random.default_rng(seed)
    seqs = {k: list(v) for k, v in aln.sequences.items()}
    n_codons = aln.codon_length
    n_sites = max(1, int(frac_sites * n_codons))
    sites = rng.choice(n_codons, size=n_sites, replace=False)

    nts = np.array(["A", "C", "G", "T"])
    for _taxon, seq_chars in seqs.items():
        for s in sites:
            if strong or rng.uniform() < 0.5:
                codon = "".join(rng.choice(nts, size=3))
                start = s * 3
                seq_chars[start : start + 3] = list(codon)
    return CodonAlignment(gene=aln.gene, sequences={k: "".join(v) for k, v in seqs.items()})


def _add_gaps_ambiguity(aln: CodonAlignment, seed: int, frac: float = 0.1) -> CodonAlignment:
    rng = np.random.default_rng(seed)
    seqs = {k: list(v) for k, v in aln.sequences.items()}
    n = len(next(iter(aln.sequences.values())))
    n_mut = max(1, int(frac * n))
    for _taxon, chars in seqs.items():
        idx = rng.choice(n, size=n_mut, replace=False)
        for i in idx:
            chars[i] = "N" if rng.uniform() < 0.5 else "-"
    return CodonAlignment(gene=aln.gene, sequences={k: "".join(v) for k, v in seqs.items()})


def run_simulation_benchmark(
    cfg: BABAPPAConfig,
    outdir: str | Path,
    seed: int,
    n_trials: int = 24,
) -> pd.DataFrame:
    out = Path(outdir)
    out.mkdir(parents=True, exist_ok=True)

    example_dir = Path(__file__).resolve().parents[1] / "data" / "examples" / "small_demo"
    tree = read_newick(example_dir / "tree.nwk")
    reference = read_codon_alignment(example_dir / "geneA.fasta", gene_name="reference")

    bundle, _bundle_path, _ = fit_null_bundle(
        alignment=reference,
        tree=tree,
        cfg=cfg,
        outdir=out / "fit_null",
        seed=seed,
        energy_family="gaussian",
        null_replicates_train=max(100, cfg.fit.null_replicates_train),
    )

    null_model = make_default_null_model(
        alignment=reference,
        kappa=cfg.null_model.kappa,
        omega=cfg.null_model.omega,
        codon_frequency_mode=cfg.null_model.codon_frequency_mode,
        rate_variation=cfg.null_model.rate_variation,
        gamma_alpha=cfg.null_model.gamma_alpha,
    )

    scenarios = [
        "type1",
        "weak_distributed",
        "sparse_strong",
        "gaps_ambiguity",
        "tree_misspecification",
    ]

    rows: list[dict[str, object]] = []
    rng = np.random.default_rng(seed + 99)
    trial_seeds = [int(x) for x in rng.integers(1, 2**31 - 1, size=n_trials)]

    for scenario in scenarios:
        for i, s in enumerate(trial_seeds):
            obs = simulate_alignment(
                tree=tree,
                model=null_model,
                n_codons=reference.codon_length,
                gene_name=f"{scenario}_{i+1}",
                seed=s,
            )

            score_tree = tree
            if scenario == "weak_distributed":
                obs = _mutate_alignment(obs, seed=s + 3, frac_sites=0.20, strong=False)
            elif scenario == "sparse_strong":
                obs = _mutate_alignment(obs, seed=s + 5, frac_sites=0.05, strong=True)
            elif scenario == "gaps_ambiguity":
                obs = _add_gaps_ambiguity(obs, seed=s + 7, frac=0.12)
            elif scenario == "tree_misspecification":
                # Create mild misspecification by inflating branch lengths.
                from copy import deepcopy

                score_tree = deepcopy(tree)

                def _inflate(node):
                    for c in node.children:
                        c.length *= 1.8
                        _inflate(c)

                _inflate(score_tree.root)

            score_df = score_genes(
                alignments={obs.gene: obs},
                tree=score_tree,
                bundle=bundle,
                cfg=cfg,
                outdir=out / "scores" / scenario / f"trial_{i+1}",
                seed=s + 11,
                jobs=1,
                resume=False,
                n_replicates_score=299,
            )
            p = float(score_df.loc[0, "p_value"])
            rows.append(
                {
                    "scenario": scenario,
                    "trial": i + 1,
                    "p_value": p,
                    "significant_0_05": p < 0.05,
                }
            )

    df = pd.DataFrame(rows)
    df.to_csv(out / "simulation_trials.tsv", sep="\t", index=False)

    summary = (
        df.groupby("scenario", as_index=False)
        .agg(
            n=("p_value", "size"),
            significant=("significant_0_05", "sum"),
            rejection_rate=("significant_0_05", "mean"),
            median_p=("p_value", "median"),
        )
        .sort_values("scenario")
    )
    summary.to_csv(out / "simulation_summary.tsv", sep="\t", index=False)
    return summary
