#include <algorithm/glue_algorithm.hpp>

#include "internal_primitive_desc.hpp"

static inline auto manually_projection(const raw_vector3d_t                    &point,
                                       const Eigen::Ref<const Eigen::Vector3d> &x_dir,
                                       const Eigen::Ref<const Eigen::Vector3d> &y_dir,
                                       const Eigen::Ref<const Eigen::Vector2d> &offset)
{
    Eigen::Vector3d p;
    vec3d_conversion(point, p);
    return Eigen::Vector2d{p.dot(x_dir), p.dot(y_dir)} + offset;
}

static inline auto manually_projection(raw_vector3d_t                         &&point,
                                       const Eigen::Ref<const Eigen::Vector3d> &x_dir,
                                       const Eigen::Ref<const Eigen::Vector3d> &y_dir,
                                       const Eigen::Ref<const Eigen::Vector2d> &offset)
{
    Eigen::Vector3d p;
    vec3d_conversion(std::move(point), p);
    return Eigen::Vector2d{p.dot(x_dir), p.dot(y_dir)} + offset;
}

namespace internal
{
void polyline::build_as_axis(const polyline_descriptor_t                       &desc,
                             const raw_vector3d_t                              &profile_ref_normal,
                             Eigen::Transform<double, 3, Eigen::AffineCompact> &axis_to_world,
                             aabb_t<>                                          &aabb)
{
    auto &matrix_handle = axis_to_world.matrix();
    vec3d_conversion(profile_ref_normal, matrix_handle.col(1));
    matrix_handle.col(1).normalize();
    vec3d_conversion(desc.reference_normal, matrix_handle.col(2));
    matrix_handle.col(2).normalize();
    matrix_handle.col(0) = matrix_handle.col(1).cross(matrix_handle.col(2));
    vec3d_conversion(desc.points[0], matrix_handle.col(3));

    aabb_t<2> profile_aabb;
    build_as_profile(desc, matrix_handle.col(2), matrix_handle.col(0), matrix_handle.col(3), profile_aabb);
    const auto &profile_aabb_min = profile_aabb.min();
    const auto &profile_aabb_max = profile_aabb.max();
    aabb                         = {
        Eigen::Vector3d{profile_aabb_min.x(), profile_aabb_min.y(), 0},
        Eigen::Vector3d{profile_aabb_max.x(), profile_aabb_max.y(), 0}
    };
    aabb.transform(axis_to_world);
}

void polyline::build_as_axis(polyline_descriptor_t                            &&desc,
                             const raw_vector3d_t                              &profile_ref_normal,
                             Eigen::Transform<double, 3, Eigen::AffineCompact> &axis_to_world,
                             aabb_t<>                                          &aabb)
{
    auto &matrix_handle = axis_to_world.matrix();
    vec3d_conversion(profile_ref_normal, matrix_handle.col(1));
    matrix_handle.col(1).normalize();
    vec3d_conversion(std::move(desc.reference_normal), matrix_handle.col(2));
    matrix_handle.col(2).normalize();
    matrix_handle.col(0) = matrix_handle.col(1).cross(matrix_handle.col(2));
    vec3d_conversion(desc.points[0], matrix_handle.col(3));

    aabb_t<2> profile_aabb;
    build_as_profile(std::move(desc), matrix_handle.col(2), matrix_handle.col(0), matrix_handle.col(3), profile_aabb);
    const auto &profile_aabb_min = profile_aabb.min();
    const auto &profile_aabb_max = profile_aabb.max();
    aabb                         = {
        Eigen::Vector3d{profile_aabb_min.x(), profile_aabb_min.y(), 0},
        Eigen::Vector3d{profile_aabb_max.x(), profile_aabb_max.y(), 0}
    };
    aabb.transform(axis_to_world);
}

void polyline::build_as_profile(const polyline_descriptor_t             &desc,
                                const Eigen::Ref<const Eigen::Vector3d> &proj_x_dir,
                                const Eigen::Ref<const Eigen::Vector3d> &proj_y_dir,
                                const Eigen::Ref<const Eigen::Vector3d> &proj_origin,
                                aabb_t<2>                               &aabb)
{
    assert((desc.is_close && desc.bulge_number == desc.point_number)
           || (!desc.is_close && desc.bulge_number + 1 == desc.point_number));

    // collect indices info
    this->thetas.resize(desc.bulge_number);
    this->start_indices.resize(desc.bulge_number + 1, uint8_t{0});
    stl_vector<uint8_t> line_start_indices(desc.bulge_number + 1), circle_start_indices(desc.bulge_number + 1);
    if (this->start_indices.size() > 16) {
        std::atomic_size_t line_start_indices_offset{0};
        std::atomic_size_t circle_start_indices_offset{0};
        std::transform_inclusive_scan(std::execution::par_unseq,
                                      counting_iterator<size_t>{},
                                      counting_iterator<size_t>{this->start_indices.size()},
                                      this->start_indices.begin() + 1,
                                      std::plus<uint8_t>{},
                                      [&](size_t index) {
                                          auto &theta = this->thetas[index];
                                          theta       = std::atan(fabs(desc.bulges[index])) * 4;
                                          if (theta <= EPSILON) {
                                              const auto offset          = line_start_indices_offset.fetch_add(1);
                                              line_start_indices[offset] = index;
                                              return uint8_t{1};
                                          } else {
                                              const auto offset = circle_start_indices_offset.fetch_add(1);
                                              return uint8_t{3};
                                          }
                                      });
    } else {
        size_t line_start_indices_offset{0};
        size_t circle_start_indices_offset{0};
        std::transform_inclusive_scan(counting_iterator<size_t>{},
                                      counting_iterator<size_t>{this->start_indices.size()},
                                      this->start_indices.begin() + 1,
                                      std::plus<uint8_t>{},
                                      [&](size_t index) {
                                          auto &theta = this->thetas[index];
                                          theta       = std::atan(fabs(desc.bulges[index])) * 4;
                                          if (theta <= EPSILON) {
                                              line_start_indices[line_start_indices_offset++] = index;
                                              return uint8_t{1};
                                          } else {
                                              circle_start_indices[circle_start_indices_offset++] = index;
                                              return uint8_t{3};
                                          }
                                      });
    }
    // i.e. iff close curve, deduce loop to simple segements by copying the first vertex to the end
    this->vertices.resize(this->start_indices.back() + static_cast<size_t>(desc.is_close));
    this->start_indices.pop_back();

    const Eigen::Vector2d offset = {-proj_x_dir.dot(proj_origin), -proj_y_dir.dot(proj_origin)};

    // first loop: process all line segment vertices
    algorithm::for_loop(size_t{}, line_start_indices.size(), [&](size_t i) {
        const auto &index       = line_start_indices[i];
        const auto &start_index = this->start_indices[index];
        auto       &v           = this->vertices[start_index];
        v                       = manually_projection(desc.points[index], proj_x_dir, proj_y_dir, offset);
        aabb.extend(v);
    });
    if (desc.is_close) this->vertices.back() = this->vertices.front();

    // second loop: process all circle vertices
    algorithm::for_loop(size_t{}, circle_start_indices.size(), [&](size_t i) {
        const auto &index           = line_start_indices[i];
        const auto &start_index     = this->start_indices[index];
        const auto &bulge           = desc.bulges[index];
        const auto &theta           = this->thetas[index];
        this->vertices[start_index] = manually_projection(desc.points[index], proj_x_dir, proj_y_dir, offset);

        const auto &a           = this->vertices[start_index];
        const auto &b           = this->vertices[start_index + 3];
        auto       &c           = this->vertices[start_index + 1];
        auto       &d           = this->vertices[start_index + 2];
        const auto  half_vec_ab = 0.5 * (a + b);
        if (theta == PI) {
            c = std::move(half_vec_ab);
        } else {
            // NOTE: bulge = tan(theta/4), so we can directly calculate tan(theta/2)
            // and no more operations are needed for identifying the circle center
            // since all direction info is already included in the sign & magnitude of bulge
            const auto inv_tan_half_theta = (1 - bulge * bulge) / (2 * bulge);
            c                             = half_vec_ab + inv_tan_half_theta * Eigen::Vector2d{half_vec_ab[1], -half_vec_ab[0]};
        }
        const auto vec_ca = a - c;
        if (bulge > 0)
            d = c + Eigen::Vector2d{vec_ca[1], -vec_ca[0]};
        else
            d = c + Eigen::Vector2d{-vec_ca[1], vec_ca[0]};

        aabb.extend(a);
        aabb.extend(c);
        aabb.extend(d);
    });
}

void polyline::build_as_profile(polyline_descriptor_t                  &&desc,
                                const Eigen::Ref<const Eigen::Vector3d> &proj_x_dir,
                                const Eigen::Ref<const Eigen::Vector3d> &proj_y_dir,
                                const Eigen::Ref<const Eigen::Vector3d> &proj_origin,
                                aabb_t<2>                               &aabb)
{
    assert((desc.is_close && desc.bulge_number == desc.point_number)
           || (!desc.is_close && desc.bulge_number + 1 == desc.point_number));

    // collect indices info
    this->thetas.resize(desc.bulge_number);
    this->start_indices.resize(desc.bulge_number + 1, uint8_t{0});
    stl_vector<uint8_t> line_start_indices(desc.bulge_number + 1), circle_start_indices(desc.bulge_number + 1);
    if (this->start_indices.size() > 16) {
        std::atomic_size_t line_start_indices_offset{0};
        std::atomic_size_t circle_start_indices_offset{0};
        std::transform_inclusive_scan(std::execution::par_unseq,
                                      counting_iterator<size_t>{},
                                      counting_iterator<size_t>{this->start_indices.size()},
                                      this->start_indices.begin() + 1,
                                      std::plus<uint8_t>{},
                                      [&](size_t index) {
                                          auto &theta = this->thetas[index];
                                          theta       = std::atan(fabs(desc.bulges[index])) * 4;
                                          if (theta <= EPSILON) {
                                              const auto offset          = line_start_indices_offset.fetch_add(1);
                                              line_start_indices[offset] = index;
                                              return uint8_t{1};
                                          } else {
                                              const auto offset = circle_start_indices_offset.fetch_add(1);
                                              return uint8_t{3};
                                          }
                                      });
    } else {
        size_t line_start_indices_offset{0};
        size_t circle_start_indices_offset{0};
        std::transform_inclusive_scan(counting_iterator<size_t>{},
                                      counting_iterator<size_t>{this->start_indices.size()},
                                      this->start_indices.begin() + 1,
                                      std::plus<uint8_t>{},
                                      [&](size_t index) {
                                          auto &theta = this->thetas[index];
                                          theta       = std::atan(fabs(desc.bulges[index])) * 4;
                                          if (theta <= EPSILON) {
                                              line_start_indices[line_start_indices_offset++] = index;
                                              return uint8_t{1};
                                          } else {
                                              circle_start_indices[circle_start_indices_offset++] = index;
                                              return uint8_t{3};
                                          }
                                      });
    }
    // i.e. iff close curve, deduce loop to simple segements by copying the first vertex to the end
    this->vertices.resize(this->start_indices.back() + static_cast<size_t>(desc.is_close));
    this->start_indices.pop_back();

    const Eigen::Vector2d offset = {-proj_x_dir.dot(proj_origin), -proj_y_dir.dot(proj_origin)};

    // first loop: process all line segment vertices
    algorithm::for_loop(size_t{}, line_start_indices.size(), [&](size_t i) {
        const auto &index       = line_start_indices[i];
        const auto &start_index = this->start_indices[index];
        auto       &v           = this->vertices[start_index];
        v                       = manually_projection(std::move(desc.points[index]), proj_x_dir, proj_y_dir, offset);
        aabb.extend(v);
    });
    if (desc.is_close) this->vertices.back() = this->vertices.front();

    // second loop: process all circle vertices
    algorithm::for_loop(size_t{}, circle_start_indices.size(), [&](size_t i) {
        const auto &index           = line_start_indices[i];
        const auto &start_index     = this->start_indices[index];
        const auto &bulge           = desc.bulges[index];
        const auto &theta           = this->thetas[index];
        this->vertices[start_index] = manually_projection(std::move(desc.points[index]), proj_x_dir, proj_y_dir, offset);

        const auto &a           = this->vertices[start_index];
        const auto &b           = this->vertices[start_index + 3];
        auto       &c           = this->vertices[start_index + 1];
        auto       &d           = this->vertices[start_index + 2];
        const auto  half_vec_ab = 0.5 * (a + b);
        if (theta == PI) {
            c = std::move(half_vec_ab);
        } else {
            // NOTE: bulge = tan(theta/4), so we can directly calculate tan(theta/2)
            // and no more operations are needed for identifying the circle center
            // since all direction info is already included in the sign & magnitude of bulge
            const auto inv_tan_half_theta = (1 - bulge * bulge) / (2 * bulge);
            c                             = half_vec_ab + inv_tan_half_theta * Eigen::Vector2d{half_vec_ab[1], -half_vec_ab[0]};
        }
        const auto vec_ca = a - c;
        if (bulge > 0)
            d = c + Eigen::Vector2d{vec_ca[1], -vec_ca[0]};
        else
            d = c + Eigen::Vector2d{-vec_ca[1], vec_ca[0]};

        aabb.extend(a);
        aabb.extend(c);
        aabb.extend(d);
    });
}

[[nodiscard]] Eigen::Vector2d polyline::calculate_tangent(double t) const
{
    assert(t >= 0 && t <= this->start_indices.size());

    double     frac;
    const auto n    = static_cast<uint32_t>(std::modf(t, &frac));
    const auto iter = this->vertices.begin() + this->start_indices[n];
    if (this->thetas[n] <= EPSILON) {
        return (*(iter + 1) - *iter).normalized();
    } else {
        const auto alpha_deriv = -std::sin(t * this->thetas[n]);
        const auto beta_deriv  = std::cos(t * this->thetas[n]);
        return (alpha_deriv * *iter + beta_deriv * *(iter + 2) - (alpha_deriv + beta_deriv) * *(iter + 1)).normalized();
    }
}

// CAUTION: always make normal points inward
[[nodiscard]] Eigen::Vector2d polyline::calculate_normal(double t) const
{
    assert(t >= 0 && t <= this->start_indices.size());

    double     frac;
    const auto n    = static_cast<uint32_t>(std::modf(t, &frac));
    const auto iter = this->vertices.begin() + this->start_indices[n];
    if (this->thetas[n] <= EPSILON) {
        // HINT: assume the plane normal (ref_normal) to be {0, 0, 1}, then the normal should be this
        const auto temp = (*(iter + 1) - *iter).normalized();
        return {temp.y(), -temp.x()};
    } else {
        const auto alpha = std::cos(t * this->thetas[n]);
        const auto beta  = std::sin(t * this->thetas[n]);
        return (alpha + beta) * *(iter + 1) - alpha * *iter - beta * *(iter + 2);
    }
}

[[nodiscard]] auto polyline::calculate_closest_param(const Eigen::Ref<const Eigen::Vector3d> &p) const
    -> std::pair<line_closest_param_t, double>
{
    struct comparable_closest_param_t {
        Eigen::Vector3d point{};
        double          t{-1.};
        double          dist{std::numeric_limits<double>::max()};
        double          dist_sign{1.};

        inline bool operator<(const comparable_closest_param_t &other) const { return dist < other.dist; }
    };

    auto line_cpm = [&](auto index) -> comparable_closest_param_t {
        const auto iter     = this->vertices.begin() + this->start_indices[index];
        const auto line_vec = *(iter + 1) - *iter;
        const auto p_dir    = p.topRows<2>() - *iter;

        // illegal solve, return null
        const auto line_vec_length_2 = line_vec.squaredNorm();
        const auto raw_t             = line_vec.dot(p_dir);
        if (raw_t < 0 || raw_t > line_vec_length_2) return {};

        const auto p_dir_length_2 = p_dir.squaredNorm();
        const auto t              = raw_t / line_vec_length_2;
        return {
            {(1 - t) * iter->x() + t * (iter + 1)->x(), (1 - t) * iter->y() + t * (iter + 1)->y(), 0},
            t,
            std::sqrt(p_dir_length_2 - raw_t * raw_t / line_vec_length_2 + p[2] * p[2])
        };
    };

    auto circle_cpm = [&](auto index) -> comparable_closest_param_t {
        const auto iter      = this->vertices.begin() + this->start_indices[index];
        const auto theta     = this->thetas[index];
        const auto base_vec1 = *iter - *(iter + 1);
        const auto base_vec2 = *(iter + 2) - *(iter + 1);
        const auto p_vec     = p.topRows<2>() - *(iter + 1);

        const auto local_x = base_vec1.dot(p_vec);
        const auto local_y = base_vec2.dot(p_vec);
        const auto phi     = std::atan2(local_y, local_x);
        // illegal solve, return null
        if (phi < 0 || phi > theta) return {};

        const auto p_vec_norm    = p_vec.norm();
        const auto r             = base_vec1.norm();
        const auto dis_on_plane  = p_vec_norm - r;
        const auto closest_point = p_vec / p_vec_norm * r + *(iter + 1);
        return {
            {closest_point.x(), closest_point.y(), 0},
            phi / theta,
            std::sqrt(p[2] * p[2] + dis_on_plane * dis_on_plane)
        };
    };

    std::vector<comparable_closest_param_t> closest_params(this->thetas.size());
    algorithm::transform(counting_iterator<size_t>(),
                         counting_iterator<size_t>(this->thetas.size()),
                         closest_params.begin(),
                         [&](size_t index) {
                             if (this->thetas[index] <= EPSILON)
                                 return line_cpm(index);
                             else
                                 return circle_cpm(index);
                         });

    const auto iter = algorithm::min_element(closest_params.begin(), closest_params.end());
    if (iter->t >= 0) { // i.e. has closest point prependicular to the line, which is the best solution
        iter->t += std::distance(closest_params.begin(), iter);
        return {
            line_closest_param_t{std::move(iter->point), iter->t, true},
            iter->dist * iter->dist_sign
        };
    } else { // in this case, the closest point can only be the vertices of the line
        algorithm::transform(counting_iterator<size_t>(),
                             counting_iterator<size_t>(this->thetas.size()),
                             closest_params.begin(),
                             [&](size_t index) {
                                 const auto iter = this->vertices.begin() + this->start_indices[index];
                                 const auto dist = std::sqrt((p.topRows<2>() - *iter).squaredNorm() + p[2] * p[2]);
                                 return comparable_closest_param_t{
                                     {iter->x(), iter->y(), 0},
                                     static_cast<double>(index),
                                     dist
                                 };
                             });
        const auto iter_ = algorithm::min_element(closest_params.begin(), closest_params.end());
        return {
            line_closest_param_t{std::move(iter->point), iter->t, false},
            iter->dist * iter->dist_sign
        };
    }
}

[[nodiscard]] bool polyline::pmc(const Eigen::Ref<const Eigen::Vector2d>&p, const line_closest_param_t& closest_param) const {
		if (closest_param.is_peak_value) {
			return (p - closest_param.point.topRows<2>()).dot(this->calculate_normal(closest_param.t)) < 0;
		}
		// the closest point is the vertex of the line
		// todo: may not be robust for rare cases
		return (p -  closest_param.point.topRows<2>()).dot(this->calculate_normal(closest_param.t)) < 0;
}

[[nodiscard]] bool polyline::isEnd(const double t) const {
		return t >= thetas.size() - EPSILON || t < EPSILON;
} 
} // namespace internal