use std::io::{self, BufWriter, Write};
use std::path::PathBuf;

use clap::Args;
use kodama::{Method, linkage};
use obifastwrite::{JsonVal, write_record};
use obikindex::{
    BasePairTally, DistanceMetric, KmerIndex, PHatEstimate, RawSnpDistanceOutput, SankoffWeights,
    SiblingAnnexStats, SnpAlignment, build_cost_matrix, c_ctx_from_p_hat, calibrate_p_hat,
    mean_substitution_cost, substitution_costs_from_tally,
};
use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick};
use tracing::info;

#[derive(clap::ValueEnum, Clone, Copy, Debug)]
pub enum MetricArg {
    Jaccard,
    Mash,
    Hamming,
    BrayCurtis,
    #[value(name = "relfreq-bray-curtis")]
    RelfreqBrayCurtis,
    Euclidean,
    #[value(name = "relfreq-euclidean")]
    RelfreqEuclidean,
    Hellinger,
    #[value(name = "hellinger-euclidean")]
    HellingerEuclidean,
}

impl From<MetricArg> for DistanceMetric {
    fn from(m: MetricArg) -> Self {
        match m {
            MetricArg::Jaccard            => DistanceMetric::Jaccard,
            MetricArg::Mash               => DistanceMetric::Mash,
            MetricArg::Hamming            => DistanceMetric::Hamming,
            MetricArg::BrayCurtis         => DistanceMetric::BrayCurtis,
            MetricArg::RelfreqBrayCurtis  => DistanceMetric::RelfreqBrayCurtis,
            MetricArg::Euclidean          => DistanceMetric::Euclidean,
            MetricArg::RelfreqEuclidean   => DistanceMetric::RelfreqEuclidean,
            MetricArg::Hellinger          => DistanceMetric::Hellinger,
            MetricArg::HellingerEuclidean => DistanceMetric::HellingerEuclidean,
        }
    }
}

#[derive(Args)]
pub struct DistanceArgs {
    /// Index directory
    pub index: PathBuf,

    /// Distance metric to compute
    #[arg(long, value_enum, default_value = "jaccard")]
    pub metric: MetricArg,

    /// Minimum count to consider a kmer present when computing Jaccard on count indexes
    #[arg(long, default_value = "1")]
    pub presence_threshold: u32,

    /// Also output the shared-kmer count matrix (CSV)
    #[arg(long)]
    pub shared_kmers: bool,

    /// Compute and write a Neighbor-Joining tree (Newick)
    #[arg(long)]
    pub nj: bool,

    /// Compute and write a UPGMA tree (Newick)
    #[arg(long)]
    pub upgma: bool,

    /// Build the sibling-count/minorant annex on this (multi-genome) index
    /// — see `docmd/theory/evolutionary_distances.md`, Step 2b. Construction
    /// only; does not by itself compute or write any statistics.
    #[arg(long)]
    pub sibling_annex: bool,

    /// Tally the sibling-count distribution (CSV) of an already-built annex
    /// (run with `--sibling-annex` first, in this invocation or an earlier
    /// one). A separate, occasional diagnostic pass — not run every time the
    /// annex itself is (re)built.
    #[arg(long)]
    pub sibling_stats: bool,

    /// Compute the raw p-distance restricted to loci that are single-copy
    /// in both genomes of each pair (an already-built sibling annex is
    /// required — run with `--sibling-annex` first, in this invocation or
    /// an earlier one). A quick way to test the central-position SNP
    /// estimator against a real index; not the full `SnpTally` design.
    #[arg(long)]
    pub raw_snp_distance: bool,

    /// Write a SNP-only pseudo-alignment (FASTA, IUPAC-coded) from an
    /// already-built sibling annex — one row per genome, one column per
    /// variable family (monomorphic families skipped), no flanking
    /// sequence. See `docmd/theory/evolutionary_distances.md`,
    /// "Multi-genome framing: family as pseudo-alignment column".
    #[arg(long)]
    pub snp: bool,

    /// Calibrate a 16-state Sankoff cost matrix and its matching
    /// pseudo-alignment from an already-built sibling annex (run with
    /// `--sibling-annex` first, in this invocation or an earlier one), for
    /// use with TNT/PhyG. See `docmd/theory/evolutionary_distances.md`,
    /// "Sankoff parsimony as the resolution of the 16-state model problem".
    #[arg(long)]
    pub sankoff: bool,

    /// Exclude genome pairs whose raw SNP ratio exceeds this value from the
    /// `p_hat` calibration pooled by `--sankoff` — a pair this close to
    /// saturation carries no information about `p_hat` and would bias it
    /// upward if pooled in (unlike a low eligible-loci count, which barely
    /// moves the pooled estimate either way — see design doc).
    #[arg(long, default_value = "0.5")]
    pub sankoff_ratio_ceiling: f64,

