// 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;
//use std::env;
use std::time::Instant;
use std::process;
use std::collections::HashMap;
use std::cmp::Ordering;

use ndarray::Array1;
use rayon::prelude::*;

// ------------------------ Correlation Functions ------------------------

fn pearson_fast(x: &Array1<f64>, y: &Array1<f64>) -> f64 {
    let n = x.len() as f64;
    let sum_x = x.sum();
    let sum_y = y.sum();
    let sum_xy = x.dot(y);
    let sum_x2 = x.dot(x);
    let sum_y2 = y.dot(y);
    let numerator = sum_xy - (sum_x * sum_y) / n;
    let denom = ((sum_x2 - sum_x*sum_x/n) * (sum_y2 - sum_y*sum_y/n)).sqrt();
    if denom.abs() < 1e-10 { 0.0 } else { numerator / denom }
}

fn rank_data_fast(data: &Array1<f64>) -> Array1<f64> {
    let mut idx: Vec<usize> = (0..data.len()).collect();
    idx.sort_by(|&i, &j| data[i].partial_cmp(&data[j]).unwrap_or(Ordering::Equal));

    let mut ranks = Array1::zeros(data.len());
    let mut start = 0;
    while start < idx.len() {
        let val = data[idx[start]];
        let mut end = start + 1;
        while end < idx.len() && (data[idx[end]] - val).abs() < 1e-10 {
            end += 1;
        }
        let rank = (start + end - 1) as f64 / 2.0 + 1.0;
        for k in start..end {
            ranks[idx[k]] = rank;
        }
        start = end;
    }
    ranks
}

fn calculate_spearman_fast(x: &Array1<f64>, y: &Array1<f64>) -> f64 {
    let rx = rank_data_fast(x);
    let ry = rank_data_fast(y);
    pearson_fast(&rx, &ry)
}

// ------------------------ P-values & stats ------------------------

fn calculate_p_value(r: f64, n: usize) -> f64 {
    if n <= 2 { return 1.0; }
    if n > 10 {
        let z = 0.5 * ((1.0 + r)/(1.0 - r)).ln() * (n as f64 - 3.0).sqrt();
        2.0 * (1.0 - normal_cdf(z.abs()))
    } else {
        let t = r * ((n as f64 - 2.0)/(1.0 - r*r)).sqrt();
        2.0 * (1.0 - student_t_cdf(t.abs(), n as f64 - 2.0))
    }.max(0.0).min(1.0)
}

fn normal_cdf(x: f64) -> f64 {
    let a1=0.254829592; let a2=-0.284496736; let a3=1.421413741; let a4=-1.453152027; let a5=1.061405429; let p=0.3275911;
    let sign = if x<0.0 {-1.0} else {1.0};
    let x = x.abs();
    let t = 1.0/(1.0+p*x);
    let y = 1.0 - (((((a5*t+a4)*t+a3)*t+a2)*t+a1)*t*(-x*x).exp());
    0.5*(1.0 + sign*y)
}

fn student_t_cdf(_t: f64, _df: f64) -> f64 { 0.5 }

// ------------------------ Non-zero correlation ------------------------

fn calculate_non_zero_correlation(x: &Array1<f64>, y: &Array1<f64>) -> (f64, f64, usize, f64, f64) {
    let mask = x.iter().zip(y.iter()).map(|(&a,&b)| a.abs()>1e-10 && b.abs()>1e-10).collect::<Vec<_>>();
    let n = mask.iter().filter(|&&b| b).count();
    if n < 3 { return (0.0,0.0,n,1.0,1.0); }

    let x_nz = Array1::from_iter(x.iter().zip(mask.iter()).filter(|&(_, &b)| b).map(|(&v,_)| v));
    let y_nz = Array1::from_iter(y.iter().zip(mask.iter()).filter(|&(_, &b)| b).map(|(&v,_)| v));

    let spearman = calculate_spearman_fast(&x_nz, &y_nz);
    let pearson = pearson_fast(&x_nz, &y_nz);
    let spearman_p = calculate_p_value(spearman, n);
    let pearson_p = calculate_p_value(pearson, n);

    (spearman, pearson, n, spearman_p, pearson_p)
}

