use csv::Reader;
use plotters::prelude::*;
use plotters::element::Polygon;
use std::{collections::HashMap, path::PathBuf};
use std::fs;
use std::path::Path;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
enum PlotType {
    FitnessEvolution,
    SuccessProbability,
}

#[derive(Debug)]
struct DataPoint {
    evaluations: u32,
    fitness: f64,
    percentage_deviation: f64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct PlotConfig {
    instances: Vec<String>,
    algorithms: Vec<String>,
    group_by_algorithm: bool,
    base_path: String,
    output_path: String,
    targets: Vec<f64>,
    plot_type: PlotType,
    average_targets: bool,
    algorithm_labels: Option<HashMap<String, String>>,
}

impl Default for PlotConfig {
    fn default() -> Self {
        Self {
            instances: vec!["eil51".to_string()],
            algorithms: vec!["ea".to_string(), "ls".to_string(), "rs".to_string()],
            group_by_algorithm: true,
            base_path: "../tsp_hw01/solutions".to_string(),
            output_path: "comparison_eil51.svg".to_string(),
            targets: vec![1.0, 5.0, 10.0],
            plot_type: PlotType::FitnessEvolution,
            average_targets: false,
            algorithm_labels: None,
        }
    }
}

fn load_optimal_cost(instance_filename: &PathBuf) -> Result<f64, Box<dyn std::error::Error>> {
    let instance_name = instance_filename
        .file_stem()
        .and_then(|s| s.to_str())
        .ok_or("Could not extract instance name")?
        .trim_end_matches(".tsp");
    println!("{:?}", instance_name);

    let solutions_path = instance_filename
        .parent().unwrap()
        .parent().unwrap()
        .parent().unwrap()
        .join("instances/solutions.txt");
    println!("{:?}", solutions_path);

    let content = std::fs::read_to_string(solutions_path)?;

    for line in content.lines() {
        let line = line.trim();
        if let Some(colon_pos) = line.find(':') {
            let name = line[..colon_pos].trim();
            if name == instance_name {
                let cost_str = line[colon_pos + 1..].trim();
                return cost_str.parse::<f64>()
                    .map_err(|e| format!("Could not parse cost '{}': {}", cost_str, e).into());
            }
        }
    }

    Err(format!("Optimal cost not found for instance '{}'", instance_name).into())
}

fn calculate_percentage_deviation(fitness: f64, optimal: f64) -> f64 {
    ((fitness - optimal) / optimal) * 100.0
}

fn get_algorithm_label(algorithm: &str, config: &PlotConfig) -> String {
    config
        .algorithm_labels
        .as_ref()
        .and_then(|labels| labels.get(algorithm))
        .map(|label| label.clone())
        .unwrap_or_else(|| algorithm.to_uppercase())
}

fn create_step_function(data: Vec<DataPoint>) -> Vec<DataPoint> {
    if data.is_empty() {
        return data;
    }

    let mut result = Vec::new();

    for i in 0..data.len() {
        let current_point = &data[i];

        // Add the actual data point (the vertical part of the step)
        result.push(DataPoint {
            evaluations: current_point.evaluations,
            fitness: current_point.fitness,
            percentage_deviation: current_point.percentage_deviation,
        });

        // If this is not the last point, add a horizontal step to the next evaluation
        if i + 1 < data.len() {
            let next_evaluation = data[i + 1].evaluations;
            // Add a point just before the next evaluation with current fitness value
            // This creates the horizontal step line
            result.push(DataPoint {
                evaluations: next_evaluation,
                fitness: current_point.fitness,
                percentage_deviation: current_point.percentage_deviation,
            });
        }
    }

    result
}

#[derive(Debug)]
struct ProbabilityPoint {
    evaluations: u32,
    probability: f64,
}

#[derive(Debug)]
struct ProbabilityPointWithDeviation {
    evaluations: u32,
    probability: f64,
    std_dev: f64,
    lower_bound: f64,
    upper_bound: f64,
}

fn calculate_success_probability(
    algorithm_data: &HashMap<String, Vec<DataPoint>>,
    target_percentage: f64,
) -> Vec<ProbabilityPoint> {
    if algorithm_data.is_empty() {
        return Vec::new();
    }

    // Collect all unique evaluation points across all runs
    let mut all_evaluations = std::collections::BTreeSet::new();
    for (_, points) in algorithm_data {
        for point in points {
            all_evaluations.insert(point.evaluations);
        }
    }

    let mut probability_points = Vec::new();

    for &evaluation in &all_evaluations {
        let total_runs = algorithm_data.len();
        let mut successful_runs = 0;

        // For each run, check if it has achieved the target at this evaluation
        for (_, points) in algorithm_data {
            // Find the best performance achieved up to this evaluation
            let best_percentage = points
                .iter()
                .filter(|p| p.evaluations <= evaluation)
                .map(|p| p.percentage_deviation)
                .fold(f64::INFINITY, f64::min);

            if best_percentage <= target_percentage {
                successful_runs += 1;
            }
        }

        let probability = successful_runs as f64 / total_runs as f64;
        probability_points.push(ProbabilityPoint {
            evaluations: evaluation,
            probability,
        });
    }

    probability_points
}

fn calculate_averaged_success_probability_with_deviation(
    algorithm_data: &HashMap<String, Vec<DataPoint>>,
    targets: &[f64],
) -> Vec<ProbabilityPointWithDeviation> {
    if algorithm_data.is_empty() || targets.is_empty() {
        return Vec::new();
    }

    // Calculate probability for each target
    let target_probabilities: Vec<Vec<ProbabilityPoint>> = targets
        .iter()
        .map(|&target| calculate_success_probability(algorithm_data, target))
        .collect();

    if target_probabilities.is_empty() {
        return Vec::new();
    }

    // Collect all unique evaluation points across all targets
    let mut all_evaluations = std::collections::BTreeSet::new();
    for prob_data in &target_probabilities {
        for point in prob_data {
            all_evaluations.insert(point.evaluations);
        }
    }

    let mut averaged_points = Vec::new();

    for &evaluation in &all_evaluations {
        let mut probabilities = Vec::new();

        // Collect probabilities across all targets for this evaluation
        for prob_data in &target_probabilities {
            if let Some(point) = prob_data.iter().find(|p| p.evaluations == evaluation) {
                probabilities.push(point.probability);
            }
        }

        if !probabilities.is_empty() {
            let mean = probabilities.iter().sum::<f64>() / probabilities.len() as f64;
            
            // Calculate standard deviation
            let variance = probabilities.iter()
                .map(|p| (p - mean).powi(2))
                .sum::<f64>() / probabilities.len() as f64;
            let std_dev = variance.sqrt();
            
            // Calculate bounds (mean ± 1 standard deviation, clamped to [0, 1])
            let lower_bound = (mean - std_dev).max(0.0);
            let upper_bound = (mean + std_dev).min(1.0);
            
            averaged_points.push(ProbabilityPointWithDeviation {
                evaluations: evaluation,
                probability: mean,
                std_dev,
                lower_bound,
                upper_bound,
            });
        }
    }

    averaged_points
}


fn read_csv_file(file_path: &Path, optimal_solution: f64) -> Result<Vec<DataPoint>, Box<dyn std::error::Error>> {
    let mut reader = Reader::from_path(file_path)?;
    let mut data = Vec::new();

    for result in reader.records() {
        let record = result?;
        if record.len() >= 2 {
            let evaluations: u32 = record[0].parse()?;
            let fitness: f64 = record[1].parse()?;
            let percentage_deviation = calculate_percentage_deviation(fitness, optimal_solution);
            data.push(DataPoint { evaluations, fitness, percentage_deviation });
        }
    }

    // Add a point at evaluation 1 with the same fitness as the first point
    if !data.is_empty() {
        let first_fitness = data[0].fitness;
        let first_percentage = data[0].percentage_deviation;
        data.insert(0, DataPoint {
            evaluations: 1,
            fitness: first_fitness,
            percentage_deviation: first_percentage
        });
    }

    Ok(data)
}

fn read_all_csv_files(directory: &Path, optimal_solution: f64) -> Result<HashMap<String, Vec<DataPoint>>, Box<dyn std::error::Error>> {
    let mut all_data = HashMap::new();

    if directory.is_dir() {
        for entry in fs::read_dir(directory)? {
            let entry = entry?;
            let path = entry.path();

            if path.extension().and_then(|s| s.to_str()) == Some("csv") {
                let file_name = path.file_stem()
                    .and_then(|s| s.to_str())
                    .unwrap_or("unknown")
                    .to_string();

                match read_csv_file(&path, optimal_solution) {
                    Ok(data) => {
                        all_data.insert(file_name, data);
                    }
                    Err(e) => {
                        eprintln!("Error reading {}: {}", path.display(), e);
                    }
                }
            }
        }
    }

    Ok(all_data)
}

#[derive(Debug)]
struct PlotData {
    data: HashMap<String, HashMap<String, Vec<DataPoint>>>,
}

fn read_plot_data(config: &PlotConfig) -> Result<PlotData, Box<dyn std::error::Error>> {
    let mut data = HashMap::new();
    let base_path = Path::new(&config.base_path);

    for algorithm in &config.algorithms {
        let mut algorithm_data = HashMap::new();

        for instance in &config.instances {
            let path = base_path.join(format!("{}/{}", algorithm, instance));

            if path.exists() {
                let optimal_solution = load_optimal_cost(&path).unwrap();

                let instance_data = read_all_csv_files(&path, optimal_solution)?;

                // Apply step function to create proper step visualization
                let mut step_data = HashMap::new();
                for (run_key, points) in instance_data {
                    let step_points = create_step_function(points);
                    step_data.insert(run_key, step_points);
                }

                algorithm_data.extend(step_data);
            }
        }

        if !algorithm_data.is_empty() {
            data.insert(algorithm.clone(), algorithm_data);
        }
    }

    Ok(PlotData { data })
}

fn get_color_palette() -> Vec<RGBColor> {
    vec![
        BLUE, RED, GREEN, CYAN, MAGENTA,
        RGBColor(255, 165, 0), // Orange
        RGBColor(128, 0, 128),  // Purple
        RGBColor(255, 192, 203), // Pink
        RGBColor(165, 42, 42),   // Brown
        RGBColor(128, 128, 128), // Gray
    ]
}

fn create_plot(plot_data: &PlotData, config: &PlotConfig) -> Result<(), Box<dyn std::error::Error>> {
    // Create plots directory if it doesn't exist
    fs::create_dir_all("plots")?;

    let output_path = Path::new("plots").join(&config.output_path);

    let root = SVGBackend::new(&output_path, (1024, 768)).into_drawing_area();
    root.fill(&WHITE)?;

    let mut min_evaluations = u32::MAX;
    let mut max_evaluations = 0u32;

    let mut min_percentage = f64::MAX;
    let mut max_percentage = f64::MIN;

    // Find min/max across all data
    for (_, algorithm_data) in &plot_data.data {
        for (_, points) in algorithm_data {
            for point in points {
                min_evaluations = min_evaluations.min(point.evaluations);
                max_evaluations = max_evaluations.max(point.evaluations);
                min_percentage = min_percentage.min(point.percentage_deviation);
                max_percentage = max_percentage.max(point.percentage_deviation);
            }
        }
    }

    println!("Percentage deviation range: {:.2}% to {:.2}%", min_percentage, max_percentage);

    // Use actual data range with some padding for better visualization
    let padding_factor = 0.1;
    let range = max_percentage - min_percentage;
    let y_min = (min_percentage - range * padding_factor).max(0.1); // Ensure y_min > 0 for log scale
    let y_max = max_percentage + range * padding_factor;

    let mut chart = ChartBuilder::on(&root)
        .margin(10)
        .x_label_area_size(50)
        .y_label_area_size(90)
        .build_cartesian_2d(
            (min_evaluations as f64..max_evaluations as f64).log_scale(),
            (y_min..y_max).log_scale(),
        )?;

    chart
        .configure_mesh()
        .x_desc("Evaluations")
        .y_desc("Percentage deviation from optimal (%)")
        .y_labels(10)
        .x_label_formatter(&|x| {
            let power = x.log10().round() as i32;
            match power {
                0 => "10⁰".to_string(),
                1 => "10¹".to_string(),
                2 => "10²".to_string(),
                3 => "10³".to_string(),
                4 => "10⁴".to_string(),
                5 => "10⁵".to_string(),
                6 => "10⁶".to_string(),
                7 => "10⁷".to_string(),
                8 => "10⁸".to_string(),
                9 => "10⁹".to_string(),
                _ => format!("10^{}", power),
            }
        })
        .y_label_formatter(&|y| {
            if *y >= 100.0 {
                format!("{:.0}%", *y)
            } else if *y >= 10.0 {
                format!("{:.1}%", *y)
            } else {
                format!("{:.2}%", *y)
            }
        })
        .x_max_light_lines(8)
        .y_max_light_lines(12)
        .axis_desc_style(("sans-serif", 18))
        .label_style(("sans-serif", 16))
        .draw()?;

    let colors = get_color_palette();
    let mut color_index = 0;
    let mut legend_added = HashMap::new();

    if config.group_by_algorithm {
        // Group by algorithm: EA, LS, RS each get one color
        for algorithm in &config.algorithms {
            if let Some(algorithm_data) = plot_data.data.get(algorithm) {
                let color = colors[color_index % colors.len()];
                color_index += 1;

                for (_, points) in algorithm_data {
                    let series = chart
                        .draw_series(LineSeries::new(
                            points.iter().map(|p| (p.evaluations as f64, p.percentage_deviation)),
                            &color,
                        ))?;

                    if !legend_added.contains_key(algorithm) {
                        let label = get_algorithm_label(algorithm, config);
                        series.label(label).legend(move |(x, y)| PathElement::new(vec![(x, y), (x + 10, y)], &color));
                        legend_added.insert(algorithm.clone(), true);
                    }
                }
            }
        }
    } else {
        // Group by instance: Each algorithm-instance combination gets different color
        for algorithm in &config.algorithms {
            if let Some(algorithm_data) = plot_data.data.get(algorithm) {
                for instance in &config.instances {
                    for (run_key, points) in algorithm_data {
                        if run_key.contains(instance) || config.instances.len() == 1 {
                            let color = colors[color_index % colors.len()];
                            color_index += 1;

                            let series = chart
                                .draw_series(LineSeries::new(
                                    points.iter().map(|p| (p.evaluations as f64, p.percentage_deviation)),
                                    &color,
                                ))?;

                            let algo_label = get_algorithm_label(algorithm, config);
                            let label = format!("{} {}", algo_label, instance);
                            if !legend_added.contains_key(&label) {
                                series.label(&label).legend(move |(x, y)| PathElement::new(vec![(x, y), (x + 10, y)], &color));
                                legend_added.insert(label, true);
                            }
                        }
                    }
                }
            }
        }
    }

    // Draw horizontal target lines
    for &target in &config.targets {
        if target >= y_min && target <= y_max {
            let series = chart
                .draw_series(LineSeries::new(
                    vec![(min_evaluations as f64, target), (max_evaluations as f64, target)],
                    BLACK.stroke_width(2),
                ))?;

            series
                .label(&format!("{}% target", target))
                .legend(move |(x, y)| PathElement::new(vec![(x, y), (x + 10, y)], BLACK));
        }
    }

    chart.configure_series_labels()
        .background_style(&WHITE)
        .border_style(&BLACK)
        .position(SeriesLabelPosition::LowerLeft)
        .draw()?;
    root.present()?;

    println!("Plot saved to: {}", output_path.display());
    Ok(())
}

fn create_success_probability_plot(plot_data: &PlotData, config: &PlotConfig) -> Result<(), Box<dyn std::error::Error>> {
    // Create plots directory if it doesn't exist
    fs::create_dir_all("plots")?;

    let output_path = Path::new("plots").join(&config.output_path);

    let root = SVGBackend::new(&output_path, (1024, 768)).into_drawing_area();
    root.fill(&WHITE)?;

    let mut min_evaluations = u32::MAX;
    let mut max_evaluations = 0u32;

    // Find min/max evaluations across all data
    for (_, algorithm_data) in &plot_data.data {
        for (_, points) in algorithm_data {
            for point in points {
                min_evaluations = min_evaluations.min(point.evaluations);
                max_evaluations = max_evaluations.max(point.evaluations);
            }
        }
    }

    let mut chart = ChartBuilder::on(&root)
        .margin(10)
        .x_label_area_size(50)
        .y_label_area_size(90)
        .build_cartesian_2d(
            (min_evaluations as f64..max_evaluations as f64).log_scale(),
            0.0f64..1.0f64,
        )?;

    chart
        .configure_mesh()
        .x_desc("Evaluations")
        .y_desc("Probability of Success")
        .y_labels(11)
        .x_label_formatter(&|x| {
            let power = x.log10().round() as i32;
            match power {
                0 => "10⁰".to_string(),
                1 => "10¹".to_string(),
                2 => "10²".to_string(),
                3 => "10³".to_string(),
                4 => "10⁴".to_string(),
                5 => "10⁵".to_string(),
                6 => "10⁶".to_string(),
                7 => "10⁷".to_string(),
                8 => "10⁸".to_string(),
                9 => "10⁹".to_string(),
                _ => format!("10^{}", power),
            }
        })
        .y_label_formatter(&|y| {
            format!("{:.1}", *y)
        })
        .x_max_light_lines(8)
        .y_max_light_lines(0)
        .axis_desc_style(("sans-serif", 18))
        .label_style(("sans-serif", 16))
        .draw()?;

    let colors = get_color_palette();
    let mut color_index = 0;

    // Plot probability curves
    if config.average_targets {
        // Plot one averaged curve per algorithm with error bars
        for algorithm in &config.algorithms {
            if let Some(algorithm_data) = plot_data.data.get(algorithm) {
                let probability_data = calculate_averaged_success_probability_with_deviation(algorithm_data, &config.targets);

                if !probability_data.is_empty() {
                    let color = colors[color_index % colors.len()];
                    color_index += 1;

                    // Create transparent confidence band
                    let transparent_color = color.mix(0.3); // 30% opacity
                    
                    // Create upper and lower bound points for the filled area
                    let upper_points: Vec<(f64, f64)> = probability_data
                        .iter()
                        .map(|p| (p.evaluations as f64, p.upper_bound))
                        .collect();
                    
                    let mut lower_points: Vec<(f64, f64)> = probability_data
                        .iter()
                        .map(|p| (p.evaluations as f64, p.lower_bound))
                        .collect();
                    
                    // Reverse the lower points to create a closed polygon
                    lower_points.reverse();
                    
                    // Combine upper and lower points to form a polygon
                    let mut polygon_points = upper_points;
                    polygon_points.extend(lower_points);
                    
                    // Draw the filled confidence band
                    if polygon_points.len() > 2 {
                        chart.draw_series(std::iter::once(Polygon::new(
                            polygon_points,
                            transparent_color.filled(),
                        )))?;
                    }

                    // Draw the main line on top of the confidence band
                    let series = chart
                        .draw_series(LineSeries::new(
                            probability_data.iter().map(|p| (p.evaluations as f64, p.probability)),
                            &color,
                        ))?;

                    let algo_label = get_algorithm_label(algorithm, config);
                    let label = format!("{} (avg ± σ)", algo_label);
                    series
                        .label(&label)
                        .legend(move |(x, y)| PathElement::new(vec![(x, y), (x + 10, y)], &color));
                }
            }
        }
    } else {
        // Plot individual curves for each algorithm and target combination
        for algorithm in &config.algorithms {
            if let Some(algorithm_data) = plot_data.data.get(algorithm) {
                for &target in &config.targets {
                    let probability_data = calculate_success_probability(algorithm_data, target);

                    if !probability_data.is_empty() {
                        let color = colors[color_index % colors.len()];
                        color_index += 1;

                        let series = chart
                            .draw_series(LineSeries::new(
                                probability_data.iter().map(|p| (p.evaluations as f64, p.probability)),
                                &color,
                            ))?;

                        let algo_label = get_algorithm_label(algorithm, config);
                        let label = format!("{} - {}%", algo_label, target);
                        series
                            .label(&label)
                            .legend(move |(x, y)| PathElement::new(vec![(x, y), (x + 10, y)], &color));
                    }
                }
            }
        }
    }

    chart.configure_series_labels().background_style(&WHITE).border_style(&BLACK).draw()?;
    root.present()?;

    println!("Probability plot saved to: {}", output_path.display());
    Ok(())
}

fn load_config(config_path: &str) -> Result<PlotConfig, Box<dyn std::error::Error>> {

    if Path::new(config_path).exists() {
        let config_str = fs::read_to_string(config_path)?;
        let config: PlotConfig = serde_json::from_str(&config_str)?;
        println!("Loaded configuration from {}", config_path);
        Ok(config)
    } else {
        let config = PlotConfig::default();

        // Create default config file
        let config_str = serde_json::to_string_pretty(&config)?;
        fs::write(config_path, config_str)?;
        println!("Created default configuration file: {}", config_path);

        Ok(config)
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = std::env::args().collect();

    if args.len() != 2 {
        eprintln!("Usage: {} <config_file.json>", args[0]);
        eprintln!("Example: {} plot_config.json", args[0]);
        std::process::exit(1);
    }

    let config_path = &args[1];
    let config = load_config(config_path)?;

    println!("Configuration:");
    println!("  Instances: {:?}", config.instances);
    println!("  Algorithms: {:?}", config.algorithms);
    println!("  Group by algorithm: {}", config.group_by_algorithm);
    println!("  Base path: {}", config.base_path);
    println!("  Output path: {}", config.output_path);
    println!("  Plot type: {:?}", config.plot_type);
    println!("  Average targets: {}", config.average_targets);
    println!();

    let base_directory = Path::new(&config.base_path);
    println!("Reading CSV files from: {}", base_directory.display());

    let plot_data = read_plot_data(&config)?;

    if plot_data.data.is_empty() {
        eprintln!("No CSV files found for any configured algorithms and instances");
        return Ok(());
    }

    for (algorithm, algorithm_data) in &plot_data.data {
        println!("Found {} CSV files for algorithm: {}", algorithm_data.len(), algorithm);
    }

    match config.plot_type {
        PlotType::FitnessEvolution => {
            create_plot(&plot_data, &config)?;
        }
        PlotType::SuccessProbability => {
            create_success_probability_plot(&plot_data, &config)?;
        }
    }

    Ok(())
}