use crate::algo::EdgesVec;
use std::collections::HashSet;
use std::ops::Mul;

#[derive(Clone, Copy)]
pub enum AntDirection {
    ToSecond,
    ToFirst,
}

pub struct Ant {
    allowed_vertices: HashSet<usize>,
    location: AntLocation,
}

pub enum AntLocation {
    OnEdge {
        edge_index: usize,
        offset: u16,
        direction: AntDirection,
    },
    OnVertex {
        vertex_index: usize,
    },
}

pub struct AntsSimulationState<'a, Rng: rand::Rng> {
    rng: Rng,
    vertices: &'a Vec<HashSet<usize>>,
    edges: &'a mut EdgesVec,
    ants: &'a mut [Ant],
}

pub fn updateState<Rng: rand::Rng>(
    state: &mut AntsSimulationState<Rng>,
    ferment_weight: f64,
    heuristic_coefficient: f64,
    q: f64,
    r: f64,
    speed: u16,
) -> bool {
    let mut finished_ants_count = 0usize;
    for ant in state.ants.iter_mut() {
        match &mut ant.location {
            AntLocation::OnEdge {
                edge_index,
                offset,
                direction,
            } => match offset.checked_add(speed) {
                Some(newOffset) => *offset = newOffset,
                None => {
                    let edge = &mut state.edges[*edge_index];
                    edge.ferment_intensity += (q / edge.length) * r;
                    let vertex_index;
                    match direction {
                        AntDirection::ToSecond => vertex_index = edge.vertex2_index,
                        AntDirection::ToFirst => vertex_index = edge.vertex1_index,
                    }
                    ant.location = AntLocation::OnVertex { vertex_index }
                }
            },
            AntLocation::OnVertex { vertex_index } => {
                let allowed_outbounds = state.vertices[*vertex_index]
                    .iter()
                    .map(|ei| (*ei, &state.edges[*ei]))
                    .map(|(ei, e)| {
                        if *vertex_index == e.vertex1_index {
                            return (ei, e, e.vertex2_index, AntDirection::ToSecond);
                        } else {
                            return (ei, e, e.vertex1_index, AntDirection::ToFirst);
                        }
                    })
                    .filter(|e| ant.allowed_vertices.contains(&e.2))
                    .map(|e| {
                        return (
                            e.0,
                            (e.2, e.3),
                            e.1.ferment_intensity.powf(ferment_weight)
                                * (1.0 / e.1.length).powf(heuristic_coefficient),
                        );
                    })
                    .collect::<Vec<(usize, (usize, AntDirection), f64)>>();

                let mut accumulator: f64 = allowed_outbounds
                    .iter()
                    .fold(0.0, |a, e| a + e.2)
                    .mul(state.rng.random::<f64>());
                let target_vertex_index = allowed_outbounds.iter().find(|e| {
                    accumulator -= e.2;
                    return accumulator <= 0.0;
                });

                match target_vertex_index {
                    None => finished_ants_count += 1,
                    Some((ei, (vi, dir), _)) => {
                        ant.allowed_vertices.remove(vi);
                        ant.location = AntLocation::OnEdge {
                            edge_index: *ei,
                            offset: 0,
                            direction: *dir
                        }
                    }
                };
            }
        }
    }

    return finished_ants_count == state.ants.len()
}