#include <iostream>
#include <vector>
#include <cmath>
#include <algorithm>
#include <chrono>
#include <thread>
#include <csignal>
#include <sqlite3.h>
#include <Eigen/Dense>

using namespace std;
using namespace Eigen;

// ===== ФЛАГ ЗАВЕРШЕНИЯ =====
volatile sig_atomic_t running = 1;

void shutdown(int signum) {
    running = 0;
}

// ===== СТРУКТУРЫ =====

struct IMUData {
    double t;
    Vector3d acc;
    Vector3d gyro;
};

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

struct State {
    VectorXd x;
    MatrixXd P;
};

// ===== ЧТЕНИЕ ИЗ SQLite =====

vector<IMUData> readIMU(sqlite3* db) {
    vector<IMUData> data;
    sqlite3_stmt* stmt;

    const char* sql = "SELECT timestamp, ax, ay, az, gx, gy, gz FROM imu_data ORDER BY timestamp";
    sqlite3_prepare_v2(db, sql, -1, &stmt, nullptr);

    while (sqlite3_step(stmt) == SQLITE_ROW) {
        IMUData d;
        d.t = sqlite3_column_double(stmt, 0);
        d.acc.x() = sqlite3_column_double(stmt, 1);
        d.acc.y() = sqlite3_column_double(stmt, 2);
        d.acc.z() = sqlite3_column_double(stmt, 3);
        d.gyro.x() = sqlite3_column_double(stmt, 4);
        d.gyro.y() = sqlite3_column_double(stmt, 5);
        d.gyro.z() = sqlite3_column_double(stmt, 6);
        data.push_back(d);
    }

    sqlite3_finalize(stmt);
    return data;
}

vector<LidarData> readLidar(sqlite3* db) {
    vector<LidarData> data;
    sqlite3_stmt* stmt;

    const char* sql = "SELECT timestamp, angle, distance_mm FROM lidar_data ORDER BY timestamp";
    sqlite3_prepare_v2(db, sql, -1, &stmt, nullptr);

    while (sqlite3_step(stmt) == SQLITE_ROW) {
        LidarData d;
        d.t = sqlite3_column_double(stmt, 0);
        d.angle = sqlite3_column_double(stmt, 1);
        d.distance = sqlite3_column_double(stmt, 2);
        data.push_back(d);
    }

    sqlite3_finalize(stmt);
    return data;
}

// ===== ЗАПИСЬ В SQLite =====

void initDB(sqlite3* db) {
    const char* sql =
        "CREATE TABLE IF NOT EXISTS trajectory ("
        "timestamp REAL, x REAL, y REAL, z REAL,"
        "roll REAL, pitch REAL, yaw REAL);"
        "DELETE FROM trajectory;"
        "CREATE TABLE IF NOT EXISTS lidar_points ("
        "x REAL, y REAL, z REAL);"
        "DELETE FROM lidar_points;";
    sqlite3_exec(db, sql, nullptr, nullptr, nullptr);
}

void writeTraj(sqlite3* db, double t, Vector3d pos, Vector3d ang) {
    sqlite3_stmt* stmt;
    const char* sql = "INSERT INTO trajectory VALUES (?,?,?,?,?,?,?)";
    sqlite3_prepare_v2(db, sql, -1, &stmt, nullptr);
    sqlite3_bind_double(stmt, 1, t);
    sqlite3_bind_double(stmt, 2, pos.x());
    sqlite3_bind_double(stmt, 3, pos.y());
    sqlite3_bind_double(stmt, 4, pos.z());
    sqlite3_bind_double(stmt, 5, ang.x());
    sqlite3_bind_double(stmt, 6, ang.y());
    sqlite3_bind_double(stmt, 7, ang.z());
    sqlite3_step(stmt);
    sqlite3_finalize(stmt);
}

void writeLidarPoint(sqlite3* db, Vector3d p) {
    sqlite3_stmt* stmt;
    const char* sql = "INSERT INTO lidar_points VALUES (?,?,?)";
    sqlite3_prepare_v2(db, sql, -1, &stmt, nullptr);
    sqlite3_bind_double(stmt, 1, p.x());
    sqlite3_bind_double(stmt, 2, p.y());
    sqlite3_bind_double(stmt, 3, p.z());
    sqlite3_step(stmt);
    sqlite3_finalize(stmt);
}

// ===== ФИЗИКА (без изменений) =====

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();
}

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;
}

// ===== ICP (без изменений) =====

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);
}

Matrix3d computeRotation(const vector<Vector3d>& src, const vector<Vector3d>& dst) {
    Vector3d c1 = Vector3d::Zero(), 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(), 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() {
    signal(SIGTERM, shutdown);
    signal(SIGINT, shutdown);

    sqlite3* db;
    sqlite3_open("../inertial_data.db", &db);
    initDB(db);

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

    vector<Vector3d> prev_scan, curr_scan;
    size_t lidar_idx = 0;
    size_t imu_idx = 0;

    cout << "Расчёты запущены..." << endl;

    while (running) {
        auto imu = readIMU(db);
        auto lidar = readLidar(db);

        // Обрабатываем только новые данные
        for (size_t i = max((size_t)1, imu_idx); 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;
                }
                for (auto& p : curr_scan)
                    writeLidarPoint(db, p);
                prev_scan = curr_scan;
                curr_scan.clear();
            }

            Vector3d pos = state.x.segment<3>(0);
            Vector3d ang = state.x.segment<3>(6);
            writeTraj(db, imu[i].t, pos, ang);
        }

        imu_idx = imu.size();

        // Ждём 200 мс перед следующим циклом
        this_thread::sleep_for(chrono::milliseconds(200));
    }

    sqlite3_close(db);
    cout << "Расчёты остановлены" << endl;
    return 0;
}