"""Core end-to-end workflow operations for BABAPPA CLI."""

from __future__ import annotations

import json
import time
import tracemalloc
from concurrent.futures import ThreadPoolExecutor
from copy import deepcopy
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd

from babappa.config import BABAPPAConfig
from babappa.features.representation import RepresentationSpec, gene_representation
from babappa.features.smoothing import smooth_matrix_with_site_seed_schedule
from babappa.io.fasta import CodonAlignment
from babappa.io.newick import PhyloTree
from babappa.model.bundle import FrozenNullBundle, make_bundle, write_bundle
from babappa.model.energy import (
    FittedGaussian,
    fit_gaussian_energy,
    fitted_from_dict,
    gene_dispersion_statistic,
)
from babappa.null.model import CodonNullModel, make_default_null_model
from babappa.null.simulator import simulate_replicates
from babappa.qc.validators import QCResult, validate_alignment_tree
from babappa.stats.contributions import summarize_site_contributions
from babappa.stats.inference import effect_size_summary, infer_pvalue
from babappa.stats.multiple_testing import adjust_pvalues
from babappa.utils.fs import ensure_dir, write_json
from babappa.utils.random import deterministic_seed, spawn_seeds


@dataclass
class GeneScoreOutput:
    gene: str
    statistic: float
    p_value: float
    n_replicates: int
    mc_resolution: float
    null_mean: float
    null_sd: float
    delta_vs_null: float
    z_vs_null: float
    warning_flags: str


def representation_spec_from_config(cfg: BABAPPAConfig) -> RepresentationSpec:
    return RepresentationSpec(
        embedding_version=cfg.representation.embedding_version,
        embed_dim=8,
        projection_dim=cfg.representation.projection_dim,
        projection_seed=cfg.representation.projection_seed,
    )


def _diag_energy_fit(samples: np.ndarray, ridge: float) -> FittedGaussian:
    mu = np.mean(samples, axis=0)
    var = np.var(samples, axis=0, ddof=1) + ridge
    precision = np.diag(1.0 / var)
    pseudo = {
        "mean": mu.tolist(),
        "precision": precision.tolist(),
        "ridge": float(ridge),
        "n_training_sites": int(samples.shape[0]),
        "score_matching_objective": float(np.mean(var)),
        "convergence_flag": "closed_form_diagonal",
        "min_eigenvalue": float(np.min(var)),
        "max_eigenvalue": float(np.max(var)),
    }
    return fitted_from_dict(pseudo)


def fit_null_bundle(
    alignment: CodonAlignment,
    tree: PhyloTree,
    cfg: BABAPPAConfig,
    outdir: str | Path,
    seed: int,
    energy_family: str = "gaussian",
    null_replicates_train: int | None = None,
) -> tuple[FrozenNullBundle, Path, pd.DataFrame]:
    out = ensure_dir(outdir)

    spec = representation_spec_from_config(cfg)
    taxa_order = tree.leaf_names()

    null_model = make_default_null_model(
        alignment=alignment,
        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,
    )

    n_train = int(null_replicates_train or cfg.fit.null_replicates_train)
    train_sim_seed = deterministic_seed(seed, "train", alignment.gene, "simulate")
    reps = simulate_replicates(
        tree=tree,
        model=null_model,
        n_codons=alignment.codon_length,
        n_reps=n_train,
        base_seed=train_sim_seed,
        gene_prefix=f"{alignment.gene}_nulltrain",
    )

    blocks = []
    smoothed_by_gene: dict[str, np.ndarray] = {}
    for rep in reps:
        site_vectors, _ = gene_representation(rep, taxa_order=taxa_order, spec=spec)
        rep_smoothed = smooth_matrix_with_site_seed_schedule(
            site_vectors,
            cfg.representation.smoothing_bandwidth,
            seed,
            "train",
            alignment.gene,
            rep.gene,
            "smooth",
        )
        smoothed_by_gene[rep.gene] = rep_smoothed
        blocks.append(rep_smoothed)
    train_mat = np.vstack(blocks)

    energy_family_norm = energy_family.lower()
    if energy_family_norm == "gaussian":
        fitted = fit_gaussian_energy(train_mat, ridge=cfg.fit.ridge)
    elif energy_family_norm in {"diag", "diagonal", "diagonal_gaussian"}:
        fitted = _diag_energy_fit(train_mat, ridge=cfg.fit.ridge)
    else:
        raise ValueError("Unsupported energy family; use 'gaussian' or 'diagonal_gaussian'")

    rep_stats = []
    for rep in reps:
        rep_vectors = smoothed_by_gene[rep.gene]
        stat, _ = gene_dispersion_statistic(rep_vectors, fitted=fitted)
        rep_stats.append({"gene": rep.gene, "statistic": stat})
    diag_df = pd.DataFrame(rep_stats)
    diag_path = out / "null_training_diagnostics.tsv"
    diag_df.to_csv(diag_path, sep="\t", index=False)

    bundle = make_bundle(
        taxa_order=taxa_order,
        representation_spec=spec.to_dict(),
        smoothing_bandwidth=cfg.representation.smoothing_bandwidth,
        null_model={
            **null_model.to_dict(),
            "energy_family": energy_family_norm,
        },
        energy_model=fitted.model,
        training={
            "n_null_replicates_train": n_train,
            "n_training_sites": int(train_mat.shape[0]),
            "training_gene": alignment.gene,
            "training_gene_codon_length": alignment.codon_length,
            "seed": int(seed),
            "seed_schedule": "deterministic_seed(base_seed, stage, gene, replicate, phase)",
        },
    )
    bundle_path = out / "null_bundle.json"
    write_bundle(bundle_path, bundle)

    return bundle, bundle_path, diag_df


