basin 1.10.0

// `!(a > b)` comparisons below are deliberate NaN-preserving ports of PRIMA's
// `~(a > b)` masks (`a <= b` would exclude NaN, changing tie-breaking); keep them.
#![allow(clippy::neg_cmp_op_on_partial_ord)]

//! COBYLA driver (PRIMA `cobylb`): the trust-region / ρ schedule loop.
//!
//! [`CobylaWork`] owns the simplex, the penalty parameter `cpen`, the
//! trust-region radius `delta`, the resolution `rho`, and the return filter.
//! [`CobylaWork::step`] runs **one** iteration of PRIMA's `do tr` loop (0, 1, or
//! 2 function evaluations) and reports a [`Transition`]; the basin
//! [`Executor`](crate::Executor) drives it and owns budget / `ρ`-floor
//! termination.

use crate::core::math::Scalar;

use super::filter::{moderatef, savefilt, selectx};
use super::geometry::{
    GeometryWork, assess_geo, geostep, setdrop_geo, setdrop_tr,
};
use super::init::initxfc;
use super::linalg::dot;
use super::model::ModelWork;
use super::trstlp::{TrstlpWork, trrad};
use super::update::{NO_DROP, UpdateWork, updatepole, updatexfc};

/// A point evaluator: maps `x` to its (raw) objective value and constraint
/// vector, propagating the problem's error type `E`.
pub(crate) type EvalFn<'a, F, E> =
    dyn FnMut(&[F]) -> Result<(F, Vec<F>), E> + 'a;

/// Outcome of one [`CobylaWork::step`].
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub(crate) enum Transition {
    /// Keep iterating.
    Continue,
    /// `ρ` was reduced this step (resolution refined).
    RhoReduced,
    /// `ρ` reached `ρ_end`: natural convergence.
    Converged,
    /// Damaging rounding broke the simplex bookkeeping; abort.
    Failed,
}

/// The penalty search may repole several times. Keep its trial simplex separate
/// from the live one so intermediate roundoff cannot alter the driver path.
struct PenaltyWork<F> {
    conmat: Vec<F>,
    cval: Vec<F>,
    fval: Vec<F>,
    sim: Vec<F>,
    simi: Vec<F>,
    model: ModelWork<F>,
    d: Vec<F>,
    valid: bool,
}

impl<F: Scalar> PenaltyWork<F> {
    fn new(n: usize, m: usize) -> Self {
        Self {
            conmat: vec![F::zero(); m * (n + 1)],
            cval: vec![F::zero(); n + 1],
            fval: vec![F::zero(); n + 1],
            sim: vec![F::zero(); n * (n + 1)],
            simi: vec![F::zero(); n * n],
            model: ModelWork::new(n, m),
            d: vec![F::zero(); n],
            valid: false,
        }
    }

    fn matches(&self, fval: &[F], conmat: &[F], simi: &[F]) -> bool {
        self.valid
            && [
                (&self.fval[..], fval),
                (&self.conmat[..], conmat),
                (&self.simi[..], simi),
            ]
            .into_iter()
            .all(|(a, b)| {
                a.iter().zip(b).all(|(&x, &y)| {
                    // Equality alone conflates signed zeros. NaNs invalidate reuse.
                    x == y && x.is_sign_negative() == y.is_sign_negative()
                })
            })
    }
}

/// Resumable COBYLA working state.
pub(crate) struct CobylaWork<F = f64> {
    n: usize,
    m: usize,
    sim: Vec<F>,    // n × (n+1)
    simi: Vec<F>,   // n × n
    fval: Vec<F>,   // n+1
    conmat: Vec<F>, // m × (n+1)
    cval: Vec<F>,   // n+1
    cpen: F,
    delta: F,
    rho: F,
    rho_end: F,
    eta1: F,
    eta2: F,
    gamma1: F,
    gamma2: F,
    gamma3: F,
    factor_alpha: F,
    factor_beta: F,
    factor_gamma: F,
    ctol: F,
    cweight: F,
    lp: TrstlpWork<F>,
    update: UpdateWork<F>,
    penalty: PenaltyWork<F>,
    geometry: GeometryWork<F>,
    // Return filter.
    maxfilt: usize,
    nfilt: usize,
    xfilt: Vec<F>,
    ffilt: Vec<F>,
    cfilt: Vec<F>,
}

