basin 3.0.0

use crate::core::math::{
    ComponentZip, NormInfinity, NormSquared, Scalar, ScaleInPlace,
};

pub(super) fn guarded_progress<F: Scalar>(
    progress: crate::Termination<F>,
    gradient: F,
    tolerance: F,
) -> crate::Termination<F> {
    use crate::{ConvergenceEvidence as E, ConvergenceTest as T, Measurement};
    let crate::Termination::Converged(progress) = progress else {
        unreachable!("only progress convergence can be guarded");
    };
    let criterion = &progress.criteria()[0];
    let (key, definition) = match criterion.test {
        T::RelativeModelReduction => (
            "lm.guarded_model_reduction",
            "Relative model reduction AND retained objective gradient infinity norm <= gradient_tolerance.",
        ),
        T::RelativeTrialStep => (
            "lm.guarded_trial_step",
            "Relative trial step AND retained objective gradient infinity norm <= gradient_tolerance.",
        ),
        T::RelativeTrustRadius => (
            "lm.guarded_trust_radius",
            "Relative trust radius AND retained objective gradient infinity norm <= gradient_tolerance.",
        ),
        _ => unreachable!("only relative progress tests can be guarded"),
    };
    let mut measurements = vec![
        Measurement {
            name: "gradient_inf".into(),
            value: gradient,
        },
        Measurement {
            name: "gradient_tolerance".into(),
            value: tolerance,
        },
    ];
    let mut add = |name: &str, value: F| {
        measurements.push(Measurement {
            name: name.into(),
            value,
        });
    };
    match &criterion.evidence {
        E::ModelReduction {
            actual,
            predicted,
            reference_cost,
            gain_ratio,
            tolerance,
        } => {
            add("actual", *actual);
            add("predicted", *predicted);
            add("reference_cost", *reference_cost);
            add("gain_ratio", *gain_ratio);
            add("progress_tolerance", *tolerance);
        }
        E::FactoredUpperBound {
            value,
            bound,
            tolerance,
            reference,
        } => {
            // Preserve exponents so the compound report cannot overflow norms.
            add("value_significand", value.significand);
            add("value_exponent", F::from_i32(value.exponent).unwrap());
            add("bound_significand", bound.significand);
            add("bound_exponent", F::from_i32(bound.exponent).unwrap());
            add("progress_tolerance", *tolerance);
            if let Some(reference) = reference {
                add("reference_significand", reference.significand);
                add(
                    "reference_exponent",
                    F::from_i32(reference.exponent).unwrap(),
                );
            }
        }
        _ => unreachable!(
            "native progress tests expose reduction or factored evidence"
        ),
    }
    crate::Termination::custom(key, definition, measurements)
}

pub(super) fn finite_unchanged_trial<V, F>(x: &V, trial: &V) -> bool
where
    F: Scalar,
    V: ComponentZip<F>,
{
    x.all_zip(trial, |base, proposed| base.is_finite() && base == proposed)
}

pub(super) fn all_finite<V, F>(v: &V) -> bool
where
    F: Scalar,
    V: ComponentZip<F>,
{
    // An infinity norm can hide NaNs on backends that reduce with scalar max.
    v.all_zip(v, |value, _| value.is_finite())
}

pub(super) fn orthogonality_converged<V, F>(
    g: &V,
    diagonal: &V,
    r: &V,
    tol: F,
) -> bool
where
    F: Scalar,
    V: Clone
        + NormInfinity<F>
        + NormSquared<F>
        + ScaleInPlace<F>
        + ComponentZip<F>,
{
    let gradient = g.norm_infinity();
    let Some((rscale, rnorm)) = norm_parts(r) else {
        return false;
    };
    if !gradient.is_finite() || !diagonal.norm_infinity().is_finite() {
        return false;
    }
    // Exact zero refers to the gradient, before any normalization can underflow.
    if gradient == F::zero() {
        return true;
    }
    if tol == F::zero() || rscale == F::zero() {
        return false;
    }
    // Compare each cosine in factored form. Even dividing by a column norm
    // first can erase a subnormal projection before residual normalization.
    g.all_zip(diagonal, |value, square| {
        if square < F::zero() || !square.is_finite() {
            return false;
        }
        if square == F::zero() {
            return value == F::zero();
        }
        product_le(&[value.abs()], &[tol, rscale, rnorm, square.sqrt()])
    })
}