def _codon_model_from_bundle(bundle: FrozenNullBundle) -> CodonNullModel:
    payload = dict(bundle.null_model)
    payload.pop("energy_family", None)
    return CodonNullModel(**payload)


def _score_one_gene_with_bundle(
    gene_name: str,
    alignment: CodonAlignment,
    tree: PhyloTree,
    bundle: FrozenNullBundle,
    cfg: BABAPPAConfig,
    seed: int,
    outdir: Path,
    n_replicates_score: int | None = None,
    two_sided: bool | None = None,
) -> GeneScoreOutput:
    spec = RepresentationSpec(**bundle.representation_spec)
    fitted = fitted_from_dict(bundle.energy_model)

    qc: QCResult = validate_alignment_tree(alignment, tree, cfg.qc)

    site_vectors, _ = gene_representation(alignment, taxa_order=bundle.taxa_order, spec=spec)
    site_vectors_smoothed = smooth_matrix_with_site_seed_schedule(
        site_vectors,
        bundle.smoothing_bandwidth,
        seed,
        "obs",
        gene_name,
        "smooth",
    )
    obs_stat, per_site = gene_dispersion_statistic(site_vectors_smoothed, fitted)

    n_reps = int(n_replicates_score or cfg.score.null_replicates_score)
    model = _codon_model_from_bundle(bundle)

    reps = simulate_replicates(
        tree=tree,
        model=model,
        n_codons=alignment.codon_length,
        n_reps=n_reps,
        base_seed=deterministic_seed(seed, "calib", gene_name, "simulate"),
        gene_prefix=f"{gene_name}_score_null",
    )
    null_stats = []
    for rep in reps:
        rep_vectors, _ = gene_representation(rep, taxa_order=bundle.taxa_order, spec=spec)
        rep_smoothed = smooth_matrix_with_site_seed_schedule(
            rep_vectors,
            bundle.smoothing_bandwidth,
            seed,
            "calib",
            gene_name,
            rep.gene,
            "smooth",
        )
        stat, _ = gene_dispersion_statistic(rep_smoothed, fitted)
        null_stats.append(stat)
    null_stats_arr = np.array(null_stats, dtype=float)

    inference = infer_pvalue(
        observed=obs_stat,
        null_stats=null_stats_arr,
        two_sided=cfg.score.two_sided if two_sided is None else two_sided,
        warning_threshold=cfg.score.mc_warning_threshold,
    )
    effects = effect_size_summary(obs_stat, null_stats_arr)

    top_sites, windows = summarize_site_contributions(per_site, top_n=10, window=10)
    site_table = pd.DataFrame(
        {
            "site": np.arange(1, len(per_site) + 1),
            "contribution": per_site,
        }
    )

    gene_out = ensure_dir(outdir / "genes" / gene_name)
    site_table.to_csv(gene_out / "site_contributions.tsv", sep="\t", index=False)
    top_sites.to_csv(gene_out / "top_sites.tsv", sep="\t", index=False)
    windows.to_csv(gene_out / "rolling_windows.tsv", sep="\t", index=False)
    write_json(
        gene_out / "null_distribution.json",
        {
            "gene": gene_name,
            "observed_statistic": float(obs_stat),
            "null_statistics": [float(x) for x in null_stats_arr],
            "inference": inference.to_dict(),
            "effect_size": effects,
            "qc": qc.to_dict(),
        },
    )

    warning_flags = sorted(set([*qc.flags, *inference.warning_flags]))
    return GeneScoreOutput(
        gene=gene_name,
        statistic=float(obs_stat),
        p_value=float(inference.p_value),
        n_replicates=int(inference.n_replicates),
        mc_resolution=float(inference.mc_resolution),
        null_mean=float(effects["null_mean"]),
        null_sd=float(effects["null_sd"]),
        delta_vs_null=float(effects["delta_vs_null"]),
        z_vs_null=float(effects["z_vs_null"]),
        warning_flags="|".join(warning_flags),
    )