// ------------------------ Direction ------------------------

fn calculate_direction(
    x: &Array1<f64>, 
    y: &Array1<f64>, 
    nzpairs: usize, 
    threshold: f64
) -> (char, f64, String, f64) {
    if nzpairs < 3 || x.len() < 3 || y.len() < 3 { 
        return ('?', 0.0, "??".to_string(), 0.0); 
    }

    let n = x.len() as f64;
    let sum_x = x.sum(); 
    let sum_y = y.sum();
    let sum_xy = x.dot(y);
    let sum_x2 = x.dot(x); 
    let sum_y2 = y.dot(y);

    let numerator = sum_xy - (sum_x * sum_y) / n;
    let denom_x = (sum_x2 - sum_x * sum_x / n).max(1e-10);
    let denom_y = (sum_y2 - sum_y * sum_y / n).max(1e-10);

    let slope_yx = numerator / denom_x;
    let slope_xy = numerator / denom_y;

    // --- Updated direction logic ---
    let yx_valid = slope_yx.abs() > slope_xy.abs() * threshold;
    let xy_valid = slope_xy.abs() > slope_yx.abs() * threshold;

    let (direction, slope) = match (yx_valid, xy_valid) {
        (true, false) => ('>', slope_yx),  // only y->x is valid
        (false, true) => ('<', slope_xy),  // only x->y is valid
        _ => ('?', 0.0),                   // both valid or both invalid
    };

    // Determine which array is perturbation and which is observation
    let (perturb, observe) = if direction == '>' { (x, y) } else { (y, x) };
    let mut positive_matches = 0; 
    let mut negative_matches = 0;
    let mut positive_effect_counts = HashMap::new();
    let mut negative_effect_counts = HashMap::new();

    for i in 0..perturb.len() {
        if perturb[i].abs() < 1e-10 || observe[i].abs() < 1e-10 { continue; }
        let delta = (0.1 * perturb[i].abs()).max(0.01);

        let val_plus = perturb[i] * (1.0 + delta);
        let effect_plus = observe[i] + slope * (val_plus - perturb[i]);
        let plus_sign = match (effect_plus - observe[i]).signum() as i32 {
            1 => '+', 
            -1 => '-', 
            _ => '0'
        };
        if (effect_plus - observe[i]).signum() == slope.signum() {
            positive_matches += 1; 
            *positive_effect_counts.entry(plus_sign).or_insert(0) += 1;
        }

        let val_minus = perturb[i] * (1.0 - delta);
        let effect_minus = observe[i] + slope * (val_minus - perturb[i]);
        let minus_sign = match (effect_minus - observe[i]).signum() as i32 {
            1 => '+', 
            -1 => '-', 
            _ => '0'
        };
        if (effect_minus - observe[i]).signum() == slope.signum() {
            negative_matches += 1; 
            *negative_effect_counts.entry(minus_sign).or_insert(0) += 1;
        }
    }

    let (dominant_perturb_sign, effect_counts) = if positive_matches >= negative_matches {
        ('+', positive_effect_counts) 
    } else { 
        ('-', negative_effect_counts) 
    };

    let dominant_effect_sign = effect_counts.iter()
        .max_by_key(|&(_, count)| count)
        .map(|(sign, _)| *sign)
        .unwrap_or('0');

    let vardirect = format!("{}/{}", dominant_perturb_sign, dominant_effect_sign);
    let roeffect = if direction == '>' { 
        slope_yx.abs() / slope_xy.abs().max(1e-10) 
    } else { 
        slope_xy.abs() / slope_yx.abs().max(1e-10) 
    };

    (direction, roeffect, vardirect, slope.abs())
}



// ------------------------ Weighted score & Fisher ------------------------

fn calculate_weighted_score(spearman:f64, pearson:f64, spearman_p:f64, pearson_p:f64, nzpairs:usize) -> f64 {
    let avg_p = (spearman_p + pearson_p)/2.0;
    let p_factor = 1.0/(1.0 + (avg_p/0.05).powi(2));
    (spearman.abs() + pearson.abs()) * p_factor * (nzpairs as f64).ln().max(1.0)
}

