#include "mbed.h"

// STD
#include <cmath>

// Mbed
#include "mbed_events.h"

// target definitions
#include "src/targets/revo_f4/pins.h"
// Mmath
#include "src/math/Utilities.h"
// Drivers
#include "src/drivers/MPU6000.h"
#include "src/drivers/stepper.h"
#include "src/drivers/servo.h"
// Control
#include "src/control/lpf.h"
#include "src/control/PID.h"
#include "src/control/ImuFusion.h"
#include "src/control/MadgwickAHRS.h"
// Serialization
#include "src/serialization/RCProtocol.h"

#define WHEEL_SIZE 0.09f

using namespace targets::revo_f4;
using namespace drivers;
using namespace control;
using namespace math;
using namespace serialization;

EventQueue queue;

// Serial port (Servo Outputs)
RawSerial serial(PA_2, PA_3, 250000);
RCProtocol RC;

// Pool to send Remote Control packages between threads
MemoryPool<RCProtocol::Packet, 16> mpool;
Queue<RCProtocol::Packet, 16> RCQueue;

// MPU setup
SPI spi(PA_7, PA_6, PA_5);   //define the SPI (mosi, miso, sclk). Default frequency is 1Mhz
mpu6000_spi imu(spi,PA_4);   //define the mpu6000 object

PwmOut ledBlue(D4);
DigitalOut ledOrg(D5);

// Stepper motorL(PC_9, PB_1, PC_8);
Stepper motorL(PC_9, PC_7, PC_8);
Stepper motorR(PB_15, PB_14, PC_6);

Servo servo(PA_0);

// Interrupt pin from Gyro to MCU
InterruptIn gyroISR(PC_4);
// Madwick filter
Madgwick madgwick;
// Timer to calculate dT
Timer timer;

// A PI controller for throttle control
controllerPI throttleControl(35.0, 3.0, 20, 60); // 35.0, 3.0, 20, 60

// A PID for angle control
controllerPID angleControl(100.0f, 200.0f, 0.5f, 4000, 2000); // 100.0f, 200.0f, 0.5f, 4000, 2000

// The control function to balance the robot
// Ran at high prio and triggered as an event by the gyro Interrupt pin
void runControl()
{
    static ImuFusion imuFusion(&imu);
    static float controlOutput(0.0f);

    static RCProtocol::Packet remote;
    osEvent evt = RCQueue.get(0);
    if (evt.status == osEventMessage) {
        RCProtocol::Packet* pPacket = (RCProtocol::Packet*)evt.value.p;
        remote = *pPacket;
        mpool.free(pPacket);

        throttleControl.setGainScaling(remote.Ki, remote.Kd);
        //throttleControl.setGainScaling(0, 0);

        //angleControl.setGainScaling(remote.Kp, remote.Ki, remote.Kd);
        angleControl.setGainScaling(remote.Kp, 0, -900);

        servo.setPosition((float)remote.Throttle/1000.0f);
    }
    

    // Calculate dT in sec
    float dT = timer.read_us()/(float)10e6;
    timer.reset();

    // Retrieve IMU angle
    madgwick.updateIMU(imu.read_rot(0),
                       imu.read_rot(1),
                       imu.read_rot(2),
                       imu.read_acc(0),
                       imu.read_acc(1), 
                       imu.read_acc(2));
    float angleX = madgwick.getPitch();

    // Reset anything left in the IMU FIFO queue
    imu.fifo_reset();

    // If the robot is above this angle we turn off motors
    static bool disabled = false;
    if (!remote.Enabled || (abs(angleX) > (float)50.0f && !disabled))
    {
        controlOutput = 0.0f; // rinse integral
        disabled = true;
        motorL.disable();
        motorR.disable();
    }
    else if (remote.Enabled && abs(angleX) < (float)10.0f && disabled)
    {
        disabled = false;
        motorL.enable();
        motorR.enable();
    }

    if (disabled) 
    { 
        throttleControl.flushIntegral(); 
        angleControl.flushIntegral();
    }

    /* --------------------------
    Calculate estimated groundspeed 
    of the robot in m/s
    -------------------------- */
    // Static variables
    static float angleXOld = 0;
    static float motorAngularVelocity = 0;
    static float motorAngularVelocityOld = 0;
    static incrementalLPF filteredGroundSpeed;

    if (disabled) 
    {
        angleXOld = 0;
        motorAngularVelocity = 0;
        motorAngularVelocityOld = 0;
        filteredGroundSpeed.clear();
    }
    
    // angular velocity - This way we get use of the madgwick filter
    float angularVelocity = (angleX - angleXOld) * dT;
    angleXOld = angleX;

    // Motor Speed - Convert from deg/s to m/s
    motorAngularVelocityOld = motorAngularVelocity;
    motorAngularVelocity = ((motorL.getSpeed() + motorR.getSpeed())/2);

    float estimatedGroundSpeed = 
        (motorAngularVelocityOld - angularVelocity)/360.0f*PI*WHEEL_SIZE;
    estimatedGroundSpeed = filteredGroundSpeed.filter(estimatedGroundSpeed);

    /* --------------------------
    Calculate the setpoint for 
    the main control (control the angle)
    through this throttle control loop (PI)
    -------------------------- */
    float throttle = remote.Throttle*0.001f;

    if (remote.Throttle < 50 && remote.Throttle > -50) {throttle = 0;}
    
    // TODO: Figure out why we have the wrong sign here. I also had to reverse motors
    // For the angle control loop. Is the madgwick filter outputting the wrong sign?
    float angleSP = -throttleControl.run(dT, estimatedGroundSpeed, throttle);

    /* --------------------------
    The last control loop. Angle (PD)
    -------------------------- */
    if (!disabled)
    {
        controlOutput = angleControl.run(dT, angleX, angleSP);

        controlOutput = constrain(controlOutput, 3000.0f);
    }

    /* --------------------------
     Lastly the steering is added
     straight on the output
    -------------------------- */
    float steering = remote.Steering*0.5;
    if (remote.Steering < 50 && remote.Steering > -50) {steering = 0;}
    motorL.setSpeed(controlOutput - steering);
    motorR.setSpeed(controlOutput + steering);

    // Blink LED at 1hz
    static int i = 0;
    if (++i > (disabled ? 25 : 100))
    {
        i = 0;
        ledOrg = !ledOrg;
    }
}

