use crate::algo::Ant;
use crate::algo::state::AntsSimulationState;
use bgtu_ai_utility::weighted_random;
use rand::Rng;
use std::collections::HashSet;
use std::ptr;

impl AntsSimulationState<'_> {
    pub fn is_ended(&self) -> bool {
        return self.ants.iter().all(|a| a.allowed_vertices.len() == 0);
    }

    pub fn update1(&mut self, rng: &mut impl Rng) {
        self._move_ants(rng);
        self._evaporate_ferments();
    }

    fn _move_ants(&mut self, rng: &mut impl Rng) {
        for ant in unsafe { (*ptr::from_mut(&mut self.ants)).iter_mut() } {
            let outbound_weights =
                self._get_outbounds_weights(ant.current_vertex, &ant.allowed_vertices);

            match weighted_random(rng, &outbound_weights, |e| e.2) {
                None => {}
                Some((ei, nvi, _)) => {
                    self.ferments[*ei] += self._calculate_ferment(*ei);
                    ant.current_vertex = *nvi;
                    ant.allowed_vertices.remove(nvi);
                }
            }
        }
    }

    fn _get_outbounds_weights(
        &self,
        location: usize,
        allowed_outbounds: &HashSet<usize>,
    ) -> Vec<(usize, usize, f64)> {
        let mut weights = Vec::with_capacity(allowed_outbounds.len());
        for ei in self.graph.vertices[location].iter().map(|ei| *ei) {
            let outbound = self.graph.edges[ei].another(location);
            if !allowed_outbounds.contains(&outbound) {
                continue;
            }
            weights.push((ei, outbound, self._calculate_weight(ei)))
        }
        return weights;
    }

    fn _calculate_weight(&self, edge: usize) -> f64 {
        let a = self.ferments[edge].powf(self.cfg.ferment_weight);
        let b = (1.0 / self.graph.edges[edge].length).powf(self.cfg.heuristic_coefficient);
        return a * b;
    }

    fn _calculate_ferment(&self, edge: usize) -> f64 {
        return self.cfg.q;
    }

    fn _evaporate_ferments(&mut self) {}
}