pub mod tsp;
pub mod graph;

use tsp::{TSPInstance, TSPRandomInitializer, SwapPerturbation};
use nalgebra::{Const, Dim, 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 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 main() -> Result<(), Box<dyn std::error::Error>> {
    // Load a TSP instance from file
    let filename = "instances/berlin52.tsp.gz";
    let instance = load_tsp_instance(filename)?;
    let dimension = instance.dimension();
    let optimal_cost = load_optimal_cost(filename)?;

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

    // 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);
    println!("Gap to optimal: {:.2} ({:.1}%)",
        result.best_candidate.fit - optimal_cost,
        ((result.best_candidate.fit - optimal_cost) / optimal_cost) * 100.0);

    Ok(())
}