// This is a personal academic project. Dear PVS-Studio, please check it.
// PVS-Studio Static Code Analyzer for C, C++, C#, and Java: https://pvs-studio.com
#pragma once

#include "MPU6050_6Axis_MotionApps20.h"
#include "Preferences.h"
#include "board_pins.h"
#include "esp_log.h"

static volatile bool _isDMPDataReady = false;

static void IRAM_ATTR _dmpInterruption() {
    _isDMPDataReady = true;
}

typedef std::function<void()> OnCalibrationFinishedCb;

class MPU {
public:
    MPU(MPU6050_6Axis_MotionApps20 *mpu) {
        assert(mpu != nullptr);
        _mpu = mpu;
    }

    ~MPU() {
        delete _mpu;
    }

    bool initialize() {
        _mpu->initialize();

        if (!_preferences.begin("imu")) {
            ESP_LOGE(_TAG, "Failed to open the preferences!");
        } else {
            _loadOffsets();
        }

        delay(250);
        if (!_mpu->testConnection()) {
            ESP_LOGE(_TAG, "MPU6050 test connection failed!");
            return false;
        }
        _mpu->setDLPFMode(MPU6050_DLPF_BW_10); // 10 Hz bandwidth
        //mpu.setFullScaleGyroRange(MPU6050_GYRO_FS_500); // set sensivity, not recommended
        //mpu.setFullScaleAccelRange(MPU6050_ACCEL_FS_8);

        if (_mpu->dmpInitialize()) {
            ESP_LOGE(_TAG, "Failed to initialize DMP!");
            return false;
        }
        _mpu->setDMPEnabled(true);
        attachInterrupt(MPU6050_INT_PIN, _dmpInterruption, RISING);
        return true;
    }

    bool tick(bool &err) {
        err = false;
        if (!_isDMPDataReady) {
            return false;
        }
        if (!_mpu->dmpGetCurrentFIFOPacket(_fifoBuffer)) {
            ESP_LOGE(_TAG, "Failed to get DMP data!");
            err = true;
            return false;
        }

        Quaternion q;
        VectorFloat gravity;
        VectorInt16 accel;
        VectorInt16 gyro;
        VectorInt16 accelReal;

        _mpu->dmpGetQuaternion(&q, _fifoBuffer);
        _mpu->dmpGetGravity(&gravity, &q);
        _mpu->dmpGetYawPitchRoll(_ypr, &q, &gravity);
        _mpu->dmpGetAccel(&accel, _fifoBuffer);
        _mpu->dmpGetGyro(&gyro, _fifoBuffer);
        _mpu->dmpGetLinearAccel(&accelReal, &accel, &gravity);
        _isDMPDataReady = false;

        _ax = accel.x / 8192.f;
        _ay = accel.y / 8192.f;
        _az = accel.z / 8192.f;

        _gx = gyro.x / 131.f;
        _gy = gyro.y / 131.f;
        _gz = gyro.z / 131.f;

        _calculateZInertial(gravity);

        if (_isCalibration) {
            if (_calibrationsIterCounter >= _calibrationIterCount) {
                _axOffset /= _calibrationsIterCounter;
                _ayOffset /= _calibrationsIterCounter;
                _azOffset /= _calibrationsIterCounter;

                _gxOffset /= _calibrationsIterCounter;
                _gyOffset /= _calibrationsIterCounter;
                _gzOffset /= _calibrationsIterCounter;

                _prOffset[0] /= _calibrationsIterCounter;
                _prOffset[1] /= _calibrationsIterCounter;

                _prOffset[0] = DEG_TO_RAD * (180 - RAD_TO_DEG * _prOffset[0]);
                _prOffset[1] = DEG_TO_RAD * (180 - RAD_TO_DEG * _prOffset[1]);
                ESP_LOGI(_TAG, "Calibration offsets: roll: %f, pitch: %f", RAD_TO_DEG * _prOffset[1],
                         RAD_TO_DEG * _prOffset[0]);

                _isCalibration = false;
                _updateOffsets();
                _calibrationsIterCounter = 0;
                if (_calibrationFinishedCb) {
                    _calibrationFinishedCb();
                }
                ESP_LOGI(_TAG, "Calibration finished!");
            } else {
                _axOffset += _ax;
                _ayOffset += _ay;
                _azOffset += _az;
                _gxOffset += _gx;
                _gyOffset += _gy;
                _gzOffset += _gz;
                _prOffset[0] += pitch() * DEG_TO_RAD;
                _prOffset[1] += roll() * DEG_TO_RAD;
                _calibrationsIterCounter++;
                ESP_LOGI(_TAG, "Angles on calibration: roll: %f, pitch: %f", roll(),
                         pitch());
            }
        }

        return true;
    }

