#define _USE_MATH_DEFINES
#include <iostream>
#include <vector>
#include <cmath>
#include <unistd.h>
#include "sqlite3.h"

using namespace std;

static const double G_MS2 = 9.81;   // g -> м/с²
static const double DEG2RAD = M_PI / 180.0;


class Kalman1D {
private:
    double x;      // положение
    double v;      // скорость
    double P[2][2]; // ковариационная матрица 2x2
    double Q;       // интенсивность шума процесса (ускорение)
    double R;       // дисперсия шума измерения
    double dt;      // временной шаг

public:
    Kalman1D(double q_accel, double r_measure, double init_x = 0.0, double init_v = 0.0, double dt0 = 0.01)
        : Q(q_accel), R(r_measure), dt(dt0), x(init_x), v(init_v)
    {
        P[0][0] = 1.0;
        P[0][1] = 0.0;
        P[1][0] = 0.0;
        P[1][1] = 100.0;
    }

    void setDt(double new_dt) { dt = new_dt; }

    void predict() {
        double new_x = x + v * dt;
        double new_v = v;

        double P00 = P[0][0] + dt * (P[1][0] + P[0][1]) + dt * dt * P[1][1];
        double P01 = P[0][1] + dt * P[1][1];
        double P10 = P[1][0] + dt * P[1][1];
        double P11 = P[1][1];

        double Qv = Q * dt;
        double Qp = Q * dt * dt * dt / 3.0;
        double Qpv = Q * dt * dt / 2.0;

        P00 += Qp;
        P01 += Qpv;
        P10 += Qpv;
        P11 += Qv;

        P[0][0] = P00; P[0][1] = P01;
        P[1][0] = P10; P[1][1] = P11;

        x = new_x;
        v = new_v;
    }

    double update(double measurement) {
        double y = measurement - x;
        double S = P[0][0] + R;
        double K_pos = P[0][0] / S;
        double K_vel = P[1][0] / S;

        x += K_pos * y;
        v += K_vel * y;

        double P00_new = (1.0 - K_pos) * P[0][0];
        double P01_new = (1.0 - K_pos) * P[0][1];
        double P11_new = P[1][1] - K_vel * P[1][0];

        P[0][0] = P00_new; P[0][1] = P01_new;
        P[1][0] = P01_new; P[1][1] = P11_new;

        return x;
    }

    double getState() const { return x; }
    double getVelocity() const { return v; }
};


// ─────────────────────────────────────────────
//  Структуры
// ─────────────────────────────────────────────
struct AccSample {
    double timestamp;
    double ax, ay, az;
};

struct GyroSample {
    double timestamp;
    double gx, gy, gz;
};

// ─────────────────────────────────────────────
//  Чтение из БД
// ─────────────────────────────────────────────

// Читаем последний timestamp и последнее состояние навигации из trajectory
double readLastTrajectoryState(sqlite3* db,
    double& px, double& py, double& pz,
    double& vx, double& vy, double& vz,
    double& roll, double& pitch, double& yaw)
{
    double last_ts = 0.0;
    sqlite3_stmt* stmt = nullptr;
    if (sqlite3_prepare_v2(db,
        "SELECT timestamp, x, y, z, vx, vy, vz, roll, pitch, yaw "
        "FROM trajectory ORDER BY timestamp DESC LIMIT 1",
        -1, &stmt, nullptr) == SQLITE_OK)
    {
        if (sqlite3_step(stmt) == SQLITE_ROW) {
            last_ts = sqlite3_column_double(stmt, 0);
            px    = sqlite3_column_double(stmt, 1);
            py    = sqlite3_column_double(stmt, 2);
            pz    = sqlite3_column_double(stmt, 3);
            vx    = sqlite3_column_double(stmt, 4);
            vy    = sqlite3_column_double(stmt, 5);
            vz    = sqlite3_column_double(stmt, 6);
            roll  = sqlite3_column_double(stmt, 7);
            pitch = sqlite3_column_double(stmt, 8);
            yaw   = sqlite3_column_double(stmt, 9);
        }
    }
    sqlite3_finalize(stmt);
    return last_ts;
}