    /// Also write <prefix>_sankoff.tnt, a ready-to-run TNT script (`proc
    /// <file>;`) for the same matrix/alignment `--sankoff` computes —
    /// recoded to TNT's default xread alphabet (0-9A-F only; TNT rejects
    /// the wider IUPAC set `--sankoff`'s own output uses unless `nstates
    /// dna` is set, which imposes TNT's own incompatible DNA encoding
    /// instead) with integer-scaled costs (TNT's smatrix/cost commands
    /// reject decimals). Implies `--sankoff`.
    #[arg(long)]
    pub tnt: bool,

    /// Also write <prefix>_sankoff.tcm and <prefix>_sankoff.pg, a
    /// custom-alphabet cost matrix and a ready-to-run PhyG script (`read`/
    /// `search`/`report`) for the same matrix/alignment `--sankoff`
    /// computes. Reuses `--sankoff`'s own `_sankoff.fasta` directly — PhyG's
    /// `tcm:` alphabet is read from the matrix file itself, so the IUPAC+`0`
    /// alphabet needs no recoding here, unlike `--tnt`. Implies `--sankoff`.
    #[arg(long)]
    pub phyg: bool,

    /// Scale factor applied before rounding real-valued costs to the
    /// integers both `--tnt`'s smatrix/cost commands and `--phyg`'s `tcm:`
    /// matrix require. Keep this small: the total tree score is this scale
    /// times the sum of per-character costs across every character (908k+
    /// for a typical run here), and there are hints in TNT's own manual
    /// that at least some of its internal accumulators are 32-bit — a large
    /// scale risks a silent integer overflow (undetectable, not just a
    /// crash) far more costly than the resolution a bigger factor would
    /// buy. Shared between `--tnt` and `--phyg` rather than split into two
    /// flags: both scale the same calibrated matrix for the same reason
    /// (integer-only cost commands), and no PhyG-specific accumulator-width
    /// constraint has actually been found to justify a different default.
    #[arg(long, default_value = "100")]
    pub sankoff_cost_scale: f64,

    /// Output prefix: <prefix>_dist.csv, <prefix>_shared.csv,
    /// <prefix>_siblings.csv, <prefix>_rawsnp.csv, <prefix>_snp.fasta,
    /// <prefix>_sankoff_matrix.csv, <prefix>_sankoff_params.yaml,
    /// <prefix>_sankoff.fasta, <prefix>_sankoff.tnt, <prefix>_sankoff.tcm,
    /// <prefix>_sankoff.pg, <prefix>_nj.nwk, <prefix>_upgma.nwk.
    /// If omitted, the distance matrix is written to stdout.
    #[arg(short, long)]
    pub output: Option<PathBuf>,
}

