//! `AntColonyTsp` — Dorigo-style Ant System for permutation problems on a
//! complete graph (TSP-style).

use rand::Rng as _;

use crate::core::candidate::Candidate;
use crate::core::objective::Direction;
use crate::core::population::Population;
use crate::core::problem::Problem;
use crate::core::result::OptimizationResult;
use crate::core::rng::rng_from_seed;
use crate::traits::Optimizer;

/// Configuration for [`AntColonyTsp`].
#[derive(Debug, Clone)]
pub struct AntColonyTspConfig {
    /// Number of ants per generation.
    pub ants: usize,
    /// Number of generations.
    pub generations: usize,
    /// Pheromone weight `α`.
    pub alpha: f64,
    /// Heuristic weight `β`.
    pub beta: f64,
    /// Pheromone evaporation rate `ρ` ∈ [0, 1].
    pub evaporation: f64,
    /// Pheromone deposit constant `Q`. Reinforcement on edge (i, j) is
    /// `Q / tour_length` for every ant whose tour uses (i, j).
    pub deposit: f64,
    /// Initial pheromone level on every edge.
    pub initial_pheromone: f64,
    /// Seed for the deterministic RNG.
    pub seed: u64,
}

impl Default for AntColonyTspConfig {
    fn default() -> Self {
        Self {
            ants: 30,
            generations: 100,
            alpha: 1.0,
            beta: 2.0,
            evaporation: 0.5,
            deposit: 1.0,
            initial_pheromone: 1.0,
            seed: 42,
        }
    }
}

/// Ant Colony Optimization for permutation-style problems on a complete graph.
///
/// `Vec<usize>` decisions only (the permutation `[0, 1, …, n_cities - 1]`).
/// Single-objective only — typically minimizing total tour length, but the
/// algorithm is direction-aware for completeness.
///
/// Each ant builds a tour by repeatedly choosing the next node with
/// probability `∝ τ_ij^α · η_ij^β` over the unvisited cities, where
/// `η_ij = 1 / distance_ij` is the heuristic desirability.
///
/// # Example
///
/// ```
/// use heuropt::prelude::*;
///
/// struct Tsp { distances: Vec<Vec<f64>> }
/// impl Problem for Tsp {
///     type Decision = Vec<usize>;
///     fn objectives(&self) -> ObjectiveSpace {
///         ObjectiveSpace::new(vec![Objective::minimize("length")])
///     }
///     fn evaluate(&self, tour: &Vec<usize>) -> Evaluation {
///         let mut len = 0.0;
///         for w in tour.windows(2) { len += self.distances[w[0]][w[1]]; }
///         len += self.distances[*tour.last().unwrap()][tour[0]];
///         Evaluation::new(vec![len])
///     }
/// }
///
/// // 5 cities laid out in a small square + center. The optimal tour
/// // is the perimeter; the diagonal is suboptimal.
/// let cities = [(0.0_f64, 0.0), (3.0, 0.0), (3.0, 3.0), (0.0, 3.0), (1.5, 1.5)];
/// let n = cities.len();
/// let mut d = vec![vec![0.0; n]; n];
/// for i in 0..n {
///     for j in 0..n {
///         let dx = cities[i].0 - cities[j].0;
///         let dy = cities[i].1 - cities[j].1;
///         d[i][j] = (dx * dx + dy * dy).sqrt();
///     }
/// }
/// let problem = Tsp { distances: d.clone() };
///
/// let mut opt = AntColonyTsp::new(AntColonyTspConfig {
///     ants: 10,
///     generations: 50,
///     alpha: 1.0,
///     beta: 5.0,
///     evaporation: 0.5,
///     deposit: 1.0,
///     initial_pheromone: 0.1,
///     seed: 42,
/// }, d);
/// let r = opt.run(&problem);
/// assert!(r.best.is_some());
/// ```
pub struct AntColonyTsp {
    /// Algorithm configuration.
    pub config: AntColonyTspConfig,
    /// Symmetric distance matrix; size `n_cities × n_cities`. Diagonal must
    /// be zero.
    pub distances: Vec<Vec<f64>>,
}

