pub mod tsp;
pub mod graph;

use tsp::{TSPInstance, TSPRandomInitializer, SwapPerturbation, ReverseSubsequencePerturbation, EdgeRecombinationCrossover};
use nalgebra::{Const, Dim, Dyn, U100};
use eoa_lib::{
    comparison::MinimizingOperator, evolution::evolution_algorithm, initializer::Initializer, local_search::local_search_first_improving, pairing::AdjacentPairing, perturbation::{CombinedPerturbation, MutationPerturbation}, replacement::BestReplacement, selection::TournamentSelection, terminating::{MaximumCyclesTerminatingCondition, NoBetterForCyclesTerminatingCondition}
};
use rand::rng;
use std::env;
use std::fs::{File, create_dir_all};
use std::io::{BufRead, BufReader, Read, Write};
use std::path::Path;
use flate2::read::GzDecoder;
use chrono::{DateTime, Local};

fn load_tsp_instance(filename: &str) -> Result<TSPInstance<Dyn>, Box<dyn std::error::Error>> {
    let file = File::open(filename)?;
    let reader: Box<dyn BufRead> = if filename.ends_with(".gz") {
        Box::new(BufReader::new(GzDecoder::new(file)))
    } else {
        Box::new(BufReader::new(file))
    };

    let mut cities = Vec::new();
    let mut in_coord_section = false;
    let mut dimension = 0;

    for line in reader.lines() {
        let line = line?.trim().to_string();

        if line.starts_with("DIMENSION") {
            dimension = line.split(':').nth(1)
                .or_else(|| line.split_whitespace().nth(1))
                .ok_or("Could not parse dimension")?
                .trim()
                .parse::<usize>()?;
        } else if line == "NODE_COORD_SECTION" {
            in_coord_section = true;
        } else if line == "EOF" {
            break;
        } else if in_coord_section && !line.is_empty() {
            let parts: Vec<&str> = line.split_whitespace().collect();
            if parts.len() >= 3 {
                let x: f64 = parts[1].parse()?;
                let y: f64 = parts[2].parse()?;
                cities.push((x, y));
            }
        }
    }

    if cities.len() != dimension {
        return Err(format!("Expected {} cities, found {}", dimension, cities.len()).into());
    }

    Ok(TSPInstance::new_dyn(cities))
}

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

    let solutions_path = std::path::Path::new(instance_filename)
        .parent()
        .unwrap_or_else(|| std::path::Path::new("."))
        .join("solutions.txt");

    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 run_evolution_algorithm(instance: &TSPInstance<Dyn>, optimal_cost: f64, base_path: &str) -> Result<(), Box<dyn std::error::Error>> {
    let mut rng = rng();
    let initializer = TSPRandomInitializer::new();
    let dimension = instance.dimension();

    // Create combined perturbation with two mutations wrapped in MutationPerturbation
    let swap_mutation = MutationPerturbation::new(Box::new(SwapPerturbation::new()), 0.5);
    let reverse_mutation = MutationPerturbation::new(Box::new(ReverseSubsequencePerturbation::new()), 0.5);
    let mut combined_perturbation = CombinedPerturbation::new(vec![
        Box::new(swap_mutation),
        Box::new(reverse_mutation),
    ]);

    // Set up other components
    let mut crossover = EdgeRecombinationCrossover::new();
    let mut selection = TournamentSelection::new(5, 0.8);
    let mut replacement = BestReplacement::new();
    let mut pairing = AdjacentPairing::new();
    let better_than_operator = MinimizingOperator::new();

    // Create initial population
    let population_size = 500;
    let initial_population = initializer.initialize(dimension, population_size, &mut rng);
    let initial_population = eoa_lib::replacement::Population::from_vec(initial_population);

    let evaluated_initial = initial_population.clone().evaluate(instance)?;
    let initial_best = evaluated_initial.best_candidate(&better_than_operator);
    println!("Initial best cost: {:.2}", initial_best.evaluation);

    // Run evolution algorithm
    let parents_count = 250;
    let result = evolution_algorithm(
        initial_population.clone(),
        parents_count,
        instance,
        &mut selection,
        &mut pairing,
        &mut crossover,
        &mut combined_perturbation,
        &mut replacement,
        &better_than_operator,
        5000, // max iterations
        &mut rng,
        |iteration, stats, _, _, _, _, perturbation, _| {
            let iters_till_end = 5000 - iteration + 1;
            let iters_since_better =
                iteration - stats.best_candidates.last().map(|c| c.iteration).unwrap_or(0);
            MutationPerturbation::apply_to_mutations(
                perturbation,
                &mut |p| {
                    p.probability = (0.5 * (1.0 + (iters_since_better as f64 / iters_till_end as f64))).min(1.0);
                }
            );
        }
    )?;

    // Plot the best solution
    let best_solution = &result.best_candidate.chromosome;
    let plot_path = format!("{}.png", base_path);
    instance.draw_solution(best_solution, &plot_path)?;

    // Save statistics to CSV
    let stats_path = format!("{}.csv", base_path);
    let mut stats_file = File::create(&stats_path)?;
    writeln!(stats_file, "fitness_evaluations,evaluation")?;

    // Calculate fitness evaluations: initial_population + iteration * offspring_count
    // offspring_count = parents_count / 2 (due to adjacent pairing)
    let offspring_count = parents_count / 2;
    let initial_population_size = initial_population.iter().count();

    for candidate in &result.stats.best_candidates {
        let fitness_evaluations = initial_population_size + candidate.iteration * offspring_count;
        writeln!(stats_file, "{},{}", fitness_evaluations, candidate.evaluated_chromosome.evaluation)?;
    }

    println!("Evolution completed in {} generations", result.iterations);
    println!("Final cost: {:.2}", result.best_candidate.evaluation);
    println!("Gap to optimal: {:.2} ({:.1}%)",
        result.best_candidate.evaluation - optimal_cost,
        ((result.best_candidate.evaluation - optimal_cost) / optimal_cost) * 100.0);

    Ok(())
}

