//! PON Builder: Rust-accelerated aggregation for Panel of Normals construction.
//!
//! This module provides high-performance sample aggregation functions for building
//! PON (Panel of Normals) models. These functions are 3-10x faster than Python
//! equivalents due to optimized memory access and parallel processing.
//!
//! # Functions
//! - `compute_gc_bias_model`: Aggregate GC bias curves from sample data
//! - `compute_fsd_baseline`: Aggregate FSD per-arm baselines
//! - `compute_wps_baseline`: Aggregate WPS per-region baselines

use pyo3::prelude::*;
use pyo3::types::PyDict;
use std::collections::HashMap;
use log::info;

/// GC bins for bias model: [0.25, 0.27, ..., 0.73]
const GC_BIN_START: f64 = 0.25;
const GC_BIN_END: f64 = 0.75;
const GC_BIN_STEP: f64 = 0.02;

/// Result structure for GC bias model
#[derive(Debug)]
pub struct GcBiasResult {
    pub gc_bins: Vec<f64>,
    pub short_expected: Vec<f64>,
    pub short_std: Vec<f64>,
    pub intermediate_expected: Vec<f64>,
    pub intermediate_std: Vec<f64>,
    pub long_expected: Vec<f64>,
    pub long_std: Vec<f64>,
}

/// Compute median of a slice
fn median(data: &mut Vec<f64>) -> f64 {
    if data.is_empty() {
        return 1.0;
    }
    data.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    let mid = data.len() / 2;
    if data.len().is_multiple_of(2) {
        (data[mid - 1] + data[mid]) / 2.0
    } else {
        data[mid]
    }
}

/// Compute standard deviation of a slice
fn std_dev(data: &[f64]) -> f64 {
    if data.len() < 2 {
        return 0.1;
    }
    let mean = data.iter().sum::<f64>() / data.len() as f64;
    let variance = data.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / (data.len() - 1) as f64;
    variance.sqrt()
}