// The corrected residual may be enormous when robust curvature is clipped.
// Normalize by the actual objective so that this cannot imply stationarity.
pub(super) fn robust_gradient_converged<V, F>(
    g: &V,
    diagonal: &V,
    cost: F,
    tol: F,
) -> bool
where
    F: Scalar,
    V: NormInfinity<F> + ComponentZip<F>,
{
    if !cost.is_finite() || cost < F::zero() || !g.norm_infinity().is_finite() {
        return false;
    }
    g.all_zip(diagonal, |value, square| {
        if !value.is_finite() || !square.is_finite() || square < F::zero() {
            return false;
        }
        if value == F::zero() {
            return true;
        }
        if tol == F::zero() || cost == F::zero() || square == F::zero() {
            return false;
        }
        product_le(
            &[value.abs()],
            &[
                tol,
                cost.sqrt(),
                F::from_f64(2.0).unwrap().sqrt(),
                square.sqrt(),
            ],
        )
    })
}

pub(super) fn relative_step_converged<V, F>(h: &V, x: &V, tol: F) -> bool
where
    F: Scalar,
    V: Clone + NormInfinity<F> + NormSquared<F> + ScaleInPlace<F>,
{
    let (Some((hscale, hnorm)), Some((xscale, xnorm))) =
        (norm_parts(h), norm_parts(x))
    else {
        return false;
    };
    if hscale == F::zero() {
        return true;
    }
    if tol == F::zero() || xscale == F::zero() {
        return false;
    }
    product_le(&[hscale, hnorm], &[tol, xscale, xnorm])
}

pub(super) fn relative_trust_radius_converged<V, F>(
    radius: F,
    x: &V,
    diagonal: &V,
    tol: F,
) -> bool
where
    F: Scalar,
    V: ComponentZip<F>,
{
    if !radius.is_finite()
        || radius < F::zero()
        || !tol.is_finite()
        || tol < F::zero()
    {
        return false;
    }
    let Some((exponent, sum)) = weighted_norm_parts(x, diagonal) else {
        return false;
    };
    if radius == F::zero() {
        return true;
    }
    if tol == F::zero() || sum == F::zero() {
        return false;
    }
    // Keep the norm's exponent separate even when the norm itself cannot fit.
    let (power, fraction) = product_parts(&[tol, sum.sqrt()]);
    product_parts(&[radius]) <= (power + exponent, fraction)
}

fn weighted_norm_parts<V, F>(x: &V, diagonal: &V) -> Option<(i32, F)>
where
    F: Scalar,
    V: ComponentZip<F>,
{
    let two = F::from_f64(2.).unwrap();
    let mut exponent = 0;
    let mut sum = F::zero();
    // Factor each weighted coordinate before summing. Separate normalizations
    // of x and D can erase every term when their large entries do not coincide.
    let finite = x.all_zip(diagonal, |value, square| {
        if !value.is_finite() || !square.is_finite() || square < F::zero() {
            return false;
        }
        if value == F::zero() || square == F::zero() {
            return true;
        }
        let (power, fraction) = product_parts(&[value.abs(), square.sqrt()]);
        if sum == F::zero() {
            exponent = power;
            sum = fraction * fraction;
        } else if power > exponent {
            let factor = two.powi(exponent - power);
            sum = (sum * factor) * factor + fraction * fraction;
            exponent = power;
        } else {
            let term = fraction * two.powi(power - exponent);
            sum = sum + term * term;
        }
        true
    });
    if !finite {
        return None;
    }
    Some((exponent, sum))
}

// Keep the norm factored: even an unsquared norm can exceed the scalar range.
fn norm_parts<V, F>(v: &V) -> Option<(F, F)>
where
    F: Scalar,
    V: Clone + NormInfinity<F> + NormSquared<F> + ScaleInPlace<F>,
{
    let scale = v.norm_infinity();
    if !scale.is_finite() {
        return None;
    }
    if scale == F::zero() {
        return Some((F::zero(), F::zero()));
    }
    let mut normalized = v.clone();
    let reciprocal = F::one() / scale;
    if reciprocal.is_finite() {
        normalized.scale_in_place(reciprocal);
    } else {
        // The reciprocal of a subnormal scale may not fit, but these factors do.
        let factor = F::one() / scale.sqrt();
        normalized.scale_in_place(factor);
        normalized.scale_in_place(factor);
    }
    Some((scale, normalized.norm_squared().sqrt()))
}

