#include "CurveFitting.h"
#include <Eigen/Dense>
#include <Eigen/Eigenvalues>
#include <cmath>

using namespace Eigen;

LineFitResult FitLine3D(const std::vector<SVzNL3DPosition>& points)
{
    LineFitResult result;
    result.success = false;
    result.error = 0.0;

    if (points.size() < 2)
        return result;

    // 计算质心
    Vector3d centroid(0, 0, 0);
    for (const auto& p : points)
    {
        centroid.x() += p.pt3D.x;
        centroid.y() += p.pt3D.y;
        centroid.z() += p.pt3D.z;
    }
    centroid /= points.size();

    // 构建协方差矩阵
    Matrix3d covariance = Matrix3d::Zero();
    for (const auto& p : points)
    {
        Vector3d pt(p.pt3D.x - centroid.x(),
                    p.pt3D.y - centroid.y(),
                    p.pt3D.z - centroid.z());
        covariance += pt * pt.transpose();
    }

    // 特征值分解，最大特征值对应的特征向量即为方向向量
    SelfAdjointEigenSolver<Matrix3d> solver(covariance);
    Vector3d direction = solver.eigenvectors().col(2);

    // 计算拟合误差
    double totalError = 0.0;
    for (const auto& p : points)
    {
        Vector3d pt(p.pt3D.x - centroid.x(),
                    p.pt3D.y - centroid.y(),
                    p.pt3D.z - centroid.z());
        double dist = (pt - pt.dot(direction) * direction).norm();
        totalError += dist * dist;
    }

    result.point.x = centroid.x();
    result.point.y = centroid.y();
    result.point.z = centroid.z();
    result.direction.x = direction.x();
    result.direction.y = direction.y();
    result.direction.z = direction.z();
    result.error = std::sqrt(totalError / points.size());
    result.success = true;

    return result;
}

ParabolaFitResult FitParabola(const std::vector<SVzNL3DPosition>& points, Plane plane)
{
    ParabolaFitResult result;
    result.success = false;
    result.error = 0.0;

    if (points.size() < 3)
        return result;

    MatrixXd A(points.size(), 3);
    VectorXd B(points.size());

    for (size_t i = 0; i < points.size(); ++i)
    {
        double u, v;
        if (plane == XY) {
            u = points[i].pt3D.x;
            v = points[i].pt3D.y;
        } else if (plane == XZ) {
            u = points[i].pt3D.x;
            v = points[i].pt3D.z;
        } else {
            u = points[i].pt3D.y;
            v = points[i].pt3D.z;
        }
        A(i, 0) = u * u;
        A(i, 1) = u;
        A(i, 2) = 1.0;
        B(i) = v;
    }

    Vector3d coeffs = A.colPivHouseholderQr().solve(B);
    result.a = coeffs(0);
    result.b = coeffs(1);
    result.c = coeffs(2);

    double totalError = 0.0;
    for (const auto& p : points)
    {
        double u, v;
        if (plane == XY) {
            u = p.pt3D.x;
            v = p.pt3D.y;
        } else if (plane == XZ) {
            u = p.pt3D.x;
            v = p.pt3D.z;
        } else {
            u = p.pt3D.y;
            v = p.pt3D.z;
        }
        double v_fit = result.a * u * u + result.b * u + result.c;
        double error = v_fit - v;
        totalError += error * error;
    }

    result.error = std::sqrt(totalError / points.size());
    result.success = true;
    return result;
}

CircleFitResult FitCircle(const std::vector<SVzNL3DPosition>& points, Plane plane)
{
    CircleFitResult result;
    result.success = false;
    result.error = 0.0;

    if (points.size() < 3)
        return result;

    MatrixXd A(points.size(), 3);
    VectorXd B(points.size());

    for (size_t i = 0; i < points.size(); ++i)
    {
        double u, v;
        if (plane == XY) {
            u = points[i].pt3D.x;
            v = points[i].pt3D.y;
        } else if (plane == XZ) {
            u = points[i].pt3D.x;
            v = points[i].pt3D.z;
        } else {
            u = points[i].pt3D.y;
            v = points[i].pt3D.z;
        }
        A(i, 0) = u;
        A(i, 1) = v;
        A(i, 2) = 1.0;
        B(i) = -(u * u + v * v);
    }

    Vector3d coeffs = A.colPivHouseholderQr().solve(B);
    double D = coeffs(0);
    double E = coeffs(1);
    double F = coeffs(2);

    double cu = -D / 2.0;
    double cv = -E / 2.0;
    result.radius = std::sqrt(D * D / 4.0 + E * E / 4.0 - F);

    if (plane == XY) {
        result.center.x = cu;
        result.center.y = cv;
        result.center.z = 0.0;
    } else if (plane == XZ) {
        result.center.x = cu;
        result.center.y = 0.0;
        result.center.z = cv;
    } else {
        result.center.x = 0.0;
        result.center.y = cu;
        result.center.z = cv;
    }

    double totalError = 0.0;
    for (const auto& p : points)
    {
        double u, v;
        if (plane == XY) {
            u = p.pt3D.x;
            v = p.pt3D.y;
        } else if (plane == XZ) {
            u = p.pt3D.x;
            v = p.pt3D.z;
        } else {
            u = p.pt3D.y;
            v = p.pt3D.z;
        }
        double du = u - cu;
        double dv = v - cv;
        double dist = std::sqrt(du * du + dv * dv);
        double error = dist - result.radius;
        totalError += error * error;
    }

    result.error = std::sqrt(totalError / points.size());
    result.success = true;
    return result;
}