impl AntColonyTsp {
    /// Construct an `AntColonyTsp`. Validates that `distances` is square
    /// and has a zero diagonal.
    pub fn new(config: AntColonyTspConfig, distances: Vec<Vec<f64>>) -> Self {
        let n = distances.len();
        assert!(
            n >= 2,
            "AntColonyTsp distances matrix must have >= 2 cities"
        );
        for (i, row) in distances.iter().enumerate() {
            assert_eq!(row.len(), n, "AntColonyTsp distances matrix must be square");
            assert_eq!(
                row[i], 0.0,
                "AntColonyTsp distance from city to itself must be 0"
            );
        }
        Self { config, distances }
    }
}

impl<P> Optimizer<P> for AntColonyTsp
where
    P: Problem<Decision = Vec<usize>> + Sync,
{
    fn run(&mut self, problem: &P) -> OptimizationResult<P::Decision> {
        assert!(self.config.ants >= 1, "AntColonyTsp ants must be >= 1");
        let objectives = problem.objectives();
        assert!(
            objectives.is_single_objective(),
            "AntColonyTsp requires exactly one objective",
        );
        let direction = objectives.objectives[0].direction;
        let n = self.distances.len();
        let mut rng = rng_from_seed(self.config.seed);

        // Heuristic desirability 1/distance, pre-raised to β. η is constant
        // for the whole run, so β is applied exactly once here instead of
        // once per ant per step inside `build_tour`.
        let eta_pow: Vec<Vec<f64>> = self
            .distances
            .iter()
            .map(|row| {
                row.iter()
                    .map(|&d| {
                        let e = if d > 0.0 { 1.0 / d } else { 0.0 };
                        e.powf(self.config.beta)
                    })
                    .collect()
            })
            .collect();

        // Pheromone matrix, plus a reused buffer holding τ pre-raised to α.
        let mut pheromone: Vec<Vec<f64>> = vec![vec![self.config.initial_pheromone; n]; n];
        let mut pheromone_pow: Vec<Vec<f64>> = vec![vec![0.0_f64; n]; n];

        let mut best_decision: Option<Vec<usize>> = None;
        let mut best_eval: Option<crate::core::evaluation::Evaluation> = None;
        let mut evaluations = 0usize;

        for _ in 0..self.config.generations {
            // τ is constant across the ant loop, so raise it to α once per
            // generation rather than once per ant per step per candidate.
            for (src, dst) in pheromone.iter().zip(pheromone_pow.iter_mut()) {
                for (&t, p) in src.iter().zip(dst.iter_mut()) {
                    *p = t.max(0.0).powf(self.config.alpha);
                }
            }

            let mut tours: Vec<Vec<usize>> = Vec::with_capacity(self.config.ants);
            let mut tour_evals: Vec<crate::core::evaluation::Evaluation> =
                Vec::with_capacity(self.config.ants);

            for _ in 0..self.config.ants {
                let start = rng.random_range(0..n);
                let tour = build_tour(n, start, &pheromone_pow, &eta_pow, &mut rng);
                let eval = problem.evaluate(&tour);
                evaluations += 1;
                tours.push(tour);
                tour_evals.push(eval);
            }

            // Update best.
            for (tour, eval) in tours.iter().zip(tour_evals.iter()) {
                let beats = match &best_eval {
                    None => true,
                    Some(b) => better_than_so(eval, b, direction),
                };
                if beats {
                    best_decision = Some(tour.clone());
                    best_eval = Some(eval.clone());
                }
            }

            // Pheromone evaporation.
            for row in pheromone.iter_mut() {
                for v in row.iter_mut() {
                    *v *= 1.0 - self.config.evaporation;
                }
            }

            // Pheromone deposit on each ant's tour.
            for (tour, eval) in tours.iter().zip(tour_evals.iter()) {
                let length = eval
                    .objectives
                    .first()
                    .copied()
                    .unwrap_or(f64::INFINITY)
                    .max(1e-12);
                let deposit = self.config.deposit / length;
                for w in tour.windows(2) {
                    let (i, j) = (w[0], w[1]);
                    pheromone[i][j] += deposit;
                    pheromone[j][i] += deposit;
                }
                // Close the loop.
                let (i, j) = (*tour.last().unwrap(), tour[0]);
                pheromone[i][j] += deposit;
                pheromone[j][i] += deposit;
            }
        }

        let best = Candidate::new(best_decision.unwrap(), best_eval.unwrap());
        let population = Population::new(vec![best.clone()]);
        let front = vec![best.clone()];
        OptimizationResult::new(
            population,
            front,
            Some(best),
            evaluations,
            self.config.generations,
        )
    }
}

