mahalanobis.h

// Copyright 2026 Julien Michot.
// SPDX-License-Identifier: Apache-2.0
 
#pragma once
 
#include <tinyopt/math.h>
#include <tinyopt/traits.h>
 
namespace tinyopt::losses {


template <typename T, typename Cov_t, typename ExportJ = std::nullptr_t>
auto MahaSquaredNorm(const T &x, const Cov_t &cov_or_var, const ExportJ &Jx_or_bool = nullptr,
                     bool add_scale = true) {
  using Scalar = typename traits::params_trait<T>::Scalar;
  constexpr Index Dims = traits::params_trait<T>::Dims;
  if constexpr (traits::is_scalar_v<T>) {  // scalar
    const T s = cov_or_var < FloatEpsilon<T>() ? T(1.0) : T(T(1.0) / cov_or_var);
    const T n2 = x * x * s;  // same as sqrt(x²*s)
    if constexpr (traits::is_bool_v<ExportJ>) {
      const T J = add_scale ? Scalar(2) * s * x : s * x;
      return std::make_pair(n2, J);
    } else if constexpr (!traits::is_nullptr_v<ExportJ>) {
      const T J = add_scale ? Scalar(2) * s * x : s * x;
      return std::make_pair(n2, T(J * Jx_or_bool));
    } else {
      return n2;
    }
  } else if constexpr (traits::is_nullptr_v<ExportJ>) {  // x is a vector, no jacobian export
    const auto cov_vars = cov_or_var.template cast<Scalar>().eval();
    if (x.cols() != 1) throw std::invalid_argument("'x' must be a vector");
    Scalar n2 = Scalar(0.0);
    if constexpr (traits::params_trait<Cov_t>::ColsAtCompileTime == 1) {  // variances
      assert(x.size() == cov_vars.size());
      n2 = x.dot(cov_vars.cwiseInverse().asDiagonal() * x);
    } else {  // covariance matrix
      if (cov_vars.cols() > 1) {
        const auto I = InvCov(cov_vars);
        if (!I.has_value())
          throw std::invalid_argument("Covariance is not invertible, make sure it is");
        n2 = x.dot(I.value() * x);
      } else if constexpr (traits::params_trait<Cov_t>::ColsAtCompileTime ==
                           Dynamic) {  // variances
        assert(x.size() == cov_vars.size());
        n2 = x.dot(cov_vars.cwiseInverse().asDiagonal() * x);
      }
    }
    return n2;
  } else {  // x is a vector, exporting jacobian
    if (x.cols() != 1) throw std::invalid_argument("'x' must be a vector");
    const auto cov_vars = cov_or_var.template cast<Scalar>().eval();
    Scalar n2 = Scalar(0.0);
    Vector<Scalar, Dims> Jt(x.size());
 
    if constexpr (traits::params_trait<Cov_t>::ColsAtCompileTime == 1) {  // variances
      assert(x.size() == cov_vars.size());
      Jt = cov_vars.cwiseInverse().asDiagonal() * x;
      n2 = x.dot(Jt);
    } else {  // covariance matrix
      if (cov_vars.cols() > 1) {
        const auto I = InvCov(cov_vars);
        if (!I.has_value())
          throw std::invalid_argument("Covariance is not invertible, make sure it is");
        Jt = I.value() * x;
        n2 = x.dot(Jt);
      } else if constexpr (traits::params_trait<Cov_t>::ColsAtCompileTime ==
                           Dynamic) {  // variances
        Jt = cov_vars.cwiseInverse().asDiagonal() * x;
        n2 = x.dot(Jt);
      }
    }
    if (add_scale) Jt *= Scalar(2);
    if constexpr (traits::is_bool_v<ExportJ>)
      return std::make_pair(n2, Jt.transpose().eval());
    else
      return std::make_pair(n2, (Jt.transpose() * Jx_or_bool).eval());
  }
}

template <typename T, typename Cov_t, typename ExportJ = std::nullptr_t>
auto MahaNorm(const T &x, const Cov_t &cov_or_var, const ExportJ &Jx_or_bool = nullptr) {
  using Scalar = typename traits::params_trait<T>::Scalar;
  using std::sqrt;
  constexpr bool add_scale = false;
  if constexpr (traits::is_nullptr_v<ExportJ>) {
    const auto &n2 = MahaSquaredNorm(x, cov_or_var, nullptr, add_scale);
    return sqrt(n2);
  } else if constexpr (traits::is_scalar_v<T>) {  // scalar
    using std::abs;
    const auto &[n2, J] = MahaSquaredNorm(x, cov_or_var, Jx_or_bool, add_scale);
    const auto n = sqrt(n2);
    const auto s = n > FloatEpsilon<Scalar>() ? n : Scalar(1);
    return std::make_pair(n, J / s);
  } else {
    const auto &[n2, J] = MahaSquaredNorm(x, cov_or_var, Jx_or_bool, add_scale);
    const auto n = sqrt(n2);
    const auto s = n > FloatEpsilon<Scalar>() ? n : Scalar(1);
    return std::make_pair(n, (J / s).eval());
  }
}

template <typename Derived, typename Cov_t, typename ExportJ = std::nullptr_t>
auto MahaWhitened(const MatrixBase<Derived> &res, const Cov_t &cov_stevs,
                  const ExportJ &Jx_or_bool = nullptr) {
  using Scalar = typename Derived::Scalar;
  if constexpr (traits::is_scalar_v<Cov_t>) {  // iso
    const auto res2 = (res / cov_stevs).eval();
    // Jacobian
    if constexpr (traits::is_nullptr_v<ExportJ>)
      return res2;
    else if constexpr (traits::is_bool_v<ExportJ>)
      return std::make_pair(res2, 1.0 / cov_stevs);
    else
      return std::make_pair(res2, (Jx_or_bool / cov_stevs).eval());
 
  } else if constexpr (traits::params_trait<Cov_t>::ColsAtCompileTime ==
                       1) {  // standard deviations
    assert(res.size() == cov_stevs.size());
    const auto res2 = (res.array() / cov_stevs.array()).matrix().eval();
    using Mat = Matrix<Scalar, Derived::RowsAtCompileTime, Derived::RowsAtCompileTime>;
 
    // Jacobian
    if constexpr (traits::is_nullptr_v<ExportJ>)
      return res2;
    else if constexpr (traits::is_bool_v<ExportJ>)
      return std::make_pair(res2, Mat(cov_stevs.cwiseInverse().asDiagonal()));
    else
      return std::make_pair(res2, (cov_stevs.cwiseInverse().asDiagonal() * Jx_or_bool).eval());
 
  } else {  // cov matrix
    const auto cov = cov_stevs.template cast<Scalar>().eval();
    {
      static constexpr Index Dims = Derived::RowsAtCompileTime;
      using Mat = Matrix<Scalar, Dims, Dims>;
      const auto chol = Eigen::SelfAdjointView<const Mat, Upper>(cov).llt();
      const auto res2 = chol.matrixL().solve(res).eval();
 
      // Jacobian
      if constexpr (traits::is_nullptr_v<ExportJ>)
        return res2;
      else if constexpr (traits::is_bool_v<ExportJ>)
        return std::make_pair(res2, (chol.matrixL().solve(Mat::Identity())).eval());
      else
        return std::make_pair(res2, (chol.matrixL().solve(Jx_or_bool)).eval());
 
    }  // TODO if constexpr (traits::params_trait<Cov_t>::ColsAtCompileTime == Dynamic)
  }
}

template <typename Derived, typename DerivedC, typename ExportJ = std::nullptr_t>
auto MahaWhitenedInfoU(const MatrixBase<Derived> &res, const MatrixBase<DerivedC> &U,
                       const ExportJ &Jx_or_bool = nullptr) {
  auto UU = U.template triangularView<Upper>();
  const auto res2 = (UU * res).eval();
  // Jacobian
  if constexpr (traits::is_nullptr_v<ExportJ>)
    return res2;
  else if constexpr (traits::is_bool_v<ExportJ>)
    return std::make_pair(res2, U.eval());
  else
    return std::make_pair(res2, (UU * Jx_or_bool).eval());
}

 
}  // namespace tinyopt::losses