koth-ms 0.10.0

use std::io::{BufWriter, Write as IoWrite};
use std::path::{Path, PathBuf};
use std::time::Instant;

use anyhow::Context;
use clap::Parser;

use koth_ms::{
    alignment::{align_runs, RunInput},
    config::{AlignConfig, OutputFormat},
    input::{check_batch_mobility_scales, discover_runs, read_features, read_hills},
    lfq::{
        quantify, quantify_guided,
        targets::{build_targets, read_identifications, write_search_outputs, TargetFormat},
        IntensityMatrix,
    },
    mem::log_mem,
    output::{build_align_report, write_align_report, AlignTiming},
};

#[derive(Parser, Debug)]
#[command(
    name = "koth_align",
    about = "Multi-run alignment and label-free quantification from koth_ff output",
    long_about = "Discovers all per-run subdirectories inside <batch_dir>, aligns retention \
                  time, m/z ppm, and ion mobility across runs, then performs Sage-style LFQ \
                  by extracting ion chromatogram grids from hill data.",
    version = koth_ms::VERSION
)]
struct Args {
    /// Directory containing one subdirectory per run.
    /// Each subdirectory must contain hills.{tsv,parquet} and features.{tsv,parquet}.
    batch_dir: PathBuf,

    /// Output directory (default: <batch_dir>/align_output)
    #[arg(short, long)]
    output: Option<PathBuf>,

    /// Alignment + LFQ configuration TOML (uses defaults if not provided)
    #[arg(short, long)]
    config: Option<PathBuf>,

    /// Generic identification target TSV (see docs/search-guided-lfq.md).
    #[arg(long, conflicts_with = "sage_psms")]
    targets: Option<PathBuf>,

    /// Sage results.sage.tsv or results.sage.parquet (rank-1, target PSMs).
    #[arg(long, conflicts_with = "targets")]
    sage_psms: Option<PathBuf>,

    /// Maximum imported identification q-value (both spectrum_q and peptide_q for Sage).
    #[arg(long, default_value_t = 0.01)]
    max_id_qvalue: f64,

    /// Extract only in runs with accepted same-run identifications; disable transfers.
    #[arg(long)]
    no_mbr: bool,

    /// Maximum exploratory extraction q-value in the peptide intensity matrix.
    #[arg(long, default_value_t = 0.01)]
    max_extraction_qvalue: f64,

    /// Ignore imported IM when search-engine coordinates are not native 1/K0.
    #[arg(long)]
    ignore_target_im: bool,

    /// Log level: error, warn, info, debug, trace
    #[arg(long, default_value = "info")]
    log_level: String,
}

