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

use clap::Args;
use kodama::{Method, linkage};
use obikindex::{DistanceMetric, KmerIndex};
use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick};
use tracing::info;

#[derive(clap::ValueEnum, Clone, Copy, Debug)]
pub enum MetricArg {
    Jaccard,
    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::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,

    /// Output prefix: <prefix>_dist.csv, <prefix>_shared.csv,
    /// <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();
    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}");
    }
}

// ── 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())
}
