<commit_before>#pragma once

#include <foundation/newton_raphson.h>
#include <foundation/types.h>

#include <iostream>

enum class ProjectionType { PERSPECTIVE, BROWN, FISHEYE, SPHERICAL, DUAL };
enum class Disto { K1 = 0, K2 = 1, K3 = 2, P1 = 3, P2 = 4, COUNT = 5 };

struct FisheyeProjection {
  template <class T>
  static Vec2<T> Forward(const Vec3<T>& point, const VecX<T>& p) {
    const T r = point.template head<2>().norm();
    const auto theta = atan2(r, point[2]);
    return Vec2<T>(theta / r * point[0], theta / r * point[1]);
  }

  template <class T>
  static Vec3<T> Backward(const Vec2<T>& point, const VecX<T>& p) {
    const T theta = point.norm();
    const auto s = tan(theta) / theta;
    return Vec3<T>(point[0] * s, point[1] * s, 1.0).normalized();
  }
};

struct PerspectiveProjection {
  template <class T>
  static Vec2<T> Forward(const Vec3<T>& point, const VecX<T>& p) {
    return Vec2<T>(point[0] / point[2], point[1] / point[2]);
  }

  template <class T>
  static Vec3<T> Backward(const Vec2<T>& point, const VecX<T>& p) {
    return Vec3<T>(point[0], point[1], 1.0).normalized();
  }
};

struct DualProjection {
  template <class T>
  static Vec2<T> Forward(const Vec3<T>& point, const VecX<T>& p) {
    const auto p_persp = PerspectiveProjection::Forward(point, p);
    const auto p_fish = FisheyeProjection::Forward(point, p);
    return p[0] * p_persp + (1.0 - p[0]) * p_fish;
  }

  template <class T>
  static Vec3<T> Backward(const Vec2<T>& point, const VecX<T>& p) {
    // Perform a bit iterations for finding theta from r
    const T r = point.norm();
    ThetaEval<T> eval_function{0, r, p[0]};
    const auto theta_refined =
        NewtonRaphson<ThetaEval<T>, 1, 1, ManualDiff<ThetaEval<T>, 1, 1>>(
            eval_function, 0, iterations);

    const auto s = tan(theta_refined) /
                   (p[0] * tan(theta_refined) + (1.0 - p[0]) * theta_refined);
    return Vec3<T>(point[0] * s, point[1] * s, 1.0).normalized();
  }

  static constexpr int iterations = 5;
  template <class T>
  struct ThetaEval {
    mutable int count;
    const T& r;
    const T& transition;
    T operator()(const T& x) const {
      return transition * tan(x) + (1.0 - transition) * x - r;
    }
    T derivative(const T& x) const {
      /* Here's some trick : use a half shorter step to prevent gross
       * overfitting on tan(x) */
      const T mult = count++ == 0 ? T(2.0) : T(1.0);
      const auto secant = 1.0 / cos(x);
      return mult * (transition * secant * secant - transition + 1);
    }
  };
};

struct SphericalProjection {
  template <class T>
  static Vec2<T> Forward(const Vec3<T>& point, const VecX<T>& p) {
    const auto lon = atan2(point[0], point[2]);
    const auto lat = atan2(-point[1], hypot(point[0], point[2]));
    return Vec2<T>(lon / (2 * M_PI), -lat / (2 * M_PI));
  }

  template <class T>
  static Vec3<T> Backward(const Vec2<T>& point, const VecX<T>& p) {
    const auto lon = point[0] * 2 * M_PI;
    const auto lat = -point[1] * 2 * M_PI;
    return Vec3<T>(cos(lat) * sin(lon), -sin(lat), cos(lat) * cos(lon));
  }
};

struct Disto24 {
  template <class T>
  static Vec2<T> Forward(const Vec2<T>& point, const VecX<T>& k) {
    const T r2 = point.squaredNorm();
    const auto distortion = Distortion(r2, k(static_cast<int>(Disto::K1)),
                                       k(static_cast<int>(Disto::K2)));
    return point * distortion;
  }

