pub mod tsp;
pub mod graph;

use tsp::{TSPInstance, TSPRandomInitializer, SwapPerturbation};
use nalgebra::{Const, Dyn};
use eoa_lib::{
    initializer::Initializer,
    local_search::local_search_first_improving,
    terminating::NoBetterForCyclesTerminatingCondition,
    comparison::MinimizingOperator,
};
use rand::rng;
use std::fs::File;
use std::io::{BufRead, BufReader, Read};
use flate2::read::GzDecoder;

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 main() -> Result<(), Box<dyn std::error::Error>> {
    // Load a TSP instance from file
    let instance = load_tsp_instance("instances/berlin52.tsp.gz")?;
    let dimension = instance.dimension();

    // Plot just the instance
    instance.plot("tsp_instance.png")?;
    println!("Created tsp_instance.png");

    // Create a random solution
    let initializer = TSPRandomInitializer::new();
    let mut rng = rng();
    let initial_solution = initializer.initialize_single(dimension, &mut rng);

    // Plot the initial solution
    instance.draw_solution(&initial_solution, "tsp_initial_solution.png")?;
    println!("Created tsp_initial_solution.png with cost: {:.2}", instance.solution_cost(&initial_solution));

    // Run local search to improve the solution
    let mut perturbation = SwapPerturbation::new();
    let mut terminating_condition = NoBetterForCyclesTerminatingCondition::new(100000);
    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
    instance.draw_solution(&result.best_candidate.pos, "tsp_improved_solution.png")?;
    println!("Created tsp_improved_solution.png with cost: {:.2}", result.best_candidate.fit);
    println!("Improvement: {:.2} -> {:.2} (cycles: {})",
        instance.solution_cost(&initial_solution),
        result.best_candidate.fit,
        result.cycles);

    Ok(())
}