pub fn run(args: DistanceArgs) {
    let idx = KmerIndex::open(&args.index).unwrap_or_else(|e| {
        eprintln!("error opening index: {e}");
        std::process::exit(1);
    });


    let labels: Vec<String> = idx.meta().genomes.iter().map(|g| g.label.clone()).collect();
    let n = labels.len();

    // ── Sibling-count/minorant annex (independent of the distance metric) ──
    // Construction (`--sibling-annex`) and stats (`--sibling-stats`) are
    // deliberately decoupled: the annex is meant to be (re)built routinely,
    // the distribution only occasionally, on demand.
    if args.sibling_annex {
        // Writes into the index directory — hold an exclusive lock for the
        // duration so a second, concurrent `--sibling-annex` run on the same
        // index can't corrupt these writes (see obisys::DirLock).
        let _lock = obisys::DirLock::acquire(&args.index).unwrap_or_else(|e| {
            eprintln!("error locking index directory {}: {e}", args.index.display());
            std::process::exit(1);
        });
        info!("building sibling-count/minorant annex");
        idx.build_sibling_annex().unwrap_or_else(|e| {
            eprintln!("error building sibling annex: {e}");
            std::process::exit(1);
        });
    }
    if args.sibling_stats {
        let stats = idx.sibling_annex_stats().unwrap_or_else(|e| {
            eprintln!("error computing sibling-annex stats: {e}");
            std::process::exit(1);
        });
        write_sibling_stats_csv(&stats, &labels, &args.output);
    }
    if args.raw_snp_distance {
        let result = idx.raw_snp_distance().unwrap_or_else(|e| {
            eprintln!("error computing raw SNP distance: {e}");
            std::process::exit(1);
        });
        write_raw_snp_distance_csv(&result, &labels, &args.output);
    }
    if args.snp {
        let alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| {
            eprintln!("error computing SNP pseudo-alignment: {e}");
            std::process::exit(1);
        });
        write_snp_fasta(&alignment, &labels, &args.output);
    }
    if args.sankoff || args.tnt || args.phyg {
        let raw = idx.raw_snp_distance().unwrap_or_else(|e| {
            eprintln!("error computing raw SNP distance: {e}");
            std::process::exit(1);
        });
        let estimate = calibrate_p_hat(&raw, args.sankoff_ratio_ceiling);
        let m = (idx.kmer_size() - 1) / 2;

        let tally = idx.base_pair_tally(&raw, args.sankoff_ratio_ceiling).unwrap_or_else(|e| {
            eprintln!("error computing base-pair tally: {e}");
            std::process::exit(1);
        });
        let sub_cost = substitution_costs_from_tally(&tally);
        let mean_sub_cost = mean_substitution_cost(&tally, &sub_cost);
        let c_ctx = c_ctx_from_p_hat(estimate.p_hat, m, mean_sub_cost);

        let weights = SankoffWeights { sub_cost, c_ctx };
        let matrix = build_cost_matrix(&weights);
        write_sankoff_matrix_csv(&matrix, &estimate, &weights, &args.output);
        write_sankoff_params(&estimate, &tally, &weights, args.sankoff_ratio_ceiling, m, mean_sub_cost, &args.output);

        let alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| {
            eprintln!("error computing SNP pseudo-alignment: {e}");
            std::process::exit(1);
        });
        write_sankoff_alignment_fasta(&alignment, &labels, &args.output);

        if args.tnt {
            write_sankoff_tnt(&matrix, &alignment, &labels, &args.output, args.sankoff_cost_scale);
        }
        if args.phyg {
            write_sankoff_phyg(&matrix, &args.output, args.sankoff_cost_scale);
        }
    }

    // `--sibling-annex`/`--sibling-stats`/`--raw-snp-distance`/`--snp`/
    // `--sankoff`/`--tnt` are their own operation, not a modifier on top of
    // a distance-metric computation — a metric was never requested by
    // asking for any of them, so there is nothing for the rest of this
    // function to compute. Not a historical accident to keep: stop here
    // rather than always also running a Jaccard (or whichever `--metric`
    // defaults to) pass and printing an unrequested matrix.
    if args.sibling_annex
        || args.sibling_stats
        || args.raw_snp_distance
        || args.snp
        || args.sankoff
        || args.tnt
        || args.phyg
    {
        return;
    }

    info!(
        "computing {:?} distances for {} genome(s)",
        args.metric, n
    );

    let need_shared = args.shared_kmers || args.nj || args.upgma;
    let result = idx
        .distance(args.metric.into(), need_shared, args.presence_threshold)
        .unwrap_or_else(|e| {
            eprintln!("error computing distances: {e}");
            std::process::exit(1);
        });

    // ── Distance matrix → CSV ─────────────────────────────────────────────────
    let write_dist_csv = |w: &mut dyn Write| {
        write!(w, "genome").unwrap();
        for g in &labels { write!(w, ",{g}").unwrap(); }
        writeln!(w).unwrap();
        for (i, g) in labels.iter().enumerate() {
            write!(w, "{g}").unwrap();
            for j in 0..n {
                write!(w, ",{:.6}", result.matrix[[i, j]]).unwrap();
            }
            writeln!(w).unwrap();
        }
    };

    match &args.output {
        Some(prefix) => {
            let path = format!("{}_dist.csv", prefix.display());
            let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
                eprintln!("error creating {path}: {e}");
                std::process::exit(1);
            }));
            write_dist_csv(&mut f);
            info!("distance matrix → {path}");
        }
        None => {
            let stdout = io::stdout();
            let mut out = BufWriter::new(stdout.lock());
            write_dist_csv(&mut out);
        }
    }

    // ── Shared-kmer matrix → CSV ──────────────────────────────────────────────
    if args.shared_kmers {
        if let Some(shared) = &result.shared_kmers {
            let path = args.output.as_ref()
                .map(|p| format!("{}_shared.csv", p.display()))
                .unwrap_or_else(|| "shared.csv".into());
            let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
                eprintln!("error creating {path}: {e}");
                std::process::exit(1);
            }));
            write!(f, "genome").unwrap();
            for g in &labels { write!(f, ",{g}").unwrap(); }
            writeln!(f).unwrap();
            for (i, g) in labels.iter().enumerate() {
                write!(f, "{g}").unwrap();
                for j in 0..n { write!(f, ",{}", shared[[i, j]]).unwrap(); }
                writeln!(f).unwrap();
            }
            info!("shared-kmer matrix → {path}");
        }
    }

    // ── NJ tree via speedytree ────────────────────────────────────────────────
    if args.nj {
        let rows: Vec<Vec<f64>> = (0..n)
            .map(|i| (0..n).map(|j| result.matrix[[i, j]]).collect())
            .collect();
        let dm = DistanceMatrix::build(rows, labels.clone()).unwrap_or_else(|e| {
            eprintln!("error building distance matrix for NJ: {e}");
            std::process::exit(1);
        });
        let tree = NeighborJoiningSolver::<Hybrid>::default(dm).solve().unwrap_or_else(|e| {
            eprintln!("error computing NJ tree: {e}");
            std::process::exit(1);
        });
        let newick = to_newick(&tree);
        let path = args.output.as_ref()
            .map(|p| format!("{}_nj.nwk", p.display()))
            .unwrap_or_else(|| "nj.nwk".into());
        std::fs::write(&path, &newick).unwrap_or_else(|e| {
            eprintln!("error writing {path}: {e}");
            std::process::exit(1);
        });
        info!("NJ tree → {path}");
    }

    // ── UPGMA tree via kodama ─────────────────────────────────────────────────
    if args.upgma {
        let mut condensed: Vec<f64> = Vec::with_capacity(n * (n - 1) / 2);
        for i in 0..n {
            for j in (i + 1)..n {
                condensed.push(result.matrix[[i, j]]);
            }
        }
        let dendro = linkage(&mut condensed, n, Method::Average);
        let newick = upgma_to_newick(&dendro, &labels);
        let path = args.output.as_ref()
            .map(|p| format!("{}_upgma.nwk", p.display()))
            .unwrap_or_else(|| "upgma.nwk".into());
        std::fs::write(&path, &newick).unwrap_or_else(|e| {
            eprintln!("error writing {path}: {e}");
            std::process::exit(1);
        });
        info!("UPGMA tree → {path}");
    }
}