fn calculate_fishers_combined_p(spearman_p:f64, pearson_p:f64) -> f64 {
    if spearman_p<=0.0 || pearson_p<=0.0 { return 0.0; }
    let chi_sq = -2.0*(spearman_p.ln() + pearson_p.ln());
    1.0 - chi_square_cdf(chi_sq, 4)
}

fn chi_square_cdf(x:f64, df:u32) -> f64 { incomplete_gamma(df as f64/2.0, x/2.0) }
fn incomplete_gamma(s:f64,x:f64) -> f64 {
    if x<=0.0 { return 0.0; }
    if x.is_infinite() { return 1.0; }
    let mut term = 1.0/s;
    let mut sum = term;
    for n in 1..100 {
        term *= x/(s+n as f64);
        sum += term;
        if term<1e-20 { break; }
    }
    sum*x.powf(s)*(-x).exp()
}

// ------------------------ Index builder (NEW) ------------------------

fn build_index_from_oall(oall_path: &str, output_prefix: &str) -> Result<(), Box<dyn Error>> {
    let mut rdr = ReaderBuilder::new().from_path(oall_path)?;

    #[derive(Default)]
    struct Stats {
        outdeg: usize,
        indeg: usize,
        out_plus: usize,
        out_minus: usize,
        in_plus: usize,
        in_minus: usize,
    }

    let mut index: HashMap<(String, String), Stats> = HashMap::new();

    for rec in rdr.records() {
        let r = rec?;

        let meta = r[0].to_string();
        let name1 = r[1].to_string();
        let name2 = r[2].to_string();
        let vardirect = r[11].chars().collect::<Vec<_>>();

        if vardirect.len() != 3 { continue; }
        let out_sign = vardirect[0];
        let in_sign = vardirect[2];

        let s1 = index.entry((meta.clone(), name1)).or_default();
        s1.outdeg += 1;
        match out_sign {
            '+' => s1.out_plus += 1,
            '-' => s1.out_minus += 1,
            _ => {}
        }

        let s2 = index.entry((meta, name2)).or_default();
        s2.indeg += 1;
        match in_sign {
            '+' => s2.in_plus += 1,
            '-' => s2.in_minus += 1,
            _ => {}
        }
    }

    let index_path = format!("{}_index.csv", output_prefix);
    let file = File::create(&index_path)?;
    let mut wtr = Writer::from_writer(BufWriter::new(file));

    wtr.write_record(&[
        "Name",
        "meta",
        "degree",
        "in_degree",
        "out_degree",
        "net_dominance",
        "outdeg_+",
        "outdeg_-",
        "indeg_+",
        "indeg_-",
    ])?;


    let mut rows = Vec::new();

    for ((meta, name), s) in index {
        let degree = s.indeg + s.outdeg;
        let net = s.outdeg as isize - s.indeg as isize;
    
        rows.push(vec![
            name,
            meta,
            degree.to_string(),
            s.indeg.to_string(),
            s.outdeg.to_string(),
            net.to_string(),
            s.out_plus.to_string(),
            s.out_minus.to_string(),
            s.in_plus.to_string(),
            s.in_minus.to_string(),
        ]);
    }
    
    // ---------- SORT INDEX ----------
    rows.sort();
    
    for row in rows {
        wtr.write_record(&row)?;
    }


    wtr.flush()?;
    println!("Index written to {}", index_path);
    Ok(())
}

fn sort_and_write_csv(
    writer: &mut csv::Writer<std::io::BufWriter<std::fs::File>>,
    mut rows: Vec<Vec<String>>,
) -> Result<(), Box<dyn std::error::Error>> {
    rows.sort(); // alphabetical, like `sort`
    for row in rows {
        writer.write_record(&row)?;
    }
    Ok(())
}

fn sort_csv_file_in_place(path: &str) -> Result<(), Box<dyn std::error::Error>> {
    let mut rdr = csv::Reader::from_path(path)?;
    let headers = rdr.headers()?.clone();

    let mut rows: Vec<Vec<String>> = rdr
        .records()
        .filter_map(|r| r.ok())
        .map(|r| r.iter().map(|s| s.to_string()).collect())
        .collect();

    rows.sort(); // alphabetical

    let file = std::fs::File::create(path)?;
    let mut wtr = csv::Writer::from_writer(std::io::BufWriter::new(file));

    wtr.write_record(&headers)?;
    for row in rows {
        wtr.write_record(&row)?;
    }

    wtr.flush()?;
    Ok(())
}