fn main() -> anyhow::Result<()> {
    let args = Args::parse();

    // The binary logs as `koth_ff`; the library crate logs as `koth_ms`.
    let log_filter = format!("koth_ff={0},koth_ms={0}", args.log_level);
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or(&log_filter)).init();

    let config = match &args.config {
        Some(p) => AlignConfig::from_toml(p)
            .with_context(|| format!("Failed to load config from {}", p.display()))?,
        None => AlignConfig::default(),
    };

    let guided = args.targets.is_some() || args.sage_psms.is_some();
    anyhow::ensure!(
        guided || (!args.no_mbr && !args.ignore_target_im),
        "--no-mbr and --ignore-target-im require --targets or --sage-psms"
    );
    for (name, q) in [
        ("--max-id-qvalue", args.max_id_qvalue),
        ("--max-extraction-qvalue", args.max_extraction_qvalue),
    ] {
        anyhow::ensure!(
            q.is_finite() && (0.0..=1.0).contains(&q),
            "{name} must be in [0, 1]"
        );
    }

    let out_dir = args
        .output
        .clone()
        .unwrap_or_else(|| args.batch_dir.join("align_output"));
    std::fs::create_dir_all(&out_dir)
        .with_context(|| format!("Failed to create output directory {}", out_dir.display()))?;

    log::info!("Batch dir:  {}", args.batch_dir.display());
    log::info!("Output dir: {}", out_dir.display());

    let total_start = Instant::now();

    // ── Discover runs ─────────────────────────────────────────────────────────
    let run_paths =
        discover_runs(&args.batch_dir).context("Failed to discover runs in batch directory")?;

    if run_paths.is_empty() {
        anyhow::bail!(
            "No runs found in '{}'. Each run must be a subdirectory containing \
             hills.{{tsv,parquet}} and features.{{tsv,parquet}}.",
            args.batch_dir.display()
        );
    }

    log::info!("Found {} runs:", run_paths.len());
    for rp in &run_paths {
        log::info!(
            "  {} → hills: {}, features: {}",
            rp.name,
            rp.hills_path.display(),
            rp.features_path.display()
        );
    }

    // ── Load runs ─────────────────────────────────────────────────────────────
    let mut runs: Vec<RunInput> = Vec::with_capacity(run_paths.len());

    // Load only FEATURES up front (small). Hills (the large intensity profiles)
    // are streamed one run at a time during LFQ so peak memory stays O(one run's
    // hills) rather than O(all runs' hills).
    for rp in &run_paths {
        let t = Instant::now();

        let features = read_features(&rp.features_path).with_context(|| {
            format!(
                "Failed to read features from {}",
                rp.features_path.display()
            )
        })?;

        log::info!(
            "Loaded '{}': {} features [{:.2?}] (hills streamed during LFQ)",
            rp.name,
            features.len(),
            t.elapsed()
        );

        // A targeted run can have hills but no complete isotope features.
        // Recover only its RT bounds; hills remain streamed during extraction.
        let rt_bounds = if guided && features.is_empty() {
            let hills = read_hills(&rp.hills_path)?;
            let lo = hills
                .iter()
                .map(|h| h.rt_start)
                .fold(f64::INFINITY, f64::min);
            let hi = hills
                .iter()
                .map(|h| h.rt_end)
                .fold(f64::NEG_INFINITY, f64::max);
            anyhow::ensure!(
                lo.is_finite() && hi > lo,
                "Run '{}' has no usable RT range",
                rp.name
            );
            Some((lo, hi))
        } else {
            None
        };
        runs.push(RunInput {
            name: rp.name.clone(),
            features,
            hills: Vec::new(),      // streamed per-run in quantify()
            scan_times: Vec::new(), // not stored in output files; interpolate hill RTs
            rt_bounds,
        });
    }

    log_mem("after loading all features");

    // Refuse to align runs whose 1/K0 values are on different scales.
    check_batch_mobility_scales(
        run_paths
            .iter()
            .zip(&runs)
            .map(|(p, r)| (p, r.features.as_slice())),
    )?;

    // ── Alignment ─────────────────────────────────────────────────────────────
    log::info!("Running alignment...");
    let t = Instant::now();
    if let Some(name) = &config.alignment.reference_run {
        anyhow::ensure!(
            runs.iter().any(|r| &r.name == name),
            "Unknown reference run: {}",
            name
        );
    }
    let alignment = align_runs(&runs, &config.alignment);
    let alignment_elapsed = t.elapsed();
    log::info!(
        "[timing] alignment: {:.2?} (reference: '{}')",
        alignment_elapsed,
        alignment.reference_name
    );

    for (name, al) in &alignment.alignments {
        log::info!(
            "  {}: {} anchors, mass drift {:.3} + {:.4}×RT ppm",
            name,
            al.n_anchors,
            al.mass_drift.intercept,
            al.mass_drift.slope
        );
    }

    // ── LFQ ───────────────────────────────────────────────────────────────────
    log::info!("Running LFQ quantification...");
    let t = Instant::now();
    let imported = match (&args.targets, &args.sage_psms) {
        (Some(path), _) => Some(read_identifications(
            path,
            TargetFormat::Generic,
            &runs,
            args.max_id_qvalue,
            args.ignore_target_im,
        )?),
        (_, Some(path)) => Some(read_identifications(
            path,
            TargetFormat::Sage,
            &runs,
            args.max_id_qvalue,
            args.ignore_target_im,
        )?),
        _ => None,
    };
    let load_hills = |i: usize| {
        // Stream this run's hills on demand; dropped inside quantify() before
        // the next run loads. Discovery already verified the file exists, so
        // a failure here means it vanished or was corrupted mid-run — fail
        // loudly rather than silently zero that run's signal, which would
        // otherwise let a corrupted intermediate silently skew the intensity
        // matrix and every downstream statistic computed from it.
        read_hills(&run_paths[i].hills_path).unwrap_or_else(|e| {
            panic!(
                "Failed to read hills from {}: {e} — aborting rather than silently \
                 contributing zero signal for this run",
                run_paths[i].hills_path.display()
            )
        })
    };
    let matrix = if let Some((ids, summary)) = imported {
        let (candidates, guidance) = build_targets(
            ids,
            summary,
            &runs,
            &alignment,
            &config.lfq,
            config.alignment.min_anchor_count,
            !args.no_mbr,
        )?;
        log::info!(
            "Search targets: {} accepted, {} rejected as ambiguous/inconsistent",
            candidates.len(),
            guidance.rejected_targets.len()
        );
        for (run, aligned) in runs.iter().zip(&guidance.aligned_runs) {
            if !aligned && !args.no_mbr {
                log::warn!(
                    "Run '{}': insufficient alignment inliers; direct extraction only",
                    run.name
                );
            }
        }
        let matrix = quantify_guided(
            &runs,
            &alignment,
            &config.lfq,
            candidates,
            guidance,
            load_hills,
        );
        write_search_outputs(
            &matrix,
            &out_dir,
            args.max_extraction_qvalue,
            config.lfq.run_tdc,
            config.lfq.search_gate_transfers_only,
        )?;
        matrix
    } else {
        quantify(&runs, &alignment, &config.lfq, load_hills)
    };
    let lfq_elapsed = t.elapsed();
    log::info!(
        "[timing] LFQ: {:.2?} ({} features × {} runs)",
        lfq_elapsed,
        matrix.n_features,
        matrix.n_runs
    );

    // ── Write outputs ─────────────────────────────────────────────────────────
    if matches!(config.output.format, OutputFormat::Parquet) {
        log::warn!(
            "Parquet output for the intensity matrix is not yet implemented; \
             wrote TSV instead."
        );
    }

    if config.output.export_long {
        let feature_paths: Vec<_> = run_paths.iter().map(|p| p.features_path.clone()).collect();
        koth_ms::output::long_matrix::write_bundle(
            &matrix,
            &runs,
            &alignment,
            &feature_paths,
            &config,
            &out_dir,
        )?;
    }

    let t = Instant::now();
    write_consensus_tsv(&matrix, &out_dir, config.output.max_qvalue)?;
    log::info!("[timing] write consensus_features.tsv: {:.2?}", t.elapsed());

    let t = Instant::now();
    write_matrix_tsv(&matrix, &out_dir)?;
    log::info!("[timing] write intensity_matrix.tsv: {:.2?}", t.elapsed());

    if config.lfq.run_tdc {
        let t = Instant::now();
        write_qvalue_matrix_tsv(&matrix, &out_dir)?;
        log::info!("[timing] write qvalue_matrix.tsv: {:.2?}", t.elapsed());
        if config.output.export_decoys {
            let t = Instant::now();
            write_decoy_matrix_tsv(&matrix, &out_dir)?;
            log::info!(
                "[timing] write decoy_intensity_matrix.tsv: {:.2?}",
                t.elapsed()
            );
        }
    }
    if config.output.export_details {
        let t = Instant::now();
        write_lfq_details_tsv(&matrix, &out_dir)?;
        log::info!("[timing] write lfq_details.tsv: {:.2?}", t.elapsed());
    }

    // Save the config used
    let cfg_path = out_dir.join("align_config.toml");
    let cfg_str = config
        .to_toml_string()
        .context("Failed to serialize config")?;
    std::fs::write(&cfg_path, cfg_str)
        .with_context(|| format!("Failed to write config to {}", cfg_path.display()))?;

    // ── Diagnostics report ────────────────────────────────────────────────────
    let total_elapsed = total_start.elapsed();
    let timing = AlignTiming {
        alignment_sec: alignment_elapsed.as_secs_f64(),
        lfq_sec: lfq_elapsed.as_secs_f64(),
        total_sec: total_elapsed.as_secs_f64(),
    };
    let t = Instant::now();
    let report = build_align_report(&runs, &alignment, &matrix, &config, timing);
    log::info!("[timing] build_align_report: {:.2?}", t.elapsed());
    let report_path = out_dir.join("align_report.json");
    let t = Instant::now();
    write_align_report(&report, &report_path).context("Failed to write align report")?;
    log::info!("[timing] write_align_report: {:.2?}", t.elapsed());

    log::info!("[timing] total: {:.2?}", total_elapsed);
    log::info!("Done. Results in {}", out_dir.display());
    Ok(())
}