// ── Family-size distribution → CSV ──────────────────────────────────────────
//
// Each row is a family (the up-to-4 k-mers sharing flanks, differing only at
// the centre), counted once — at its minorant — regardless of how many of
// its members are observed. Family size 1..4 (not "sibling count" 0..3):
// see `docmd/theory/evolutionary_distances.md`, "Definitions".

fn write_sibling_stats_csv(stats: &SiblingAnnexStats, labels: &[String], output: &Option<PathBuf>) {
    // One row per genome (4 columns, family size 1-4: number of families of
    // that size for which the genome carries at least one member), plus a
    // `global` row — the actual deduplicated family-size histogram
    // (`stats.counts`), NOT a sum of the per-genome columns (a family shared
    // by several genomes would otherwise be counted once per genome it
    // appears in, inflating the total beyond the real family count).
    let path = output.as_ref()
        .map(|p| format!("{}_siblings.csv", p.display()))
        .unwrap_or_else(|| "siblings.csv".into());
    let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
        eprintln!("error creating {path}: {e}");
        std::process::exit(1);
    }));
    writeln!(f, "genome,1,2,3,4").unwrap();
    for (label, counts) in labels.iter().zip(stats.per_genome.iter()) {
        writeln!(f, "{label},{},{},{},{}", counts[0], counts[1], counts[2], counts[3]).unwrap();
    }
    writeln!(
        f, "global,{},{},{},{}",
        stats.counts[0], stats.counts[1], stats.counts[2], stats.counts[3],
    ).unwrap();
    let total: u64 = stats.counts.iter().sum();
    info!("family-size distribution → {path} (total {total} famil{})",
        if total == 1 { "y" } else { "ies" });
}

// ── Raw single-copy SNP distance → CSV ──────────────────────────────────────
//
// p_hat[i,j] = snp[i,j] / (snp[i,j] + shared[i,j]) over loci single-copy in
// both i and j — see `RawSnpDistanceOutput` / `KmerIndex::raw_snp_distance`.
// A single file: the distance matrix, with an eligible-loci count alongside
// each value so a 0/0 pair (no eligible locus at all) is distinguishable
// from a genuinely identical pair.

fn write_raw_snp_distance_csv(result: &RawSnpDistanceOutput, labels: &[String], output: &Option<PathBuf>) {
    let path = output.as_ref()
        .map(|p| format!("{}_rawsnp.csv", p.display()))
        .unwrap_or_else(|| "rawsnp.csv".into());
    let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
        eprintln!("error creating {path}: {e}");
        std::process::exit(1);
    }));
    let n = labels.len();
    write!(f, "genome").unwrap();
    for g in labels { write!(f, ",{g}").unwrap(); }
    writeln!(f).unwrap();
    for (i, g) in labels.iter().enumerate() {
        write!(f, "{g}").unwrap();
        for j in 0..n {
            let snp = result.snp[[i, j]];
            let shared = result.shared[[i, j]];
            let eligible = snp + shared;
            if eligible == 0 {
                write!(f, ",NA").unwrap();
            } else {
                write!(f, ",{:.6}", snp as f64 / eligible as f64).unwrap();
            }
        }
        writeln!(f).unwrap();
    }
    info!("raw single-copy SNP distance matrix → {path}");
}

// ── SNP-only pseudo-alignment → FASTA ───────────────────────────────────────
//
// One record per genome, IUPAC-coded, no flanking sequence — see
// `SnpAlignment` / `KmerIndex::snp_pseudo_alignment`. Uses the project's
// existing FASTA writer (`obifastwrite::write_record`) rather than
// hand-rolling one.