#[cfg(feature = "async")]
impl AntColonyTsp {
    /// Async version of [`Optimizer::run`] — drives evaluations through
    /// the user-chosen async runtime. Available only with the `async`
    /// feature.
    ///
    /// `concurrency` bounds in-flight evaluations per generation.
    pub async fn run_async<P>(
        &mut self,
        problem: &P,
        concurrency: usize,
    ) -> OptimizationResult<Vec<usize>>
    where
        P: crate::core::async_problem::AsyncProblem<Decision = Vec<usize>>,
    {
        use crate::algorithms::parallel_eval_async::evaluate_batch_async;

        assert!(self.config.ants >= 1, "AntColonyTsp ants must be >= 1");
        let objectives = problem.objectives();
        assert!(
            objectives.is_single_objective(),
            "AntColonyTsp requires exactly one objective",
        );
        let direction = objectives.objectives[0].direction;
        let n = self.distances.len();
        let mut rng = rng_from_seed(self.config.seed);

        let eta_pow: Vec<Vec<f64>> = self
            .distances
            .iter()
            .map(|row| {
                row.iter()
                    .map(|&d| {
                        let e = if d > 0.0 { 1.0 / d } else { 0.0 };
                        e.powf(self.config.beta)
                    })
                    .collect()
            })
            .collect();

        let mut pheromone: Vec<Vec<f64>> = vec![vec![self.config.initial_pheromone; n]; n];
        let mut pheromone_pow: Vec<Vec<f64>> = vec![vec![0.0_f64; n]; n];

        let mut best_decision: Option<Vec<usize>> = None;
        let mut best_eval: Option<crate::core::evaluation::Evaluation> = None;
        let mut evaluations = 0usize;

        for _ in 0..self.config.generations {
            for (src, dst) in pheromone.iter().zip(pheromone_pow.iter_mut()) {
                for (&t, p) in src.iter().zip(dst.iter_mut()) {
                    *p = t.max(0.0).powf(self.config.alpha);
                }
            }

            let mut tours: Vec<Vec<usize>> = Vec::with_capacity(self.config.ants);
            for _ in 0..self.config.ants {
                let start = rng.random_range(0..n);
                let tour = build_tour(n, start, &pheromone_pow, &eta_pow, &mut rng);
                tours.push(tour);
            }

            let cands = evaluate_batch_async(problem, tours.clone(), concurrency).await;
            evaluations += cands.len();
            let tour_evals: Vec<crate::core::evaluation::Evaluation> =
                cands.into_iter().map(|c| c.evaluation).collect();

            for (tour, eval) in tours.iter().zip(tour_evals.iter()) {
                let beats = match &best_eval {
                    None => true,
                    Some(b) => better_than_so(eval, b, direction),
                };
                if beats {
                    best_decision = Some(tour.clone());
                    best_eval = Some(eval.clone());
                }
            }

            for row in pheromone.iter_mut() {
                for v in row.iter_mut() {
                    *v *= 1.0 - self.config.evaporation;
                }
            }

            for (tour, eval) in tours.iter().zip(tour_evals.iter()) {
                let length = eval
                    .objectives
                    .first()
                    .copied()
                    .unwrap_or(f64::INFINITY)
                    .max(1e-12);
                let deposit = self.config.deposit / length;
                for w in tour.windows(2) {
                    let (i, j) = (w[0], w[1]);
                    pheromone[i][j] += deposit;
                    pheromone[j][i] += deposit;
                }
                let (i, j) = (*tour.last().unwrap(), tour[0]);
                pheromone[i][j] += deposit;
                pheromone[j][i] += deposit;
            }
        }

        let best = Candidate::new(best_decision.unwrap(), best_eval.unwrap());
        let population = Population::new(vec![best.clone()]);
        let front = vec![best.clone()];
        OptimizationResult::new(
            population,
            front,
            Some(best),
            evaluations,
            self.config.generations,
        )
    }
}