/// Compute GC bias model from sample data.
/// 
/// Aggregates GC-binned fragment counts across samples and computes
/// median expected coverage and std per GC bin.
/// 
/// # Arguments
/// * `all_gc_data` - List of dicts, each with keys: gc, short, intermediate, long (as arrays)
/// 
/// # Returns
/// * Dict with gc_bins, short_expected, short_std, intermediate_expected, intermediate_std, long_expected, long_std
/// 
/// # Performance
/// 3-10x faster than Python numpy/pandas aggregation
#[pyfunction]
pub fn compute_gc_bias_model(py: Python<'_>, all_gc_data: Vec<pyo3::Bound<'_, PyDict>>) -> PyResult<PyObject> {
    info!("PON Builder: Computing GC bias model from {} samples", all_gc_data.len());
    
    // Generate GC bins
    let num_bins = ((GC_BIN_END - GC_BIN_START) / GC_BIN_STEP) as usize;
    let gc_bins: Vec<f64> = (0..num_bins)
        .map(|i| GC_BIN_START + i as f64 * GC_BIN_STEP)
        .collect();
    
    // Aggregation structures
    let mut short_by_gc: Vec<Vec<f64>> = vec![Vec::new(); num_bins];
    let mut intermediate_by_gc: Vec<Vec<f64>> = vec![Vec::new(); num_bins];
    let mut long_by_gc: Vec<Vec<f64>> = vec![Vec::new(); num_bins];
    
    // Process each sample
    for sample in all_gc_data.iter() {
        // Extract arrays
        let gc_arr: Vec<f64> = sample.get_item("gc")
            .ok().and_then(|v| v.and_then(|item| item.extract().ok()))
            .unwrap_or_default();
        let short_arr: Vec<f64> = sample.get_item("short")
            .ok().and_then(|v| v.and_then(|item| item.extract().ok()))
            .unwrap_or_default();
        let intermediate_arr: Vec<f64> = sample.get_item("intermediate")
            .ok().and_then(|v| v.and_then(|item| item.extract().ok()))
            .unwrap_or_default();
        let long_arr: Vec<f64> = sample.get_item("long")
            .ok().and_then(|v| v.and_then(|item| item.extract().ok()))
            .unwrap_or_default();
        
        if gc_arr.is_empty() {
            continue;
        }
        
        // Compute means for normalization
        let short_mean: f64 = short_arr.iter().filter(|x| !x.is_nan()).sum::<f64>() 
            / short_arr.iter().filter(|x| !x.is_nan()).count().max(1) as f64;
        let intermediate_mean: f64 = intermediate_arr.iter().filter(|x| !x.is_nan()).sum::<f64>() 
            / intermediate_arr.iter().filter(|x| !x.is_nan()).count().max(1) as f64;
        let long_mean: f64 = long_arr.iter().filter(|x| !x.is_nan()).sum::<f64>() 
            / long_arr.iter().filter(|x| !x.is_nan()).count().max(1) as f64;
        
        // Bin by GC
        for (i, &gc_val) in gc_arr.iter().enumerate() {
            if gc_val.is_nan() || gc_val <= 0.0 || !(GC_BIN_START..GC_BIN_END).contains(&gc_val) {
                continue;
            }
            
            let bin_idx = ((gc_val - GC_BIN_START) / GC_BIN_STEP) as usize;
            if bin_idx >= num_bins {
                continue;
            }
            
            // Normalized values
            if i < short_arr.len() && !short_arr[i].is_nan() && short_mean > 0.0 {
                short_by_gc[bin_idx].push(short_arr[i] / short_mean);
            }
            if i < intermediate_arr.len() && !intermediate_arr[i].is_nan() && intermediate_mean > 0.0 {
                intermediate_by_gc[bin_idx].push(intermediate_arr[i] / intermediate_mean);
            }
            if i < long_arr.len() && !long_arr[i].is_nan() && long_mean > 0.0 {
                long_by_gc[bin_idx].push(long_arr[i] / long_mean);
            }
        }
    }
    
    // Compute median and std per bin
    let mut short_expected: Vec<f64> = Vec::with_capacity(num_bins);
    let mut short_std: Vec<f64> = Vec::with_capacity(num_bins);
    let mut intermediate_expected: Vec<f64> = Vec::with_capacity(num_bins);
    let mut intermediate_std: Vec<f64> = Vec::with_capacity(num_bins);
    let mut long_expected: Vec<f64> = Vec::with_capacity(num_bins);
    let mut long_std: Vec<f64> = Vec::with_capacity(num_bins);
    
    for i in 0..num_bins {
        short_expected.push(median(&mut short_by_gc[i].clone()));
        short_std.push(std_dev(&short_by_gc[i]));
        intermediate_expected.push(median(&mut intermediate_by_gc[i].clone()));
        intermediate_std.push(std_dev(&intermediate_by_gc[i]));
        long_expected.push(median(&mut long_by_gc[i].clone()));
        long_std.push(std_dev(&long_by_gc[i]));
    }
    
    // Log summary
    let total_samples: usize = short_by_gc.iter().map(|v| v.len()).sum();
    info!("PON Builder: GC model computed: {} bins, {} data points", num_bins, total_samples);
    
    // Build result dict
    let result = PyDict::new(py);
    result.set_item("gc_bins", gc_bins)?;
    result.set_item("short_expected", short_expected)?;
    result.set_item("short_std", short_std)?;
    result.set_item("intermediate_expected", intermediate_expected)?;
    result.set_item("intermediate_std", intermediate_std)?;
    result.set_item("long_expected", long_expected)?;
    result.set_item("long_std", long_std)?;
    
    Ok(result.into())
}