    /* getters */
    float accZInertial() const {
        return _AccZInertial;
    }

    float yaw() const {
        return degrees(_ypr[0]) + 180.f;
    }

    float pitch() const {
        if (_isCalibration) {
            return 360.f - (degrees(_ypr[1]) + 180);
        } else {
            return 360.f - (degrees(_ypr[1]) + 180) + degrees(_prOffset[0]);
        }
    }

    float roll() const {
        if (_isCalibration) {
            return 360.f - (degrees(_ypr[2]) + 180);
        } else {
            return 360.f - (degrees(_ypr[2] + _prOffset[1]) + 180) + degrees(_prOffset[1]);
        }
    }

    float ax() const {
        if (_isCalibration) {
            return _ax;
        } else {
            return _ax - _axOffset;
        }
    }

    float ay() const {
        if (_isCalibration) {
            return _ay;
        } else {
            return _ay - _ayOffset;
        }
    }

    float az() const {
        if (_isCalibration) {
            return _az;
        } else {
            return _az - _azOffset;
        }
    }

    float gravityZ() const {
        return _gravity;
    }

    float gx() const {
        if (_isCalibration) {
            return _gx;
        } else {
            return _gx - _gxOffset;
        }
    }

    float gy() const {
        if (_isCalibration) {
            return _gy;
        } else {
            return _gy - _gyOffset;
        }
    }

    float gz() const {
        if (_isCalibration) {
            return _gz;
        } else {
            return _gz - _gzOffset;
        }
    }

    void startCalibration() {
        _isCalibration = true;
    }

    void onCalibrationFinished(OnCalibrationFinishedCb callback) {
        _calibrationFinishedCb = callback;
    }

private:
    MPU6050_6Axis_MotionApps20 *_mpu = nullptr;
    float _ypr[3] = {0};
    float _ax = 0, _ay = 0, _az = 0;
    float _gx = 0, _gy = 0, _gz = 0;
    float _gravity = 0;
    float _AccZInertial = 0;

    float _axOffset = 0, _ayOffset = 0, _azOffset = 0;
    float _gxOffset = 0, _gyOffset = 0, _gzOffset = 0;
    float _prOffset[2] = {0}; // yaw isn't used

    uint8_t _fifoBuffer[45] = {0};
    Preferences _preferences;

    const char *_TAG = "MPU6050 module";
    static constexpr const int _calibrationIterCount = 250;

    bool _isCalibration = false;
    int _calibrationsIterCounter = 0;
    OnCalibrationFinishedCb _calibrationFinishedCb = nullptr;

    void _calculateZInertial(VectorFloat &gravity) {
        const float anglePitch = _ypr[1];
        const float angleRoll = _ypr[2];
        _gravity = gravity.z;

        _AccZInertial = -sin(anglePitch) * _ax + cos(anglePitch) * sin(angleRoll) * _ay
                        + cos(anglePitch) * cos(angleRoll) * _az;
        _AccZInertial = (_AccZInertial - 1) * gravity.z * 100.f;
    }

    void _loadOffsets() {
        _axOffset = _preferences.getFloat("ax", 0.f);
        _ayOffset = _preferences.getFloat("ax", 0.f);
        _azOffset = _preferences.getFloat("az", 0.f);
        _gxOffset = _preferences.getFloat("gx", 0.f);
        _gyOffset = _preferences.getFloat("gy", 0.f);
        _gzOffset = _preferences.getFloat("gz", 0.f);
        _prOffset[0] = _preferences.getFloat("pitch", 0.f);
        _prOffset[1] = _preferences.getFloat("roll", 0.f);
        ESP_LOGI(_TAG, "Offsets loaded: pitch: %f, roll %f", _prOffset[0], _prOffset[1]);
    }

    void _updateOffsets() {
        _preferences.putFloat("ax", _axOffset);
        _preferences.putFloat("ay", _ayOffset);
        _preferences.putFloat("az", _azOffset);
        _preferences.putFloat("gx", _gxOffset);
        _preferences.putFloat("gy", _gyOffset);
        _preferences.putFloat("gz", _gzOffset);
        _preferences.putFloat("pitch", _prOffset[0]);
        _preferences.putFloat("roll", _prOffset[1]);
        ESP_LOGI(_TAG, "Offsets saved.");
    }
};