metalforge 0.3.0

//! NSGA-II — Non-dominated Sorting Genetic Algorithm II (Deb, Pratap, Agarwal &
//! Meyarivan 2002, IEEE Trans. Evol. Comput. 6(2), 182–197).
//!
//! The reference elitist multi-objective evolutionary algorithm. Each
//! generation it ranks the combined parent+offspring population into Pareto
//! fronts (fast non-dominated sorting), measures diversity within each front
//! (crowding distance), and selects the next generation by `(rank, crowding)`.
//! Offspring come from binary-tournament selection, simulated binary crossover
//! (SBX), and polynomial mutation — the canonical real-coded operators.
//!
//! Returns a [`ParetoFront`] (the final non-dominated set), not a single
//! solution. Deterministic for a given seed. Closes the multi-objective gap the
//! sibling engine Anvil left pending.

use super::sample;
use crate::problem::MultiProblem;
use crate::rng::Rng;
use crate::solution::{MultiSolution, ParetoFront};
use crate::termination::Termination;

/// NSGA-II configuration.
#[derive(Debug, Clone, Copy)]
pub struct NsgaII {
    /// Population size (forced even and ≥ 2: offspring are bred in pairs).
    pub pop_size: usize,
    /// SBX distribution index `η_c`: larger ⇒ children closer to parents.
    pub crossover_eta: f64,
    /// Probability of applying SBX to a parent pair.
    pub crossover_prob: f64,
    /// Polynomial-mutation distribution index `η_m`.
    pub mutation_eta: f64,
    /// Per-variable mutation probability; `None` uses `1/dim` (Deb's default).
    pub mutation_prob: Option<f64>,
    /// RNG seed; same seed + same problem + same budget ⇒ same result.
    pub seed: u64,
}

impl Default for NsgaII {
    fn default() -> Self {
        NsgaII {
            pop_size: 100,
            crossover_eta: 15.0,
            crossover_prob: 0.9,
            mutation_eta: 20.0,
            mutation_prob: None,
            seed: 42,
        }
    }
}

impl NsgaII {
    /// Approximates the Pareto front of `problem` within the evaluation budget
    /// of `term` (its `target` is ignored — meaningless for a front).
    pub fn optimize(&self, problem: &dyn MultiProblem, term: &Termination) -> ParetoFront {
        crate::problem::validate_multi(problem)
            .unwrap_or_else(|e| panic!("NsgaII: invalid problem: {e}"));
        let bounds = problem.bounds();
        let dim = bounds.len();
        let n = self.pop_size.max(2) + (self.pop_size & 1); // even, ≥ 2
        let pm = self.mutation_prob.unwrap_or(1.0 / dim as f64);
        let mut rng = Rng::new(self.seed);

        let eval = |x: &[f64]| -> Vec<f64> {
            let mut o = problem.objectives(x);
            for v in &mut o {
                if !v.is_finite() {
                    *v = f64::INFINITY; // infeasible ⇒ dominated by everything
                }
            }
            o
        };

        // Initial population (the budget binds here too).
        let mut pop: Vec<MultiSolution> = Vec::with_capacity(n);
        let mut evaluations = 0;
        for _ in 0..n {
            if evaluations >= term.max_evaluations {
                break;
            }
            let x = sample(bounds, &mut rng);
            let objectives = eval(&x);
            evaluations += 1;
            pop.push(MultiSolution { x, objectives });
        }
        if pop.is_empty() {
            return ParetoFront {
                solutions: Vec::new(),
                evaluations,
            };
        }
        let full_pop = pop.len() == n;

        while full_pop && evaluations < term.max_evaluations {
            let (rank, crowd) = rank_and_crowd(&pop);

            // Breed n offspring via tournament + SBX + polynomial mutation,
            // stopping mid-generation when the budget is spent.
            let mut offspring: Vec<MultiSolution> = Vec::with_capacity(n);
            while offspring.len() < n && evaluations < term.max_evaluations {
                let p1 = tournament(&rank, &crowd, &mut rng);
                let p2 = tournament(&rank, &crowd, &mut rng);
                let (mut c1, mut c2) = sbx(
                    &pop[p1].x,
                    &pop[p2].x,
                    bounds,
                    self.crossover_eta,
                    self.crossover_prob,
                    &mut rng,
                );
                mutate(&mut c1, bounds, self.mutation_eta, pm, &mut rng);
                mutate(&mut c2, bounds, self.mutation_eta, pm, &mut rng);

                let o1 = eval(&c1);
                evaluations += 1;
                offspring.push(MultiSolution {
                    x: c1,
                    objectives: o1,
                });
                if offspring.len() < n && evaluations < term.max_evaluations {
                    let o2 = eval(&c2);
                    evaluations += 1;
                    offspring.push(MultiSolution {
                        x: c2,
                        objectives: o2,
                    });
                }
            }

            // Elitist environmental selection over parents ∪ offspring.
            let mut union = pop;
            union.extend(offspring);
            pop = environmental_selection(union, n);
        }

        // Final front: the rank-0 set of the last population.
        let fronts = fast_non_dominated_sort(&pop);
        let solutions = fronts[0].iter().map(|&i| pop[i].clone()).collect();
        ParetoFront {
            solutions,
            evaluations,
        }
    }
}

