// Copyright (c) 2026 Ranjith N
// SPDX-License-Identifier: MIT OR Apache-2.0

use std::collections::HashMap;
use std::error::Error;
use std::fs::File;
use std::io::{BufRead, BufReader};
use csv::Writer;

// Compute geometric mean of nonzero positive values
fn geometric_mean(values: &[f64]) -> f64 {
    let positives: Vec<f64> = values.iter().cloned().filter(|&v| v > 0.0).collect();
    if positives.is_empty() {
        return 1.0;
    }
    let log_sum: f64 = positives.iter().map(|x| x.ln()).sum();
    (log_sum / positives.len() as f64).exp()
}

// Extract taxon by optional column index
fn extract_taxon_by_column(taxon_str: &str, col_index: Option<usize>) -> String {
    if let Some(idx) = col_index {
        let parts: Vec<&str> = taxon_str.split('|').collect();
        if idx < parts.len() {
            return parts[idx].trim().to_string();
        } else {
            return String::new();
        }
    }
    taxon_str.to_string()
}

pub fn run(args: Vec<String>) -> Result<(), Box<dyn Error>> {
    if args.len() < 5 {
        eprintln!(
            "Usage: {} -i <input.tsv> -o <output.csv> [-c <1|2|3|4|5|6|7>] [-f clr|abun|log1p|logret]",
            args[0]
        );
        std::process::exit(1);
    }

    let mut input_file = String::new();
    let mut output_file = String::new();
    let mut col_index: Option<usize> = None;
    let mut function = String::from("clr"); // default

    let mut i = 1;
    while i < args.len() {
        match args[i].as_str() {
            "-i" => {
                input_file = args.get(i + 1).cloned().unwrap_or_default();
                i += 1;
            }
            "-o" => {
                output_file = args.get(i + 1).cloned().unwrap_or_default();
                i += 1;
            }
            "-c" => {
                let val: usize = args.get(i + 1).and_then(|v| v.parse().ok()).unwrap_or(0);
                if val >= 1 {
                    col_index = Some(val - 1);
                }
                i += 1;
            }
            "-f" => {
                function = args.get(i + 1).cloned().unwrap_or_default().to_lowercase();
                i += 1;
            }
            _ => {}
        }
        i += 1;
    }

    if input_file.is_empty() || output_file.is_empty() {
        eprintln!("❌ Missing required arguments.");
        std::process::exit(1);
    }

    // Open file
    let file = File::open(&input_file)?;
    let reader = BufReader::new(file);

    // Read header
    let mut lines = reader.lines();
    let header_line = lines.next().ok_or("Empty file")??;
    let mut header_fields = header_line.split('\t');
    header_fields.next(); // skip Taxon
    let sample_names: Vec<String> = header_fields.map(|s| s.to_string()).collect();
    let n_samples = sample_names.len();

    let mut aggregated: HashMap<String, Vec<f64>> = HashMap::new();

    // Read data
    for line in lines {
        let line = line?;
        if line.trim().is_empty() {
            continue;
        }
        let fields: Vec<&str> = line.split('\t').collect();
        if fields.len() < 2 {
            continue;
        }

        let taxon_name = extract_taxon_by_column(fields[0], col_index);
        if taxon_name.is_empty() {
            continue;
        }

        let entry = aggregated
            .entry(taxon_name)
            .or_insert(vec![0.0; n_samples]);

        let limit = std::cmp::min(fields.len() - 1, n_samples);
        for i in 0..limit {
            let val: f64 = fields[i + 1].parse().unwrap_or(0.0);
            entry[i] += val;
        }
    }

    // =========================
    // Normalization
    // =========================
    match function.as_str() {
        "clr" => {
            for col in 0..n_samples {
                let col_vals: Vec<f64> = aggregated.values().map(|v| v[col]).collect();
                let gm = geometric_mean(&col_vals);
                for v in aggregated.values_mut() {
                    v[col] = if v[col] > 0.0 {
                        (v[col] / gm).ln()
                    } else {
                        0.0
                    };
                }
            }
        }

        "log1p" => {
            for v in aggregated.values_mut() {
                for x in v.iter_mut() {
                    *x = (*x + 1.0).ln();
                }
            }
        }

        "logret" => {
            for v in aggregated.values_mut() {
                let mut prev = v[0];
                v[0] = 0.0;
                for i in 1..v.len() {
                    let curr = v[i];
                    if curr > 0.0 && prev > 0.0 {
                        v[i] = (curr / prev).ln();
                    } else {
                        v[i] = 0.0;
                    }
                    prev = curr;
                }
            }
        }

        "abun" => {}

        _ => {
            eprintln!("❌ Unknown function '{}'. Allowed: clr, abun, log1p, logret", function);
            std::process::exit(1);
        }
    }

    // Sort taxa
    let mut taxa_sorted: Vec<_> = aggregated.iter().collect();
    taxa_sorted.sort_by(|a, b| a.0.cmp(b.0));

    // Write CSV
    let mut wtr = Writer::from_path(&output_file)?;
    let mut header = vec!["Name".to_string()];
    header.extend(sample_names);
    wtr.write_record(&header)?;

    for (taxon, values) in taxa_sorted {
        let mut row = vec![taxon.clone()];
        row.extend(values.iter().map(|v| v.to_string()));
        wtr.write_record(&row)?;
    }

    wtr.flush()?;
    println!("✅ '{}' normalized data written to '{}'", function, output_file);

    Ok(())
}

