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

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

// Состояние системы
Vector3d pos    = Vector3d::Zero();  // позиция (м)
Vector3d vel    = Vector3d::Zero();  // скорость (м/с)
Vector3d angles = Vector3d::Zero();  // углы roll, pitch, yaw (рад)

// Bias акселерометра (в единицах g).
// При калибровке измеряется начальное значение,
// затем автоматически уточняется в каждую остановку (ZUPT).
double ax_bias = 0.0;
double ay_bias = 0.0;
double az_bias = 1.0;  // ~1g — гравитация по Z когда датчик горизонтальный

// Параметры ZUPT
const double ZUPT_ACC_THR  = 0.08;  // отклонение нормы от 1g
const double ZUPT_GYRO_THR = 0.05;  // рад/с (~3 °/с)
const int    ZUPT_WINDOW   = 10;    // сэмплов подряд (~120 мс)

// Коэффициент комплементарного фильтра
// 0.98 = 98% гироскоп (быстро, но дрейфует)
// 0.02 = 2%  акселерометр (медленно, но абсолютное значение)
const double CF_ALPHA = 0.98;

// ─────────────────────────────────────────────
// Чтение БД
// ─────────────────────────────────────────────
vector<IMUData> readIMU(sqlite3* db, double after_ts) {
    vector<IMUData> data;
    sqlite3_stmt* stmt;
    sqlite3_prepare_v2(db,
        "SELECT timestamp, ax, ay, az, gx, gy, gz "
        "FROM imu_data WHERE timestamp > ? ORDER BY timestamp",
        -1, &stmt, nullptr);
    sqlite3_bind_double(stmt, 1, after_ts);
    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, double after_ts) {
    vector<LidarData> data;
    sqlite3_stmt* stmt;
    sqlite3_prepare_v2(db,
        "SELECT timestamp, angle, distance_mm "
        "FROM lidar_data WHERE timestamp > ? ORDER BY timestamp",
        -1, &stmt, nullptr);
    sqlite3_bind_double(stmt, 1, after_ts);
    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;
}

// ─────────────────────────────────────────────
// Запись в БД
// ─────────────────────────────────────────────
void initDB(sqlite3* db) {
    sqlite3_exec(db,
        "PRAGMA journal_mode=WAL;"
        "PRAGMA synchronous=NORMAL;"
        "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;",
        nullptr, nullptr, nullptr);
}

void writeTraj(sqlite3* db, double t) {
    sqlite3_stmt* stmt;
    sqlite3_prepare_v2(db,
        "INSERT INTO trajectory VALUES (?,?,?,?,?,?,?)",
        -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, angles.x());
    sqlite3_bind_double(stmt, 6, angles.y());
    sqlite3_bind_double(stmt, 7, angles.z());
    sqlite3_step(stmt);
    sqlite3_finalize(stmt);
}

void writeLidarPoint(sqlite3* db, Vector3d p) {
    sqlite3_stmt* stmt;
    sqlite3_prepare_v2(db, "INSERT INTO lidar_points VALUES (?,?,?)",
                       -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);
}

// ─────────────────────────────────────────────
// Матрица поворота (roll, pitch, yaw)
// ─────────────────────────────────────────────
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();
}

Vector3d lidarToWorld(const LidarData& l) {
    double angle = l.angle * M_PI / 180.0;
    double dist  = l.distance / 1000.0;
    Vector3d local(dist * cos(angle), dist * sin(angle), 0);
    return pos + rotationMatrix(angles[0], angles[1], angles[2]) * local;
}