void serialISR()
{
    // Receive
    while (serial.readable())
    {
        bool newPackage = RC.appendByte(serial.getc());

        if (newPackage)
        {
            RCProtocol::Packet* pPacket = mpool.alloc();
            *pPacket = RC.read();
            RCQueue.put(pPacket);
        }
    }
    // Sleep some?
    //Thread::wait(1);
    
    // Transmit 
    // serial.putc()...
}

// This context just pulses the blue LED
void pulseLedContext()
{
    while (true)
    {
        float in, out;
  
        for (in = 0; in < 2*PI; in = in + PI/50)
        {
            out = sin(in) + 0.5;
            // We cannot set pulsewidth or period due to
            // sharing timer with motor outputs. "write"
            // instead sets the dutycycle in range 0-1
            ledBlue.write(out);
            Thread::wait(10);
        }
    }
}

// main() runs in its own thread
int main() {

    // MPU startup at 100hz
    if(imu.init(10,BITS_DLPF_CFG_188HZ)){
        //printf("\nCouldn't initialize MPU6000 via SPI!");
    }  

    wait(0.1);
    // Set IMU scale
    int accScale = imu.set_acc_scale(BITS_FS_16G);
    wait(0.1);
    int gyroScale = imu.set_gyro_scale(BITS_FS_2000DPS);
    wait(0.1);

    // Calibrate and Trim acc & gyro
    bool calibrationResult = imu.resetOffset_gyro();
    imu.calib_acc(0);
    calibrationResult = imu.resetOffset_acc();

    // Setup Madwick filter
    madgwick.begin(100);

    /*-------------- Visible start sequence ------------*/
    // Start sweeping the arm
    //servo.sweep(0.0, 1, 2);

    // Enable motor controllers (Will power the motors with no movement)
    // motorL.enable();
    // motorR.enable();
    motorL.disable();
    motorR.disable();
    motorL.setDirection(-1);
    motorR.setDirection(1);

    // Servo Nod to tell us that we are done
    //servo.nod();

    //servo.setPosition(-0.2);

    /*********************** Start all threads last **********************/
    
    // Start the pulsing blue led
    ledBlue.period_ms(10); // Any period. This will be overwritten by motors
    Thread ledPulseThread;
    //ledPulseThread.start(callback(&pulseLedContext));
    Thread::wait(100);

    // Serial in / out
    serial.set_dma_usage_rx(DMA_USAGE_ALWAYS);
    serial.attach(&serialISR);
    Thread::wait(100);

    // ISR won't be called if the serial is not emptied first.
    serialISR();

    // Enable/Activate the Gyro interrupt
    imu.enableInterrupt();

    // Start the timer used by the control loop
    timer.start(); // Used to calc dT

    // Create realtime eventhandler for control loop
    Thread eventThread(osPriorityHigh);
    eventThread.start(callback(&queue, &EventQueue::dispatch_forever));

    // Attach gyro interrupt to add a control event
    gyroISR.rise(queue.event(&runControl));

    wait(osWaitForever);
}