impl<F: Scalar> CobylaWork<F> {
    /// Build and seed the simplex from the start point `x0`, returning the work
    /// plus the initial reported `(x, f)`. `eval(x) -> (f, constr)` evaluates
    /// the (raw) objective and constraint vector.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn try_init<E>(
        x0: Vec<F>,
        m: usize,
        rho_beg: F,
        rho_end: F,
        eval: &mut EvalFn<F, E>,
    ) -> Result<(Self, Vec<F>, F), E> {
        let n = x0.len();
        let zero = F::zero();

        // Evaluate the start point first (moderated).
        let (f0_raw, c0_raw) = eval(&x0)?;
        let f0 = moderatef(f0_raw);
        let constr0: Vec<F> = c0_raw.iter().map(|&v| moderatef(v)).collect();

        let out = initxfc(n, m, &x0, f0, &constr0, rho_beg, eval)?;

        let eta1 = F::from_f64(0.1).unwrap();
        let eta2 = F::from_f64(0.7).unwrap();
        let gamma1 = F::from_f64(0.5).unwrap();
        let gamma2 = F::from_f64(2.0).unwrap();
        let gamma3 = F::one().max(
            (F::from_f64(0.75).unwrap() * gamma2)
                .min(F::from_f64(1.5).unwrap()),
        );
        let ctol = F::epsilon().sqrt();
        let cweight = F::from_f64(1.0e8).unwrap();
        let cpenmin = F::epsilon();

        let cpen = cpenmin.max(F::from_f64(1.0e3).unwrap().min(fcratio(
            &out.conmat,
            &out.fval,
            n,
            m,
        )));

        let maxfilt = 2000usize;
        let mut work = Self {
            n,
            m,
            sim: out.sim,
            simi: out.simi,
            fval: out.fval,
            conmat: out.conmat,
            cval: out.cval,
            cpen,
            delta: rho_beg,
            rho: rho_beg,
            rho_end,
            eta1,
            eta2,
            gamma1,
            gamma2,
            gamma3,
            factor_alpha: F::from_f64(0.25).unwrap(),
            factor_beta: F::from_f64(2.1).unwrap(),
            factor_gamma: F::from_f64(0.5).unwrap(),
            ctol,
            cweight,
            lp: TrstlpWork::new(n, m),
            update: UpdateWork::new(n),
            penalty: PenaltyWork::new(n, m),
            geometry: GeometryWork::new(n),
            maxfilt,
            nfilt: 0,
            xfilt: Vec::new(),
            ffilt: Vec::new(),
            cfilt: Vec::new(),
        };

        // Seed the filter from the initial simplex vertices.
        for j in 0..=n {
            let mut x = vec![zero; n];
            for r in 0..n {
                x[r] = if j < n {
                    work.sim[r + j * n] + work.sim[r + n * n]
                } else {
                    work.sim[r + n * n]
                };
            }
            work.save_to_filter(&x, work.fval[j], work.cval[j]);
        }

        let (bx, bf) = work.best();
        Ok((work, bx, bf))
    }

    /// Current trust-region resolution `ρ`.
    pub(crate) fn rho(&self) -> F {
        self.rho
    }

    fn pole(&self) -> &[F] {
        &self.sim[self.n * self.n..]
    }

    fn save_to_filter(&mut self, x: &[F], f: F, cstrv: F) {
        savefilt(
            x,
            f,
            cstrv,
            self.n,
            self.ctol,
            self.cweight,
            self.maxfilt,
            &mut self.nfilt,
            &mut self.xfilt,
            &mut self.ffilt,
            &mut self.cfilt,
        );
    }

    /// The filter-selected incumbent `(x, f)` COBYLA would return now.
    pub(crate) fn best(&self) -> (Vec<F>, F) {
        let (x, f) = self.best_ref();
        (x.to_vec(), f)
    }

    /// Borrow the incumbent so the public solver can reuse its parameter buffer.
    pub(crate) fn best_ref(&self) -> (&[F], F) {
        let kopt = selectx(
            &self.ffilt[..self.nfilt],
            &self.cfilt[..self.nfilt],
            self.cpen.max(self.cweight),
            self.ctol,
        );
        let x = &self.xfilt[kopt * self.n..(kopt + 1) * self.n];
        (x, self.ffilt[kopt])
    }

    /// Evaluate `(f, constr, cstrv)` at `x`, moderated.
    fn eval_moderated<E>(
        eval: &mut EvalFn<F, E>,
        x: &[F],
    ) -> Result<(F, Vec<F>, F), E> {
        let (f_raw, mut constr) = eval(x)?;
        let f = moderatef(f_raw);
        for v in &mut constr {
            *v = moderatef(*v);
        }
        let cstrv = constr
            .iter()
            .cloned()
            .fold(F::zero(), F::max)
            .max(F::zero());
        Ok((f, constr, cstrv))
    }

    /// One iteration of PRIMA's `do tr` loop.
    pub(crate) fn step<E>(
        &mut self,
        eval: &mut EvalFn<F, E>,
    ) -> Result<Transition, E> {
        let n = self.n;
        let m = self.m;
        let zero = F::zero();
        let one = F::one();

        // (1) Increase CPEN if needed so PREREM > 0, then re-pole.
        self.cpen = self.get_cpen();
        if !updatepole(
            self.cpen,
            &mut self.conmat,
            &mut self.cval,
            &mut self.fval,
            &mut self.sim,
            &mut self.simi,
            n,
            m,
            &mut self.update,
        ) {
            return Ok(Transition::Failed);
        }

        let adequate_geo = assess_geo(
            self.delta,
            self.factor_alpha,
            self.factor_beta,
            &self.sim,
            &self.simi,
            n,
        );

        // (2) Linear models and the trust-region step.
        // Reuse the penalty search's LP only when its model inputs survived
        // repoling exactly. Multiple pole changes or inverse repair can differ.
        if !self.penalty.matches(&self.fval, &self.conmat, &self.simi) {
            self.penalty
                .model
                .build(&self.fval, &self.conmat, &self.simi);
            let model = &self.penalty.model;
            self.penalty.d.copy_from_slice(
                self.lp.solve(&model.a, &model.b, self.delta, &model.g),
            );
        }
        let g = &self.penalty.model.g;
        let a = &self.penalty.model.a;
        let d = self.penalty.d.clone();
        let dnorm = self.delta.min(dot(&d, &d).sqrt());
        let shortd = dnorm < F::from_f64(0.1).unwrap() * self.rho;

        let preref = -dot(&d, g);
        // prerec = cval[pole] − max(0, maxᵢ(conmat[i,pole] + (Aᵀd)[i])).
        let mut lin_cv = zero;
        for i in 0..m {
            let adi: F = (0..n).map(|l| d[l] * a[l + i * n]).sum();
            lin_cv = lin_cv.max(self.conmat[i + n * m] + adi);
        }
        lin_cv = lin_cv.max(zero);
        let prerec = self.cval[n] - lin_cv;
        let prerem = preref + self.cpen * prerec;
        let trfail = !(prerem
            > F::from_f64(1.0e-5).unwrap()
                * self.cpen.min(one)
                * self.rho
                * self.rho);

        let mut ratio = -one;
        let mut jdrop_tr = NO_DROP;

        if shortd || trfail {
            self.delta = F::from_f64(0.1).unwrap() * self.delta;
            if self.delta <= self.gamma3 * self.rho {
                self.delta = self.rho;
            }
        } else {
            let pole = self.pole();
            let x: Vec<F> = (0..n).map(|r| pole[r] + d[r]).collect();
            let (f, constr, cstrv) = Self::eval_moderated(eval, &x)?;
            self.save_to_filter(&x, f, cstrv);

            let actrem = (self.fval[n] + self.cpen * self.cval[n])
                - (f + self.cpen * cstrv);
            ratio = redrat(actrem, prerem, self.eta1);
            self.delta = trrad(
                self.delta,
                dnorm,
                self.eta1,
                self.eta2,
                self.gamma1,
                self.gamma2,
                ratio,
            );
            if self.delta <= self.gamma3 * self.rho {
                self.delta = self.rho;
            }
            let ximproved = actrem > zero;
            jdrop_tr = setdrop_tr(
                ximproved,
                &d,
                self.delta,
                self.rho,
                &self.sim,
                &self.simi,
                n,
                &mut self.geometry,
            );
            if !updatexfc(
                jdrop_tr,
                &constr,
                self.cpen,
                cstrv,
                &d,
                f,
                &mut self.conmat,
                &mut self.cval,
                &mut self.fval,
                &mut self.sim,
                &mut self.simi,
                n,
                m,
                &mut self.update,
            ) {
                return Ok(Transition::Failed);
            }
        }

        // (3) Decide geometry / ρ-reduction.
        let bad_trstep =
            shortd || trfail || ratio <= zero || jdrop_tr == NO_DROP;
        let improve_geo = bad_trstep && !adequate_geo;
        let reduce_rho =
            bad_trstep && adequate_geo && self.delta.max(dnorm) <= self.rho;

        if improve_geo
            && !assess_geo(
                self.delta,
                self.factor_alpha,
                self.factor_beta,
                &self.sim,
                &self.simi,
                n,
            )
        {
            let jdrop_geo = setdrop_geo(
                self.delta,
                self.factor_alpha,
                self.factor_beta,
                &self.sim,
                &self.simi,
                n,
            );
            if jdrop_geo == NO_DROP {
                return Ok(Transition::Failed);
            }
            let d = geostep(
                jdrop_geo,
                &self.conmat,
                self.cpen,
                self.delta,
                &self.fval,
                self.factor_gamma,
                &self.simi,
                n,
                m,
                &mut self.penalty.model,
            );
            let pole = self.pole();
            let x: Vec<F> = (0..n).map(|r| pole[r] + d[r]).collect();
            let (f, constr, cstrv) = Self::eval_moderated(eval, &x)?;
            self.save_to_filter(&x, f, cstrv);
            if !updatexfc(
                jdrop_geo,
                &constr,
                self.cpen,
                cstrv,
                &d,
                f,
                &mut self.conmat,
                &mut self.cval,
                &mut self.fval,
                &mut self.sim,
                &mut self.simi,
                n,
                m,
                &mut self.update,
            ) {
                return Ok(Transition::Failed);
            }
        }

        if reduce_rho {
            if self.rho <= self.rho_end {
                // Converged. Refine with the pending short step if it was never
                // evaluated (PRIMA's end-of-run trust-region tail).
                if shortd {
                    let pole = self.pole();
                    let x: Vec<F> = (0..n).map(|r| pole[r] + d[r]).collect();
                    let (f, _, cstrv) = Self::eval_moderated(eval, &x)?;
                    self.save_to_filter(&x, f, cstrv);
                }
                return Ok(Transition::Converged);
            }
            self.delta = (F::from_f64(0.5).unwrap() * self.rho)
                .max(redrho(self.rho, self.rho_end));
            self.rho = redrho(self.rho, self.rho_end);
            self.cpen = F::epsilon().max(self.cpen.min(fcratio(
                &self.conmat,
                &self.fval,
                n,
                m,
            )));
            if !updatepole(
                self.cpen,
                &mut self.conmat,
                &mut self.cval,
                &mut self.fval,
                &mut self.sim,
                &mut self.simi,
                n,
                m,
                &mut self.update,
            ) {
                return Ok(Transition::Failed);
            }
            return Ok(Transition::RhoReduced);
        }

        Ok(Transition::Continue)
    }

    /// Increase `cpen` so the predicted merit reduction is positive (PRIMA
    /// `getcpen`). Works on local copies of the simplex bookkeeping.
    fn get_cpen(&mut self) -> F {
        let n = self.n;
        let m = self.m;
        let zero = F::zero();
        let realmax = F::max_value();
        let two = F::from_f64(2.0).unwrap();

        let mut cpen = self.cpen;
        let PenaltyWork {
            conmat,
            cval,
            fval,
            sim,
            simi,
            model,
            d,
            valid,
        } = &mut self.penalty;
        conmat.copy_from_slice(&self.conmat);
        cval.copy_from_slice(&self.cval);
        fval.copy_from_slice(&self.fval);
        sim.copy_from_slice(&self.sim);
        simi.copy_from_slice(&self.simi);
        *valid = false;

        for _ in 0..(n + 1) {
            if !updatepole(
                cpen,
                conmat,
                cval,
                fval,
                sim,
                simi,
                n,
                m,
                &mut self.update,
            ) {
                break;
            }
            model.build(fval, conmat, simi);
            let a = &model.a;
            d.copy_from_slice(self.lp.solve(a, &model.b, self.delta, &model.g));
            *valid = true;
            let preref = -dot(d, &model.g);
            let mut lin_cv = zero;
            for i in 0..m {
                let adi: F = (0..n).map(|l| d[l] * a[l + i * n]).sum();
                lin_cv = lin_cv.max(conmat[i + n * m] + adi);
            }
            lin_cv = lin_cv.max(zero);
            let prerec = cval[n] - lin_cv;
            if !(prerec > zero && preref < zero) {
                break;
            }
            cpen = cpen.max((-two * (preref / prerec)).min(realmax));
            if super::update::findpole(cpen, cval, fval, n) == n {
                break;
            }
        }
        cpen
    }
}

