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

use csv::{ReaderBuilder, Writer};
use std::error::Error;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::time::Instant;
use std::process;
use std::collections::HashMap;

use ndarray::Array1;

// Powerlaw crate
use powerlaw::dist::exponential::Exponential;
use powerlaw::dist::pareto;

// ------------------------ Math Helpers ------------------------

fn safe_log(v: f64) -> f64 {
    (v.max(1e-12)).ln()  // safer log for numerical stability
}

fn variance(values: &Array1<f64>) -> f64 {
    if values.is_empty() { return 0.0; }
    let mean = values.mean().unwrap_or(0.0);
    let mean_sq = values.mapv(|v| v * v).mean().unwrap_or(0.0);
    (mean_sq - mean * mean).max(0.0)
}

// ------------------------ Regression Helpers ------------------------

fn linear_regression(xs: &[f64], ys: &[f64]) -> (f64, f64, f64) {
    let n = xs.len() as f64;
    if n == 0.0 { return (0.0, 0.0, 0.0); }

    let sum_x = xs.iter().sum::<f64>();
    let sum_y = ys.iter().sum::<f64>();
    let sum_xx = xs.iter().map(|v| v*v).sum::<f64>();
    let sum_xy = xs.iter().zip(ys).map(|(x,y)| x*y).sum::<f64>();
    let mean_y = sum_y / n;

    let denom = (n*sum_xx - sum_x*sum_x).max(1e-12);
    let slope = (n*sum_xy - sum_x*sum_y) / denom;
    let intercept = (sum_y - slope*sum_x)/n;

    let ss_tot: f64 = ys.iter().map(|y| (y-mean_y).powi(2)).sum();
    let ss_res: f64 = xs.iter().zip(ys).map(|(x,y)| (y - (slope*x + intercept)).powi(2)).sum();
    let r2 = if ss_tot>0.0 {1.0 - ss_res/ss_tot} else {0.0};

    (slope, intercept, r2)
}

// ------------------------ Powerlaw Fit Helper ------------------------

fn fit_powerlaw(data: &[f64], group: &str) -> Vec<String> {
    let mut output_lines = Vec::new();

    if data.is_empty() {
        output_lines.push(format!("# No valid data for group '{}', skipping powerlaw fit.", group));
        return output_lines;
    }

    let mut data_vec = data.to_vec();

    // Step 1: Pareto Type I
    let (x_mins, alphas) = pareto::find_alphas_fast(&mut data_vec);
    let pareto_fit = pareto::gof(&data_vec, &x_mins, &alphas);

    // Step 2: Exponential
    let exp_dist = Exponential::from_fitment(&data_vec, &pareto_fit);

    // Step 3: Summary
    output_lines.push(format!("# --- Powerlaw fits for group '{}' ---", group));
    output_lines.push(format!("# Pareto Type I: x_min = {}, alpha = {}", pareto_fit.x_min, pareto_fit.alpha));
    output_lines.push(format!("# Exponential fit: lambda = {}", exp_dist.lambda));

    output_lines
}

// ------------------------ CATEGORY HELPERS ------------------------

fn categorize_prevalence(prevalence: f64) -> &'static str {
    match prevalence {
        p if p < 40.0 => "rare",
        p if p < 70.0 => "moderate",
        _ => "common",
    }
}

fn categorize_deviation(log_var: f64, log_taylor: f64) -> &'static str {
    let diff = log_var - log_taylor;
    match diff {
        d if d < -0.35 => "morestable",
        d if d <= 0.35 => "stable",
        _ => "lessstable",
    }
}

// ------------------------ MAIN ------------------------