vector<AccSample> readAcc(sqlite3* db, double last_timestamp) {
    vector<AccSample> data;
    sqlite3_stmt* stmt = nullptr;
    sqlite3_prepare_v2(db,
        "SELECT timestamp, ax, ay, az "
        "FROM filtered_acc_data WHERE timestamp > ? ORDER BY timestamp ASC",
        -1, &stmt, nullptr);
    sqlite3_bind_double(stmt, 1, last_timestamp);
    while (sqlite3_step(stmt) == SQLITE_ROW) {
        AccSample s;
        s.timestamp = sqlite3_column_double(stmt, 0);
        s.ax = sqlite3_column_double(stmt, 1);
        s.ay = sqlite3_column_double(stmt, 2);
        s.az = sqlite3_column_double(stmt, 3);
        data.push_back(s);
    }
    sqlite3_finalize(stmt);
    return data;
}

vector<GyroSample> readGyro(sqlite3* db, double last_timestamp) {
    vector<GyroSample> data;
    sqlite3_stmt* stmt = nullptr;
    sqlite3_prepare_v2(db,
        "SELECT timestamp, gx, gy, gz "
        "FROM filtered_gyro_data WHERE timestamp > ? ORDER BY timestamp ASC",
        -1, &stmt, nullptr);
    sqlite3_bind_double(stmt, 1, last_timestamp);
    while (sqlite3_step(stmt) == SQLITE_ROW) {
        GyroSample s;
        s.timestamp = sqlite3_column_double(stmt, 0);
        s.gx = sqlite3_column_double(stmt, 1);
        s.gy = sqlite3_column_double(stmt, 2);
        s.gz = sqlite3_column_double(stmt, 3);
        data.push_back(s);
    }
    sqlite3_finalize(stmt);
    return data;
}

// ─────────────────────────────────────────────
//  Запись в БД
// ─────────────────────────────────────────────
void createTable(sqlite3* db) {
    char* err = nullptr;
    sqlite3_exec(db,
        "CREATE TABLE IF NOT EXISTS trajectory ("
        "timestamp REAL, "
        "x REAL, y REAL, z REAL, "
        "vx REAL, vy REAL, vz REAL, "
        "roll REAL, pitch REAL, yaw REAL)",
        nullptr, nullptr, &err);
}

void writePoint(sqlite3_stmt* stmt,
    double ts,
    double x, double y, double z,
    double vx, double vy, double vz,
    double roll, double pitch, double yaw)
{
    sqlite3_reset(stmt);
    sqlite3_bind_double(stmt, 1, ts);
    sqlite3_bind_double(stmt, 2, x);
    sqlite3_bind_double(stmt, 3, y);
    sqlite3_bind_double(stmt, 4, z);
    sqlite3_bind_double(stmt, 5, vx);
    sqlite3_bind_double(stmt, 6, vy);
    sqlite3_bind_double(stmt, 7, vz);
    sqlite3_bind_double(stmt, 8, roll);
    sqlite3_bind_double(stmt, 9, pitch);
    sqlite3_bind_double(stmt, 10, yaw);
    sqlite3_step(stmt);
}