/// Ratio between the typical change of `F` and that of the constraints
/// (Powell 1994 eq. 12–13). PRIMA `fcratio`.
fn fcratio<F: Scalar>(conmat: &[F], fval: &[F], n: usize, m: usize) -> F {
    let zero = F::zero();
    let half = F::from_f64(0.5).unwrap();
    let np = n + 1;
    let mut r = zero;
    if m == 0 {
        return r;
    }
    // cmin/cmax over -conmat (convert to Powell's c ≥ 0 form).
    let cmin: Vec<F> = (0..m)
        .map(|i| {
            (0..np)
                .map(|j| -conmat[i + j * m])
                .fold(F::infinity(), F::min)
        })
        .collect();
    let cmax: Vec<F> = (0..m)
        .map(|i| {
            (0..np)
                .map(|j| -conmat[i + j * m])
                .fold(F::neg_infinity(), F::max)
        })
        .collect();
    let fmin = fval.iter().cloned().fold(F::infinity(), F::min);
    let fmax = fval.iter().cloned().fold(F::neg_infinity(), F::max);
    let any = (0..m).any(|i| cmin[i] < half * cmax[i]);
    if any && fmin < fmax {
        let denom = (0..m)
            .filter(|&i| cmin[i] < half * cmax[i])
            .map(|i| cmax[i].max(zero) - cmin[i])
            .fold(F::infinity(), F::min);
        r = (fmax - fmin) / denom;
    }
    r
}

