#include <array>
#include <cmath>
#include <iomanip>
#include <sstream>
#include <vector>
#include <cmath>

#include "camera.hpp"

// =================================================================================================
// =================================================================================================

template <size_t M, size_t N>
static const std::string print_matrix(const std::array<std::array<float, N>, M> &matrix)
{
    std::ostringstream out;
    out << "[";
    for (size_t j = 0; j < M; ++j)
    {
        out << "[";
        for (size_t i = 0; i < N; ++i)
        {
            out << matrix[j][i];
            if (i < N - 1)
            {
                out << ", ";
            }
        }
        out << "]";
        if (j < M - 1)
        {
            out << ", ";
        }
    }
    out << "]";
    return out.str();
}

// =================================================================================================
// =================================================================================================

std::string Camera::to_string() const
{
    std::ostringstream out;
    out << std::fixed << std::setprecision(6);

    out << "{";
    out << "'name': '" << name << "', ";

    out << "'K': " << print_matrix(K) << ", ";

    out << "'DC': [";
    for (size_t i = 0; i < DC.size(); ++i)
    {
        out << DC[i];
        if (i < DC.size() - 1)
            out << ", ";
    }
    out << "], ";

    out << "'R': " << print_matrix(R) << ", ";
    out << "'T': " << print_matrix(T) << ", ";

    out << "'width': " << width << ", ";
    out << "'height': " << height << ", ";
    out << "'type': " << type;

    out << "}";
    return out.str();
}

// =================================================================================================

std::ostream &operator<<(std::ostream &out, const Camera &cam)
{
    out << cam.to_string();
    return out;
}

// =================================================================================================
// =================================================================================================

CameraInternal::CameraInternal(const Camera &cam)
{
    this->cam = cam;
    this->invR = transpose3x3(cam.R);

    // Camera center:
    // C = -(Rᵀ * t) = -(Rᵀ * (R * (T * -1))) = -(Rᵀ * (R * -T)) = -(Rᵀ * -R * T) = -(-T) = T
    this->center = {cam.T[0][0], cam.T[1][0], cam.T[2][0]};

    // Undistort camera matrix
    // As with the undistortion, the own implementation avoids some overhead compared to OpenCV
    if (cam.type == "fisheye")
    {
        newK = calc_optimal_camera_matrix_fisheye(1.0, {cam.width, cam.height});
    }
    else
    {
        newK = calc_optimal_camera_matrix_pinhole(1.0, {cam.width, cam.height});
    }
    this->invK = invert3x3(newK);
}

// =================================================================================================

std::array<std::array<float, 3>, 3> CameraInternal::transpose3x3(
    const std::array<std::array<float, 3>, 3> &M)
{
    return {{{M[0][0], M[1][0], M[2][0]},
             {M[0][1], M[1][1], M[2][1]},
             {M[0][2], M[1][2], M[2][2]}}};
}

// =================================================================================================

std::array<std::array<float, 3>, 3> CameraInternal::invert3x3(
    const std::array<std::array<float, 3>, 3> &M)
{
    // Compute the inverse using the adjugate method
    // See: https://scicomp.stackexchange.com/a/29206

    std::array<std::array<float, 3>, 3> adj = {
        {{
             M[1][1] * M[2][2] - M[1][2] * M[2][1],
             M[0][2] * M[2][1] - M[0][1] * M[2][2],
             M[0][1] * M[1][2] - M[0][2] * M[1][1],
         },
         {
             M[1][2] * M[2][0] - M[1][0] * M[2][2],
             M[0][0] * M[2][2] - M[0][2] * M[2][0],
             M[0][2] * M[1][0] - M[0][0] * M[1][2],
         },
         {
             M[1][0] * M[2][1] - M[1][1] * M[2][0],
             M[0][1] * M[2][0] - M[0][0] * M[2][1],
             M[0][0] * M[1][1] - M[0][1] * M[1][0],
         }}};

    float det = M[0][0] * adj[0][0] + M[0][1] * adj[1][0] + M[0][2] * adj[2][0];
    if (std::fabs(det) < 1e-6f)
    {
        return {{{0.0, 0.0, 0.0}, {0.0, 0.0, 0.0}, {0.0, 0.0, 0.0}}};
    }

    float idet = 1.0f / det;
    std::array<std::array<float, 3>, 3> inv = {
        {{
             adj[0][0] * idet,
             adj[0][1] * idet,
             adj[0][2] * idet,
         },
         {
             adj[1][0] * idet,
             adj[1][1] * idet,
             adj[1][2] * idet,
         },
         {
             adj[2][0] * idet,
             adj[2][1] * idet,
             adj[2][2] * idet,
         }}};

    return inv;
}