/// Compute FSD baseline from sample TSV files.
/// 
/// Reads FSD TSV files, aggregates counts per arm and size bin,
/// computes mean and std across samples.
/// 
/// # Arguments
/// * `fsd_paths` - List of paths to FSD.tsv files
/// 
/// # Returns
/// * Dict mapping arm -> {size_bins, expected, std}
#[pyfunction]
pub fn compute_fsd_baseline(py: Python<'_>, fsd_paths: Vec<String>) -> PyResult<PyObject> {
    use std::fs::File;
    use std::io::{BufRead, BufReader};
    
    info!("PON Builder: Computing FSD baseline from {} samples", fsd_paths.len());
    
    // Aggregation: arm -> size_bin -> list of values
    let mut arm_data: HashMap<String, HashMap<i32, Vec<f64>>> = HashMap::new();
    
    for path in &fsd_paths {
        let file = match File::open(path) {
            Ok(f) => f,
            Err(e) => {
                info!("PON Builder: Skipping {}: {}", path, e);
                continue;
            }
        };
        let reader = BufReader::new(file);
        let mut lines = reader.lines();
        
        // Parse header
        let header_line = match lines.next() {
            Some(Ok(line)) => line,
            _ => continue,
        };
        let headers: Vec<&str> = header_line.split('\t').collect();
        
        // Find bin columns (format: "65-69", "70-74", etc.)
        let bin_indices: Vec<(usize, i32)> = headers.iter().enumerate()
            .filter_map(|(i, h)| {
                if let Some(dash_pos) = h.find('-') {
                    if let Ok(start) = h[..dash_pos].parse::<i32>() {
                        return Some((i, start));
                    }
                }
                None
            })
            .collect();
        
        // Process data rows
        for line_result in lines {
            let line = match line_result {
                Ok(l) => l,
                Err(_) => continue,
            };
            
            let fields: Vec<&str> = line.split('\t').collect();
            if fields.is_empty() {
                continue;
            }
            
            let arm = fields[0].to_string();
            
            for (col_idx, size) in &bin_indices {
                let value: f64 = fields.get(*col_idx)
                    .and_then(|v| v.parse().ok())
                    .unwrap_or(0.0);
                
                arm_data
                    .entry(arm.clone())
                    .or_default()
                    .entry(*size)
                    .or_default()
                    .push(value);
            }
        }
    }
    
    // Compute mean and std per arm/bin
    let result = PyDict::new(py);
    
    for (arm, size_data) in &arm_data {
        let arm_dict = PyDict::new(py);
        let mut size_bins: Vec<i32> = size_data.keys().copied().collect();
        size_bins.sort();
        
        let mut expected: Vec<f64> = Vec::new();
        let mut std_vals: Vec<f64> = Vec::new();
        
        for size in &size_bins {
            let values = size_data.get(size).unwrap();
            let mean = values.iter().sum::<f64>() / values.len().max(1) as f64;
            expected.push(mean);
            std_vals.push(std_dev(values));
        }
        
        arm_dict.set_item("size_bins", size_bins)?;
        arm_dict.set_item("expected", expected)?;
        arm_dict.set_item("std", std_vals)?;
        
        result.set_item(arm, arm_dict)?;
    }
    
    info!("PON Builder: FSD baseline computed: {} arms", arm_data.len());
    
    Ok(result.into())
}

/// Extract 200-element WPS vector from a ListArray column.
/// Returns None if the column is not a List type or extraction fails.
fn extract_wps_vector(col: &dyn arrow::array::Array, row: usize) -> Option<Vec<f32>> {
    use arrow::array::{Float32Array, ListArray, Array};
    
    if let Some(list_arr) = col.as_any().downcast_ref::<ListArray>() {
        if row < list_arr.len() && !list_arr.is_null(row) {
            let inner = list_arr.value(row);
            if let Some(f32_arr) = inner.as_any().downcast_ref::<Float32Array>() {
                // Extract all values from the list
                let vec: Vec<f32> = (0..f32_arr.len())
                    .map(|i| if f32_arr.is_null(i) { 0.0 } else { f32_arr.value(i) })
                    .collect();
                return Some(vec);
            }
        }
    }
    None
}

/// Compute element-wise mean of multiple vectors.
/// All vectors must have the same length.
fn element_wise_mean(vectors: &[Vec<f32>]) -> Vec<f32> {
    if vectors.is_empty() {
        return Vec::new();
    }
    let len = vectors[0].len();
    let n = vectors.len() as f32;
    
    (0..len)
        .map(|i| {
            let sum: f32 = vectors.iter().map(|v| v.get(i).copied().unwrap_or(0.0)).sum();
            sum / n
        })
        .collect()
}

/// Compute element-wise standard deviation of multiple vectors.
fn element_wise_std(vectors: &[Vec<f32>], means: &[f32]) -> Vec<f32> {
    if vectors.len() < 2 || means.is_empty() {
        return means.iter().map(|_| 0.1_f32).collect();
    }
    let n = vectors.len() as f32;
    
    means.iter().enumerate()
        .map(|(i, &mean)| {
            let variance: f32 = vectors.iter()
                .map(|v| {
                    let val = v.get(i).copied().unwrap_or(0.0);
                    (val - mean).powi(2)
                })
                .sum::<f32>() / (n - 1.0);
            variance.sqrt().max(0.01) // Minimum std to avoid divide by zero
        })
        .collect()
}