/// Trust-region reduction ratio with Inf/NaN handling (PRIMA `redrat`).
fn redrat<F: Scalar>(ared: F, pred: F, rshrink: F) -> F {
    let realmax = F::max_value();
    let half = F::from_f64(0.5).unwrap();
    if ared.is_nan() {
        -realmax
    } else if pred.is_nan() || pred <= F::zero() {
        if ared > F::zero() {
            half * rshrink
        } else {
            -realmax
        }
    } else if pred.is_infinite()
        && pred > F::zero()
        && ared.is_infinite()
        && ared > F::zero()
    {
        F::one()
    } else if pred.is_infinite()
        && pred > F::zero()
        && ared.is_infinite()
        && ared < F::zero()
    {
        -realmax
    } else {
        ared / pred
    }
}

/// Reduce `ρ` toward `ρ_end` (PRIMA `redrho`, shared schedule).
fn redrho<F: Scalar>(rho_in: F, rhoend: F) -> F {
    let tenth = F::from_f64(0.1).unwrap();
    let ratio = rho_in / rhoend;
    if ratio > F::from_f64(250.0).unwrap() {
        tenth * rho_in
    } else if ratio <= F::from_f64(16.0).unwrap() {
        rhoend
    } else {
        ratio.sqrt() * rhoend
    }
}