// =================================================================================================

void CameraInternal::undistort_point_pinhole(std::array<float, 3> &p, const std::vector<float> &k)
{
    // Following: cv::cvUndistortPointsInternal
    // Uses only the distortion coefficients: [k1, k2, p1, p2, k3]
    // https://github.com/opencv/opencv/blob/4.x/modules/calib3d/src/undistort.dispatch.cpp#L432

    float x0 = p[0];
    float y0 = p[1];
    float x = x0;
    float y = y0;

    // Iteratively refine the estimate for the undistorted point.
    int max_iterations = 5;
    for (int iter = 0; iter < max_iterations; ++iter)
    {
        float r2 = x * x + y * y;
        double icdist = 1.0 / (1 + ((k[4] * r2 + k[1]) * r2 + k[0]) * r2);
        if (icdist < 0)
        {
            x = x0;
            y = y0;
            break;
        }

        float deltaX = 2 * k[2] * x * y + k[3] * (r2 + 2 * x * x);
        float deltaY = k[2] * (r2 + 2 * y * y) + 2 * k[3] * x * y;
        x = (x0 - deltaX) * icdist;
        y = (y0 - deltaY) * icdist;
    }

    p[0] = x;
    p[1] = y;
}

// =================================================================================================

void CameraInternal::undistort_point_fisheye(std::array<float, 3> &p, const std::vector<float> &k)
{
    // Following: cv::fisheye::undistortPoints
    // Uses only the distortion coefficients: [k1, k2, k3, k4]
    // https://github.com/opencv/opencv/blob/4.x/modules/calib3d/src/fisheye.cpp#L429

    float theta_d = std::sqrt(p[0] * p[0] + p[1] * p[1]);
    float pi_half = M_PI * 0.5;
    theta_d = std::min(std::max(-pi_half, theta_d), pi_half);

    if (theta_d < 1e-6)
    {
        return;
    }

    float scale = 0.0;
    float theta = theta_d;

    int max_iterations = 5;
    for (int iter = 0; iter < max_iterations; ++iter)
    {
        float theta2 = theta * theta;
        float theta4 = theta2 * theta2;
        float theta6 = theta4 * theta2;
        float theta8 = theta4 * theta4;

        float k0_theta2 = k[0] * theta2;
        float k1_theta4 = k[1] * theta4;
        float k2_theta6 = k[2] * theta6;
        float k3_theta8 = k[3] * theta8;

        float theta_fix = (theta * (1 + k0_theta2 + k1_theta4 + k2_theta6 + k3_theta8) - theta_d) /
                          (1 + 3 * k0_theta2 + 5 * k1_theta4 + 7 * k2_theta6 + 9 * k3_theta8);
        theta = theta - theta_fix;

        if (std::fabs(theta_fix) < 1e-6)
        {
            break;
        }
    }

    scale = std::tan(theta) / theta_d;
    p[0] *= scale;
    p[1] *= scale;
}

// =================================================================================================