// Compare nonnegative finite products without overflowing or rounding a tiny
// tolerance to zero. Binary exponents remain integers until after comparison.
fn product_le<F: Scalar>(left: &[F], right: &[F]) -> bool {
    crate::BinaryValue::product_le(left, right)
}

fn product_parts<F: Scalar>(factors: &[F]) -> (i32, F) {
    let value = crate::BinaryValue::product(factors);
    (value.exponent, value.significand)
}

fn binary_product<F: Scalar>(factors: &[F]) -> crate::BinaryValue<F> {
    crate::BinaryValue::product(factors)
}

pub(super) fn step_evidence<V, F>(
    h: &V,
    x: &V,
    tolerance: F,
) -> crate::ConvergenceEvidence<F>
where
    F: Scalar,
    V: Clone + NormInfinity<F> + NormSquared<F> + ScaleInPlace<F>,
{
    try_step_evidence(h, x, tolerance)
        .expect("a passing predicate has finite evidence")
}

pub(super) fn try_step_evidence<V, F>(
    h: &V,
    x: &V,
    tolerance: F,
) -> Option<crate::ConvergenceEvidence<F>>
where
    F: Scalar,
    V: Clone + NormInfinity<F> + NormSquared<F> + ScaleInPlace<F>,
{
    let (hs, hn) = norm_parts(h)?;
    let (xs, xn) = norm_parts(x)?;
    if [hs, hn, xs, xn, tolerance]
        .iter()
        .any(|value| !value.is_finite() || *value < F::zero())
    {
        return None;
    }
    Some(crate::ConvergenceEvidence::FactoredUpperBound {
        value: binary_product(&[hs, hn]),
        bound: binary_product(&[tolerance, xs, xn]),
        tolerance,
        reference: Some(binary_product(&[xs, xn])),
    })
}

pub(super) fn radius_evidence<V, F>(
    radius: F,
    x: &V,
    d: &V,
    tolerance: F,
) -> crate::ConvergenceEvidence<F>
where
    F: Scalar,
    V: ComponentZip<F>,
{
    try_radius_evidence(radius, x, d, tolerance)
        .expect("a passing predicate has finite evidence")
}

pub(super) fn try_radius_evidence<V, F>(
    radius: F,
    x: &V,
    d: &V,
    tolerance: F,
) -> Option<crate::ConvergenceEvidence<F>>
where
    F: Scalar,
    V: ComponentZip<F>,
{
    if !radius.is_finite()
        || radius < F::zero()
        || !tolerance.is_finite()
        || tolerance < F::zero()
    {
        return None;
    }
    let (exponent, sum) = weighted_norm_parts(x, d)?;
    let mut bound = binary_product(&[tolerance, sum.sqrt()]);
    bound.exponent += exponent;
    Some(crate::ConvergenceEvidence::FactoredUpperBound {
        value: binary_product(&[radius]),
        bound,
        tolerance,
        reference: Some(crate::BinaryValue {
            significand: sum.sqrt(),
            exponent,
        }),
    })
}

pub(super) fn orthogonality_evidence<V, F>(
    g: &V,
    d: &V,
    r: &V,
    robust_cost: Option<F>,
    tolerance: F,
) -> crate::ConvergenceEvidence<F>
where
    F: Scalar,
    V: Clone
        + NormInfinity<F>
        + NormSquared<F>
        + ScaleInPlace<F>
        + ComponentZip<F>,
{
    try_orthogonality_evidence(g, d, r, robust_cost, tolerance)
        .expect("a passing predicate has finite evidence")
}