fn write_snp_fasta(alignment: &SnpAlignment, labels: &[String], output: &Option<PathBuf>) {
    let path = output.as_ref()
        .map(|p| format!("{}_snp.fasta", p.display()))
        .unwrap_or_else(|| "snp.fasta".into());
    let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
        eprintln!("error creating {path}: {e}");
        std::process::exit(1);
    }));
    let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0);
    for (label, seq) in labels.iter().zip(alignment.sequences.iter()) {
        write_record(seq, label, &[("n_sites", JsonVal::Num(n_sites as u64))], &mut f).unwrap_or_else(|e| {
            eprintln!("error writing {path}: {e}");
            std::process::exit(1);
        });
    }
    info!("SNP pseudo-alignment → {path} ({n_sites} site{})",
        if n_sites == 1 { "" } else { "s" });
}

// ── Sankoff pseudo-alignment → FASTA ────────────────────────────────────────
//
// Same data as `--snp`'s pseudo-alignment (`SnpAlignment`/
// `snp_pseudo_alignment`), re-coded so its symbols match the accompanying
// `--sankoff-matrix` output exactly: `0` for the empty/absent state instead
// of `-`, which TNT/PhyG would otherwise read as their own gap character
// rather than our "family absent" state.

fn write_sankoff_alignment_fasta(alignment: &SnpAlignment, labels: &[String], output: &Option<PathBuf>) {
    let path = output.as_ref()
        .map(|p| format!("{}_sankoff.fasta", p.display()))
        .unwrap_or_else(|| "sankoff.fasta".into());
    let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
        eprintln!("error creating {path}: {e}");
        std::process::exit(1);
    }));
    let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0);
    for (label, seq) in labels.iter().zip(alignment.sequences.iter()) {
        let recoded: Vec<u8> = seq.iter().map(|&b| if b == b'-' { b'0' } else { b }).collect();
        write_record(&recoded, label, &[("n_sites", JsonVal::Num(n_sites as u64))], &mut f).unwrap_or_else(|e| {
            eprintln!("error writing {path}: {e}");
            std::process::exit(1);
        });
    }
    info!("Sankoff pseudo-alignment → {path} ({n_sites} site{})",
        if n_sites == 1 { "" } else { "s" });
}

// ── Sankoff cost matrix → CSV ────────────────────────────────────────────────
//
// 16 states indexed by bitmask (bit 0=A, 1=C, 2=G, 3=T; state 0 is `∅`),
// matching the convention already used for `--snp`'s IUPAC-coded output and
// for the external TNT/PhyG scripts this feeds. Calibration report (p_hat,
// its variance, how many pairs/loci went into it, the resulting c_ctx) goes
// to the log, not the CSV, since it's a run-level fact, not per-cell data.

// IUPAC ambiguity code per state (same mapping as `siblings::iupac_code`,
// already used for `--snp`'s pseudo-alignment — a biologist reads "R" as
// "A or G" without needing this file's convention explained), with `0`
// standing in for the empty state (`-` would collide with TNT/PhyG's own
// gap/range syntax). Bit order: 0=A, 1=C, 2=G, 3=T. This project's
// canonical alphabet — `write_sankoff_tnt` below recodes it to TNT's own
// default alphabet at the adapter boundary, rather than using it here.
const STATE_SYMBOL: [char; 16] = [
    '0', 'A', 'C', 'M', 'G', 'R', 'S', 'V', 'T', 'W', 'Y', 'H', 'K', 'D', 'B', 'N',
];

fn write_sankoff_matrix_csv(
    matrix: &[[f64; 16]; 16],
    estimate: &PHatEstimate,
    weights: &SankoffWeights,
    output: &Option<PathBuf>,
) {
    info!(
        p_hat = format_args!("{:.6}", estimate.p_hat),
        variance = format_args!("{:.3e}", estimate.variance),
        n_pairs_included = estimate.n_pairs_included,
        n_loci_total = estimate.n_loci_total,
        c_ctx = format_args!("{:.4}", weights.c_ctx),
        "Sankoff matrix calibration"
    );

    let path = output.as_ref()
        .map(|p| format!("{}_sankoff_matrix.csv", p.display()))
        .unwrap_or_else(|| "sankoff_matrix.csv".into());
    let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
        eprintln!("error creating {path}: {e}");
        std::process::exit(1);
    }));
    write!(f, "state").unwrap();
    for sym in STATE_SYMBOL { write!(f, ",{sym}").unwrap(); }
    writeln!(f).unwrap();
    for (s, row) in matrix.iter().enumerate() {
        write!(f, "{}", STATE_SYMBOL[s]).unwrap();
        for cost in row { write!(f, ",{cost:.4}").unwrap(); }
        writeln!(f).unwrap();
    }
    info!("Sankoff cost matrix → {path}");
}

// ── Sankoff calibration parameters → YAML report ────────────────────────────
//
// Everything `--sankoff` estimates from real data, in one durable,
// machine-readable file: `p_hat` and its variance (with how many pairs/loci
// went into it), the derived `c_ctx`, and the base-pair substitution tally
// (raw counts, not just the derived costs) — kept for the same reason raw
// counts are kept anywhere else in this project: costs are a modelling
// choice built *from* the counts, and reproducing/re-deriving them later
// needs the counts, not just their current derived value. Structured (YAML,
// not an ad hoc key=value text file) so R/Python/etc. can load it directly
// rather than re-parsing free text.

