use anyhow::Result;
use clap::Parser;
use ml::data_loaders::parquet_utils::load_parquet_data;

#[derive(Parser, Debug)]
#[command(author, version, about = "Validate 54-feature dataset quality", long_about = None)]
struct Args {
    #[arg(long, help = "Path to parquet file")]
    parquet_file: String,

    #[arg(long, default_value = "1000", help = "Number of samples to analyze")]
    num_samples: usize,

    #[arg(long, default_value = "0.95", help = "Correlation threshold for redundancy detection")]
    correlation_threshold: f64,

    #[arg(long, default_value = "0.01", help = "Variance threshold for low-signal detection")]
    variance_threshold: f64,
}

fn main() -> Result<()> {
    let args = Args::parse();
    println!("54-Feature Dataset Quality Validation");
    println!("======================================\n");
    println!("Loading data from: {}", args.parquet_file);
    println!("Analyzing {} samples", args.num_samples);
    println!("Correlation threshold: {}", args.correlation_threshold);
    println!("Variance threshold: {}\n", args.variance_threshold);

    // Load data using production parquet_utils (returns Vec<Vec<f32>>)
    let warmup_bars = 50; // Standard warmup for technical indicators
    let feature_vectors = load_parquet_data(std::path::Path::new(&args.parquet_file), warmup_bars)?;

    println!("Dataset loaded:");
    println!("  Total feature vectors: {}", feature_vectors.len());
    if !feature_vectors.is_empty() {
        println!("  Feature dimensions: {}\n", feature_vectors[0].len());
    }

    let num_samples = args.num_samples.min(feature_vectors.len());
    println!("Analyzing {} samples...", num_samples);

    // Convert to column-oriented matrix (transpose)
    let num_features = if feature_vectors.is_empty() { 0 } else { feature_vectors[0].len() };
    let mut feature_matrix: Vec<Vec<f64>> = vec![Vec::new(); num_features];

    for i in 0..num_samples {
        let feature_vec = &feature_vectors[i];

        // Sanity check
        if feature_vec.len() != num_features {
            eprintln!("WARNING: Sample {} has {} features, expected {}", i, feature_vec.len(), num_features);
            continue;
        }

        // Append to feature matrix (convert f32 to f64)
        for (j, val) in feature_vec.iter().enumerate() {
            feature_matrix[j].push(*val as f64);
        }

        if (i + 1) % 100 == 0 {
            print!(".");
            std::io::Write::flush(&mut std::io::stdout())?;
        }
    }
    println!("\n");

    // Analyze each feature
    println!("Analyzing individual features...\n");
    println!("{:<8} {:<20} {:<15} {:<15} {:<15}", "Index", "Type", "Mean", "Std Dev", "Status");
    println!("{:-<75}", "");

    let mut constant_features = Vec::new();
    let mut low_variance_features = Vec::new();
    let mut healthy_features = Vec::new();
    let mut ofi_zeros = Vec::new();

    // Determine feature boundaries based on actual feature count
    // Expected: 128 features total (54 populated + 74 padding)
    // - Base features (0-45): 46
    // - OFI placeholders (46-53): 8 (zeros until MBP-10 integrated)
    // - Padding (54-127): 74 (zeros for future expansion)
    let base_end = 46;
    let ofi_end = 54;
    let padding_end = 128;

    for (idx, values) in feature_matrix.iter().enumerate() {
        let (feature_type, should_be_zero) = if idx < base_end {
            ("Base", false)
        } else if idx < ofi_end {
            ("OFI Placeholder", true)
        } else {
            ("Padding", true)
        };

        let mean = values.iter().sum::<f64>() / values.len() as f64;
        let variance = values.iter()
            .map(|x| (x - mean).powi(2))
            .sum::<f64>() / values.len() as f64;
        let std_dev = variance.sqrt();

        let status = if std_dev < 1e-10 {
            if should_be_zero {
                ofi_zeros.push(idx);
                "EXPECTED ZERO"
            } else {
                constant_features.push(idx);
                "CONSTANT (BAD)"
            }
        } else if variance < args.variance_threshold && !should_be_zero {
            low_variance_features.push(idx);
            "LOW VARIANCE"
        } else if should_be_zero {
            healthy_features.push(idx);
            "UNEXPECTED NON-ZERO"
        } else {
            healthy_features.push(idx);
            "HEALTHY"
        };

        println!("{:<8} {:<20} {:<15.6} {:<15.6} {:<15}",
                 idx, feature_type, mean, std_dev, status);
    }

    // Correlation analysis (only for base features 0-45)
    println!("\n\nCorrelation Analysis (Base Features 0-{})...", base_end - 1);
    println!("==========================================\n");

    let mut high_correlations = Vec::new();

    for i in 0..base_end {
        for j in (i + 1)..base_end {
            let corr = calculate_correlation(&feature_matrix[i], &feature_matrix[j]);
            if corr.abs() > args.correlation_threshold {
                high_correlations.push((i, j, corr));
            }
        }
        if (i + 1) % 10 == 0 {
            print!(".");
            std::io::Write::flush(&mut std::io::stdout())?;
        }
    }
    println!();

    if high_correlations.is_empty() {
        println!("No high correlations found (threshold: {})", args.correlation_threshold);
    } else {
        println!("High correlation pairs (>{}):", args.correlation_threshold);
        println!("{:<15} {:<15} {:<15}", "Feature A", "Feature B", "Correlation");
        println!("{:-<45}", "");
        for (i, j, corr) in &high_correlations {
            println!("{:<15} {:<15} {:<15.4}", i, j, corr);
        }
    }

    // Generate summary report
    println!("\n\n54-Feature Dataset Quality Report (Padded to 128)");
    println!("==================================================\n");
    println!("Total Features: {} (54 populated + 74 padding)", num_features);
    println!("  - Base Features (0-45): 46");
    println!("  - OFI Placeholders (46-53): 8");
    println!("  - Padding (54-127): 74\n");

    // Count constant features in base region only
    let base_constant = constant_features.iter().filter(|&&x| x < base_end).count();
    let base_low_var = low_variance_features.iter().filter(|&&x| x < base_end).count();
    let base_healthy = healthy_features.iter().filter(|&&x| x < base_end).count();

    println!("Base Feature Analysis (0-45):");
    println!("  - Constant features: {} (should be 0)", base_constant);
    println!("  - Low variance (<{}): {} (should be 0)", args.variance_threshold, base_low_var);
    println!("  - High correlation pairs (>{}): {} (should be 0)", args.correlation_threshold, high_correlations.len());
    println!("  - Healthy features: {}/46\n", base_healthy);

    let ofi_zeros_count = ofi_zeros.iter().filter(|&&x| x >= base_end && x < ofi_end).count();
    println!("OFI Placeholder Analysis (46-53):");
    println!("  - All zeros: {} (should be YES until MBP-10 integrated)\n",
             if ofi_zeros_count == 8 { "YES" } else { "NO" });

    let padding_zeros_count = ofi_zeros.iter().filter(|&&x| x >= ofi_end).count();
    println!("Padding Analysis (54-127):");
    println!("  - All zeros: {} (should be YES)\n",
             if padding_zeros_count == 74 { "YES" } else { "NO" });

    // Feature-specific details (only for base features 0-45)
    let base_constant_list: Vec<usize> = constant_features.iter().filter(|&&x| x < base_end).copied().collect();
    let base_low_var_list: Vec<usize> = low_variance_features.iter().filter(|&&x| x < base_end).copied().collect();

    if !base_constant_list.is_empty() {
        println!("Constant Base Features (0-45): {:?}", base_constant_list);
    }
    if !base_low_var_list.is_empty() {
        println!("Low Variance Base Features (0-45): {:?}", base_low_var_list);
    }

    // Detailed correlation matrix for suspicious pairs
    if !high_correlations.is_empty() {
        println!("\nDetailed Correlation Analysis:");
        println!("{:-<80}", "");
        for (i, j, corr) in &high_correlations {
            let var_i = calculate_variance(&feature_matrix[*i]);
            let var_j = calculate_variance(&feature_matrix[*j]);
            println!("Feature {} vs Feature {}: correlation={:.4}, variance_{}={:.4}, variance_{}={:.4}",
                     i, j, corr, i, var_i, j, var_j);
        }
    }

    // Verdict
    println!("\n{:-<80}", "");
    let verdict = if base_constant == 0
        && base_low_var == 0
        && high_correlations.is_empty()
        && ofi_zeros_count == 8
        && padding_zeros_count == 74 {
        "CLEAN - Dataset is high quality with no redundancy"
    } else if base_constant > 5 || high_correlations.len() > 10 {
        "NOISY - Significant issues detected requiring cleanup"
    } else {
        "ACCEPTABLE - Some minor issues but usable for training"
    };

    println!("VERDICT: {}", verdict);
    println!("{:-<80}", "");

    // Top 10 most variable features (signal strength) - base features only
    println!("\nTop 10 Base Features by Variance (Signal Strength):");
    println!("{:-<50}", "");
    let mut variance_pairs: Vec<(usize, f64)> = (0..base_end)
        .map(|i| (i, calculate_variance(&feature_matrix[i])))
        .collect();
    variance_pairs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

    println!("{:<8} {:<15} {:<15}", "Index", "Variance", "Feature Type");
    println!("{:-<50}", "");
    for (idx, var) in variance_pairs.iter().take(10) {
        println!("{:<8} {:<15.6} {:<15}", idx, var, "Base");
    }

    Ok(())
}

fn calculate_correlation(x: &[f64], y: &[f64]) -> f64 {
    let n = x.len() as f64;
    let mean_x = x.iter().sum::<f64>() / n;
    let mean_y = y.iter().sum::<f64>() / n;

    let cov: f64 = x.iter()
        .zip(y.iter())
        .map(|(xi, yi)| (xi - mean_x) * (yi - mean_y))
        .sum::<f64>() / n;

    let var_x = x.iter().map(|xi| (xi - mean_x).powi(2)).sum::<f64>() / n;
    let var_y = y.iter().map(|yi| (yi - mean_y).powi(2)).sum::<f64>() / n;

    if var_x < 1e-10 || var_y < 1e-10 {
        return 0.0;
    }

    cov / (var_x.sqrt() * var_y.sqrt())
}

fn calculate_variance(values: &[f64]) -> f64 {
    let n = values.len() as f64;
    let mean = values.iter().sum::<f64>() / n;
    values.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / n
}