// ─────────────────────────────────────────────
// predict() — формулы с доски + комплементарный фильтр
//
// Формулы преподавателя (для всех трёх осей OX, OY, OZ):
//   V(nT) = V((n-1)T) + a·T
//   X(nT) = V((n-1)T)·T + a·T²/2
//   φ     = ω·T
//
// Комплементарный фильтр — решает проблему наклона:
//   angle = α·(angle + gyro·dt) + (1-α)·angle_from_accel
//   Гироскоп отслеживает быстрые повороты.
//   Акселерометр корректирует накопленный дрейф.
//
// ZUPT — когда стоим: замораживаем всё, уточняем bias.
// ─────────────────────────────────────────────
void predict(const IMUData& imu, double dt) {
    // Угловые скорости (рад/с)
    Vector3d gyro = imu.gyro * M_PI / 180.0;

    // ── Комплементарный фильтр ориентации ──────────────
    // Углы наклона из акселерометра (абсолютные, но шумные)
    double roll_acc  = atan2(imu.acc.y(), imu.acc.z());
    double pitch_acc = atan2(-imu.acc.x(),
                              sqrt(imu.acc.y() * imu.acc.y() +
                                   imu.acc.z() * imu.acc.z()));

    // Комбинируем: быстрая часть (гироскоп) + медленная коррекция (акселерометр)
    // φ = α·(φ + ω·T) + (1-α)·φ_акселерометр
    angles.x() = CF_ALPHA * (angles.x() + gyro.x() * dt) + (1.0 - CF_ALPHA) * roll_acc;
    angles.y() = CF_ALPHA * (angles.y() + gyro.y() * dt) + (1.0 - CF_ALPHA) * pitch_acc;
    angles.z() += gyro.z() * dt;  // yaw — только гироскоп (нет абсолютного источника)

    // ── Линейное ускорение в мировой системе ───────────
    // Шаг 1: вычитаем bias датчика и переводим в м/с²
    Vector3d acc_sensor(
        (imu.acc.x() - ax_bias) * 9.81,
        (imu.acc.y() - ay_bias) * 9.81,
        imu.acc.z() * 9.81   // Z: bias не вычитаем, гравитация убирается через матрицу поворота
    );

    // Шаг 2: поворачиваем в мировую систему координат
    Matrix3d R = rotationMatrix(angles.x(), angles.y(), angles.z());
    Vector3d acc_world = R * acc_sensor;

    // Шаг 3: вычитаем гравитацию (она направлена вниз = +Z в мировой системе)
    acc_world.z() -= 9.81;

    // ── ZUPT: детектор покоя ────────────────────────────
    double acc_mag  = imu.acc.norm();
    double gyro_mag = gyro.norm();
    static int zupt_counter = 0;

    if (abs(acc_mag - 1.0) < ZUPT_ACC_THR && gyro_mag < ZUPT_GYRO_THR) {
        zupt_counter++;
    } else {
        zupt_counter = 0;
    }

    if (zupt_counter >= ZUPT_WINDOW) {
        // ── СТОИМ ──────────────────────────────────────
        // Замораживаем скорость
        vel = Vector3d::Zero();
        // Уточняем bias по всем трём осям (медленное скользящее среднее)
        ax_bias = 0.99 * ax_bias + 0.01 * imu.acc.x();
        ay_bias = 0.99 * ay_bias + 0.01 * imu.acc.y();
        az_bias = 0.99 * az_bias + 0.01 * imu.acc.z();
        // pos не меняем
    } else {
        // ── ДВИЖЕМСЯ ───────────────────────────────────
        // Формулы с доски для OX, OY, OZ:
        // X(nT) = V((n-1)T)·T + a·T²/2
        // V(nT) = V((n-1)T) + a·T
        pos += vel * dt + 0.5 * acc_world * dt * dt;
        vel += acc_world * dt;
    }
}

// ─────────────────────────────────────────────
// 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; }
    KDNode* near_ = target[node->axis] < node->point[node->axis] ? node->left  : node->right;
    KDNode* far_  = target[node->axis] < node->point[node->axis] ? node->right : node->left;
    nearestSearch(near_, target, best, best_dist);
    double diff = target[node->axis] - node->point[node->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;
    }
}