/// Compute WPS baseline from sample parquet files (Vector v2.0 format).
/// 
/// Reads WPS parquet files, aggregates wps_nuc and wps_tf 200-bin vectors per region,
/// computes element-wise mean and std vectors across all samples.
/// 
/// # Arguments
/// * `wps_paths` - List of paths to WPS.parquet files
/// 
/// # Returns
/// * Dict mapping region_id -> {wps_nuc_mean: [200], wps_nuc_std: [200], wps_tf_mean: [200], wps_tf_std: [200]}
/// 
/// # Schema
/// v2.0 vector format enables position-specific z-score computation and shape correlation analysis.
/// 
/// # Performance
/// 3-10x faster than Python pandas concat + groupby
#[pyfunction]
pub fn compute_wps_baseline(py: Python<'_>, wps_paths: Vec<String>) -> PyResult<PyObject> {
    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
    use arrow::array::{StringArray, Array};
    use std::fs::File;
    
    info!("PON Builder: Computing WPS vector baseline from {} samples", wps_paths.len());
    
    // Aggregation: region_id -> list of 200-element vectors
    let mut region_nuc_vectors: HashMap<String, Vec<Vec<f32>>> = HashMap::new();
    let mut region_tf_vectors: HashMap<String, Vec<Vec<f32>>> = HashMap::new();
    
    let mut files_processed = 0;
    let mut regions_found = 0;
    
    for path in &wps_paths {
        let file = match File::open(path) {
            Ok(f) => f,
            Err(e) => {
                info!("PON Builder: Skipping {}: {}", path, e);
                continue;
            }
        };
        
        let reader = match ParquetRecordBatchReaderBuilder::try_new(file) {
            Ok(builder) => match builder.build() {
                Ok(r) => r,
                Err(e) => {
                    info!("PON Builder: Failed to read parquet {}: {}", path, e);
                    continue;
                }
            },
            Err(e) => {
                info!("PON Builder: Failed to open parquet {}: {}", path, e);
                continue;
            }
        };
        
        files_processed += 1;
        
        for batch_result in reader {
            let batch = match batch_result {
                Ok(b) => b,
                Err(_) => continue,
            };
            
            let schema = batch.schema();
            
            // Find required columns
            let region_idx = schema.index_of("region_id")
                .or_else(|_| schema.index_of("group_id"))
                .ok();
            let nuc_idx = schema.index_of("wps_nuc").ok();
            let tf_idx = schema.index_of("wps_tf").ok();
            
            if region_idx.is_none() {
                continue;
            }
            
            let region_col = batch.column(region_idx.unwrap())
                .as_any().downcast_ref::<StringArray>();
            
            if region_col.is_none() {
                continue;
            }
            
            let region_array = region_col.unwrap();
            
            for i in 0..region_array.len() {
                let region_id = region_array.value(i).to_string();
                regions_found += 1;
                
                // Extract wps_nuc vector (expecting List<Float32>[200])
                if let Some(idx) = nuc_idx {
                    if let Some(vec) = extract_wps_vector(batch.column(idx).as_ref(), i) {
                        region_nuc_vectors.entry(region_id.clone())
                            .or_default()
                            .push(vec);
                    }
                }
                
                // Extract wps_tf vector
                if let Some(idx) = tf_idx {
                    if let Some(vec) = extract_wps_vector(batch.column(idx).as_ref(), i) {
                        region_tf_vectors.entry(region_id)
                            .or_default()
                            .push(vec);
                    }
                }
            }
        }
    }
    
    info!("PON Builder: Processed {} files, found {} region entries", files_processed, regions_found);
    info!("PON Builder: Unique regions: {} nuc, {} tf", region_nuc_vectors.len(), region_tf_vectors.len());
    
    // Compute element-wise mean and std per region
    let result = PyDict::new(py);
    
    for (region_id, nuc_vectors) in &region_nuc_vectors {
        let region_dict = PyDict::new(py);
        
        // Compute WPS-Nuc mean/std vectors
        let nuc_mean = element_wise_mean(nuc_vectors);
        let nuc_std = element_wise_std(nuc_vectors, &nuc_mean);
        region_dict.set_item("wps_nuc_mean", nuc_mean.iter().map(|&x| x as f64).collect::<Vec<f64>>())?;
        region_dict.set_item("wps_nuc_std", nuc_std.iter().map(|&x| x as f64).collect::<Vec<f64>>())?;
        
        // Compute WPS-TF mean/std vectors
        if let Some(tf_vectors) = region_tf_vectors.get(region_id) {
            let tf_mean = element_wise_mean(tf_vectors);
            let tf_std = element_wise_std(tf_vectors, &tf_mean);
            region_dict.set_item("wps_tf_mean", tf_mean.iter().map(|&x| x as f64).collect::<Vec<f64>>())?;
            region_dict.set_item("wps_tf_std", tf_std.iter().map(|&x| x as f64).collect::<Vec<f64>>())?;
        } else {
            // Empty vectors if no TF data
            region_dict.set_item("wps_tf_mean", Vec::<f64>::new())?;
            region_dict.set_item("wps_tf_std", Vec::<f64>::new())?;
        }
        
        // Add sample count for diagnostics
        region_dict.set_item("n_samples", nuc_vectors.len())?;
        
        result.set_item(region_id, region_dict)?;
    }
    
    info!("PON Builder: WPS vector baseline computed: {} regions", region_nuc_vectors.len());
    
    Ok(result.into())
}