// ------------------------ 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> [-t <threshold>] ...", args[0]);
        eprintln!("Usage: {} -i <input.csv> -m <metadata.csv> -mc <meta_column> [-mp <min_percent>] [-t <threshold>] ...", args[0]);

        process::exit(1);
    }

    let input_file = &args[2];
    let mut meta_file = None; let mut meta_column = None;
    let mut selected_condition: Option<String> = None;
    let mut output_prefix = "META".to_string();
    let mut min_coc = 0;
    let mut s_minval: Option<f64> = None;
    let mut p_minval: Option<f64> = None;
    let mut s_maxp: Option<f64> = None;
    let mut p_maxp: Option<f64> = None;
    let mut threshold: f64 = 1.5;

    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; },
            "-mc" => {
                let raw = args.get(i+1).cloned().expect("Missing -mc argument");
                // Check if condition is specified: "Column:Condition"
                if let Some((col, cond)) = raw.split_once(':') {
                    meta_column = Some(col.to_string());
                    // store selected metadata condition
                    selected_condition = Some(cond.to_string());
                } else {
                    meta_column = Some(raw);
                    selected_condition = None;
                }
                i += 2;
            }

            "-o"=>{ output_prefix=args.get(i+1).cloned().unwrap_or(output_prefix.clone()); i+=2; },
            "-mp"=>{ min_coc=args.get(i+1).and_then(|s| s.parse::<usize>().ok()).unwrap_or(0); i+=2; },
            "-Sminval"=>{ s_minval=args.get(i+1).and_then(|s| s.parse::<f64>().ok()); i+=2; },
            "-Pminval"=>{ p_minval=args.get(i+1).and_then(|s| s.parse::<f64>().ok()); i+=2; },
            "-Smaxp"=>{ s_maxp=args.get(i+1).and_then(|s| s.parse::<f64>().ok()); i+=2; },
            "-Pmaxp"=>{ p_maxp=args.get(i+1).and_then(|s| s.parse::<f64>().ok()); i+=2; },
            "-t"=>{ threshold=args.get(i+1).and_then(|s| s.parse::<f64>().ok()).unwrap_or(1.5); 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");

    let mut rdr = ReaderBuilder::new().from_path(input_file)?;
    let headers = rdr.headers()?.clone();

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

    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::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 user specified one condition: skip all others
        if let Some(ref wanted) = selected_condition {
            if &group != wanted {
                continue;
            }
        }
    
        if let Some(idx) = headers.iter().position(|h| h==&id) {
            meta_map.entry(group).or_insert_with(Vec::new).push(idx - 1);
        }
    }


    let combined_header = vec!["meta","Name1","Name2","Spearman","Pearson","nzpairs","spearman_p","pearson_p","dir","Roeffect","slope","vardirect","weighted_c_score","fishers_c_pval"];
    let mut combined_writer = None;
    if selected_condition.is_none() {
        let combined_output = format!("{}_oall.csv", output_prefix);
        let combined_file = File::create(&combined_output)?;
        let mut writer = Writer::from_writer(BufWriter::new(combined_file));
        writer.write_record(&combined_header)?;
        combined_writer = Some(writer);
    }


    for (group_name, col_indices) in meta_map {
        if let Some(ref wanted) = selected_condition {
            if &group_name != wanted {
                continue;
            }
        }

        let group_output = format!("{}_{}.csv", output_prefix, group_name);
        let output_file = File::create(&group_output)?;
        let mut writer = Writer::from_writer(BufWriter::new(output_file));
        //writer.write_record(&combined_header[1..])?;
        writer.write_record(&combined_header)?;

        let entries_arc = &entries;
        let col_indices_arc = &col_indices;
        let group_name_arc = &group_name;

        let total_reps = col_indices_arc.len();

        let group_records: Vec<Vec<String>> = (0..entries.len()).into_par_iter().flat_map(|i| {
            let mut rows = Vec::new();
            for j in (i+1)..entries_arc.len() {
                let (name1, vals1) = &entries_arc[i];
                let (name2, vals2) = &entries_arc[j];
                let x_vals = Array1::from(col_indices_arc.iter().map(|&idx| vals1[idx]).collect::<Vec<_>>());
                let y_vals = Array1::from(col_indices_arc.iter().map(|&idx| vals2[idx]).collect::<Vec<_>>());
                let (spearman, pearson, nzpairs, spearman_p, pearson_p) = calculate_non_zero_correlation(&x_vals, &y_vals);

                // ---------- NEW PERCENTAGE MIN-COC ----------
                let min_required = ((total_reps as f64) * (min_coc as f64 / 100.0)).ceil() as usize;
                if nzpairs < min_required { continue; }

                if let Some(smin) = s_minval { if spearman.abs()<smin { continue; } }
                if let Some(pmin) = p_minval { if pearson.abs()<pmin { continue; } }
                if let Some(smaxp) = s_maxp { if spearman_p>smaxp { continue; } }
                if let Some(pmaxp) = p_maxp { if pearson_p>pmaxp { continue; } }

                let (mut direction, roeffect, mut vardirect, slope) = calculate_direction(&x_vals, &y_vals, nzpairs, threshold);
                if direction=='?' || vardirect=="??" { continue; }

                let mut n1=name1.clone(); 
                let mut n2=name2.clone();
                let s=spearman; let p=pearson; let sp=spearman_p; let pp=pearson_p;


                if direction == '<' {
                    std::mem::swap(&mut n1, &mut n2);
                
                    // flip vardirect like "+/-" -> "-/+"
                    let parts: Vec<char> = vardirect.chars().collect();
                    if parts.len() == 3 && parts[1] == '/' {
                        vardirect = format!("{}/{}", parts[2], parts[0]);
                    }
                
                    direction = '>';
                }


                let weighted_c_score = calculate_weighted_score(s,p,sp,pp,nzpairs);
                let fishers_c_pval = calculate_fishers_combined_p(sp,pp);

                // ---------- ROUND ALL FLOATS TO 4 DECIMALS ----------
                let row = vec![
                    group_name_arc.clone(),
                    n1,
                    n2,
                    format!("{:.4}", s),
                    format!("{:.4}", p),
                    nzpairs.to_string(),
                    format!("{:.4}", sp),
                    format!("{:.4}", pp),
                    direction.to_string(),
                    format!("{:.4}", roeffect),
                    format!("{:.4}", slope),
                    vardirect,
                    format!("{:.4}", weighted_c_score),
                    format!("{:.4}", fishers_c_pval),
                ];

                rows.push(row);
            }
            rows
        }).collect();

        let mut sorted_records = group_records.clone();
        //sorted_records.sort_by(|a,b| b[14].parse::<f64>().unwrap_or(0.0).partial_cmp(&a[14].parse::<f64>().unwrap_or(0.0)).unwrap_or(Ordering::Equal));
        sorted_records.sort_by(|a,b| b[12].parse::<f64>().unwrap_or(0.0).partial_cmp(&a[12].parse::<f64>().unwrap_or(0.0)).unwrap_or(Ordering::Equal));

        // --- WRITE PER-GROUP CSV (SORTED, HEADER PRESERVED) ---
        sort_and_write_csv(&mut writer, sorted_records.clone())?;
        
        // --- STILL WRITE TO _oall.csv (UNSORTED FOR NOW) ---
        if let Some(ref mut cw) = combined_writer {
            for row in &sorted_records {
                cw.write_record(row)?;
            }
        }

        println!("Finished group: {}", group_name);
    }

    if let Some(ref mut cw) = combined_writer {
        cw.flush()?;
    
    // ---------- SORT _oall.csv ----------
        let oall_path = format!("{}_oall.csv", output_prefix);
        sort_csv_file_in_place(&oall_path)?;
    
        println!("Combined output (oall) written.");
    
    // ---------- BUILD INDEX FROM SORTED OALL ----------
        build_index_from_oall(&oall_path, &output_prefix)?;
    } else {
        println!("Condition detected → Skipping combined oall output.");
    }


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