#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <cmath>
#include <algorithm>
#include <Eigen/Dense>

using namespace std;
using namespace Eigen;

// СТРУКТУРЫ

struct IMUData {
    double t;
    Vector3d acc;   // g
    Vector3d gyro;  // deg/s
};

struct LidarData {
    double t;
    double angle;
    double distance; // mm
};

struct State {
    VectorXd x;
    MatrixXd P;
};

// ЧТЕНИЕ CSV

vector<IMUData> readIMU(const string& filename) {
    vector<IMUData> data;
    ifstream file(filename);
    string line;

    getline(file, line);

    while (getline(file, line)) {
        stringstream ss(line);
        string val;

        IMUData d;

        getline(ss, val, ','); d.t = stod(val);
        getline(ss, val, ','); d.acc.x() = stod(val);
        getline(ss, val, ','); d.acc.y() = stod(val);
        getline(ss, val, ','); d.acc.z() = stod(val);
        getline(ss, val, ','); d.gyro.x() = stod(val);
        getline(ss, val, ','); d.gyro.y() = stod(val);
        getline(ss, val, ','); d.gyro.z() = stod(val);

        data.push_back(d);
    }

    return data;
}

vector<LidarData> readLidar(const string& filename) {
    vector<LidarData> data;
    ifstream file(filename);
    string line;

    getline(file, line);

    while (getline(file, line)) {
        stringstream ss(line);
        string val;

        LidarData d;

        getline(ss, val, ','); d.t = stod(val);
        getline(ss, val, ','); d.angle = stod(val);
        getline(ss, val, ','); d.distance = stod(val);

        data.push_back(d);
    }

    return data;
}

// ПОВОРОТ

Matrix3d rotationMatrix(double roll, double pitch, double yaw) {
    AngleAxisd rx(roll, Vector3d::UnitX());
    AngleAxisd ry(pitch, Vector3d::UnitY());
    AngleAxisd rz(yaw, Vector3d::UnitZ());
    return (rz * ry * rx).toRotationMatrix();
}

// IMU PREDICT

void predict(State& s, const IMUData& imu, double dt) {
    VectorXd& x = s.x;

    Vector3d pos = x.segment<3>(0);
    Vector3d vel = x.segment<3>(3);
    Vector3d angles = x.segment<3>(6);

    Vector3d acc = imu.acc * 9.81;
    Vector3d gyro = imu.gyro * M_PI / 180.0;

    Matrix3d R = rotationMatrix(angles[0], angles[1], angles[2]);

    Vector3d g(0, 0, 9.81);
    Vector3d acc_world = R * acc - g;

    vel += acc_world * dt;
    pos += vel * dt + 0.5 * acc_world * dt * dt;
    angles += gyro * dt;

    x.segment<3>(0) = pos;
    x.segment<3>(3) = vel;
    x.segment<3>(6) = angles;

    MatrixXd F = MatrixXd::Identity(15, 15);
    F.block<3, 3>(0, 3) = Matrix3d::Identity() * dt;

    MatrixXd Q = MatrixXd::Identity(15, 15) * 0.01;

    s.P = F * s.P * F.transpose() + Q;
}

// ЛИДАР → МИР

Vector3d lidarToWorld(const State& s, const LidarData& l) {
    double angle = l.angle * M_PI / 180.0;
    double dist = l.distance / 1000.0;

    Vector3d local;
    local << dist * cos(angle), dist* sin(angle), 0;

    Vector3d pos = s.x.segment<3>(0);
    Vector3d ang = s.x.segment<3>(6);

    Matrix3d R = rotationMatrix(ang[0], ang[1], ang[2]);

    return pos + R * local;
}

// KD-TREE

struct KDNode {
    Vector3d point;
    KDNode* left;
    KDNode* right;
    int axis;

    KDNode(Vector3d p, int ax) : point(p), left(nullptr), right(nullptr), axis(ax) {}
};

KDNode* buildKDTree(vector<Vector3d>& pts, int depth = 0) {
    if (pts.empty()) return nullptr;

    int axis = depth % 3;

    sort(pts.begin(), pts.end(),
        [axis](const Vector3d& a, const Vector3d& b) {
            return a[axis] < b[axis];
        });

    int mid = pts.size() / 2;

    KDNode* node = new KDNode(pts[mid], axis);

    vector<Vector3d> left(pts.begin(), pts.begin() + mid);
    vector<Vector3d> right(pts.begin() + mid + 1, pts.end());

    node->left = buildKDTree(left, depth + 1);
    node->right = buildKDTree(right, depth + 1);

    return node;
}

