distances.h

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

template <typename TA, typename TB, typename ExportJ = std::nullptr_t>
auto Euclidean(const TA &a, const TB &b, const ExportJ &export_jac = nullptr) {
  if constexpr (traits::is_nullptr_v<ExportJ>) {
    return losses::L2(a - b);
  } else if constexpr (traits::is_scalar_v<TA>) {
    const auto &[d, J] = losses::L2(a - b, export_jac);
    return std::make_tuple(d, J, -J);
  } else {
    const auto &[d, J] = losses::L2(a - b, export_jac);
    return std::make_tuple(d, J, (-J).eval());
  }
}
 
template <typename TA, typename TB, typename ExportJ = std::nullptr_t>
auto L2(const TA &a, const TB &b, const ExportJ &export_jac = nullptr) {
  return Euclidean(a, b, export_jac);
}

template <typename TA, typename TB, typename ExportJ = std::nullptr_t>
auto Manhattan(const TA &a, const TB &b, const ExportJ &export_jac = nullptr) {
  if constexpr (traits::is_nullptr_v<ExportJ>) {
    return losses::L1(a - b);
  } else if constexpr (traits::is_scalar_v<TA>) {
    const auto &[d, J] = losses::L1(a - b, export_jac);
    return std::make_tuple(d, J, -J);
  } else {
    const auto &[d, J] = losses::L1(a - b, export_jac);
    return std::make_tuple(d, J, (-J).eval());
  }
}
 
template <typename TA, typename TB, typename ExportJ = std::nullptr_t>
auto L1(const TA &a, const TB &b, const ExportJ &export_jac = nullptr) {
  return Manhattan(a, b, export_jac);
}

template <typename TA, typename TB, typename ExportJ = std::nullptr_t>
auto Linf(const TA &a, const TB &b, const ExportJ &export_jac = nullptr) {
  if constexpr (traits::is_nullptr_v<ExportJ>) {
    return losses::Linf(a - b);
  } else if constexpr (traits::is_scalar_v<TA>) {
    const auto &[d, J] = losses::Linf(a - b, export_jac);
    return std::make_tuple(d, J, -J);
  } else {
    const auto &[d, J] = losses::Linf(a - b, export_jac);
    return std::make_tuple(d, J, (-J).eval());
  }
}

template <typename TA, typename TB, typename ExportJ = std::nullptr_t>
auto Cosine(const TA &a, const TB &b, const ExportJ & = nullptr) {
  if constexpr (traits::is_nullptr_v<ExportJ>) {
    return TA(0);
  } else if constexpr (traits::is_scalar_v<TA>) {
    return std::make_tuple(TA(0), TA(0), TA(0));
  } else {  // Vectors
    using Scalar = std::decay_t<typename TA::Scalar>;
    constexpr double eps2 = FloatEpsilon2<Scalar>();
    constexpr Index Dims = traits::params_trait<TA>::Dims;
    using Vec = RowVector<Scalar, Dims>;
    const auto a_norm = a.norm();
    const auto b_norm = b.norm();
    if (a_norm * b_norm < eps2) {
      return std::make_tuple(Scalar(0), Vec::Zero(a.size()), Vec::Zero(a.size()));
    } else {
      const auto ab = a.dot(b);
      const auto d = (ab / (a_norm * b_norm));
      Vec Ja = (b / (a_norm * b_norm)) - (ab * a / (a_norm * a_norm * a_norm * b_norm));
      Vec Jb = (a / (a_norm * b_norm)) - (ab * b / (a_norm * b_norm * b_norm));
      return std::make_tuple(d, Ja, Jb);
    }
  }
}

template <typename TA, typename TB, typename Cov_t, typename ExportJ = std::nullptr_t>
auto MahaNorm(const TA &a, const TB &b, const Cov_t &cov_or_var,
              const ExportJ &export_jac = nullptr) {
  if constexpr (traits::is_nullptr_v<ExportJ>) {
    return losses::MahaNorm(a - b, cov_or_var);
  } else if constexpr (traits::is_scalar_v<TA>) {
    const auto &[d, J] = losses::MahaNorm(a - b, cov_or_var, export_jac);
    return std::make_tuple(d, J, -J);
  } else {
    const auto &[d, J] = losses::MahaNorm(a - b, cov_or_var, export_jac);
    return std::make_tuple(d, J, (-J).eval());
  }
}
 
}  // namespace tinyopt::distances