//! RDKit accuracy comparison.
//!
//! Reads scripts/rdkit_ref_properties.tsv (generated by gen_rdkit_reference.py),
//! computes the same properties with chematic, and prints MAE/RMSE/Pearson-r
//! for MW, LogP, TPSA.
//!
//! Also reads scripts/rdkit_ref_tanimoto.tsv (50×50 ECFP4 matrix from RDKit)
//! and computes Spearman rank correlation for ECFP4 Tanimoto similarity.
//!
//! Usage:
//!   cargo run -p chematic-chem --example rdkit_compare

use chematic_chem::descriptors::{
    hba_count, hbd_count, heavy_atom_count, logp_crippen, molecular_weight, tpsa,
};
use chematic_fp::ecfp::{EcfpConfig, ecfp};
use chematic_smiles::parse;

use std::fs;

fn pearson(xs: &[f64], ys: &[f64]) -> f64 {
    let n = xs.len() as f64;
    let mx = xs.iter().sum::<f64>() / n;
    let my = ys.iter().sum::<f64>() / n;
    let cov: f64 = xs
        .iter()
        .zip(ys)
        .map(|(x, y)| (x - mx) * (y - my))
        .sum::<f64>()
        / n;
    let sx = (xs.iter().map(|x| (x - mx).powi(2)).sum::<f64>() / n).sqrt();
    let sy = (ys.iter().map(|y| (y - my).powi(2)).sum::<f64>() / n).sqrt();
    if sx < 1e-12 || sy < 1e-12 {
        return f64::NAN;
    }
    cov / (sx * sy)
}

fn mae(xs: &[f64], ys: &[f64]) -> f64 {
    xs.iter().zip(ys).map(|(x, y)| (x - y).abs()).sum::<f64>() / xs.len() as f64
}

fn rmse(xs: &[f64], ys: &[f64]) -> f64 {
    (xs.iter().zip(ys).map(|(x, y)| (x - y).powi(2)).sum::<f64>() / xs.len() as f64).sqrt()
}