  template <class T>
  static Vec2<T> Backward(const Vec2<T>& point, const VecX<T>& k) {
    /* Beware if you use Backward together with autodiff. You'll need to remove
     * the line below, otherwise, derivatives won't be propagated */
    const T rd = point.norm();<commit_msg>fix: get rid of more eigen's stack<commit_after>#pragma once

#include <foundation/newton_raphson.h>
#include <foundation/types.h>

#include <iostream>

enum class ProjectionType { PERSPECTIVE, BROWN, FISHEYE, SPHERICAL, DUAL };
enum class Disto { K1 = 0, K2 = 1, K3 = 2, P1 = 3, P2 = 4, COUNT = 5 };

struct FisheyeProjection {
  template <class T>
  static Vec2<T> Forward(const Vec3<T>& point, const VecX<T>& p) {
    const T r = point.template head<2>().norm();
    const auto theta = atan2(r, point[2]);
    return Vec2<T>(theta / r * point[0], theta / r * point[1]);
  }

  template <class T>
  static Vec3<T> Backward(const Vec2<T>& point, const VecX<T>& p) {
    const T theta = point.norm();
    const auto s = tan(theta) / theta;
    return Vec3<T>(point[0] * s, point[1] * s, 1.0).normalized();
  }
};

struct PerspectiveProjection {
  template <class T>
  static Vec2<T> Forward(const Vec3<T>& point, const VecX<T>& p) {
    return Vec2<T>(point[0] / point[2], point[1] / point[2]);
  }

  template <class T>
  static Vec3<T> Backward(const Vec2<T>& point, const VecX<T>& p) {
    return Vec3<T>(point[0], point[1], 1.0).normalized();
  }
};

struct DualProjection {
  template <class T>
  static Vec2<T> Forward(const Vec3<T>& point, const VecX<T>& p) {
    const auto p_persp = PerspectiveProjection::Forward(point, p);
    const auto p_fish = FisheyeProjection::Forward(point, p);
    return p[0] * p_persp + (1.0 - p[0]) * p_fish;
  }

  template <class T>
  static Vec3<T> Backward(const Vec2<T>& point, const VecX<T>& p) {
    // Perform a bit iterations for finding theta from r
    const T r = point.norm();
    ThetaEval<T> eval_function{0, r, p[0]};
    const auto theta_refined =
        NewtonRaphson<ThetaEval<T>, 1, 1, ManualDiff<ThetaEval<T>, 1, 1>>(
            eval_function, 0, iterations);

    const auto s = tan(theta_refined) /
                   (p[0] * tan(theta_refined) + (1.0 - p[0]) * theta_refined);
    return Vec3<T>(point[0] * s, point[1] * s, 1.0).normalized();
  }

  static constexpr int iterations = 5;
  template <class T>
  struct ThetaEval {
    mutable int count;
    const T& r;
    const T& transition;
    T operator()(const T& x) const {
      return transition * tan(x) + (1.0 - transition) * x - r;
    }
    T derivative(const T& x) const {
      /* Here's some trick : use a half shorter step to prevent gross
       * overfitting on tan(x) */
      const T mult = count++ == 0 ? T(2.0) : T(1.0);
      const auto secant = 1.0 / cos(x);
      return mult * (transition * secant * secant - transition + 1);
    }
  };
};

struct SphericalProjection {
  template <class T>
  static Vec2<T> Forward(const Vec3<T>& point, const VecX<T>& p) {
    const auto lon = atan2(point[0], point[2]);
    const auto lat = atan2(-point[1], hypot(point[0], point[2]));
    return Vec2<T>(lon / (2 * M_PI), -lat / (2 * M_PI));
  }

  template <class T>
  static Vec3<T> Backward(const Vec2<T>& point, const VecX<T>& p) {
    const auto lon = point[0] * 2 * M_PI;
    const auto lat = -point[1] * 2 * M_PI;
    return Vec3<T>(cos(lat) * sin(lon), -sin(lat), cos(lat) * cos(lon));
  }
};

struct Disto24 {
  template <class T>
  static Vec2<T> Forward(const Vec2<T>& point, const VecX<T>& k) {
    const T r2 = point.squaredNorm();
    const auto distortion = Distortion(r2, k(static_cast<int>(Disto::K1)),
                                       k(static_cast<int>(Disto::K2)));
    return Vec2<T>(point[0] / distortion, point[1] / distortion);
  }

  template <class T>
  static Vec2<T> Backward(const Vec2<T>& point, const VecX<T>& k) {
    /* Beware if you use Backward together with autodiff. You'll need to remove
     * the line below, otherwise, derivatives won't be propagated */
    const T rd = point.norm();
{'fraction_non_alphanumeric': 0.15576407506702414, 'fraction_numerical': 0.028284182305630026, 'mean_word_length': 3.1312292358803986, 'pattern_counts': {'":': 0, '<': 122, '<?xml version=': 0, '>': 122, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7377500', 'n_tokens_mistral': 2850, 'n_tokens_neox': 2715, 'n_words': 795}