// ── Output writers ────────────────────────────────────────────────────────────

/// `{:.6}` of the feature's ion mobility, or an empty string when it is exactly
/// zero (absent — e.g. Orbitrap runs carry no IM). Shared by every writer that
/// emits the `im` metadata column so the blank-cell rule can't drift.
fn feature_im_str(matrix: &IntensityMatrix, feat: usize) -> String {
    if matrix.feature_im[feat] != 0.0 {
        format!("{:.6}", matrix.feature_im[feat])
    } else {
        String::new()
    }
}

/// Write the 8 shared per-feature metadata columns (no leading or trailing tab).
/// The caller writes the header and appends the per-run cells; this keeps the
/// column set/format identical across the intensity/q-value/decoy matrices.
fn write_feature_metadata_cols(
    f: &mut BufWriter<std::fs::File>,
    matrix: &IntensityMatrix,
    feat: usize,
) -> std::io::Result<()> {
    write!(
        f,
        "{:.6}\t{:.6}\t{}\t{:.6}\t{}\t{:.6}\t{}\t{}",
        matrix.feature_mass[feat],
        matrix.feature_mz[feat],
        matrix.feature_charge[feat],
        matrix.feature_rt[feat],
        feature_im_str(matrix, feat),
        matrix.feature_combined_score[feat],
        matrix.feature_seed_run[feat],
        matrix.feature_n_contributing_runs[feat],
    )
}