// ─────────────────────────────────────────────
// Калибровка: датчик лежит горизонтально, неподвижно.
// Измеряем начальный bias по всем трём осям.
// ─────────────────────────────────────────────
void calibrate(sqlite3* db) {
    cout << "Калибровка (держите датчик горизонтально, неподвижно)..." << endl;

    double sum_ax = 0, sum_ay = 0, sum_az = 0;
    int n = 0;
    double last_ts = 0;
    bool seen_data = false;
    auto start = chrono::steady_clock::now();
    auto first_data_time = chrono::steady_clock::time_point();

    while (running) {
        auto imu = readIMU(db, last_ts);
        for (auto& d : imu) {
            if (!seen_data) {
                seen_data = true;
                first_data_time = chrono::steady_clock::now();
                cout << "Данные IMU получены, собираю калибровку..." << endl;
            }
            sum_ax += d.acc.x();
            sum_ay += d.acc.y();
            sum_az += d.acc.z();
            n++;
            last_ts = d.t;
        }

        if (seen_data) {
            auto elapsed = chrono::steady_clock::now() - first_data_time;
            if (elapsed >= chrono::seconds(2) && n >= 30) break;
        } else {
            auto waited = chrono::steady_clock::now() - start;
            if (waited >= chrono::seconds(60)) {
                cout << "IMU не отвечает. Использую стандартный bias." << endl;
                return;
            }
        }
        this_thread::sleep_for(chrono::milliseconds(100));
    }

    if (n > 0) {
        ax_bias = sum_ax / n;
        ay_bias = sum_ay / n;
        az_bias = sum_az / n;
        cout << "Калибровка завершена. Bias: ax=" << ax_bias
             << "g, ay=" << ay_bias << "g, az=" << az_bias
             << "g (по " << n << " сэмплам)" << endl;
        cout << "Bias будет уточняться автоматически во время остановок." << endl;
    }
}

// ─────────────────────────────────────────────
// main
// ─────────────────────────────────────────────
int main() {
    signal(SIGTERM, shutdown);
    signal(SIGINT,  shutdown);

    sqlite3* db;
    char exe_path[1024];
    ssize_t len = readlink("/proc/self/exe", exe_path, sizeof(exe_path)-1);
    exe_path[len] = '\0';
    string db_path = string(dirname(exe_path)) + "/../inertial_data.db";
    sqlite3_open(db_path.c_str(), &db);
    initDB(db);

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

    double last_imu_ts   = 0;
    double last_lidar_ts = 0;
    {
        sqlite3_stmt* stmt;
        sqlite3_prepare_v2(db, "SELECT MAX(timestamp) FROM imu_data",
                           -1, &stmt, nullptr);
        if (sqlite3_step(stmt) == SQLITE_ROW)
            last_imu_ts = sqlite3_column_double(stmt, 0);
        sqlite3_finalize(stmt);

        sqlite3_prepare_v2(db, "SELECT MAX(timestamp) FROM lidar_data",
                           -1, &stmt, nullptr);
        if (sqlite3_step(stmt) == SQLITE_ROW)
            last_lidar_ts = sqlite3_column_double(stmt, 0);
        sqlite3_finalize(stmt);
    }

    vector<Vector3d> prev_scan, curr_scan;
    vector<LidarData> pending_lidar;

    while (running) {
        auto imu       = readIMU(db, last_imu_ts);
        auto new_lidar = readLidar(db, last_lidar_ts);
        if (!new_lidar.empty()) last_lidar_ts = new_lidar.back().t;
        pending_lidar.insert(pending_lidar.end(),
                             new_lidar.begin(), new_lidar.end());

        for (size_t i = 0; i < imu.size(); i++) {
            double prev_ts = (i == 0) ? last_imu_ts : imu[i-1].t;
            double dt = imu[i].t - prev_ts;
            if (dt <= 0 || dt > 0.1) continue;

            predict(imu[i], dt);

            while (!pending_lidar.empty() &&
                   pending_lidar.front().t <= imu[i].t) {
                curr_scan.push_back(lidarToWorld(pending_lidar.front()));
                pending_lidar.erase(pending_lidar.begin());
            }

            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);
                    pos += t;
                }
                for (auto& p : curr_scan) writeLidarPoint(db, p);
                prev_scan = curr_scan;
                curr_scan.clear();
            }

            writeTraj(db, imu[i].t);
        }

        if (!imu.empty()) last_imu_ts = imu.back().t;
        this_thread::sleep_for(chrono::milliseconds(200));
    }

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