fn main() {
    // Find the tsv relative to the workspace root (two levels up from crates/chematic-chem)
    let tsv_path = {
        let p = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
            .join("../../scripts/rdkit_ref_properties.tsv");
        p.canonicalize().unwrap_or(p)
    };

    let content = fs::read_to_string(&tsv_path)
        .unwrap_or_else(|e| panic!("Cannot read {}: {e}", tsv_path.display()));

    let mut lines = content.lines();
    let _header = lines.next(); // skip header

    let mut names: Vec<String> = Vec::new();
    let mut rdkit_mw: Vec<f64> = Vec::new();
    let mut rdkit_lp: Vec<f64> = Vec::new();
    let mut rdkit_tp: Vec<f64> = Vec::new();
    let mut rdkit_hac: Vec<f64> = Vec::new();
    let mut rdkit_hbd: Vec<f64> = Vec::new();
    let mut rdkit_hba: Vec<f64> = Vec::new();
    let mut chem_mw: Vec<f64> = Vec::new();
    let mut chem_lp: Vec<f64> = Vec::new();
    let mut chem_tp: Vec<f64> = Vec::new();
    let mut chem_hac: Vec<f64> = Vec::new();
    let mut chem_hbd: Vec<f64> = Vec::new();
    let mut chem_hba: Vec<f64> = Vec::new();
    let mut skipped: Vec<(String, String)> = Vec::new();

    for line in lines {
        let fields: Vec<&str> = line.split('\t').collect();
        if fields.len() < 8 {
            continue;
        }
        let name = fields[0];
        let smi = fields[1];
        let rw: f64 = fields[2].parse().unwrap();
        let rl: f64 = fields[3].parse().unwrap();
        let rt: f64 = fields[4].parse().unwrap();
        let rhac: f64 = fields[5].parse().unwrap();
        let rhbd: f64 = fields[6].parse().unwrap();
        let rhba: f64 = fields[7].parse().unwrap();

        match parse(smi) {
            Ok(mol) => {
                names.push(name.to_string());
                rdkit_mw.push(rw);
                rdkit_lp.push(rl);
                rdkit_tp.push(rt);
                rdkit_hac.push(rhac);
                rdkit_hbd.push(rhbd);
                rdkit_hba.push(rhba);
                chem_mw.push(molecular_weight(&mol));
                chem_lp.push(logp_crippen(&mol));
                chem_tp.push(tpsa(&mol));
                chem_hac.push(heavy_atom_count(&mol) as f64);
                chem_hbd.push(hbd_count(&mol) as f64);
                chem_hba.push(hba_count(&mol) as f64);
            }
            Err(e) => {
                skipped.push((name.to_string(), format!("{e}")));
            }
        }
    }

    let n = names.len();
    println!("=== chematic vs RDKit property comparison ===");
    println!("Molecules parsed:  {n}");
    println!("Molecules skipped: {}", skipped.len());
    if !skipped.is_empty() {
        for (name, err) in &skipped {
            println!("  [skip] {name}: {err}");
        }
    }
    println!();

    // ── Per-molecule table (worst outliers) ──────────────────────────────────
    let mut diffs_mw: Vec<(f64, &str)> = chem_mw
        .iter()
        .zip(rdkit_mw.iter())
        .zip(names.iter())
        .map(|((c, r), n)| ((c - r).abs(), n.as_str()))
        .collect();
    diffs_mw.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap());

    println!(
        "{:<30} {:>8} {:>8} {:>7}   {:>8} {:>8} {:>7}   {:>8} {:>8} {:>7}",
        "name", "rdMW", "chMW", "ΔMW", "rdLogP", "chLogP", "ΔLogP", "rdTPSA", "chTPSA", "ΔTPSA"
    );
    println!("{}", "-".repeat(110));
    for &(_, name) in diffs_mw.iter().take(25) {
        let i = names.iter().position(|n| n == name).unwrap();
        let dmw = chem_mw[i] - rdkit_mw[i];
        let dlp = chem_lp[i] - rdkit_lp[i];
        let dtp = chem_tp[i] - rdkit_tp[i];
        println!(
            "{:<30} {:>8.3} {:>8.3} {:>+7.3}   {:>8.3} {:>8.3} {:>+7.3}   {:>8.2} {:>8.2} {:>+7.2}",
            name,
            rdkit_mw[i],
            chem_mw[i],
            dmw,
            rdkit_lp[i],
            chem_lp[i],
            dlp,
            rdkit_tp[i],
            chem_tp[i],
            dtp
        );
    }
    println!();

    // ── Summary statistics ───────────────────────────────────────────────────
    println!("=== Summary statistics (n={n}) ===");
    println!(
        "{:<8} {:>8} {:>8} {:>10}",
        "prop", "MAE", "RMSE", "Pearson-r"
    );
    println!("{}", "-".repeat(40));
    for (label, chem, rdkit) in [
        ("MW", &chem_mw[..], &rdkit_mw[..]),
        ("LogP", &chem_lp[..], &rdkit_lp[..]),
        ("TPSA", &chem_tp[..], &rdkit_tp[..]),
        ("HAC", &chem_hac[..], &rdkit_hac[..]),
        ("HBD", &chem_hbd[..], &rdkit_hbd[..]),
        ("HBA", &chem_hba[..], &rdkit_hba[..]),
    ] {
        println!(
            "{:<8} {:>8.4} {:>8.4} {:>10.4}",
            label,
            mae(chem, rdkit),
            rmse(chem, rdkit),
            pearson(chem, rdkit)
        );
    }

    // ── TSV output for downstream plotting ──────────────────────────────────
    let out_path = {
        let p = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
            .join("../../scripts/chematic_vs_rdkit.tsv");
        p.to_string_lossy().into_owned()
    };
    let mut out = String::from(
        "name\trdMW\tchMW\trdLogP\tchLogP\trdTPSA\tchTPSA\trdHAC\tchHAC\trdHBD\tchHBD\trdHBA\tchHBA\n",
    );
    for i in 0..n {
        out += &format!(
            "{}\t{:.4}\t{:.4}\t{:.4}\t{:.4}\t{:.4}\t{:.4}\t{}\t{}\t{}\t{}\t{}\t{}\n",
            names[i],
            rdkit_mw[i],
            chem_mw[i],
            rdkit_lp[i],
            chem_lp[i],
            rdkit_tp[i],
            chem_tp[i],
            rdkit_hac[i] as u32,
            chem_hac[i] as u32,
            rdkit_hbd[i] as u32,
            chem_hbd[i] as u32,
            rdkit_hba[i] as u32,
            chem_hba[i] as u32,
        );
    }
    fs::write(&out_path, out).ok();
    println!("\nFull property table written to: {out_path}");

    // ── ECFP4 Tanimoto comparison (50×50) ────────────────────────────────────
    let tanimoto_tsv_path = {
        let p = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
            .join("../../scripts/rdkit_ref_tanimoto.tsv");
        p.canonicalize().unwrap_or(p)
    };

    if let Ok(tanimoto_content) = fs::read_to_string(&tanimoto_tsv_path) {
        println!("\n=== ECFP4 Tanimoto comparison (50-molecule subset) ===");
        // Parse the header to get molecule names in order.
        let mut tani_lines = tanimoto_content.lines();
        let header_line = tani_lines.next().unwrap_or("");
        let fp_names: Vec<&str> = header_line.split('\t').skip(1).collect();

        // ---- Re-read properties TSV to get SMILES ----
        let prop_content = fs::read_to_string(&tsv_path).unwrap();
        let mut smiles_map: std::collections::HashMap<String, String> =
            std::collections::HashMap::new();
        for line in prop_content.lines().skip(1) {
            let f: Vec<&str> = line.split('\t').collect();
            if f.len() >= 2 {
                smiles_map.insert(f[0].to_string(), f[1].to_string());
            }
        }

        let fps50: Vec<Option<chematic_fp::bitvec::BitVec2048>> = fp_names
            .iter()
            .map(|name| {
                smiles_map.get(*name).and_then(|smi| {
                    parse(smi).ok().map(|mol| {
                        ecfp(
                            &mol,
                            &EcfpConfig {
                                radius: 2,
                                nbits: 2048,
                                use_chirality: false,
                                use_double_fold: false,
                            },
                        )
                    })
                })
            })
            .collect();

        // Compute chematic Tanimoto matrix and compare with RDKit.
        let mut rdkit_tani_flat: Vec<f64> = Vec::new();
        let mut chem_tani_flat: Vec<f64> = Vec::new();
        let n50 = fp_names.len();

        for (i, row_line) in tanimoto_content.lines().skip(1).enumerate() {
            if i >= n50 {
                break;
            }
            let fields: Vec<&str> = row_line.split('\t').collect();
            for j in 0..n50 {
                if i == j {
                    continue;
                } // skip self-similarity
                let rd_t: f64 = fields
                    .get(j + 1)
                    .and_then(|s| s.parse().ok())
                    .unwrap_or(0.0);
                let ch_t: f64 = match (&fps50[i], &fps50[j]) {
                    (Some(a), Some(b)) => a.tanimoto(b),
                    _ => continue,
                };
                rdkit_tani_flat.push(rd_t);
                chem_tani_flat.push(ch_t);
            }
        }

        let m = rdkit_tani_flat.len();
        if m > 0 {
            // Spearman rank correlation
            fn rank(xs: &[f64]) -> Vec<f64> {
                let mut order: Vec<usize> = (0..xs.len()).collect();
                order.sort_by(|&a, &b| xs[a].partial_cmp(&xs[b]).unwrap());
                let mut ranks = vec![0.0f64; xs.len()];
                let mut i = 0;
                while i < order.len() {
                    let mut j = i;
                    while j + 1 < order.len() && xs[order[j + 1]] == xs[order[i]] {
                        j += 1;
                    }
                    let avg_rank = (i + j) as f64 / 2.0 + 1.0;
                    for k in i..=j {
                        ranks[order[k]] = avg_rank;
                    }
                    i = j + 1;
                }
                ranks
            }

            let r_ranks = rank(&rdkit_tani_flat);
            let c_ranks = rank(&chem_tani_flat);
            let spearman = pearson(&r_ranks, &c_ranks);

            println!("Pairs (off-diagonal):   {m}");
            println!(
                "MAE  (Tanimoto):         {:.4}",
                mae(&chem_tani_flat, &rdkit_tani_flat)
            );
            println!(
                "RMSE (Tanimoto):         {:.4}",
                rmse(&chem_tani_flat, &rdkit_tani_flat)
            );
            println!(
                "Pearson-r (Tanimoto):    {:.4}",
                pearson(&chem_tani_flat, &rdkit_tani_flat)
            );
            println!("Spearman-r (Tanimoto):   {:.4}", spearman);

            // Distribution of |Δ Tanimoto|
            let mut abs_diffs: Vec<f64> = chem_tani_flat
                .iter()
                .zip(&rdkit_tani_flat)
                .map(|(c, r)| (c - r).abs())
                .collect();
            abs_diffs.sort_by(|a, b| a.partial_cmp(b).unwrap());
            let p90 = abs_diffs[(abs_diffs.len() as f64 * 0.90) as usize];
            let p99 = abs_diffs[(abs_diffs.len() as f64 * 0.99) as usize];
            println!("|ΔTanimoto| p90 = {p90:.4},  p99 = {p99:.4}");
        } else {
            println!("(no pairs computed — check tanimoto TSV)");
        }
    } else {
        println!(
            "\n[skip] ECFP4 comparison: {} not found",
            tanimoto_tsv_path.display()
        );
    }
}