/// Write a features × runs matrix TSV: the shared metadata header + per-feature
/// metadata columns, then one cell per run rendered by `cell`. Used by the
/// intensity, q-value and decoy-intensity matrix writers, which differ only in
/// the per-cell value/format and the log label.
fn write_cell_matrix<F>(
    matrix: &IntensityMatrix,
    out_dir: &Path,
    filename: &str,
    log_label: &str,
    mut cell: F,
) -> anyhow::Result<()>
where
    F: FnMut(&mut BufWriter<std::fs::File>, &IntensityMatrix, usize, usize) -> std::io::Result<()>,
{
    let path = out_dir.join(filename);
    let file = std::fs::File::create(&path)?;
    let mut f = BufWriter::with_capacity(1 << 20, file);

    write!(
        f,
        "massCalib\tmz\tcharge\trtApex\tim\tcombined_score\tseed_run\tn_contributing_runs"
    )?;
    for name in &matrix.run_names {
        write!(f, "\t{}", name)?;
    }
    writeln!(f)?;

    for feat in 0..matrix.n_features {
        write_feature_metadata_cols(&mut f, matrix, feat)?;
        for run in 0..matrix.n_runs {
            cell(&mut f, matrix, feat, run)?;
        }
        writeln!(f)?;
    }

    f.flush()?;
    log::info!("Wrote {} to {}", log_label, path.display());
    Ok(())
}

/// Write `consensus_features.tsv` — one row per consensus feature.
fn write_consensus_tsv(
    matrix: &IntensityMatrix,
    out_dir: &Path,
    max_qvalue: f64,
) -> anyhow::Result<()> {
    let path = out_dir.join("consensus_features.tsv");
    let file = std::fs::File::create(&path)?;
    let mut f = BufWriter::with_capacity(1 << 20, file);

    writeln!(
        f,
        "massCalib\tmz\tcharge\trtApex\tim\tcombined_score\tseed_run\tn_contributing_runs\tn_runs_detected\tgroup_score\tgroup_qvalue"
    )?;

    for feat in 0..matrix.n_features {
        let n_detected = (0..matrix.n_runs)
            .filter(|&run| {
                matrix.intensity(feat, run) > 0.0 && matrix.q_value(feat, run) <= max_qvalue
            })
            .count();

        writeln!(
            f,
            "{:.6}\t{:.6}\t{}\t{:.6}\t{}\t{:.6}\t{}\t{}\t{}\t{:.6}\t{:.6}",
            matrix.feature_mass[feat],
            matrix.feature_mz[feat],
            matrix.feature_charge[feat],
            matrix.feature_rt[feat],
            feature_im_str(matrix, feat),
            matrix.feature_combined_score[feat],
            matrix.feature_seed_run[feat],
            matrix.feature_n_contributing_runs[feat],
            n_detected,
            matrix.consensus[feat].group_score,
            matrix.consensus[feat].group_qvalue,
        )?;
    }

    f.flush()?;
    log::info!("Wrote consensus features to {}", path.display());
    Ok(())
}

