metalforge 0.1.1

//! NSGA-III — reference-point many-objective optimization (Deb & Jain 2014,
//! IEEE Trans. Evol. Comput. 18(4), 577–601).
//!
//! NSGA-III shares NSGA-II's elitist skeleton — fast non-dominated sorting, SBX
//! crossover, polynomial mutation, parent+offspring environmental selection —
//! but replaces crowding distance, which degrades past ~3 objectives, with
//! diversity preservation against a fixed set of **reference points** spread on
//! the unit simplex (Das–Dennis). After sorting, the splitting front is filled
//! by *niching*: objectives are adaptively normalized (ideal point + hyperplane
//! intercepts), each candidate is associated to its nearest reference line, and
//! members are picked to even out the per-reference niche counts. This keeps a
//! well-distributed front in 3+ objectives where NSGA-II clusters.
//!
//! Returns a [`ParetoFront`]. Deterministic for a given seed.

use super::nsga2::{fast_non_dominated_sort, mutate, sbx};
use crate::problem::MultiProblem;
use crate::rng::Rng;
use crate::solution::{MultiSolution, ParetoFront};
use crate::termination::Termination;

/// NSGA-III configuration.
#[derive(Debug, Clone, Copy)]
pub struct NsgaIII {
    /// Population size (forced even and ≥ 2).
    pub pop_size: usize,
    /// Das–Dennis divisions `p` for the reference-point set. The number of
    /// reference points is `C(M + p − 1, p)` for `M` objectives.
    pub divisions: usize,
    /// SBX distribution index `η_c`.
    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`.
    pub mutation_prob: Option<f64>,
    /// RNG seed; same seed + same problem + same budget ⇒ same result.
    pub seed: u64,
}

impl Default for NsgaIII {
    fn default() -> Self {
        NsgaIII {
            pop_size: 92,
            divisions: 12,
            crossover_eta: 30.0,
            crossover_prob: 1.0,
            mutation_eta: 20.0,
            mutation_prob: None,
            seed: 42,
        }
    }
}

impl NsgaIII {
    /// Approximates the Pareto front of `problem` within the evaluation budget
    /// of `term` (its `target` is ignored).
    pub fn optimize(&self, problem: &dyn MultiProblem) -> ParetoFront {
        self.optimize_within(problem, &Termination::budget(self.pop_size * 250))
    }

    /// Like [`NsgaIII::optimize`] but with an explicit evaluation budget.
    pub fn optimize_within(&self, problem: &dyn MultiProblem, term: &Termination) -> ParetoFront {
        let bounds = problem.bounds();
        let dim = bounds.len();
        let m = problem.n_objectives();
        let n = self.pop_size.max(2) + (self.pop_size & 1);
        let pm = self.mutation_prob.unwrap_or(1.0 / dim as f64);
        let mut rng = Rng::new(self.seed);
        let refs = das_dennis(m, self.divisions.max(1));

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

        let mut pop: Vec<MultiSolution> = Vec::with_capacity(n);
        let mut evaluations = 0;
        for _ in 0..n {
            let x: Vec<f64> = bounds
                .iter()
                .map(|&(lo, hi)| rng.uniform_in(lo, hi))
                .collect();
            let objectives = eval(&x);
            evaluations += 1;
            pop.push(MultiSolution { x, objectives });
        }

        while evaluations < term.max_evaluations {
            // Random mating (selection pressure comes from environmental step).
            let mut offspring: Vec<MultiSolution> = Vec::with_capacity(n);
            while offspring.len() < n {
                let a = rng.index(n);
                let b = rng.index(n);
                let (mut c1, mut c2) = sbx(
                    &pop[a].x,
                    &pop[b].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 {
                    let o2 = eval(&c2);
                    evaluations += 1;
                    offspring.push(MultiSolution {
                        x: c2,
                        objectives: o2,
                    });
                }
            }

            let mut union = pop;
            union.extend(offspring);
            pop = environmental_selection(union, n, m, &refs, &mut rng);
        }

        let fronts = fast_non_dominated_sort(&pop);
        let solutions = fronts[0].iter().map(|&i| pop[i].clone()).collect();
        ParetoFront {
            solutions,
            evaluations,
        }
    }
}

/// Reference-point niching selection of the best `n` from `union`.
fn environmental_selection(
    union: Vec<MultiSolution>,
    n: usize,
    m: usize,
    refs: &[Vec<f64>],
    rng: &mut Rng,
) -> Vec<MultiSolution> {
    let fronts = fast_non_dominated_sort(&union);

    // Accumulate whole fronts until the next would overflow; that one splits.
    let mut complete: Vec<usize> = Vec::new();
    let mut last: Vec<usize> = Vec::new();
    for front in fronts {
        if complete.len() + front.len() <= n {
            complete.extend(front);
        } else {
            last = front;
            break;
        }
    }
    if complete.len() == n || last.is_empty() {
        return take(&union, &complete);
    }

    // St = complete ∪ last; normalize objectives and associate to references.
    let n_complete = complete.len();
    let st: Vec<usize> = complete
        .iter()
        .copied()
        .chain(last.iter().copied())
        .collect();
    let normalized = normalize(&union, &st, m);
    let (assoc_ref, assoc_dist) = associate(&normalized, refs);

    // Niche counts from the already-selected (complete) members.
    let mut niche = vec![0usize; refs.len()];
    for &r in assoc_ref.iter().take(n_complete) {
        niche[r] += 1;
    }

    // Iteratively fill from the splitting front, evening out niche counts.
    let need = n - n_complete;
    let mut chosen = vec![false; st.len()]; // over st positions; only last range used
    let mut excluded = vec![false; refs.len()];
    let mut picked = 0;
    while picked < need {
        // Reference with the smallest niche count among non-excluded.
        let j = (0..refs.len())
            .filter(|&r| !excluded[r])
            .min_by_key(|&r| niche[r])
            .expect("a reference remains while members are needed");

        // Candidates in the splitting front associated to j, not yet chosen.
        let mut best_pos: Option<usize> = None;
        let mut best_dist = f64::INFINITY;
        let mut count = 0usize;
        for pos in n_complete..st.len() {
            if chosen[pos] || assoc_ref[pos] != j {
                continue;
            }
            count += 1;
            if assoc_dist[pos] < best_dist {
                best_dist = assoc_dist[pos];
                best_pos = Some(pos);
            }
        }

        match best_pos {
            None => excluded[j] = true, // no candidates for this reference
            Some(nearest) => {
                // Empty niche → take the nearest; otherwise a random candidate.
                let pos = if niche[j] == 0 {
                    nearest
                } else {
                    let pick = rng.index(count);
                    (n_complete..st.len())
                        .filter(|&p| !chosen[p] && assoc_ref[p] == j)
                        .nth(pick)
                        .unwrap()
                };
                chosen[pos] = true;
                niche[j] += 1;
                picked += 1;
            }
        }
    }

    let mut result = take(&union, &complete);
    for (pos, &is_chosen) in chosen.iter().enumerate() {
        if is_chosen {
            result.push(union[st[pos]].clone());
        }
    }
    result
}

/// Clones the `union` members at the given indices.
fn take(union: &[MultiSolution], idx: &[usize]) -> Vec<MultiSolution> {
    idx.iter().map(|&i| union[i].clone()).collect()
}

/// Adaptive normalization (Deb & Jain 2014, §IV-A): translate by the ideal
/// point, find the per-axis extreme points by an achievement scalarizing
/// function, fit the hyperplane intercepts, and scale each objective into
/// `[0, ~1]`. Returns the normalized objective vector for each `st` member.
fn normalize(union: &[MultiSolution], st: &[usize], m: usize) -> Vec<Vec<f64>> {
    // Ideal point (per-objective minimum over St).
    let mut ideal = vec![f64::INFINITY; m];
    for &i in st {
        for (idl, &o) in ideal.iter_mut().zip(&union[i].objectives) {
            *idl = idl.min(o);
        }
    }
    let trans: Vec<Vec<f64>> = st
        .iter()
        .map(|&i| (0..m).map(|j| union[i].objectives[j] - ideal[j]).collect())
        .collect();

    // Extreme point per axis: minimizes ASF with a near-axis weight vector.
    let mut extreme = vec![0usize; m];
    for (j, e) in extreme.iter_mut().enumerate() {
        let mut best = f64::INFINITY;
        for (p, t) in trans.iter().enumerate() {
            let asf = (0..m)
                .map(|d| {
                    let w = if d == j { 1.0 } else { 1e-6 };
                    t[d] / w
                })
                .fold(f64::NEG_INFINITY, f64::max);
            if asf < best {
                best = asf;
                *e = p;
            }
        }
    }

    // Intercepts a_j from the plane through the M extreme points:
    // for each extreme point e, sum_j e_j * b_j = 1, then a_j = 1/b_j.
    let mat: Vec<Vec<f64>> = extreme.iter().map(|&p| trans[p].clone()).collect();
    let intercepts = gaussian_solve(&mat, &vec![1.0; m])
        .map(|b| {
            b.iter()
                .map(|&bj| if bj.abs() > 1e-12 { 1.0 / bj } else { f64::NAN })
                .collect::<Vec<_>>()
        })
        .filter(|a: &Vec<f64>| a.iter().all(|&v| v.is_finite() && v > 1e-9));

    // Fallback: per-axis maximum of the translated objectives.
    let intercepts = intercepts.unwrap_or_else(|| {
        (0..m)
            .map(|j| {
                trans
                    .iter()
                    .map(|t| t[j])
                    .fold(f64::NEG_INFINITY, f64::max)
                    .max(1e-9)
            })
            .collect()
    });

    trans
        .iter()
        .map(|t| (0..m).map(|j| t[j] / intercepts[j]).collect())
        .collect()
}

/// Associates each normalized member to its nearest reference line (smallest
/// perpendicular distance), returning the reference index and distance per
/// member.
fn associate(normalized: &[Vec<f64>], refs: &[Vec<f64>]) -> (Vec<usize>, Vec<f64>) {
    // Precompute reference directions' squared norms.
    let ref_norm2: Vec<f64> = refs.iter().map(|r| r.iter().map(|v| v * v).sum()).collect();
    let mut idx = vec![0usize; normalized.len()];
    let mut dist = vec![0.0f64; normalized.len()];
    for (s, point) in normalized.iter().enumerate() {
        let mut best = f64::INFINITY;
        let mut best_r = 0;
        for (r, refp) in refs.iter().enumerate() {
            // Perpendicular distance from `point` to the line through `refp`.
            let dot: f64 = point.iter().zip(refp).map(|(a, b)| a * b).sum();
            let t = dot / ref_norm2[r].max(1e-12);
            let d2: f64 = point
                .iter()
                .zip(refp)
                .map(|(a, b)| {
                    let proj = a - t * b;
                    proj * proj
                })
                .sum();
            if d2 < best {
                best = d2;
                best_r = r;
            }
        }
        idx[s] = best_r;
        dist[s] = best.sqrt();
    }
    (idx, dist)
}

/// Das–Dennis structured reference points: all points on the unit simplex with
/// coordinates that are multiples of `1/p` (sum to 1).
fn das_dennis(m: usize, p: usize) -> Vec<Vec<f64>> {
    let mut out = Vec::new();
    let mut cur = vec![0usize; m];
    das_dennis_rec(0, p, m, p, &mut cur, &mut out);
    out
}

fn das_dennis_rec(
    pos: usize,
    left: usize,
    m: usize,
    p: usize,
    cur: &mut [usize],
    out: &mut Vec<Vec<f64>>,
) {
    if pos == m - 1 {
        cur[pos] = left;
        out.push(cur.iter().map(|&v| v as f64 / p as f64).collect());
        return;
    }
    for i in 0..=left {
        cur[pos] = i;
        das_dennis_rec(pos + 1, left - i, m, p, cur, out);
    }
}

/// Solves the linear system `A x = b` by Gaussian elimination with partial
/// pivoting. Returns `None` if `A` is (numerically) singular.
#[allow(clippy::needless_range_loop)] // explicit row/col indices read clearer here
fn gaussian_solve(a: &[Vec<f64>], b: &[f64]) -> Option<Vec<f64>> {
    let n = b.len();
    let mut m: Vec<Vec<f64>> = (0..n)
        .map(|i| a[i].iter().copied().chain(std::iter::once(b[i])).collect())
        .collect();
    for col in 0..n {
        // Partial pivot.
        let mut piv = col;
        for r in col + 1..n {
            if m[r][col].abs() > m[piv][col].abs() {
                piv = r;
            }
        }
        if m[piv][col].abs() < 1e-12 {
            return None;
        }
        m.swap(col, piv);
        let d = m[col][col];
        for r in 0..n {
            if r == col {
                continue;
            }
            let factor = m[r][col] / d;
            for c in col..=n {
                m[r][c] -= factor * m[col][c];
            }
        }
    }
    Some((0..n).map(|i| m[i][n] / m[i][i]).collect())
}

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

    #[test]
    fn das_dennis_counts_and_sums() {
        // M=3, p=4 ⇒ C(6,2) = 15 points, each summing to 1.
        let pts = das_dennis(3, 4);
        assert_eq!(pts.len(), 15);
        for pt in &pts {
            assert!((pt.iter().sum::<f64>() - 1.0).abs() < 1e-9);
        }
        // M=2, p=12 ⇒ 13 points.
        assert_eq!(das_dennis(2, 12).len(), 13);
    }

    #[test]
    fn gaussian_solve_identity_and_singular() {
        let a = vec![vec![2.0, 0.0], vec![0.0, 4.0]];
        let x = gaussian_solve(&a, &[2.0, 4.0]).unwrap();
        assert!((x[0] - 1.0).abs() < 1e-9 && (x[1] - 1.0).abs() < 1e-9);
        let singular = vec![vec![1.0, 1.0], vec![1.0, 1.0]];
        assert!(gaussian_solve(&singular, &[1.0, 1.0]).is_none());
    }
}