def score_genes(
    alignments: dict[str, CodonAlignment],
    tree: PhyloTree,
    bundle: FrozenNullBundle,
    cfg: BABAPPAConfig,
    outdir: str | Path,
    seed: int,
    jobs: int,
    resume: bool,
    n_replicates_score: int | None = None,
    two_sided: bool | None = None,
) -> pd.DataFrame:
    out = ensure_dir(outdir)
    gene_names = sorted(alignments.keys())
    seeds = spawn_seeds(seed, len(gene_names))

    results: list[GeneScoreOutput] = []

    def run_one(args: tuple[str, int]) -> GeneScoreOutput:
        gene, gene_seed = args
        gene_out = out / "genes" / gene / "null_distribution.json"
        if resume and gene_out.exists():
            payload = json.loads(gene_out.read_text(encoding="utf-8"))
            effects = payload["effect_size"]
            infer = payload["inference"]
            qc_flags = payload["qc"]["flags"]
            warning_flags = sorted(set([*qc_flags, *infer["warning_flags"]]))
            return GeneScoreOutput(
                gene=gene,
                statistic=float(payload["observed_statistic"]),
                p_value=float(infer["p_value"]),
                n_replicates=int(infer["n_replicates"]),
                mc_resolution=float(infer["mc_resolution"]),
                null_mean=float(effects["null_mean"]),
                null_sd=float(effects["null_sd"]),
                delta_vs_null=float(effects["delta_vs_null"]),
                z_vs_null=float(effects["z_vs_null"]),
                warning_flags="|".join(warning_flags),
            )
        return _score_one_gene_with_bundle(
            gene_name=gene,
            alignment=alignments[gene],
            tree=tree,
            bundle=bundle,
            cfg=cfg,
            seed=gene_seed,
            outdir=out,
            n_replicates_score=n_replicates_score,
            two_sided=two_sided,
        )

    tasks = list(zip(gene_names, seeds, strict=True))
    if jobs > 1 and len(tasks) > 1:
        with ThreadPoolExecutor(max_workers=jobs) as ex:
            for result in ex.map(run_one, tasks):
                results.append(result)
    else:
        for task in tasks:
            results.append(run_one(task))

    rows = [
        {
            "gene": r.gene,
            "statistic": r.statistic,
            "p_value": r.p_value,
            "n_replicates": r.n_replicates,
            "mc_resolution": r.mc_resolution,
            "null_mean": r.null_mean,
            "null_sd": r.null_sd,
            "delta_vs_null": r.delta_vs_null,
            "z_vs_null": r.z_vs_null,
            "warning_flags": r.warning_flags,
        }
        for r in results
    ]
    df = pd.DataFrame(rows).sort_values("p_value").reset_index(drop=True)
    df.to_csv(out / "scores.tsv", sep="\t", index=False)
    return df


def adjust_scores_table(scores: pd.DataFrame, method: str = "bh") -> pd.DataFrame:
    out = scores.copy()
    out["q_value"] = adjust_pvalues(out["p_value"].tolist(), method=method)
    out = out.sort_values(["q_value", "p_value"]).reset_index(drop=True)
    return out