/// True if `a` Pareto-dominates `b`: no worse in every objective, strictly
/// better in at least one (minimization).
pub(crate) fn dominates(a: &[f64], b: &[f64]) -> bool {
    let mut strictly_better = false;
    for (x, y) in a.iter().zip(b) {
        if x > y {
            return false;
        }
        if x < y {
            strictly_better = true;
        }
    }
    strictly_better
}

/// Fast non-dominated sort (Deb et al. 2002, §III-A). Returns the fronts as
/// lists of population indices, best (rank 0) first.
pub(crate) fn fast_non_dominated_sort(pop: &[MultiSolution]) -> Vec<Vec<usize>> {
    let n = pop.len();
    let mut dominated: Vec<Vec<usize>> = vec![Vec::new(); n]; // who each p dominates
    let mut dom_count = vec![0usize; n]; // how many dominate p
    let mut fronts: Vec<Vec<usize>> = vec![Vec::new()];

    for p in 0..n {
        for q in 0..n {
            if p == q {
                continue;
            }
            if dominates(&pop[p].objectives, &pop[q].objectives) {
                dominated[p].push(q);
            } else if dominates(&pop[q].objectives, &pop[p].objectives) {
                dom_count[p] += 1;
            }
        }
        if dom_count[p] == 0 {
            fronts[0].push(p);
        }
    }

    let mut i = 0;
    while !fronts[i].is_empty() {
        let mut next = Vec::new();
        for &p in &fronts[i] {
            for &q in &dominated[p] {
                dom_count[q] -= 1;
                if dom_count[q] == 0 {
                    next.push(q);
                }
            }
        }
        i += 1;
        fronts.push(next);
    }
    fronts.pop(); // drop the trailing empty front
    fronts
}