pub(super) fn try_orthogonality_evidence<V, F>(
    g: &V,
    d: &V,
    r: &V,
    robust_cost: Option<F>,
    tolerance: F,
) -> Option<crate::ConvergenceEvidence<F>>
where
    F: Scalar,
    V: Clone
        + NormInfinity<F>
        + NormSquared<F>
        + ScaleInPlace<F>
        + ComponentZip<F>,
{
    let (scale, norm) = if let Some(cost) = robust_cost {
        if !cost.is_finite() || cost < F::zero() {
            return None;
        }
        (cost.sqrt(), F::from_f64(2.).unwrap().sqrt())
    } else {
        norm_parts(r)?
    };
    if [scale, norm, tolerance]
        .iter()
        .any(|value| !value.is_finite() || *value < F::zero())
    {
        return None;
    }
    if !g.all_zip(d, |value, square| {
        value.is_finite()
            && square.is_finite()
            && square >= F::zero()
            && (value == F::zero()
                || (scale > F::zero()
                    && norm > F::zero()
                    && square > F::zero()))
    }) {
        return None;
    }
    let mut maximum = crate::BinaryValue {
        significand: F::zero(),
        exponent: 0,
    };
    g.all_zip(d, |value, square| {
        if value != F::zero() {
            let (np, nf) = product_parts(&[value.abs()]);
            let (dp, df) = product_parts(&[scale, norm, square.sqrt()]);
            let (rp, rf) = product_parts(&[nf / df]);
            let exponent = np - dp + rp;
            if maximum.significand == F::zero()
                || (exponent, rf) > (maximum.exponent, maximum.significand)
            {
                maximum = crate::BinaryValue {
                    significand: rf,
                    exponent,
                };
            }
        }
        true
    });
    Some(crate::ConvergenceEvidence::FactoredUpperBound {
        value: maximum,
        bound: binary_product(&[tolerance]),
        tolerance,
        reference: Some(binary_product(&[scale, norm])),
    })
}

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

    #[test]
    fn failing_diagnostics_keep_finite_evidence_and_omit_invalid_operands() {
        assert!(try_step_evidence(&vec![1.], &vec![2.], 0.1).is_some());
        assert!(try_step_evidence(&vec![f64::NAN], &vec![2.], 0.1).is_none());
        assert!(
            try_radius_evidence(f64::INFINITY, &vec![2.], &vec![1.], 0.1)
                .is_none()
        );
        assert!(
            try_orthogonality_evidence(
                &vec![1.],
                &vec![0.],
                &vec![1.],
                None,
                0.1
            )
            .is_none()
        );
        assert!(
            try_orthogonality_evidence(
                &vec![1.],
                &vec![1.],
                &vec![1.],
                None,
                0.1
            )
            .is_some()
        );
    }

    macro_rules! extreme_weighted_norms {
        ($name:ident, $scalar:ty, $power:expr) => {
            #[test]
            fn $name() {
                let two: $scalar = 2.;
                let large = two.powi($power);
                let small = two.powi(-$power);
                assert!(relative_trust_radius_converged(
                    5.,
                    &vec![3., -4.],
                    &vec![4., 4.],
                    0.5
                ));
                assert!(!relative_trust_radius_converged(
                    6.,
                    &vec![3., -4.],
                    &vec![4., 4.],
                    0.5
                ));
                // The weighted norm itself overflows or underflows, but
                // multiplying by the tolerance gives a representable bound.
                for (x, d, tol, bound) in [
                    (large, large, small, two.powi($power / 2)),
                    (small, small, large, two.powi(-$power / 2)),
                ] {
                    assert!(relative_trust_radius_converged(
                        bound,
                        &vec![x],
                        &vec![d],
                        tol
                    ));
                    assert!(!relative_trust_radius_converged(
                        two * bound,
                        &vec![x],
                        &vec![d],
                        tol
                    ));
                }
                // Independently normalizing x and D would lose both terms.
                let x = vec![large, small];
                let d = vec![small, large];
                let bound = two.powi($power / 2);
                assert!(relative_trust_radius_converged(
                    bound / two,
                    &x,
                    &d,
                    1.
                ));
                assert!(!relative_trust_radius_converged(
                    two * bound,
                    &x,
                    &d,
                    1.
                ));
                let tiny = <$scalar>::from_bits(1);
                assert!(relative_trust_radius_converged(
                    tiny,
                    &vec![tiny],
                    &vec![1.],
                    1.
                ));
                assert!(!relative_trust_radius_converged(
                    tiny,
                    &vec![large],
                    &vec![large],
                    0.
                ));
                assert!(relative_trust_radius_converged(
                    0.,
                    &vec![0.],
                    &vec![1.],
                    0.
                ));
                assert!(!relative_trust_radius_converged(
                    tiny,
                    &vec![0.],
                    &vec![1.],
                    large
                ));
                assert!(!relative_trust_radius_converged(
                    tiny,
                    &vec![large],
                    &vec![0.],
                    large
                ));
                for bad in [
                    <$scalar>::NAN,
                    <$scalar>::INFINITY,
                    <$scalar>::NEG_INFINITY,
                ] {
                    assert!(!relative_trust_radius_converged(
                        bad,
                        &vec![1.],
                        &vec![1.],
                        1.
                    ));
                    assert!(!relative_trust_radius_converged(
                        0.,
                        &vec![1., bad],
                        &vec![1., 1.],
                        0.
                    ));
                    assert!(!relative_trust_radius_converged(
                        0.,
                        &vec![1., 1.],
                        &vec![1., bad],
                        0.
                    ));
                    assert!(!relative_trust_radius_converged(
                        0.,
                        &vec![1.],
                        &vec![1.],
                        bad
                    ));
                }
                assert!(!relative_trust_radius_converged(
                    -1.,
                    &vec![1.],
                    &vec![1.],
                    1.
                ));
                assert!(!relative_trust_radius_converged(
                    0.,
                    &vec![1.],
                    &vec![-1.],
                    1.
                ));
                assert!(!relative_trust_radius_converged(
                    0.,
                    &vec![1.],
                    &vec![1.],
                    -1.
                ));
            }
        };
    }
    extreme_weighted_norms!(f32_extreme_weighted_norms, f32, 100);
    extreme_weighted_norms!(f64_extreme_weighted_norms, f64, 800);

    macro_rules! extreme_norms {
        ($name:ident, $scalar:ty) => {
            #[test]
            fn $name() {
                let tiny = <$scalar>::from_bits(1);
                let huge = <$scalar>::MAX;
                for scale in [tiny, huge] {
                    let h = vec![scale, scale];
                    let x = vec![scale, 0.];
                    assert!(!relative_step_converged(&h, &x, 1.));
                    assert!(relative_step_converged(&h, &x, 2.));
                    assert!(!relative_step_converged(&h, &x, 0.));
                    assert!(relative_step_converged(&vec![0., 0.], &x, 0.));
                }
                // The ratio underflows, but a nonzero step still fails exact zero.
                assert!(!relative_step_converged(&vec![tiny], &vec![huge], 0.));
                assert!(relative_step_converged(
                    &vec![tiny],
                    &vec![huge],
                    tiny
                ));
                // The right-hand norm overflows, but the finite bound is below one.
                assert!(!relative_step_converged(
                    &vec![1., 0.],
                    &vec![huge, huge],
                    tiny
                ));
            }
        };
    }
    extreme_norms!(f64_extreme_norms, f64);
    extreme_norms!(f32_extreme_norms, f32);

    #[test]
    fn nonfinite_inputs_cannot_satisfy_stopping_checks() {
        for value in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
            let bad = vec![value, 0.];
            let zero = vec![0., 0.];
            let finite = vec![1., 1.];
            for tolerance in [0., 1., f64::MAX] {
                assert!(!relative_step_converged(&bad, &finite, tolerance));
                assert!(!relative_step_converged(&zero, &bad, tolerance));
                assert!(!orthogonality_converged(
                    &bad, &finite, &finite, tolerance
                ));
                assert!(!orthogonality_converged(
                    &zero, &bad, &finite, tolerance
                ));
                assert!(!orthogonality_converged(
                    &zero, &finite, &bad, tolerance
                ));
            }
        }
    }

    #[test]
    fn zero_columns_and_exact_orthogonality() {
        let diagonal = vec![1., 0.];
        assert!(!orthogonality_converged(
            &vec![1., 0.],
            &diagonal,
            &vec![1., 0.],
            0.5
        ));
        assert!(orthogonality_converged(
            &vec![1., 0.],
            &diagonal,
            &vec![1., 0.],
            1.
        ));
        assert!(orthogonality_converged(
            &vec![0., 0.],
            &diagonal,
            &vec![0., 1.],
            0.
        ));
        assert!(!orthogonality_converged(
            &vec![1e-200, 0.],
            &vec![1e200, 0.],
            &vec![1e100, 1.],
            0.
        ));
    }

    #[test]
    fn subnormal_projection_does_not_become_orthogonal() {
        // J = [2, -1] and r = [tiny, tiny] give g = tiny and cosine 1/sqrt(10).
        let tiny = f64::from_bits(1);
        assert!(!orthogonality_converged(
            &vec![tiny],
            &vec![5.],
            &vec![tiny, tiny],
            0.25
        ));
        assert!(orthogonality_converged(
            &vec![tiny],
            &vec![5.],
            &vec![tiny, tiny],
            0.5
        ));
    }
}