fn build_tour(
    n: usize,
    start: usize,
    pheromone_pow: &[Vec<f64>],
    eta_pow: &[Vec<f64>],
    rng: &mut crate::core::rng::Rng,
) -> Vec<usize> {
    let mut tour = Vec::with_capacity(n);
    let mut visited = vec![false; n];
    tour.push(start);
    visited[start] = true;

    for _ in 1..n {
        let current = *tour.last().unwrap();
        // Build a probability vector over the unvisited candidates. Both
        // matrices are already raised to α / β by the caller, so the per-
        // candidate weight is a single multiply — no `powf` in the hot loop.
        let probs: Vec<(usize, f64)> = (0..n)
            .filter(|&j| !visited[j])
            .map(|j| {
                let p = pheromone_pow[current][j] * eta_pow[current][j];
                (j, p)
            })
            .collect();
        let total: f64 = probs.iter().map(|(_, p)| *p).sum();
        let next = if total > 0.0 {
            let r: f64 = rng.random::<f64>() * total;
            let mut acc = 0.0;
            let mut chosen = probs.last().unwrap().0;
            for (j, p) in &probs {
                acc += *p;
                if r <= acc {
                    chosen = *j;
                    break;
                }
            }
            chosen
        } else {
            // Degenerate case: pheromone × heuristic is 0 for every
            // unvisited city. Fall back to uniform random.
            let &(j, _) = probs.choose_uniform(rng);
            j
        };
        let _ = probs;
        tour.push(next);
        visited[next] = true;
    }
    tour
}

trait ChooseUniform<T> {
    fn choose_uniform(&self, rng: &mut crate::core::rng::Rng) -> &T;
}
impl<T> ChooseUniform<T> for [T] {
    fn choose_uniform(&self, rng: &mut crate::core::rng::Rng) -> &T {
        &self[rng.random_range(0..self.len())]
    }
}

fn better_than_so(
    a: &crate::core::evaluation::Evaluation,
    b: &crate::core::evaluation::Evaluation,
    direction: Direction,
) -> bool {
    match (a.is_feasible(), b.is_feasible()) {
        (true, false) => true,
        (false, true) => false,
        (false, false) => a.constraint_violation < b.constraint_violation,
        (true, true) => match direction {
            Direction::Minimize => a.objectives[0] < b.objectives[0],
            Direction::Maximize => a.objectives[0] > b.objectives[0],
        },
    }
}

