#include "M5Bala.h"

#include "M5Stack.h"

#define MPU9250_ID 0x71
#define MPU6050_ID 0x68

M5Bala::M5Bala() { wire = &Wire; }

M5Bala::M5Bala(TwoWire &w) { wire = &w; }

void M5Bala::begin() {
    // IMU
    imu_CalcInit();

    // Motor
    setMotor(0, 0);

    // PID param
    K1 = 60;
    K2 = 24;
    K5 = 0;
    K3 = 8.5;
    K4 = 5.2;
    // K5 = 8;

    loop_interval  = 0;
    left_offset    = 0;
    right_offset   = 0;
    forward_offset = 0;

    for (int i = 0; i < 128; i++) {
        imu_update();
    }
    pitch = imu_getAngleX();
}

uint8_t M5Bala::i2c_readByte(uint8_t address, uint8_t subAddress) {
    uint8_t data;
    M5.I2C.readByte(address, subAddress, &data);
    return data;  // Return data read from slave register
}

void M5Bala::setMotor(int16_t pwm0, int16_t pwm1) {
    // Value range
    int16_t m0 = constrain(pwm0, -255, 255);
    int16_t m1 = constrain(pwm1, -255, 255);

    // Dead zone
    if (((m0 > 0) && (m0 < DEAD_ZONE)) || ((m0 < 0) && (m0 > -DEAD_ZONE)))
        m0 = 0;
    if (((m1 > 0) && (m1 < DEAD_ZONE)) || ((m1 < 0) && (m1 > -DEAD_ZONE)))
        m1 = 0;

    // Same value
    static int16_t pre_m0, pre_m1;
    if ((m0 == pre_m0) && (m1 == pre_m1)) return;
    pre_m0 = m0;
    pre_m1 = m1;

    uint8_t i2c_out_buff[4];
    i2c_out_buff[0] = m0 & 0xff;
    i2c_out_buff[1] = (m0 >> 8) & 0xff;
    i2c_out_buff[2] = (m1 >> 0) & 0xff;
    i2c_out_buff[3] = (m1 >> 8) & 0xff;
    M5.I2C.writeBytes(M5GO_WHEEL_ADDR, MOTOR_CTRL_ADDR, i2c_out_buff, 4);
}

void M5Bala::readEncder() {
    static float _speed_input0 = 0, _speed_input1 = 0;
    int16_t rx_buf[2];
    M5.I2C.readBytes(M5GO_WHEEL_ADDR, ENCODER_ADDR, 4, (uint8_t *)rx_buf);

    // filter
    _speed_input0 *= 0.9;
    _speed_input0 += 0.1 * rx_buf[0];
    _speed_input1 *= 0.9;
    _speed_input1 += 0.1 * rx_buf[1];

    speed_input0 = constrain((int16_t)(-_speed_input0), -255, 255);
    speed_input1 = constrain((int16_t)(_speed_input1), -255, 255);
}

void M5Bala::PIDCompute() {
    static float last_angle;
    static int16_t last_wheel;
    float angle, angle_velocity;
    int16_t wheel, wheel_velocity;
    int16_t torque, speed_diff, speed_diff_adjust;

    angle          = pitch;
    angle_velocity = angle - last_angle;
    last_angle     = angle;

    wheel = (speed_input0 + speed_input1) /
            2; /* wheel = read_encoder()-profiler() */
    wheel_velocity = wheel - last_wheel;
    last_wheel     = wheel;

    torque = (angle_velocity * K1) + (angle * K2) + (wheel_velocity * K3) +
             (wheel * K4);

    speed_diff        = (speed_input0 - speed_input1);
    speed_diff_adjust = (K5 * speed_diff);

    pwm_out0 = torque - speed_diff_adjust;
    pwm_out1 = torque;
    pwm_out0 = constrain(pwm_out0, -255, 255);
    pwm_out1 = constrain(pwm_out1, -255, 255);
}

void M5Bala::run() {
    if (micros() >= loop_interval) {
        loop_interval = micros() + 10000;

        // Attitude sample
        imu_update();
        pitch = imu_getAngleX() + angle_offset;

        // Car down
        if (abs(pitch) > 45) {
            setMotor(0, 0);
            return;
        }

        // Encoder sample
        readEncder();

        // PID Compute
        PIDCompute();

        // Motor out
        setMotor(pwm_out0 + forward_offset + left_offset,
                 pwm_out1 + forward_offset + right_offset);
    }
}

void M5Bala::stop() {
    left_offset  = 0;
    right_offset = 0;
}

void M5Bala::move(int16_t speed, uint16_t duration) {
    forward_offset = -speed;
    if (duration != 0) {
        delay(duration);
        stop();
    }
}

void M5Bala::turn(int16_t speed, uint16_t duration) {
    if (speed > 0) {
        left_offset  = speed;
        right_offset = 0;

    } else if (speed < 0) {
        left_offset  = 0;
        right_offset = -speed;
    }

    if (duration != 0) {
        delay(duration);
        stop();
    }
}

void M5Bala::rotate(int16_t speed, uint16_t duration) {
    left_offset  = speed;
    right_offset = -speed;

    if (duration != 0) {
        delay(duration);
        stop();
    }
}