#[derive(serde::Serialize)]
struct SankoffSubstitution {
    pair: String,
    count: u64,
    cost: f64,
}

#[derive(serde::Serialize)]
struct SankoffParamsReport {
    ratio_ceiling: f64,
    flank_length_m: usize,
    p_hat: f64,
    p_hat_variance: f64,
    n_pairs_included: usize,
    n_loci_total: u64,
    /// Weighted-average substitution cost — see `c_ctx_from_p_hat`'s docs:
    /// this is what turns the raw expected mutation *count* behind `c_ctx`
    /// into an actual cost (not every mutation is worth a flat `1`).
    mean_sub_cost: f64,
    c_ctx: f64,
    substitutions: Vec<SankoffSubstitution>,
}

fn write_sankoff_params(
    estimate: &PHatEstimate,
    tally: &BasePairTally,
    weights: &SankoffWeights,
    ratio_ceiling: f64,
    m: usize,
    mean_sub_cost: f64,
    output: &Option<PathBuf>,
) {
    const BASE_LETTER: [char; 4] = ['A', 'C', 'G', 'T'];

    let mut substitutions = Vec::with_capacity(6);
    for a in 0..4 {
        for b in (a + 1)..4 {
            substitutions.push(SankoffSubstitution {
                pair: format!("{}/{}", BASE_LETTER[a], BASE_LETTER[b]),
                count: tally.counts[a][b],
                cost: weights.sub_cost[a][b],
            });
        }
    }
    let report = SankoffParamsReport {
        ratio_ceiling,
        flank_length_m: m,
        p_hat: estimate.p_hat,
        p_hat_variance: estimate.variance,
        n_pairs_included: estimate.n_pairs_included,
        n_loci_total: estimate.n_loci_total,
        mean_sub_cost,
        c_ctx: weights.c_ctx,
        substitutions,
    };

    let path = output.as_ref()
        .map(|p| format!("{}_sankoff_params.yaml", p.display()))
        .unwrap_or_else(|| "sankoff_params.yaml".into());
    let f = std::fs::File::create(&path).unwrap_or_else(|e| {
        eprintln!("error creating {path}: {e}");
        std::process::exit(1);
    });
    serde_yaml::to_writer(f, &report).unwrap_or_else(|e| {
        eprintln!("error writing {path}: {e}");
        std::process::exit(1);
    });
    info!("Sankoff calibration parameters → {path}");
}

// ── Sankoff cost matrix + alignment → ready-to-run TNT script ──────────────
//
// TNT's *default* xread reader only accepts its own 0-9A-F alphabet (see
// its manual: "up to 16 states are allowed by xread, using symbols 0-9 ...
// and A-F") — the wider IUPAC set `STATE_SYMBOL` uses is rejected as an
// "alien symbol" unless `nstates dna` is set, which imposes TNT's own fixed
// DNA encoding instead, incompatible with a custom smatrix. And TNT's
// `smatrix`/`cost` commands reject decimal costs ("found symbol . when
// reading transformation costs"). So: recode to TNT's alphabet and
// integer-scale the costs here, at this adapter's boundary, rather than
// degrading the project's own canonical (IUPAC, real-valued) output.

const TNT_STATE_SYMBOL: [char; 16] = [
    '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F',
];