fn run_local_search(instance: &TSPInstance<Dyn>, optimal_cost: f64, base_path: &str) -> Result<(), Box<dyn std::error::Error>> {
    let mut rng = rng();
    let initializer = TSPRandomInitializer::new();
    let dimension = instance.dimension();

    // Create a random initial solution
    let initial_solution = initializer.initialize_single(dimension, &mut rng);
    println!("Initial cost: {:.2}", instance.solution_cost(&initial_solution));

    // Run local search
    let mut perturbation = ReverseSubsequencePerturbation::new();
    let mut terminating_condition = MaximumCyclesTerminatingCondition::new(250 * 5000 + 500);
    let better_than_operator = MinimizingOperator::new();

    let result = local_search_first_improving(
        instance,
        &mut terminating_condition,
        &mut perturbation,
        &better_than_operator,
        &initial_solution,
        &mut rng,
    )?;

    // Plot the improved solution
    let plot_path = format!("{}.png", base_path);
    instance.draw_solution(&result.best_candidate.pos, &plot_path)?;

    // Save statistics to CSV
    let stats_path = format!("{}.csv", base_path);
    let mut stats_file = File::create(&stats_path)?;
    writeln!(stats_file, "fitness_evaluations,evaluation")?;
    // For local search, each cycle = 1 fitness evaluation
    for candidate in result.stats.candidates() {
        writeln!(stats_file, "{},{}", candidate.cycle, candidate.fit)?;
    }

    writeln!(stats_file, "{},{}", result.cycles, result.stats.candidates().iter().last().unwrap().fit)?;

    println!("Local search completed in {} cycles", result.cycles);
    println!("Final cost: {:.2}", result.best_candidate.fit);
    println!("Gap to optimal: {:.2} ({:.1}%)",
        result.best_candidate.fit - optimal_cost,
        ((result.best_candidate.fit - optimal_cost) / optimal_cost) * 100.0);

    Ok(())
}

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

    if args.len() != 3 {
        eprintln!("Usage: {} <instance_name> <algorithm>", args[0]);
        eprintln!("  instance_name: e.g., kroA100, berlin52, eil51");
        eprintln!("  algorithm: ea (evolution algorithm) or ls (local search)");
        std::process::exit(1);
    }

    let instance_name = &args[1];
    let algorithm = &args[2];

    // Load TSP instance
    let filename = format!("instances/{}.tsp.gz", instance_name);
    let instance = load_tsp_instance(&filename)?;
    let dimension = instance.dimension();
    let optimal_cost = load_optimal_cost(&filename)?;

    println!("Loaded {} with {} cities (optimal cost: {})", instance_name, dimension.value(), optimal_cost);

    // Create directory structure: algorithm/instance_name/
    let output_dir = format!("solutions/{}/{}", algorithm, instance_name);
    create_dir_all(&output_dir)?;

    // Generate timestamp for filename
    let now: DateTime<Local> = Local::now();
    let timestamp = now.format("%Y-%m-%d_%H-%M-%S");
    let solution_base_path = format!("{}/{}", output_dir, timestamp);

    // Run the specified algorithm
    match algorithm.as_str() {
        "ea" => {
            println!("Running Evolution Algorithm...");
            run_evolution_algorithm(&instance, optimal_cost, &solution_base_path)?;
        },
        "ls" => {
            println!("Running Local Search...");
            run_local_search(&instance, optimal_cost, &solution_base_path)?;
        },
        _ => {
            eprintln!("Unknown algorithm: {}. Use 'ea' or 'ls'", algorithm);
            std::process::exit(1);
        }
    }

    println!("Created {}.png and {}.csv", solution_base_path, solution_base_path);
    Ok(())
}