use std::collections::HashSet;
use rand::Rng;
use clap::Parser;

#[derive(Parser, Debug)]
#[command(about, long_about=None)]
struct Args{
    #[arg(short='n', long="ants")]
    num_ants: usize,

    #[arg(short, long)]
    iterations: usize,

    #[arg(short, long)]
    alpha: f64,

    #[arg(short, long)]
    beta: f64,

    #[arg(short, long)]
    evaporation: f64,

    #[arg(short, long)]
    q_value: f64,

    #[clap()]
    edges: Vec<String>,
}

#[derive(Debug)]
struct Edge {
    from: String,
    to: String,
    distance: f64,
    pheromone: f64,
}

fn parse_edges(edge_strs: &[String]) -> Vec<Edge> {
    let mut edges = Vec::new();
    for edge_str in edge_strs {
        let parts: Vec<&str> = edge_str.split(['=', '-']).collect();
        if parts.len() != 3 {
            panic!("Invalid edge format: {}", edge_str);
        }
        
        let distance = parts[2].parse::<f64>()
            .expect("Distance must be a number");
            
        edges.push(Edge {
            from: parts[0].to_string(),
            to: parts[1].to_string(),
            distance,
            pheromone: 1.0,
        });
        
        edges.push(Edge {
            from: parts[1].to_string(),
            to: parts[0].to_string(),
            distance,
            pheromone: 1.0,
        });
    }
    edges
}

fn construct_solution(
    edges: &[Edge],
    vertices: &Vec<String>,
    alpha: f64,
    beta: f64,
) -> Vec<String> {
    let mut rng = rand::thread_rng();
    // Текущая вершина (стартовая)
    let mut current = vertices.iter().nth(0).unwrap().clone();
    // Текущий построенный путь
    let mut path = vec![current.clone()];
    // Список непосещенных вершин
    let mut unvisited: HashSet<_> = vertices.iter().cloned().filter(|v| v != &current).collect();

    // Пока существуют непосещенные вершины
    while !unvisited.is_empty() {
        // Список доступных ребер, на которые мы можем пойти из текущей вершины
        let available_edges: Vec<_> = edges.iter()
            .filter(|e| e.from == current && unvisited.contains(&e.to))
            .collect();

        // Если нет доступных ребер
        if available_edges.is_empty() {
            break;
        }

        // Сумма ощущаемых феромонов
        let total: f64 = available_edges.iter()
            .map(|e| (e.pheromone.powf(alpha)) * ((1.0 / e.distance).powf(beta)))
            .sum();

        // Вероятность пойти на каждое ребро
        let probabilities: Vec<f64> = available_edges.iter()
            .map(|e| (e.pheromone.powf(alpha)) * ((1.0 / e.distance).powf(beta)) / total)
            .collect();

        let mut cumsum = 0.0;
        let rand_val = rng.gen::<f64>();
        // Выбираем следующее ребро
        let selected = available_edges.iter().zip(probabilities.iter())
            .find(|(_, &p)| {
                cumsum += p;
                rand_val <= cumsum
            })
            .map(|(e, _)| e)
            .unwrap();

        // Переходим на следующую вершину
        current = selected.to.clone();
        // Сохраняем новую текущую вершину в путь
        path.push(current.clone());
        // Удаляем посещенную вершину из списка непосещенных
        unvisited.remove(&current);
    }

    path
}

fn calculate_path_length(path: &[String], edges: &[Edge]) -> f64 {
    // Считаем длину пути
    path.windows(2)
        .map(|window| {
            let from = &window[0];
            let to = &window[1];
            edges.iter()
                .find(|e| e.from == *from && e.to == *to)
                .map(|e| e.distance)
                .unwrap()
        })
        .sum()
}

fn update_pheromones(
    edges: &mut [Edge],
    solutions: &[(Vec<String>, f64)],
    evaporation: f64, 
    q_value: f64,
) {
    // Испаряем часть феромонов со всех рёбер графа
    for edge in edges.iter_mut() {
        edge.pheromone *= 1.0 - evaporation;
    }

    // Для каждого найденного пути добавляем новые феромоны
    for (path, length) in solutions {
        // Вычисляем количество феромонов для отложения
        // Чем короче путь, тем больше феромонов откладывается
        let deposit = q_value / length;

        // Проходим по всем парам последовательных вершин в пути
        path.windows(2).for_each(|window| {
            let from = &window[0];
            let to = &window[1];

            // Добавляем феромоны на прямое ребро (from -> to)
            if let Some(edge) = edges.iter_mut().find(|e| e.from == *from && e.to == *to) {
                edge.pheromone += deposit;
            }
            // Добавляем феромоны на обратное ребро (to -> from), 
            // так как граф неориентированный
            if let Some(edge) = edges.iter_mut().find(|e| e.from == *to && e.to == *from) {
                edge.pheromone += deposit;
            }
        });
    }
}

fn main() {
    // Парсим аргументы командной строки
    let opt = Args::parse();
    let mut edges = parse_edges(&opt.edges);
    
    // Собираем все вершины
    let vertices: Vec<String> = edges.iter()
        .flat_map(|e| vec![e.from.clone(), e.to.clone()])
        .collect();

    let mut best_solution = None;
    let mut best_length = f64::INFINITY;

    // Основной цикл итераций
    for _ in 0..opt.iterations {
        let mut solutions = Vec::new();

        for _ in 0..opt.num_ants {
            // Строим путь
            let path = construct_solution(&edges, &vertices, opt.alpha, opt.beta);
            // Находим длину пути
            let length = calculate_path_length(&path, &edges);
            // Сохраняем решение
            solutions.push((path.clone(), length));

            if length < best_length {
                best_length = length;
                best_solution = Some(path);
            }
        }

        update_pheromones(&mut edges, &solutions, opt.evaporation, opt.q_value);
    }

    // Если решение найдено
    if let Some(_path) = best_solution {
        let mut path = _path.clone();
        // Добавляем путь до стартовой вершины
        path.push(vertices.iter().nth(0).clone().unwrap().to_string());
        best_length += edges.iter().filter(|e| e.from == path.first().unwrap().to_string() && e.to == _path.last().unwrap().to_string()).nth(0).unwrap().distance;
        println!("Best path: {}", path.join(" -> "));

        let mut sum = 0f64;

        // Считаем сумму
        for edgenum in 1..path.len(){
            let prev = path[edgenum - 1].clone();
            let cver = path[edgenum].clone();
            let edge = &edges.iter().filter(|e| e.from == prev && e.to == cver).last().expect("");
            sum += edge.distance;
            println!("{}-{}={}\t(pheromone:\t{:.3}). \tCumulative length:\t{}", edge.from, edge.to, edge.distance, edge.pheromone, sum);
        }

        println!("Total path length:\t.\t.\t.\t.\t.\t{}", best_length);
    } else {
        println!("No solution found.");
    }
}