std::array<std::array<float, 3>, 3> CameraInternal::calc_optimal_camera_matrix_fisheye(
    float balance, std::pair<int, int> new_size)
{
    // Following: cv::fisheye::estimateNewCameraMatrixForUndistortRectify
    // https://github.com/opencv/opencv/blob/4.x/modules/calib3d/src/fisheye.cpp#L630

    float fov_scale = 1.0;
    float w = static_cast<float>(cam.width);
    float h = static_cast<float>(cam.height);
    balance = std::min(std::max(balance, 0.0f), 1.0f);

    // Define four key points at the middle of each edge
    std::vector<std ::array<float, 2>> pts = {
        {w * 0.5f, 0.0},
        {w, h * 0.5f},
        {w * 0.5f, h},
        {0.0, h * 0.5f}};

    // Extract intrinsic parameters
    float fx = cam.K[0][0];
    float fy = cam.K[1][1];
    float cx = cam.K[0][2];
    float cy = cam.K[1][2];

    // Undistort the edge points
    for (auto &pt : pts)
    {
        std::array<float, 3> p_normed = {(pt[0] - cx) / fx, (pt[1] - cy) / fy, 1.0};
        undistort_point_fisheye(p_normed, cam.DC);
        pt[0] = p_normed[0];
        pt[1] = p_normed[1];
    }

    // Compute center mass of the undistorted edge points
    float sum_x = 0.0, sum_y = 0.0;
    for (const auto &pt : pts)
    {
        sum_x += pt[0];
        sum_y += pt[1];
    }
    float cn_x = sum_x / pts.size();
    float cn_y = sum_y / pts.size();

    // Convert to identity ratio
    float aspect_ratio = fx / fy;
    cn_y *= aspect_ratio;
    for (auto &pt : pts)
        pt[1] *= aspect_ratio;

    // Find the bounding box of the undistorted points
    float minx = std::numeric_limits<float>::max();
    float miny = std::numeric_limits<float>::max();
    float maxx = -std::numeric_limits<float>::max();
    float maxy = -std::numeric_limits<float>::max();
    for (const auto &pt : pts)
    {
        minx = std::min(minx, pt[0]);
        maxx = std::max(maxx, pt[0]);
        miny = std::min(miny, pt[1]);
        maxy = std::max(maxy, pt[1]);
    }

    // Calculate candidate focal lengths
    float f1 = (w * 0.5) / (cn_x - minx);
    float f2 = (w * 0.5) / (maxx - cn_x);
    float f3 = (h * 0.5 * aspect_ratio) / (cn_y - miny);
    float f4 = (h * 0.5 * aspect_ratio) / (maxy - cn_y);
    float fmin = std::min({f1, f2, f3, f4});
    float fmax = std::max({f1, f2, f3, f4});

    // Blend the candidate focal lengths
    float f_val = balance * fmin + (1.0f - balance) * fmax;
    if (fov_scale > 0.0f)
        f_val /= fov_scale;

    // Compute new intrinsic parameters
    float new_fx = f_val;
    float new_fy = f_val;
    float new_cx = -cn_x * f_val + (w * 0.5);
    float new_cy = -cn_y * f_val + (h * aspect_ratio * 0.5);

    // Restore aspect ratio
    new_fy /= aspect_ratio;
    new_cy /= aspect_ratio;

    // Optionally scale parameters to new new image size
    if (new_size.first > 0 && new_size.second > 0)
    {
        float rx = static_cast<float>(new_size.first) / w;
        float ry = static_cast<float>(new_size.second) / h;
        new_fx *= rx;
        new_fy *= ry;
        new_cx *= rx;
        new_cy *= ry;
    }

    std::array<std::array<float, 3>, 3> newK = {
        {{new_fx, 0.0, new_cx},
         {0.0, new_fy, new_cy},
         {0.0, 0.0, 1.0}}};
    return newK;
}

// =================================================================================================