/// Write `intensity_matrix.tsv` — features × runs, all intensities unfiltered.
/// Use `qvalue_matrix.tsv` to apply an FDR threshold downstream.
fn write_matrix_tsv(matrix: &IntensityMatrix, out_dir: &Path) -> anyhow::Result<()> {
    write_cell_matrix(
        matrix,
        out_dir,
        "intensity_matrix.tsv",
        "intensity matrix",
        |f, matrix, feat, run| {
            let v = matrix.intensity(feat, run);
            if v > 0.0 {
                write!(f, "\t{:.5e}", v)
            } else {
                write!(f, "\t0")
            }
        },
    )
}

/// Write `qvalue_matrix.tsv` — same layout as intensity_matrix but cells are
/// TDC q-values (1.0 = no signal / TDC not run, 0.0 = perfect score).
/// Use this file to apply an FDR threshold to intensity_matrix.tsv downstream.
fn write_qvalue_matrix_tsv(matrix: &IntensityMatrix, out_dir: &Path) -> anyhow::Result<()> {
    write_cell_matrix(
        matrix,
        out_dir,
        "qvalue_matrix.tsv",
        "q-value matrix",
        |f, matrix, feat, run| write!(f, "\t{:.4}", matrix.q_value(feat, run)),
    )
}

/// Write `decoy_intensity_matrix.tsv` — mirrors `intensity_matrix.tsv` but
/// each cell holds the decoy-pass intensity for that (feature, run).
fn write_decoy_matrix_tsv(matrix: &IntensityMatrix, out_dir: &Path) -> anyhow::Result<()> {
    write_cell_matrix(
        matrix,
        out_dir,
        "decoy_intensity_matrix.tsv",
        "decoy intensity matrix",
        |f, matrix, feat, run| {
            let v = matrix.decoy_intensity(feat, run);
            if v > 0.0 {
                write!(f, "\t{:.5e}", v)
            } else {
                write!(f, "\t0")
            }
        },
    )
}