//fn main() -> Result<(), Box<dyn Error>> {
pub fn run(args: Vec<String>) -> Result<(), Box<dyn Error>> {
    let start_time = Instant::now();
    let args = args;
//    let args: Vec<String> = env::args().collect();

    if args.len() < 4 || args[1] != "-i" {
        eprintln!("Usage: {} -i <input.csv> -m <metadata.csv> -mc <meta_column> [-o <prefix>] [-mp <min_prevalence>]", args[0]);
        process::exit(1);
    }

    let input_file = &args[2];
    let mut meta_file = None;
    let mut meta_column = None;
    let mut output_prefix = "META".to_string();
    let mut min_prevalence: Option<f64> = None;

    let mut i = 3;
    while i < args.len() {
        match args[i].as_str() {
            "-m" => { meta_file = args.get(i+1).cloned(); i += 2; }
            "-mc" => { meta_column = args.get(i+1).cloned(); i += 2; }
            "-o" => { output_prefix = args.get(i+1).cloned().unwrap_or(output_prefix.clone()); i += 2; }
            "-mp" => { min_prevalence = args.get(i+1).and_then(|v| v.parse::<f64>().ok()); i += 2; }
            _ => { eprintln!("Unknown arg {}", args[i]); process::exit(1); }
        }
    }

    let meta_file = meta_file.expect("Metadata file required");
    let meta_column = meta_column.expect("Metadata column required");

    // ---------------- Load input table ----------------
    let mut rdr = ReaderBuilder::new().from_path(input_file)?;
    let headers = rdr.headers()?.clone();

    let mut rows = Vec::new();
    for result in rdr.records() {
        let record = result?;
        let name = record[0].to_string();
        let values = record.iter().skip(1)
            .map(|s| s.parse::<f64>().unwrap_or(0.0))
            .collect::<Vec<_>>();
        rows.push((name, values));
    }

    // ---------------- Load Metadata ----------------
    let mut meta_rdr = ReaderBuilder::new().from_path(meta_file)?;
    let meta_headers = meta_rdr.headers()?.clone();
    let meta_col_idx = meta_headers.iter().position(|h| h == &meta_column)
        .expect("Metadata column not found");

    let mut meta_map = HashMap::<String, Vec<usize>>::new();
    for result in meta_rdr.into_records() {
        let record = result?;
        let id = record[0].to_string();
        let group = record[meta_col_idx].trim().to_string();

        if let Some(idx) = headers.iter().position(|h| h == &id) {
            meta_map.entry(group).or_default().push(idx - 1);
        }
    }

    // ---------------- Per-group CSV ----------------
    let out_header = vec![
        "meta","Name",
        "Sum","Mean","Var",
        "Prevalence",
        "logSum","logMean","logVar",
        "classic_taylor","log_taylor",
        "presence-cat","deviation-cat"
    ];

    for (group, col_indices) in meta_map {
        let group_output = format!("{}_{}.csv", output_prefix, group);
        let output_file = File::create(&group_output)?;
        let buf_writer = BufWriter::new(output_file);
        let mut writer = Writer::from_writer(buf_writer);

        writer.write_record(&out_header)?;

        let mut sum_column_for_fit = Vec::new();
        let mut log_means = Vec::new();
        let mut log_vars = Vec::new();

        // Store species rows temporarily
        let mut species_rows: Vec<(Vec<String>, f64, f64, f64, f64, f64)> = Vec::new();

        // First: compute per-species stats
        for (name, vals) in &rows {
            let subset = col_indices.iter().map(|&i| vals[i]).collect::<Vec<_>>();
            let arr = Array1::from(subset.clone());

            let sum = arr.sum();
            if sum == 0.0 { continue; }

            let mean = arr.mean().unwrap_or(0.0);
            let var = variance(&arr);
            let log_sum = safe_log(sum);
            let log_mean = safe_log(mean);
            let log_var = safe_log(var);

            let nonzero_count = subset.iter().filter(|&&v| v > 0.0).count() as f64;
            let prevalence = (nonzero_count / subset.len() as f64) * 100.0;

            // Apply minimum prevalence filter if specified
            if let Some(mp) = min_prevalence {
                if prevalence < mp {
                    continue;
                }
            }

            sum_column_for_fit.push(sum);
            log_means.push(log_mean);
            log_vars.push(log_var);

            species_rows.push((
                vec![
                    group.clone(),
                    name.clone(),
                    format!("{:.4}", sum),
                    format!("{:.4}", mean),
                    format!("{:.4}", var),
                    format!("{:.4}", prevalence),
                    format!("{:.4}", log_sum),
                    format!("{:.4}", log_mean),
                    format!("{:.4}", log_var),
                    "".to_string(),
                    "".to_string(),
                    "".to_string(),  // placeholder for presence-cat
                    "".to_string(),  // placeholder for deviation-cat
                ],
                mean, var, log_mean, log_var, prevalence
            ));
        }

        // ---------------- Taylor's Law ----------------
        let (tl_slope, tl_intercept, tl_r2) =
            linear_regression(&log_means, &log_vars);
        let tl_a = tl_intercept.exp();

        // ---------------- Fill species-specific expected columns and categories ----------------
        species_rows.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        for (row, mean, _var, log_mean, log_var, prevalence) in &species_rows {
            let classic_exp = tl_a * mean.powf(tl_slope);
            let log_exp = safe_log(tl_a) + tl_slope * (*log_mean);

            let mut row = row.clone();
            row[9] = format!("{:.4}", classic_exp);
            row[10] = format!("{:.4}", log_exp);

            // --- NEW: Categorize prevalence and deviation ---
            row[11] = categorize_prevalence(*prevalence).to_string();
            row[12] = categorize_deviation(*log_var, log_exp).to_string();

            writer.write_record(&row)?;
        }

        // ---------------- Write comments ----------------
        let mut buf_writer = writer.into_inner()?;
        writeln!(buf_writer, "# Total samples/replicates:\t{}", col_indices.len())?;
        writeln!(buf_writer, "# --- Taylor's Law (Variance ~ Mean^β) ---")?;
        writeln!(buf_writer, "# beta (slope):\t\t{}", tl_slope)?;
        writeln!(buf_writer, "# alpha (intercept):\t{}", tl_intercept)?;
        writeln!(buf_writer, "# a (scale = exp(alpha)):\t{}", tl_a)?;
        writeln!(buf_writer, "# R^2:\t\t\t{}", tl_r2)?;

        let pl_comments = fit_powerlaw(&sum_column_for_fit, &group);
        for line in pl_comments {
            writeln!(buf_writer, "{}", line)?;
        }

        buf_writer.flush()?;
        println!("Finished group: {}", group);
    }

    println!("Done! Processing took {:.2?}", start_time.elapsed());
    Ok(())
}