std::array<std::array<float, 3>, 3> CameraInternal::calc_optimal_camera_matrix_pinhole(
    float alpha, std::pair<int, int> new_size)
{
    // Following: cv::getOptimalNewCameraMatrix
    // https://github.com/opencv/opencv/blob/4.x/modules/calib3d/src/calibration_base.cpp#L1565

    bool center_principal_point = false;
    bool use_pix_roi = false;
    float w = static_cast<float>(cam.width);
    float h = static_cast<float>(cam.height);
    alpha = std::min(std::max(alpha, 0.0f), 1.0f);

    if (center_principal_point || use_pix_roi)
    {
        // Not implemented
        exit(1);
    }

    // Define key points
    // Calculate only the contour points of the image, and use less points,
    // the edges and centers should be enough if the camera has no strange distortions
    const size_t N = 3;
    std::vector<std ::array<float, 2>> pts;
    pts.reserve(4 * (N - 1));
    for (size_t y = 0; y < N; ++y)
    {
        for (size_t x = 0; x < N; ++x)
        {
            if (x != 0 && x != N - 1 && y != 0 && y != N - 1)
            {
                continue;
            }
            pts.push_back({x * (w - 1) / (N - 1), y * (h - 1) / (N - 1)});
        }
    }

    // Extract intrinsic parameters
    float fx = cam.K[0][0];
    float fy = cam.K[1][1];
    float cx = cam.K[0][2];
    float cy = cam.K[1][2];

    // Undistort the key points
    for (auto &pt : pts)
    {
        std::array<float, 3> p_normed = {(pt[0] - cx) / fx, (pt[1] - cy) / fy, 1.0};
        undistort_point_pinhole(p_normed, cam.DC);
        pt[0] = p_normed[0];
        pt[1] = p_normed[1];
    }

    // Get inscribed and circumscribed rectangles in normalized coordinates
    float iX0 = -std::numeric_limits<float>::max();
    float iX1 = std::numeric_limits<float>::max();
    float iY0 = -std::numeric_limits<float>::max();
    float iY1 = std::numeric_limits<float>::max();
    float oX0 = std::numeric_limits<float>::max();
    float oX1 = -std::numeric_limits<float>::max();
    float oY0 = std::numeric_limits<float>::max();
    float oY1 = -std::numeric_limits<float>::max();
    size_t k = 0;
    for (size_t y = 0; y < N; ++y)
    {
        for (size_t x = 0; x < N; ++x)
        {
            if (x != 0 && x != N - 1 && y != 0 && y != N - 1)
            {
                continue;
            }

            auto &pt = pts[k];
            k += 1;

            if (x == 0)
            {
                oX0 = std::min(oX0, pt[0]);
                iX0 = std::max(iX0, pt[0]);
            }
            if (x == N - 1)
            {
                oX1 = std::max(oX1, pt[0]);
                iX1 = std::min(iX1, pt[0]);
            }
            if (y == 0)
            {
                oY0 = std::min(oY0, pt[1]);
                iY0 = std::max(iY0, pt[1]);
            }
            if (y == N - 1)
            {
                oY1 = std::max(oY1, pt[1]);
                iY1 = std::min(iY1, pt[1]);
            }
        }
    }
    float inner_width = iX1 - iX0;
    float inner_height = iY1 - iY0;
    float outer_width = oX1 - oX0;
    float outer_height = oY1 - oY0;

    // Projection mapping inner rectangle to viewport
    float fx0 = (new_size.first - 1) / inner_width;
    float fy0 = (new_size.second - 1) / inner_height;
    float cx0 = -fx0 * iX0;
    float cy0 = -fy0 * iY0;

    // Projection mapping outer rectangle to viewport
    float fx1 = (new_size.first - 1) / outer_width;
    float fy1 = (new_size.second - 1) / outer_height;
    float cx1 = -fx1 * oX0;
    float cy1 = -fy1 * oY0;

    // Interpolate between the two optimal projections
    float new_fx = fx0 * (1 - alpha) + fx1 * alpha;
    float new_fy = fy0 * (1 - alpha) + fy1 * alpha;
    float new_cx = cx0 * (1 - alpha) + cx1 * alpha;
    float new_cy = cy0 * (1 - alpha) + cy1 * alpha;

    std::array<std::array<float, 3>, 3> newK = {
        {{new_fx, 0.0, new_cx},
         {0.0, new_fy, new_cy},
         {0.0, 0.0, 1.0}}};
    return newK;
}