void nearestSearch(KDNode* node, const Vector3d& target,
    Vector3d& best, double& best_dist) {
    if (!node) return;

    double d = (node->point - target).squaredNorm();
    if (d < best_dist) {
        best_dist = d;
        best = node->point;
    }

    int axis = node->axis;

    KDNode* near = target[axis] < node->point[axis] ? node->left : node->right;
    KDNode* far = target[axis] < node->point[axis] ? node->right : node->left;

    nearestSearch(near, target, best, best_dist);

    double diff = target[axis] - node->point[axis];
    if (diff * diff < best_dist) {
        nearestSearch(far, target, best, best_dist);
    }
}

// ICP

Matrix3d computeRotation(const vector<Vector3d>& src,
    const vector<Vector3d>& dst) {

    Vector3d c1 = Vector3d::Zero();
    Vector3d c2 = Vector3d::Zero();

    for (size_t i = 0;i < src.size();i++) {
        c1 += src[i];
        c2 += dst[i];
    }

    c1 /= src.size();
    c2 /= dst.size();

    Matrix3d H = Matrix3d::Zero();

    for (size_t i = 0;i < src.size();i++) {
        H += (src[i] - c1) * (dst[i] - c2).transpose();
    }

    JacobiSVD<Matrix3d> svd(H, ComputeFullU | ComputeFullV);
    return svd.matrixV() * svd.matrixU().transpose();
}

Vector3d computeTranslation(const vector<Vector3d>& src,
    const vector<Vector3d>& dst,
    const Matrix3d& R) {

    Vector3d c1 = Vector3d::Zero();
    Vector3d c2 = Vector3d::Zero();

    for (size_t i = 0;i < src.size();i++) {
        c1 += src[i];
        c2 += dst[i];
    }

    c1 /= src.size();
    c2 /= dst.size();

    return c2 - R * c1;
}

vector<Vector3d> downsample(const vector<Vector3d>& pts, int step = 5) {
    vector<Vector3d> out;
    for (size_t i = 0;i < pts.size();i += step)
        out.push_back(pts[i]);
    return out;
}

void ICP_fast(vector<Vector3d>& src,
    vector<Vector3d>& dst,
    Matrix3d& R, Vector3d& t) {

    R = Matrix3d::Identity();
    t = Vector3d::Zero();

    vector<Vector3d> dst_copy = dst;
    KDNode* tree = buildKDTree(dst_copy);

    for (int iter = 0; iter < 10; iter++) {

        vector<Vector3d> matched;

        for (auto& p : src) {
            Vector3d best;
            double best_dist = 1e9;

            nearestSearch(tree, p, best, best_dist);

            if (best_dist < 0.5)
                matched.push_back(best);
        }

        if (matched.size() < 10) break;

        Matrix3d dR = computeRotation(src, matched);
        Vector3d dt = computeTranslation(src, matched, dR);

        for (auto& p : src)
            p = dR * p + dt;

        R = dR * R;
        t = dR * t + dt;
    }
}

// MAIN

int main() {
    auto imu = readIMU("../imu_module/imu.csv");
    auto lidar = readLidar("../lidar_module/lidar.csv");

    State state;
    state.x = VectorXd::Zero(15);
    state.P = MatrixXd::Identity(15, 15) * 0.1;

    ofstream traj("trajectory.csv");
    traj << "timestamp,x,y,z,roll,pitch,yaw\n";

    ofstream lidar_out("lidar_points.csv");
    lidar_out << "x,y,z\n";

    vector<Vector3d> prev_scan, curr_scan;

    size_t lidar_idx = 0;

    for (size_t i = 1;i < imu.size();i++) {

        double dt = imu[i].t - imu[i - 1].t;
        if (dt <= 0 || dt > 0.1) continue;

        predict(state, imu[i], dt);

        while (lidar_idx < lidar.size() &&
            lidar[lidar_idx].t <= imu[i].t) {

            curr_scan.push_back(lidarToWorld(state, lidar[lidar_idx]));
            lidar_idx++;
        }

        if (curr_scan.size() > 200) {

            if (!prev_scan.empty()) {

                vector<Vector3d> src = downsample(curr_scan, 5);
                vector<Vector3d> dst = downsample(prev_scan, 5);

                Matrix3d R;
                Vector3d t;

                ICP_fast(src, dst, R, t);

                state.x.segment<3>(0) += t;
            }

            prev_scan = curr_scan;
            curr_scan.clear();
        }

        Vector3d pos = state.x.segment<3>(0);
        Vector3d ang = state.x.segment<3>(6);

        traj << imu[i].t << ","
            << pos.x() << ","
            << pos.y() << ","
            << pos.z() << ","
            << ang.x() << ","
            << ang.y() << ","
            << ang.z() << "\n";
    }

    for (auto& p : prev_scan) {
        lidar_out << p.x() << ","
            << p.y() << ","
            << p.z() << "\n";
    }

    cout << "Done\n";
    return 0;
}