constexpr float PI = 3.14159265358979323846f;

void undistort_points_pinhole_cpp(
    const std::vector<float> &K,
    const std::vector<float> &D,
    int num_points,
    std::vector<float> &unit_vector)
{
    std::vector<float> distorted(num_points * 2, 0.f);
    std::vector<float> undistorted(num_points * 2, 0.f);

    // Generate distorted points (same as original)
    for (int y = 0; y < 48; ++y)
    {
        for (int x = 0; x < 120; ++x)
        {
            int idx = (y * 120 + x) * 2;
            distorted[idx] = (x + 0.5f);
            distorted[idx + 1] = (y + 0.5f);
        }
    }

    unit_vector.resize(num_points * 3);

    float fx = K[0];
    float fy = K[4];
    float cx = K[2];
    float cy = K[5];

    float k1 = D[0], k2 = D[1], p1 = D[2], p2 = D[3], k3 = D[4];

    for (int idx = 0; idx < num_points; ++idx)
    {
        int idx_two = idx * 2;
        int idx_three = idx * 3;

        float u = distorted[idx_two + 0];
        float v = distorted[idx_two + 1];

        //---------------------------------------------------------
        // Step 1: Normalize distorted coords
        //---------------------------------------------------------
        float xd = (u - cx) / fx;
        float yd = (v - cy) / fy;

        //---------------------------------------------------------
        // Step 2: Iteratively solve undistorted x,y (Brown–Conrady)
        //---------------------------------------------------------
        float x = xd;
        float y = yd;

        for (int iter = 0; iter < 5; iter++)
        {
            float r2 = x * x + y * y;
            float r4 = r2 * r2;
            float r6 = r4 * r2;

            float radial = 1 + k1 * r2 + k2 * r4 + k3 * r6;

            float x_tangential = 2 * p1 * x * y + p2 * (r2 + 2 * x * x);
            float y_tangential = p1 * (r2 + 2 * y * y) + 2 * p2 * x * y;

            float x_est = (xd - x_tangential) / radial;
            float y_est = (yd - y_tangential) / radial;

            x = x_est;
            y = y_est;
        }

        //---------------------------------------------------------
        // Step 3: Convert to undistorted pixel coords
        //---------------------------------------------------------
        float u_und = fx * x + cx;
        float v_und = fy * y + cy;

        undistorted[idx_two + 0] = u_und;
        undistorted[idx_two + 1] = v_und;

        //---------------------------------------------------------
        // Step 4: Compute camera-frame unit vector
        //---------------------------------------------------------
        float X = x;
        float Y = y;
        float Z = 1.0f;

        float norm = std::sqrt(X * X + Y * Y + Z * Z);

        //---------------------------------------------------------
        // Optional ellipse mask (unchanged)
        //---------------------------------------------------------
        // float ra = 1.f / (cx - 3);
        // float rb = 1.f / (cy + 4);

        // float mask = ((u - cx)*(u - cx)*ra*ra +
        //               (v - cy)*(v - cy)*rb*rb - 1.f) < 0.f ? 1.f : 0.f;
        float mask = 1;
        unit_vector[idx_three + 0] = mask * (X / norm);
        unit_vector[idx_three + 1] = mask * (Y / norm);
        unit_vector[idx_three + 2] = mask * (Z / norm);
    }
}

void tof_to_xyz(
    const std::vector<float> &unit_vector,
    const std::vector<float> &tof,
    int num_points,
    std::vector<float> &points3D)
{
    points3D.resize(num_points * 3);

    for (int i = 0; i < num_points; ++i)
    {
        int idx3 = i * 3;
        float d = tof[i];

        points3D[idx3 + 0] = unit_vector[idx3 + 0] * d;
        points3D[idx3 + 1] = unit_vector[idx3 + 1] * d;
        points3D[idx3 + 2] = unit_vector[idx3 + 2] * d;
    }
}

void depth_to_xyz(
    const std::vector<float> &unit_vector,
    const std::vector<float> &depth,
    int num_points,
    std::vector<float> &points3D)
{
    points3D.resize(num_points * 3);

    for (int i = 0; i < num_points; ++i)
    {
        int idx3 = i * 3;
        float d = depth[i];

        points3D[idx3 + 0] = unit_vector[idx3 + 0]/ unit_vector[idx3 + 2] * d;
        points3D[idx3 + 1] = unit_vector[idx3 + 1]/ unit_vector[idx3 + 2] * d;
        points3D[idx3 + 2] = d;
    }
}

int main()
{
    int WIDTH = 40;
    int HEIGHT = 30;
    int num_points = WIDTH * HEIGHT;

    std::vector<float> unit_vector(num_points * 3, 0.f);
    std::vector<float> depth(num_points , 0.f);
    std::vector<float> points3D(num_points * 3 , 0.f); // 改成 *3

    // 内参
    float fx = 37;    
    float fy = 37;
    float cx = 20;
    float cy = 15;

    // 畸变参数
    float k1 = 0.f;
    float k2 = 0.f;
    float p1 = 0.f;
    float p2 = 0.f;
    float k3 = 0.f;

    std::vector<float> K = {fx , 0.f, cx , 0.f, fy , cy , 0.f, 0.f, 1.f};
    std::vector<float> D = {k1, k2, p1, p2, k3};

    // 生成 unit_vector
    undistort_points_pinhole_cpp(K, D, num_points, unit_vector);

    // 将 depth 转成相机坐标系
    depth_to_xyz(unit_vector, depth, num_points, points3D);

    return 0;
}