/// Compute Region MDS baseline from sample TSV files.
/// 
/// Reads MDS.gene.tsv files, aggregates per-gene MDS values across samples,
/// computes mean and std for each gene.
/// 
/// # Arguments
/// * `mds_paths` - List of paths to MDS.gene.tsv files
/// 
/// # Returns
/// * Dict mapping gene -> {mds_mean, mds_std, mds_e1_mean, mds_e1_std, n_samples}
/// 
/// # Performance
/// Optimized for aggregating MDS across hundreds of samples
#[pyfunction]
pub fn compute_region_mds_baseline(py: Python<'_>, mds_paths: Vec<String>) -> PyResult<PyObject> {
    use std::fs::File;
    use std::io::{BufRead, BufReader};
    
    info!("PON Builder: Computing Region MDS baseline from {} samples", mds_paths.len());
    
    // Aggregation: gene -> (mds_values, mds_e1_values)
    let mut gene_mds: HashMap<String, Vec<f64>> = HashMap::new();
    let mut gene_mds_e1: HashMap<String, Vec<f64>> = HashMap::new();
    
    for path in &mds_paths {
        let file = match File::open(path) {
            Ok(f) => f,
            Err(e) => {
                info!("PON Builder: Skipping {}: {}", path, e);
                continue;
            }
        };
        let reader = BufReader::new(file);
        let mut lines = reader.lines();
        
        // Parse header to find column indices
        let header_line = match lines.next() {
            Some(Ok(line)) => line,
            _ => continue,
        };
        let headers: Vec<&str> = header_line.split('\t').collect();
        
        // Find column indices
        let gene_idx = headers.iter().position(|&h| h == "gene").unwrap_or(0);
        let mds_mean_idx = headers.iter().position(|&h| h == "mds_mean");
        let mds_e1_idx = headers.iter().position(|&h| h == "mds_e1");
        
        // Process data rows
        for line_result in lines {
            let line = match line_result {
                Ok(l) => l,
                Err(_) => continue,
            };
            
            let fields: Vec<&str> = line.split('\t').collect();
            if fields.is_empty() {
                continue;
            }
            
            let gene = fields.get(gene_idx).unwrap_or(&"").to_string();
            if gene.is_empty() {
                continue;
            }
            
            // mds_mean
            if let Some(idx) = mds_mean_idx {
                if let Some(val_str) = fields.get(idx) {
                    if let Ok(val) = val_str.parse::<f64>() {
                        if !val.is_nan() && val > 0.0 {
                            gene_mds.entry(gene.clone()).or_default().push(val);
                        }
                    }
                }
            }
            
            // mds_e1
            if let Some(idx) = mds_e1_idx {
                if let Some(val_str) = fields.get(idx) {
                    if let Ok(val) = val_str.parse::<f64>() {
                        if !val.is_nan() && val > 0.0 {
                            gene_mds_e1.entry(gene).or_default().push(val);
                        }
                    }
                }
            }
        }
    }
    
    // Compute mean and std per gene
    let result = PyDict::new(py);
    
    for (gene, mds_values) in &gene_mds {
        let gene_dict = PyDict::new(py);
        
        // MDS mean stats
        let mds_mean = mds_values.iter().sum::<f64>() / mds_values.len().max(1) as f64;
        let mds_std = std_dev(mds_values);
        gene_dict.set_item("mds_mean", mds_mean)?;
        gene_dict.set_item("mds_std", mds_std)?;
        
        // MDS E1 stats
        if let Some(e1_values) = gene_mds_e1.get(gene) {
            let e1_mean = e1_values.iter().sum::<f64>() / e1_values.len().max(1) as f64;
            let e1_std = std_dev(e1_values);
            gene_dict.set_item("mds_e1_mean", e1_mean)?;
            gene_dict.set_item("mds_e1_std", e1_std)?;
        } else {
            gene_dict.set_item("mds_e1_mean", mds_mean)?;
            gene_dict.set_item("mds_e1_std", mds_std)?;
        }
        
        gene_dict.set_item("n_samples", mds_values.len())?;
        
        result.set_item(gene, gene_dict)?;
    }
    
    info!("PON Builder: Region MDS baseline computed: {} genes", gene_mds.len());
    
    Ok(result.into())
}