/// Write `lfq_details.tsv` — long-format file with one row per
/// (consensus feature, run, target/decoy pass) containing the integration
/// stats, RT diagnostics, and observed monoisotopic m/z + ion mobility.
/// `q_value` is filled for target rows only (decoys get NaN since TDC ranks
/// targets against decoys, not the other way round).
fn write_lfq_details_tsv(matrix: &IntensityMatrix, out_dir: &Path) -> anyhow::Result<()> {
    let path = out_dir.join("lfq_details.tsv");
    let file = std::fs::File::create(&path)?;
    let mut f = BufWriter::with_capacity(1 << 20, file);

    // Column dictionary (all LFQ-stage, at the apex column — NOT the feature
    // finder's chromatographic cosine_score). The two isotope signals:
    //   spectral_bhattacharyya = observed-vs-theoretical isotope PATTERN match
    //                            (penalises expected-but-missing peaks)
    //   coelution              = cosine between the matched isotopes' XIC traces
    //                            across RT (do they co-elute?)
    //   hybrid_score           = (rt × intensity × spectral_bhattacharyya × coelution)^¼
    //   seed_combined_score    = seed feature's combined_score (feature-stage)
    writeln!(
        f,
        "feature_idx\tmassCalib\tmz\tcharge\trefRtApex\trefIm\tseed_combined_score\tseed_run\
         \tn_contributing_runs\trun_name\tis_decoy\tis_mbr\tintensity\thybrid_score\tspectral_bhattacharyya\
         \trt_score\tint_score\tcoelution\
         \tn_isotopes_found\texpected_rt\tapex_rt\trt_diff\tpeak_width_rt\
         \tobserved_mz\tppm_error\tobserved_im\tim_delta\tq_value\townership_status\towned_samples\texcluded_samples\tcompeting_feature\tpreceding_signal_fraction"
    )?;

    let fmt_f = |v: f64| {
        if v.is_finite() {
            format!("{:.6}", v)
        } else {
            String::from("NaN")
        }
    };

    for entry in &matrix.entries {
        let feat = entry.feature_idx;
        let run = entry.run_idx;
        let intensity_str = if entry.intensity > 0.0 {
            format!("{:.5e}", entry.intensity)
        } else {
            String::from("0")
        };

        // Note: q-values are computed per (feature, run) in target-decoy
        // competition. We surface the target's q-value on the decoy row too
        // for convenience when filtering — same value for both rows in a pair.
        let q_value = if matrix.q_values.is_empty() {
            f64::NAN
        } else {
            matrix.q_value(feat, run)
        };

        // ppm error vs the centre m/z used to build this row's XIC grid.
        // For targets that is the alignment-predicted m/z; for decoys it is
        // the shifted decoy m/z — so target and decoy residuals are on the
        // same "distance from prediction" scale and stay comparable for
        // downstream rescorers.
        let ppm_error = if entry.observed_mz.is_finite() && entry.expected_mz > 0.0 {
            (entry.observed_mz - entry.expected_mz) / entry.expected_mz * 1e6
        } else {
            f64::NAN
        };

        // Per-row IM delta vs the centre IM used to build the grid (same
        // target/decoy symmetry as ppm_error). NaN when either side missing.
        let im_delta = if entry.observed_im.is_finite() && entry.expected_im != 0.0 {
            entry.observed_im - entry.expected_im
        } else {
            f64::NAN
        };

        let rt_diff = if entry.apex_rt.is_finite() && entry.expected_rt.is_finite() {
            entry.apex_rt - entry.expected_rt
        } else {
            f64::NAN
        };

        writeln!(
            f,
            "{feat}\t{mass:.6}\t{mz:.6}\t{charge}\t{rt:.6}\t{im}\t{seed_combined:.6}\t{seed_run}\
             \t{ncont}\t{run_name}\t{decoy}\t{is_mbr}\t{intensity}\t{hybrid:.6}\t{bhatt:.6}\
             \t{rt_score:.6}\t{int_score:.6}\t{coelution:.6}\
             \t{nslots}\t{exp_rt}\t{apex_rt}\t{rtdiff}\t{pw}\
             \t{obs_mz}\t{ppm}\t{obs_im}\t{im_delta}\t{qv}\t{ownership}\t{owned}\t{excluded}\t{competitor}\t{preceding:.6}",
            feat = feat,
            mass = matrix.feature_mass[feat],
            mz = matrix.feature_mz[feat],
            charge = matrix.feature_charge[feat],
            rt = matrix.feature_rt[feat],
            im = feature_im_str(matrix, feat),
            seed_combined = matrix.feature_combined_score[feat],
            seed_run = matrix.feature_seed_run[feat],
            ncont = matrix.feature_n_contributing_runs[feat],
            run_name = matrix.run_names[run],
            decoy = entry.is_decoy,
            is_mbr = entry.is_mbr,
            intensity = intensity_str,
            hybrid = entry.hybrid_score,
            bhatt = entry.spectral_bhattacharyya,
            rt_score = entry.rt_score,
            int_score = entry.int_score,
            coelution = entry.coelution,
            nslots = entry.n_isotopes_found,
            exp_rt = fmt_f(entry.expected_rt),
            apex_rt = fmt_f(entry.apex_rt),
            rtdiff = fmt_f(rt_diff),
            pw = fmt_f(entry.peak_width_rt),
            obs_mz = fmt_f(entry.observed_mz),
            ppm = fmt_f(ppm_error),
            obs_im = fmt_f(entry.observed_im),
            im_delta = fmt_f(im_delta),
            qv = fmt_f(q_value),
            ownership = entry.ownership_status,
            preceding = entry.preceding_signal_fraction,
            owned = entry.owned_samples,
            excluded = entry.excluded_samples,
            competitor = entry.competing_feature.map(|i| i.to_string()).unwrap_or_default(),
        )?;
    }

    f.flush()?;
    log::info!(
        "Wrote LFQ details ({} rows) to {}",
        matrix.entries.len(),
        path.display()
    );
    Ok(())
}