def _scenario_name(prefix: str, value: object) -> str:
    return f"{prefix}={value}"


def run_sensitivity_analysis(
    alignments: dict[str, CodonAlignment],
    tree: PhyloTree,
    base_bundle: FrozenNullBundle,
    cfg: BABAPPAConfig,
    outdir: str | Path,
    seed: int,
    bandwidths: list[float],
    kappas: list[float],
    replicate_counts: list[int],
    energy_families: list[str],
    codon_frequency_modes: list[str],
) -> pd.DataFrame:
    out = ensure_dir(outdir)
    rows: list[dict[str, object]] = []

    scenarios: list[tuple[str, dict[str, object]]] = [("baseline", {})]
    scenarios.extend([(_scenario_name("bandwidth", b), {"bandwidth": b}) for b in bandwidths])
    scenarios.extend([(_scenario_name("kappa", k), {"kappa": k}) for k in kappas])
    scenarios.extend([(_scenario_name("nrep", n), {"nrep": n}) for n in replicate_counts])
    scenarios.extend([(_scenario_name("energy", e), {"energy": e}) for e in energy_families])
    scenarios.extend(
        [(_scenario_name("codon_freq", c), {"codon_frequency_mode": c}) for c in codon_frequency_modes]
    )

    scenario_seeds = spawn_seeds(seed, len(scenarios))

    for (scenario_name, overrides), scenario_seed in zip(scenarios, scenario_seeds, strict=True):
        sub_cfg = deepcopy(cfg)

        if "kappa" in overrides:
            sub_cfg.null_model.kappa = float(overrides["kappa"])
        if "bandwidth" in overrides:
            sub_cfg.representation.smoothing_bandwidth = float(overrides["bandwidth"])
        if "codon_frequency_mode" in overrides:
            sub_cfg.null_model.codon_frequency_mode = str(overrides["codon_frequency_mode"])

        nrep_override = int(overrides["nrep"]) if "nrep" in overrides else None
        energy_family = str(overrides["energy"]) if "energy" in overrides else str(base_bundle.null_model.get("energy_family", "gaussian"))

        for gene, aln in alignments.items():
            # For scenario-specific model changes, recalibrate using null simulations only.
            if overrides:
                scenario_bundle, _, _ = fit_null_bundle(
                    alignment=aln,
                    tree=tree,
                    cfg=sub_cfg,
                    outdir=out / "scenario_bundles" / scenario_name / gene,
                    seed=scenario_seed,
                    energy_family=energy_family,
                    null_replicates_train=max(50, min(200, cfg.fit.null_replicates_train)),
                )
            else:
                scenario_bundle = base_bundle

            one = _score_one_gene_with_bundle(
                gene_name=gene,
                alignment=aln,
                tree=tree,
                bundle=scenario_bundle,
                cfg=sub_cfg,
                seed=scenario_seed + 13,
                outdir=out / "scenario_scores" / scenario_name,
                n_replicates_score=nrep_override,
                two_sided=None,
            )
            rows.append(
                {
                    "gene": gene,
                    "scenario": scenario_name,
                    "statistic": one.statistic,
                    "p_value": one.p_value,
                    "warning_flags": one.warning_flags,
                }
            )

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

    stable_rows = []
    for gene, sub in df.groupby("gene"):
        baseline = sub.loc[sub["scenario"] == "baseline", "p_value"]
        if baseline.empty:
            continue
        b = float(baseline.iloc[0])
        stable = bool(((sub["p_value"] < 0.05) == (b < 0.05)).all())
        stable_rows.append({"gene": gene, "baseline_p": b, "significance_stable": stable})
    pd.DataFrame(stable_rows).to_csv(out / "sensitivity_stability.tsv", sep="\t", index=False)
    return df


def profile_callable(fn, *args, **kwargs) -> tuple[Any, dict[str, float]]:
    tracemalloc.start()
    start = time.perf_counter()
    result = fn(*args, **kwargs)
    elapsed = time.perf_counter() - start
    _, peak = tracemalloc.get_traced_memory()
    tracemalloc.stop()
    return result, {"runtime_seconds": float(elapsed), "peak_memory_mb": float(peak / (1024 * 1024))}