impl crate::traits::AlgorithmInfo for AntColonyTsp {
    fn name(&self) -> &'static str {
        "Ant Colony"
    }
    fn full_name(&self) -> &'static str {
        "Ant Colony System for TSP"
    }
    fn seed(&self) -> Option<u64> {
        Some(self.config.seed)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::evaluation::Evaluation;
    use crate::core::objective::{Objective, ObjectiveSpace};

    /// A 5-city ring problem: cities placed at `(cos(2πi/5), sin(2πi/5))`.
    /// Optimal tour length: 2·5·sin(π/5) ≈ 5.878 (a regular pentagon).
    struct RingTsp {
        distances: Vec<Vec<f64>>,
    }
    impl RingTsp {
        fn new(n: usize) -> Self {
            use std::f64::consts::PI;
            let pts: Vec<(f64, f64)> = (0..n)
                .map(|i| {
                    let a = 2.0 * PI * (i as f64) / (n as f64);
                    (a.cos(), a.sin())
                })
                .collect();
            let distances = (0..n)
                .map(|i| {
                    (0..n)
                        .map(|j| {
                            let (xi, yi) = pts[i];
                            let (xj, yj) = pts[j];
                            ((xi - xj).powi(2) + (yi - yj).powi(2)).sqrt()
                        })
                        .collect()
                })
                .collect();
            Self { distances }
        }
    }
    impl Problem for RingTsp {
        type Decision = Vec<usize>;

        fn objectives(&self) -> ObjectiveSpace {
            ObjectiveSpace::new(vec![Objective::minimize("tour_length")])
        }

        fn evaluate(&self, tour: &Vec<usize>) -> Evaluation {
            let n = tour.len();
            let mut total = 0.0;
            for w in tour.windows(2) {
                total += self.distances[w[0]][w[1]];
            }
            total += self.distances[tour[n - 1]][tour[0]];
            Evaluation::new(vec![total])
        }
    }

    /// Trivial single-objective problem to test the multi-objective panic.
    struct DummyMo;
    impl Problem for DummyMo {
        type Decision = Vec<usize>;

        fn objectives(&self) -> ObjectiveSpace {
            ObjectiveSpace::new(vec![Objective::minimize("a"), Objective::minimize("b")])
        }

        fn evaluate(&self, _tour: &Vec<usize>) -> Evaluation {
            Evaluation::new(vec![0.0, 0.0])
        }
    }

    #[test]
    fn finds_near_optimum_on_5_city_ring() {
        let problem = RingTsp::new(5);
        let mut opt = AntColonyTsp::new(
            AntColonyTspConfig {
                ants: 10,
                generations: 30,
                alpha: 1.0,
                beta: 3.0,
                evaporation: 0.5,
                deposit: 1.0,
                initial_pheromone: 1.0,
                seed: 1,
            },
            problem.distances.clone(),
        );
        let r = opt.run(&problem);
        let best = r.best.unwrap();
        // Optimal pentagon perimeter ≈ 5.878. ACO should hit close.
        assert!(
            best.evaluation.objectives[0] < 5.95,
            "got tour length = {}",
            best.evaluation.objectives[0],
        );
    }

    #[test]
    fn deterministic_with_same_seed() {
        let problem = RingTsp::new(5);
        let cfg = AntColonyTspConfig {
            ants: 8,
            generations: 10,
            alpha: 1.0,
            beta: 2.0,
            evaporation: 0.5,
            deposit: 1.0,
            initial_pheromone: 1.0,
            seed: 99,
        };
        let mut a = AntColonyTsp::new(cfg.clone(), problem.distances.clone());
        let mut b = AntColonyTsp::new(cfg, problem.distances.clone());
        let ra = a.run(&problem);
        let rb = b.run(&problem);
        assert_eq!(
            ra.best.unwrap().evaluation.objectives,
            rb.best.unwrap().evaluation.objectives,
        );
    }

    #[test]
    #[should_panic(expected = "exactly one objective")]
    fn multi_objective_panics() {
        let mut opt = AntColonyTsp::new(
            AntColonyTspConfig::default(),
            vec![vec![0.0, 1.0], vec![1.0, 0.0]],
        );
        let _ = opt.run(&DummyMo);
    }

    // ---- Mutation-test pinned helpers --------------------------------------

    use crate::core::objective::Direction;
    use crate::core::rng::rng_from_seed;

    /// Raise every matrix entry to `p` — mirrors the α / β pre-raising the
    /// `run` loop now does before calling `build_tour`.
    fn raise(m: &[Vec<f64>], p: f64) -> Vec<Vec<f64>> {
        m.iter()
            .map(|row| row.iter().map(|&v| v.powf(p)).collect())
            .collect()
    }

    /// `better_than_so` follows the feasibility-first / objective-second
    /// tournament rule. Pin each of the four feasibility-cross-product
    /// branches so the `<` and `>` comparisons cannot flip silently.
    #[test]
    fn better_than_so_feasible_beats_infeasible() {
        let mut a = Evaluation::new(vec![10.0]);
        a.constraint_violation = 0.0; // feasible
        let mut b = Evaluation::new(vec![1.0]);
        b.constraint_violation = 1.0; // infeasible
        assert!(better_than_so(&a, &b, Direction::Minimize));
        assert!(!better_than_so(&b, &a, Direction::Minimize));
    }

    #[test]
    fn better_than_so_two_infeasible_compares_violation() {
        let mut a = Evaluation::new(vec![0.0]);
        a.constraint_violation = 0.5;
        let mut b = Evaluation::new(vec![0.0]);
        b.constraint_violation = 1.0;
        // a has smaller constraint_violation → "better".
        assert!(better_than_so(&a, &b, Direction::Minimize));
        assert!(!better_than_so(&b, &a, Direction::Minimize));
    }

    #[test]
    fn better_than_so_two_feasible_compares_objective_under_min() {
        let a = Evaluation::new(vec![1.0]); // feasible (default cv=0)
        let b = Evaluation::new(vec![2.0]); // feasible
        assert!(better_than_so(&a, &b, Direction::Minimize));
        assert!(!better_than_so(&b, &a, Direction::Minimize));
    }

    #[test]
    fn better_than_so_two_feasible_compares_objective_under_max() {
        let a = Evaluation::new(vec![2.0]);
        let b = Evaluation::new(vec![1.0]);
        assert!(better_than_so(&a, &b, Direction::Maximize));
        assert!(!better_than_so(&b, &a, Direction::Maximize));
    }

    #[test]
    fn better_than_so_equal_objectives_neither_strictly_better() {
        let a = Evaluation::new(vec![1.0]);
        let b = Evaluation::new(vec![1.0]);
        // Equal objectives → strict `<` is false both directions.
        assert!(!better_than_so(&a, &b, Direction::Minimize));
        assert!(!better_than_so(&b, &a, Direction::Minimize));
    }

    /// `build_tour` must produce a permutation of `[0..n)` starting at the
    /// given start city. Pin both invariants across many seeds.
    #[test]
    fn build_tour_is_permutation_starting_at_start() {
        let n = 6;
        let pher = raise(&vec![vec![1.0; n]; n], 1.0);
        let eta = raise(&vec![vec![1.0; n]; n], 2.0);
        for seed in 0..20 {
            for start in 0..n {
                let mut rng = rng_from_seed(seed);
                let tour = build_tour(n, start, &pher, &eta, &mut rng);
                assert_eq!(tour.len(), n);
                assert_eq!(tour[0], start, "tour must start at the given city");
                let mut sorted = tour.clone();
                sorted.sort();
                let expected: Vec<usize> = (0..n).collect();
                assert_eq!(sorted, expected, "tour must visit every city exactly once");
            }
        }
    }

    /// With a high `beta` and a heuristic that strongly prefers the next
    /// city, `build_tour` chooses that next city with near-certainty.
    /// Pins the heuristic-weighting arithmetic.
    #[test]
    fn build_tour_follows_strong_heuristic() {
        let n = 4;
        let pher = raise(&vec![vec![1.0; n]; n], 1.0);
        // Heuristic strongly favors city (i+1) % n: 1000x preferred.
        let mut eta = vec![vec![1.0; n]; n];
        for i in 0..n {
            eta[i][(i + 1) % n] = 1000.0;
        }
        let eta = raise(&eta, 5.0);
        let mut rng = rng_from_seed(0);
        let tour = build_tour(n, 0, &pher, &eta, &mut rng);
        // With beta=5 and 1000× heuristic, the path 0→1→2→3 has overwhelming
        // probability.
        assert_eq!(tour, vec![0, 1, 2, 3]);
    }

    /// `build_tour` with zero alpha + zero beta degenerates to uniform
    /// random over unvisited cities; the result is still a permutation.
    #[test]
    fn build_tour_zero_weights_still_produces_permutation() {
        let n = 5;
        let pher = raise(&vec![vec![1.0; n]; n], 0.0);
        let eta = raise(&vec![vec![1.0; n]; n], 0.0);
        let mut rng = rng_from_seed(42);
        let tour = build_tour(n, 2, &pher, &eta, &mut rng);
        // With alpha=beta=0, every term is 1.0 so the result is uniform but
        // still a permutation.
        assert_eq!(tour.len(), n);
        assert_eq!(tour[0], 2);
        let mut sorted = tour.clone();
        sorted.sort();
        assert_eq!(sorted, vec![0, 1, 2, 3, 4]);
    }
}