// ─────────────────────────────────────────────
//  main
// ─────────────────────────────────────────────
int main()
{
    sqlite3* db = nullptr;
    if (sqlite3_open("/home/grant/inertial_control/inertial_data.db", &db) != SQLITE_OK) {
        cerr << "Не удалось открыть БД" << endl;
        return 1;
    }

    sqlite3_busy_timeout(db, 5000);
    sqlite3_exec(db, "PRAGMA journal_mode=WAL;", nullptr, nullptr, nullptr);

    createTable(db);

    // Начальные условия — читаем из последней точки trajectory
    double px = 0.0, py = 0.0, pz = 0.0;
    double vx = 0.0, vy = 0.0, vz = 0.0;
    double roll = 0.0, pitch = 0.0, yaw = 0.0;
    double last_ts = readLastTrajectoryState(db, px, py, pz, vx, vy, vz, roll, pitch, yaw);

    // Инициализация фильтров Калмана один раз при старте
    Kalman1D kalmanAx(0.01, 0.1, 0.0, 0.0);
    Kalman1D kalmanAy(0.01, 0.1, 0.0, 0.0);
    Kalman1D kalmanAz(0.01, 0.1, 0.0, 0.0);
    Kalman1D kalmanGx(0.001, 0.01, 0.0, 0.0);
    Kalman1D kalmanGy(0.001, 0.01, 0.0, 0.0);
    Kalman1D kalmanGz(0.001, 0.01, 0.0, 0.0);

    sqlite3_stmt* ins = nullptr;
    sqlite3_prepare_v2(db,
        "INSERT INTO trajectory "
        "(timestamp,x,y,z,vx,vy,vz,roll,pitch,yaw) "
        "VALUES (?,?,?,?,?,?,?,?,?,?)",
        -1, &ins, nullptr);

    // Бесконечный цикл — читаем только новые данные каждые 3 секунды
    while (true) {
        last_ts = readLastTrajectoryState(db, px, py, pz, vx, vy, vz, roll, pitch, yaw);

        auto accData = readAcc(db, last_ts);
        auto gyroData = readGyro(db, last_ts);

        if (accData.empty() || gyroData.empty()) {
            sleep(3);
            continue;
        }

        size_t N = min(accData.size(), gyroData.size());
        cout << "Новых: acc=" << accData.size() << " gyro=" << gyroData.size() << endl;

        sqlite3_exec(db, "BEGIN TRANSACTION;", nullptr, nullptr, nullptr);
        size_t written = 0;

        for (size_t i = 1; i < N; i++) {
            double dt = 0.01;
            if (dt <= 0.0 || dt > 0.5) dt = 0.01;

            kalmanAx.setDt(dt);
            kalmanAy.setDt(dt);
            kalmanAz.setDt(dt);
            kalmanGx.setDt(dt);
            kalmanGy.setDt(dt);
            kalmanGz.setDt(dt);

            // Предсказание и обновление для акселерометра
            kalmanAx.predict(); kalmanAx.update(accData[i].ax);
            kalmanAy.predict(); kalmanAy.update(accData[i].ay);
            kalmanAz.predict(); kalmanAz.update(accData[i].az);

            // Предсказание и обновление для гироскопа
            kalmanGx.predict(); kalmanGx.update(gyroData[i].gx);
            kalmanGy.predict(); kalmanGy.update(gyroData[i].gy);
            kalmanGz.predict(); kalmanGz.update(gyroData[i].gz);

            // Получаем отфильтрованные значения
            double ax_f = kalmanAx.getState();
            double ay_f = kalmanAy.getState();
            double az_f = kalmanAz.getState();
            double gx_f = kalmanGx.getState();
            double gy_f = kalmanGy.getState();
            double gz_f = kalmanGz.getState();

            // Перевод в м/с²
            double ax_ms2 = ax_f * G_MS2;
            double ay_ms2 = ay_f * G_MS2;
            double az_ms2 = az_f * G_MS2;

            // Интегрирование углов (используем отфильтрованные угловые скорости)
            roll += gx_f * dt;
            pitch += gy_f * dt;
            yaw += gz_f * dt;

            double Gx = G_MS2 * sin(pitch * DEG2RAD);
            double Gy = G_MS2 * cos(pitch * DEG2RAD) * sin(roll * DEG2RAD) + fmod(sin(yaw * DEG2RAD) * G_MS2, sqrt(2));
            double Gz = G_MS2 * cos(pitch * DEG2RAD) * cos(roll * DEG2RAD) - fmod(sin(yaw * DEG2RAD) * G_MS2, sqrt(2));

            ax_ms2 -= Gx;
            ay_ms2 += Gy;
            az_ms2 -= Gz;

            // Интегрирование скорости: v += a * dt
            vx += ax_ms2 * dt;
            vy += ay_ms2 * dt;
            vz += az_ms2 * dt;

            // Интегрирование позиции: p += v * dt
            px += vx * dt;
            py += vy * dt;
            pz += vz * dt;

            // Запись
            writePoint(ins,
                accData[i].timestamp,
                px, py, pz,
                vx, vy, vz,
                roll, pitch, yaw);
            written++;
        }

        sqlite3_exec(db, "COMMIT;", nullptr, nullptr, nullptr);

        cout << "Готово. Записано " << written << " точек в trajectory." << endl;
        cout << "Позиция в метрах, углы в градусах." << endl;

        sleep(3);
    }

    sqlite3_finalize(ins);
    sqlite3_close(db);
    return 0;
}