fn write_sankoff_tnt(
    matrix: &[[f64; 16]; 16],
    alignment: &SnpAlignment,
    labels: &[String],
    output: &Option<PathBuf>,
    cost_scale: f64,
) {
    let path = output.as_ref()
        .map(|p| format!("{}_sankoff.tnt", p.display()))
        .unwrap_or_else(|| "sankoff.tnt".into());
    let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
        eprintln!("error creating {path}: {e}");
        std::process::exit(1);
    }));

    // IUPAC-ish symbol -> bitmask, to translate the alignment (which uses
    // `STATE_SYMBOL`, `-` already normalised to `0` by `snp_pseudo_alignment`
    // callers) into TNT's alphabet without re-deriving state indices.
    let mut iupac_to_state = [0u8; 128];
    for (state, &sym) in STATE_SYMBOL.iter().enumerate() {
        iupac_to_state[sym as usize] = state as u8;
    }

    let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0);
    writeln!(f, "xread").unwrap();
    writeln!(f, "'obikmer central-position SNP families, calibrated Sankoff 16-state encoding'").unwrap();
    writeln!(f, "{n_sites} {}", labels.len()).unwrap();
    for (label, seq) in labels.iter().zip(alignment.sequences.iter()) {
        write!(f, "{label} ").unwrap();
        for &b in seq {
            let b = if b == b'-' { b'0' } else { b };
            let state = iupac_to_state[b as usize];
            write!(f, "{}", TNT_STATE_SYMBOL[state as usize]).unwrap();
        }
        writeln!(f).unwrap();
    }
    writeln!(f, ";\n").unwrap();

    let scaled_matrix = scaled_metric_matrix(matrix, cost_scale);

    writeln!(f, "smatrix =0 (family16)").unwrap();
    for i in 0..16 {
        for j in (i + 1)..16 {
            writeln!(f, "{}/{} {}", TNT_STATE_SYMBOL[i], TNT_STATE_SYMBOL[j], scaled_matrix[i][j]).unwrap();
        }
    }
    writeln!(f, ";\n").unwrap();

    writeln!(f, "ccode ( 0.{} ;", n_sites - 1).unwrap();
    writeln!(f, "smatrix +0 0.{} ;", n_sites - 1).unwrap();
    writeln!(f).unwrap();

    // Basename only (not the full `path`/`output` prefix): TNT's natural
    // workflow is to `cd` into the output directory before `proc`-ing the
    // script, and an absolute path here would break if that directory is
    // later moved or copied elsewhere.
    let tre_name = output.as_ref()
        .and_then(|p| p.file_name())
        .map(|n| format!("{}_sankoff.tre", n.to_string_lossy()))
        .unwrap_or_else(|| "sankoff.tre".into());

    // TNT's plain command parser has no comment syntax of its own — `/* */`
    // and `[ ]` are only recognised inside the (separately-enabled) macro
    // scripting language, and fail with "No command!" here otherwise
    // (verified against this file with the local TNT binary). `quote` is
    // the closest working equivalent: it prints free text and does not
    // otherwise affect parsing, so it doubles as an explanation of the
    // defaults below when the script is run. `;` ends a `quote` block like
    // any other TNT command, so the text itself must avoid semicolons.
    writeln!(f, "quote").unwrap();
    writeln!(f, "Default search below (edit or delete this block to run your own strategy):").unwrap();
    writeln!(f, "  hold N     : size of TNT's tree buffer (how many equally-parsimonious").unwrap();
    writeln!(f, "               trees it keeps in memory at once), 20 is a small, fast").unwrap();
    writeln!(f, "               default, raise it if mult reports it had to drop trees.").unwrap();
    writeln!(f, "  mult       : traditional search (random addition sequences followed by").unwrap();
    writeln!(f, "               TBR branch-swapping, TNT's own default replication count),").unwrap();
    writeln!(f, "               a reasonable first-pass strategy on this data's memory").unwrap();
    writeln!(f, "               footprint, xmult's ratchet/drift/tree-fusion buffers ran").unwrap();
    writeln!(f, "               this out of RAM at TNT's default mxram on this dataset.").unwrap();
    writeln!(f, "  export - F : write the trees held in the buffer to file F, in").unwrap();
    writeln!(f, "               TNT/Hennig86 format ('-' means trees, as opposed to data).").unwrap();
    writeln!(f, ";").unwrap();
    writeln!(f, "hold 20;").unwrap();
    writeln!(f, "mult;").unwrap();
    writeln!(f, "export - {tre_name};").unwrap();

    info!(
        "TNT script → {path} (costs scaled x{cost_scale:.0}, runs a default `hold 20; mult;` \
         search and writes trees to {tre_name} in TNT's working directory — edit the trailing \
         comment block in the script to change this)\n\
         Run it with:\n  \
         printf 'proc {path};\\nquit;\\n' | tnt\n\
         (or start `tnt` interactively and type `proc {path};`)"
    );
}

/// Scale `matrix` by `cost_scale` and round to integers (TNT's smatrix/cost
/// commands reject decimals), then take the *metric closure* of the result
/// (Floyd-Warshall over the 16 states again, on the now-integer values).
///
/// `matrix` is already a metric in its real-valued form (it's a
/// shortest-path closure itself — see `build_cost_matrix`), but rounding
/// each cell independently can still break the triangle inequality: e.g.
/// two real costs of `1.734` each round to `173`, summing to `346`, while
/// their own real sum `3.468` rounds to `347` — TNT then reports "triangle
/// inequality violated ... Fixed" and silently substitutes its own
/// corrected value. Re-closing after rounding makes that correction
/// explicit and reproducible here instead, rather than left implicit and
/// TNT-version-dependent.
fn scaled_metric_matrix(matrix: &[[f64; 16]; 16], cost_scale: f64) -> [[i64; 16]; 16] {
    let mut m = [[0i64; 16]; 16];
    for i in 0..16 {
        for j in 0..16 {
            m[i][j] = (matrix[i][j] * cost_scale).round() as i64;
        }
    }
    for k in 0..16 {
        for i in 0..16 {
            for j in 0..16 {
                let via = m[i][k] + m[k][j];
                if via < m[i][j] {
                    m[i][j] = via;
                }
            }
        }
    }
    m
}

