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, Clone)]
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>>,
    #[serde(default = "default_show_std_dev")]
    show_std_dev: bool,
    #[serde(default = "default_average_runs")]
    average_runs: bool,
}

fn default_show_std_dev() -> bool {
    true
}

fn default_average_runs() -> bool {
    false
}

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,
            show_std_dev: true,
            average_runs: false,
        }
    }
}

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, Clone)]
struct ProbabilityPoint {
    evaluations: u32,
    probability: f64,
}

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

#[derive(Debug, Clone)]
struct DataPointWithDeviation {
    evaluations: u32,
    fitness: f64,
    percentage_deviation: 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();
    }

    // Find the maximum evaluation point across all data
    let max_evaluation = algorithm_data
        .values()
        .flat_map(|points| points.iter())
        .map(|p| p.evaluations)
        .max()
        .unwrap_or(0);

    // Extend all data to the maximum evaluation point
    let mut extended_algorithm_data = HashMap::new();
    for (run_key, points) in algorithm_data {
        let mut extended_points = points.clone();
        if !extended_points.is_empty() {
            extended_points.sort_by_key(|point| point.evaluations);
            let last_point = extended_points.last().unwrap();
            if last_point.evaluations < max_evaluation {
                extended_points.push(DataPoint {
                    evaluations: max_evaluation,
                    fitness: last_point.fitness,
                    percentage_deviation: last_point.percentage_deviation,
                });
            }
        }
        extended_algorithm_data.insert(run_key.clone(), extended_points);
    }

    // Collect all unique evaluation points across all runs
    let mut all_evaluations = std::collections::BTreeSet::new();
    for (_, points) in &extended_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 = extended_algorithm_data.len();
        let mut successful_runs = 0;

        // For each run, check if it has achieved the target at this evaluation
        for (_, points) in &extended_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],
    max_evaluations: u32,
) -> Vec<ProbabilityPointWithDeviation> {
    if algorithm_data.is_empty() || targets.is_empty() {
        return Vec::new();
    }

    // Calculate probability for each target with global max evaluation extension
    let target_probabilities: Vec<Vec<ProbabilityPoint>> = targets
        .iter()
        .map(|&target| {
            // First extend data to global max for this target calculation
            let mut extended_algorithm_data = HashMap::new();
            for (run_key, points) in algorithm_data {
                let mut extended_points = points.clone();
                if !extended_points.is_empty() {
                    extended_points.sort_by_key(|point| point.evaluations);
                    let last_point = extended_points.last().unwrap();
                    if last_point.evaluations < max_evaluations {
                        extended_points.push(DataPoint {
                            evaluations: max_evaluations,
                            fitness: last_point.fitness,
                            percentage_deviation: last_point.percentage_deviation,
                        });
                    }
                }
                extended_algorithm_data.insert(run_key.clone(), extended_points);
            }
            calculate_success_probability(&extended_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 calculate_averaged_fitness_with_deviation(
    algorithm_data: &HashMap<String, Vec<DataPoint>>,
    max_evaluations: u32,
) -> Vec<DataPointWithDeviation> {
    if algorithm_data.is_empty() {
        return Vec::new();
    }

    // Extend all data to the global maximum evaluation point
    let mut extended_algorithm_data = HashMap::new();
    for (run_key, points) in algorithm_data {
        let mut extended_points = points.clone();
        if !extended_points.is_empty() {
            extended_points.sort_by_key(|point| point.evaluations);
            let last_point = extended_points.last().unwrap();
            if last_point.evaluations < max_evaluations {
                extended_points.push(DataPoint {
                    evaluations: max_evaluations,
                    fitness: last_point.fitness,
                    percentage_deviation: last_point.percentage_deviation,
                });
            }
        }
        extended_algorithm_data.insert(run_key.clone(), extended_points);
    }

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

    let mut averaged_points = Vec::new();

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

        // Collect fitness and percentage values at this evaluation point from all runs
        for (_, points) in &extended_algorithm_data {
            // Find the best (minimum) percentage deviation 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);

            // Find the corresponding fitness value
            let best_fitness = points
                .iter()
                .filter(|p| p.evaluations <= evaluation)
                .min_by(|a, b| a.percentage_deviation.partial_cmp(&b.percentage_deviation).unwrap())
                .map(|p| p.fitness)
                .unwrap_or(f64::INFINITY);

            if best_percentage != f64::INFINITY {
                fitness_values.push(best_fitness);
                percentage_values.push(best_percentage);
            }
        }

        if !fitness_values.is_empty() {
            // Calculate means
            let mean_fitness = fitness_values.iter().sum::<f64>() / fitness_values.len() as f64;
            let mean_percentage = percentage_values.iter().sum::<f64>() / percentage_values.len() as f64;

            // Calculate standard deviation for percentage deviation
            let variance = percentage_values.iter()
                .map(|p| (p - mean_percentage).powi(2))
                .sum::<f64>() / percentage_values.len() as f64;
            let std_dev = variance.sqrt();

            // Calculate bounds (mean ± 1 standard deviation, clamped to reasonable values)
            let lower_bound = (mean_percentage - std_dev).max(0.1); // Keep above 0.1% for log scale
            let upper_bound = mean_percentage + std_dev;

            averaged_points.push(DataPointWithDeviation {
                evaluations: evaluation,
                fitness: mean_fitness,
                percentage_deviation: mean_percentage,
                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", 24))
        .label_style(("sans-serif", 24))
        .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;

                if config.average_runs {
                    // Calculate averaged fitness data with deviation
                    let averaged_data = calculate_averaged_fitness_with_deviation(algorithm_data, max_evaluations);

                    if !averaged_data.is_empty() {
                        // Conditionally draw standard deviation bands
                        if config.show_std_dev {
                            // Create transparent confidence band
                            let transparent_color = color.mix(0.15); // 15% opacity

                            // Create upper and lower bound points for the filled area
                            let upper_points: Vec<(f64, f64)> = averaged_data
                                .iter()
                                .map(|p| (p.evaluations as f64, p.upper_bound))
                                .collect();

                            let mut lower_points: Vec<(f64, f64)> = averaged_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 averaged line on top of the confidence band
                        let series = chart
                            .draw_series(LineSeries::new(
                                averaged_data.iter().map(|p| (p.evaluations as f64, p.percentage_deviation)),
                                &color,
                            ))?;

                        if !legend_added.contains_key(algorithm) {
                            let label = get_algorithm_label(algorithm, config);
                            let final_label = if config.show_std_dev {
                                format!("{} (avg ± σ)", label)
                            } else {
                                format!("{} (avg)", label)
                            };
                            series.label(final_label).legend(move |(x, y)| PathElement::new(vec![(x, y), (x + 10, y)], &color));
                            legend_added.insert(algorithm.clone(), true);
                        }
                    }
                } else {
                    // Original individual run plotting
                    for (_, points) in algorithm_data {
                        let mut extended_points = points.clone();

                        // Extend data to max evaluation inline
                        if !extended_points.is_empty() {
                            extended_points.sort_by_key(|point| point.evaluations);
                            let last_point = extended_points.last().unwrap();
                            if last_point.evaluations < max_evaluations {
                                extended_points.push(DataPoint {
                                    evaluations: max_evaluations,
                                    fitness: last_point.fitness,
                                    percentage_deviation: last_point.percentage_deviation,
                                });
                            }
                        }

                        let series = chart
                            .draw_series(LineSeries::new(
                                extended_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 mut extended_points = points.clone();

                            // Extend data to max evaluation inline
                            if !extended_points.is_empty() {
                                extended_points.sort_by_key(|point| point.evaluations);
                                let last_point = extended_points.last().unwrap();
                                if last_point.evaluations < max_evaluations {
                                    extended_points.push(DataPoint {
                                        evaluations: max_evaluations,
                                        fitness: last_point.fitness,
                                        percentage_deviation: last_point.percentage_deviation,
                                    });
                                }
                            }

                            let series = chart
                                .draw_series(LineSeries::new(
                                    extended_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", 24))
        .label_style(("sans-serif", 24))
        .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, max_evaluations);

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

                    // Conditionally draw standard deviation bands
                    if config.show_std_dev {
                        // Create transparent confidence band
                        let transparent_color = color.mix(0.15); // 15% 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 = if config.show_std_dev {
                        format!("{} (avg ± σ)", algo_label)
                    } else {
                        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)
        .position(SeriesLabelPosition::UpperLeft)
        .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(())
}