/// Crowding distance for the solutions in one front (Deb et al. 2002, §III-B):
/// the sum, over objectives, of normalized gaps to the nearest neighbors;
/// boundary solutions get infinity to preserve the extremes.
pub(crate) fn crowding_distance(front: &[usize], pop: &[MultiSolution]) -> Vec<f64> {
    let l = front.len();
    let mut dist = vec![0.0; l];
    if l == 0 {
        return dist;
    }
    let m = pop[front[0]].objectives.len();
    for obj in 0..m {
        // Sort the front by this objective.
        let mut order: Vec<usize> = (0..l).collect();
        order.sort_by(|&a, &b| {
            pop[front[a]].objectives[obj]
                .partial_cmp(&pop[front[b]].objectives[obj])
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        dist[order[0]] = f64::INFINITY;
        dist[order[l - 1]] = f64::INFINITY;
        let fmin = pop[front[order[0]]].objectives[obj];
        let fmax = pop[front[order[l - 1]]].objectives[obj];
        let span = fmax - fmin;
        if span <= 0.0 || !span.is_finite() {
            continue;
        }
        for k in 1..l - 1 {
            let prev = pop[front[order[k - 1]]].objectives[obj];
            let next = pop[front[order[k + 1]]].objectives[obj];
            dist[order[k]] += (next - prev) / span;
        }
    }
    dist
}

/// Computes each solution's Pareto rank and crowding distance.
fn rank_and_crowd(pop: &[MultiSolution]) -> (Vec<usize>, Vec<f64>) {
    let fronts = fast_non_dominated_sort(pop);
    let mut rank = vec![0usize; pop.len()];
    let mut crowd = vec![0.0; pop.len()];
    for (r, front) in fronts.iter().enumerate() {
        let cd = crowding_distance(front, pop);
        for (k, &idx) in front.iter().enumerate() {
            rank[idx] = r;
            crowd[idx] = cd[k];
        }
    }
    (rank, crowd)
}

/// Binary tournament on `(rank, crowding)`: lower rank wins; ties go to the
/// larger crowding distance (more isolated ⇒ better for diversity).
fn tournament(rank: &[usize], crowd: &[f64], rng: &mut Rng) -> usize {
    let a = rng.index(rank.len());
    let b = rng.index(rank.len());
    if rank[a] < rank[b] || (rank[a] == rank[b] && crowd[a] > crowd[b]) {
        a
    } else {
        b
    }
}

/// Picks the best `n` from `union` by filling whole fronts, then truncating the
/// splitting front by descending crowding distance.
fn environmental_selection(union: Vec<MultiSolution>, n: usize) -> Vec<MultiSolution> {
    let fronts = fast_non_dominated_sort(&union);
    let mut selected: Vec<usize> = Vec::with_capacity(n);
    for front in &fronts {
        if selected.len() + front.len() <= n {
            selected.extend_from_slice(front);
        } else {
            // Partial front: keep the most isolated solutions.
            let cd = crowding_distance(front, &union);
            let mut order: Vec<usize> = (0..front.len()).collect();
            order.sort_by(|&a, &b| {
                cd[b]
                    .partial_cmp(&cd[a])
                    .unwrap_or(std::cmp::Ordering::Equal)
            });
            for &k in order.iter().take(n - selected.len()) {
                selected.push(front[k]);
            }
            break;
        }
    }

    // Materialize, preserving `union` ownership for the chosen indices.
    let mut chosen: Vec<Option<MultiSolution>> = union.into_iter().map(Some).collect();
    selected
        .into_iter()
        .map(|i| chosen[i].take().expect("each index selected once"))
        .collect()
}

/// Simulated Binary Crossover (Deb & Agrawal 1995) on a parent pair, applied
/// per variable with probability ½ inside an overall `crossover_prob`.
pub(crate) fn sbx(
    x1: &[f64],
    x2: &[f64],
    bounds: &[(f64, f64)],
    eta: f64,
    crossover_prob: f64,
    rng: &mut Rng,
) -> (Vec<f64>, Vec<f64>) {
    let mut c1 = x1.to_vec();
    let mut c2 = x2.to_vec();
    if rng.uniform() > crossover_prob {
        return (c1, c2);
    }
    for d in 0..bounds.len() {
        let (lo, hi) = bounds[d];
        // Cross this variable half the time, only if the parents differ.
        if rng.uniform() > 0.5 || (x1[d] - x2[d]).abs() < 1e-14 {
            continue;
        }
        let (y1, y2) = (x1[d].min(x2[d]), x1[d].max(x2[d]));
        let rand = rng.uniform();
        let exp = 1.0 / (eta + 1.0);

        // Lower child.
        let beta = 1.0 + 2.0 * (y1 - lo) / (y2 - y1);
        let alpha = 2.0 - beta.powf(-(eta + 1.0));
        let betaq = beta_q(rand, alpha, exp);
        let a = 0.5 * ((y1 + y2) - betaq * (y2 - y1));

        // Upper child.
        let beta = 1.0 + 2.0 * (hi - y2) / (y2 - y1);
        let alpha = 2.0 - beta.powf(-(eta + 1.0));
        let betaq = beta_q(rand, alpha, exp);
        let b = 0.5 * ((y1 + y2) + betaq * (y2 - y1));

        let (a, b) = (a.clamp(lo, hi), b.clamp(lo, hi));
        // Random assignment of the two children to the two offspring.
        if rng.uniform() <= 0.5 {
            c1[d] = b;
            c2[d] = a;
        } else {
            c1[d] = a;
            c2[d] = b;
        }
    }
    (c1, c2)
}

/// The spread factor `β_q` of the SBX scheme.
#[inline]
fn beta_q(rand: f64, alpha: f64, exp: f64) -> f64 {
    if rand <= 1.0 / alpha {
        (rand * alpha).powf(exp)
    } else {
        (1.0 / (2.0 - rand * alpha)).powf(exp)
    }
}

/// Polynomial mutation (Deb & Goyal 1996), per variable with probability `pm`.
pub(crate) fn mutate(x: &mut [f64], bounds: &[(f64, f64)], eta: f64, pm: f64, rng: &mut Rng) {
    for d in 0..bounds.len() {
        if rng.uniform() > pm {
            continue;
        }
        let (lo, hi) = bounds[d];
        let span = hi - lo;
        if span <= 0.0 {
            continue;
        }
        let y = x[d];
        let delta1 = (y - lo) / span;
        let delta2 = (hi - y) / span;
        let rnd = rng.uniform();
        let mut_pow = 1.0 / (eta + 1.0);
        let deltaq = if rnd <= 0.5 {
            let xy = 1.0 - delta1;
            let val = 2.0 * rnd + (1.0 - 2.0 * rnd) * xy.powf(eta + 1.0);
            val.powf(mut_pow) - 1.0
        } else {
            let xy = 1.0 - delta2;
            let val = 2.0 * (1.0 - rnd) + 2.0 * (rnd - 0.5) * xy.powf(eta + 1.0);
            1.0 - val.powf(mut_pow)
        };
        x[d] = (y + deltaq * span).clamp(lo, hi);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn dominance_is_correct() {
        assert!(dominates(&[1.0, 1.0], &[2.0, 2.0]));
        assert!(dominates(&[1.0, 2.0], &[1.0, 3.0])); // equal in one, better in other
        assert!(!dominates(&[1.0, 3.0], &[2.0, 1.0])); // trade-off, neither dominates
        assert!(!dominates(&[1.0, 1.0], &[1.0, 1.0])); // equal ⇒ no domination
    }

    #[test]
    fn non_dominated_sort_splits_fronts() {
        // Three points: A and B are mutually non-dominated (front 0); C is
        // dominated by both (front 1).
        let pop = vec![
            MultiSolution {
                x: vec![],
                objectives: vec![1.0, 3.0],
            },
            MultiSolution {
                x: vec![],
                objectives: vec![3.0, 1.0],
            },
            MultiSolution {
                x: vec![],
                objectives: vec![4.0, 4.0],
            },
        ];
        let fronts = fast_non_dominated_sort(&pop);
        assert_eq!(fronts.len(), 2);
        assert_eq!(fronts[0].len(), 2);
        assert_eq!(fronts[1], vec![2]);
    }
}