// ── Sankoff cost matrix → PhyG custom-alphabet TCM + ready-to-run script ────
//
// PhyG's `tcm:STRING` format needs no alphabet recoding, unlike `--tnt`:
// its parser reads the alphabet straight from the tcm file's own first
// line, so `--sankoff`'s own `_sankoff.fasta` (already IUPAC+`0`) is reused
// as-is via `prefasta:`. PhyG auto-adds its own indel/gap state as an
// (n+1)-th row/column of the tcm — inert here since the alignment already
// encodes absence as an ordinary state (`0`), never as `-` (see
// `write_sankoff_alignment_fasta`'s own comment on why, and the RAxML-era
// bug that motivated it). The gap row/column below reuses `matrix[i][0]`/
// `matrix[0][j]` (cost to/from `∅`) as the closest principled value for a
// state that, in practice, is never actually triggered.

fn write_sankoff_phyg(matrix: &[[f64; 16]; 16], output: &Option<PathBuf>, cost_scale: f64) {
    let scaled_matrix = scaled_metric_matrix(matrix, cost_scale);

    let basename = |suffix: &str| -> String {
        output.as_ref()
            .and_then(|p| p.file_name())
            .map(|n| format!("{}{suffix}", n.to_string_lossy()))
            .unwrap_or_else(|| format!("sankoff{suffix}"))
    };
    let full_path = |suffix: &str| -> String {
        output.as_ref()
            .map(|p| format!("{}{suffix}", p.display()))
            .unwrap_or_else(|| format!("sankoff{suffix}"))
    };

    let tcm_path = full_path("_sankoff.tcm");
    let mut f = BufWriter::new(std::fs::File::create(&tcm_path).unwrap_or_else(|e| {
        eprintln!("error creating {tcm_path}: {e}");
        std::process::exit(1);
    }));
    let alphabet_line = STATE_SYMBOL.iter().map(|c| c.to_string()).collect::<Vec<_>>().join(" ");
    writeln!(f, "{alphabet_line}").unwrap();
    for i in 0..16 {
        let mut row: Vec<i64> = (0..16).map(|j| scaled_matrix[i][j]).collect();
        row.push(scaled_matrix[i][0]); // gap column: same cost as to/from ∅
        writeln!(f, "{}", row.iter().map(|v| v.to_string()).collect::<Vec<_>>().join(" ")).unwrap();
    }
    let mut gap_row: Vec<i64> = (0..16).map(|j| scaled_matrix[0][j]).collect();
    gap_row.push(0);
    writeln!(f, "{}", gap_row.iter().map(|v| v.to_string()).collect::<Vec<_>>().join(" ")).unwrap();
    info!("PhyG TCM → {tcm_path}");

    let pg_path = full_path("_sankoff.pg");
    let mut f = BufWriter::new(std::fs::File::create(&pg_path).unwrap_or_else(|e| {
        eprintln!("error creating {pg_path}: {e}");
        std::process::exit(1);
    }));
    let fasta_name = basename("_sankoff.fasta");
    let tcm_name = basename("_sankoff.tcm");
    let tre_name = basename("_sankoff.tre");
    writeln!(f, "read(prefasta:\"{fasta_name}\", tcm:\"{tcm_name}\")").unwrap();
    writeln!(f, "search(seconds:300, instances:4)").unwrap();
    writeln!(f, "report(\"{tre_name}\", graphs, newick, overwrite)").unwrap();

    let pg_dir = std::path::Path::new(&pg_path).parent()
        .filter(|d| !d.as_os_str().is_empty())
        .map(|d| d.display().to_string())
        .unwrap_or_else(|| ".".into());
    let pg_name = std::path::Path::new(&pg_path).file_name()
        .map(|n| n.to_string_lossy().into_owned())
        .unwrap_or_else(|| pg_path.clone());
    info!(
        "PhyG script → {pg_path} (costs scaled x{cost_scale:.0}, runs a default 300s/4-instance \
         search and writes trees to {tre_name})\n\
         Run it with:\n  \
         cd {pg_dir} && phyg {pg_name}\n\
         (`phyg` must run from that directory — `read()`/`report()` in the script use relative \
         file names)"
    );
}

// ── UPGMA Newick from kodama dendrogram ───────────────────────────────────────

fn upgma_to_newick(dendro: &kodama::Dendrogram<f64>, names: &[String]) -> String {
    let n = names.len();
    // node_labels[i]: Newick subtree string for node i (leaves 0..n, internals n..)
    let mut labels: Vec<String> = names.to_vec();
    // height of each node: leaves = 0, internal = dissimilarity/2
    let mut heights: Vec<f64> = vec![0.0; 2 * n - 1];

    for (k, step) in dendro.steps().iter().enumerate() {
        let new_node = n + k;
        let h = step.dissimilarity / 2.0;
        heights[new_node] = h;
        let c1 = step.cluster1;
        let c2 = step.cluster2;
        let bl1 = (h - heights[c1]).max(0.0);
        let bl2 = (h - heights[c2]).max(0.0);
        labels.push(format!(
            "({label1}:{bl1:.6},{label2}:{bl2:.6})",
            label1 = labels[c1],
            label2 = labels[c2],
        ));
    }

    format!("{};", labels.last().unwrap())
}
