//
// Created by emon on 02/22/21.
//

#include "cfJacobian.h"
// #include <math.h>
// #include "arm_math.h"
// #include "cfassert.h"
#include "deck.h"
#include "system.h"
#include "log.h"

#include "FreeRTOS.h"
#include "task.h"
#include "debug.h"



#define LEARNING_RATE 0.01
#define STEPS 50 // 500 is pretty good 1000

struct ab {
    double a, b; // a = scaling, b = offset
};

float Ag = 1.2, f=50, thetaG = 3.14/6;
float chkObs[50] = {0};

float getCost(struct ab current_coefs);
struct ab get_loss_gradient(struct ab current_coefs, double learning_rate);
struct ab doOneIter(struct ab current_coefs, double learning_rate);

// kalmanCoreData_Emon KFCoreEmonData;
// Axis3f accEmon, gyroEmon;
// float dtEmon, errorEmon, cmdThrustEmon;
// float stdMeasNoise = 0.25;
// bool quadIsFlyingEmon;
// flowMeasurement_t flowEmon;


static void ekfOptiEngineTask(){
    systemWaitStart();
    DEBUG_PRINT("\tekfOptiEngine l38\n");
    // kalmanCoreInitEmon(&KFCoreEmonData);
    DEBUG_PRINT("\tekfOptiEngine l40\n");
    vTaskDelay(M2T(2000));
    // for (int i=0; i<50; i++){ // This is your observed check. Either define it here or save it at a header file as constant
    //     chkObs[i] = Ag*sin(2*3.14*f/50*i + thetaG);
    // }
    while(1){
        struct ab guess = {.a = 1.5, .b = 0};
        DEBUG_PRINT("FINISHED. from ekfOptiEngine\n");
        vTaskDelay(M2T(250));

        for (int i=0;i<STEPS;i++){
            doOneSim(&KFCoreEmonData, cmdThrustEmon, &accEmon, &gyroEmon, dtEmon, quadIsFlyingEmon);
            // guess = doOneIter(guess, LEARNING_RATE);
            // DEBUG_PRINT("i: %d, A: %f, theta: %f\n", i, guess.a, guess.b);
            // if (i%1 ==0){
            vTaskDelay(M2T(250));
            // }
        }
    }
}



float getCost(struct ab current_coefs){ // This is where you need to have the whole evaluation
    float chk[50] = {0}, sum = 0;
    for (int i=0; i<50; i++){
        chk[i] = current_coefs.a*sin(2*3.14*f/50*i + current_coefs.b); // This is the function/model
        sum = sum + (chk[i]-chkObs[i])*(chk[i]-chkObs[i]);
    }
    return sum;
}

struct ab get_loss_gradient(struct ab current_coefs, double learning_rate){
    struct ab grad;

    struct ab coef = current_coefs;
    float c2 = getCost(coef);
    float h = coef.a * 0.005 + 0.001;
    coef.a += h;
    float c1 = getCost(coef);
    grad.a = (c1-c2)/h;

    coef = current_coefs;
    h = coef.b * 0.005  + 0.001;
    coef.b += h;
    c1 = getCost(coef);
    grad.b = (c1-c2)/h;

    return grad;
}

struct ab doOneIter(struct ab current_coefs, double learning_rate){
    struct ab new_coefs;
    struct ab grad = get_loss_gradient(current_coefs, learning_rate);
    new_coefs.a = current_coefs.a - learning_rate * grad.a;
    new_coefs.b = current_coefs.b - learning_rate * grad.b;

    return new_coefs;
}

// /*
//  * get_loss_gradient
//  * --------------------
//  * Computes gradient in given location (defined by current_coefs).
//  * 
//  * data: training dataset
//  * n: size of the dataset
//  * current_coefs: regression coefficients
//  * 
//  * returns: gradient
//  */
// struct ab get_loss_gradient(double **data, int n, struct ab current_coefs) {
//     double x, y_true, a_grad = 0, b_grad = 0;
//     struct ab grad;

//     for (int i = 0; i < n; i++) {
//         // x = *(*(data + 0) + i);
//         // y_true = *(*(data + 1) + i);

//         // a_grad += -x * (y_true - (current_coefs.a * x + current_coefs.b));
//         // b_grad += -(y_true - (current_coefs.a * x + current_coefs.b));

//         // record the observed data. This is chk_{obs}
//         // get plant check/equation, and run for c2 = f(x)
//         // get plant check/equation, and run for c1 = f(x + [1,0] * h)
//         // grad1 = (c1 - c2)/h
//         // similar step for grad2


//         // get c1 = f(x + [1,0] * h), c2 = f(x)
//         // get grad.a += (c1 - c2)/h

//         // get c1 = f(x + [0,1] * h), c2 = f(x)
//         // get grad.b += (c1 - c2)/h
//     }

//     grad.a = (a_grad) / n;
//     grad.b = (b_grad) / n;

//     return grad;
// }

// /*
//  * do_step
//  * --------------------
//  * Moves one step in the gradient direction.
//  * 
//  * data: training dataset
//  * n: size of the dataset
//  * current_coefs: regression coefficients
//  * learning_rate: step size
//  * 
//  * returns: new coefficients
//  */
// struct ab do_step(double **data, int n, struct ab current_coefs, double learning_rate) {
//     struct ab grad, new_coefs;

//     grad = get_loss_gradient(data, n, current_coefs);
//     new_coefs.a = current_coefs.a - learning_rate * grad.a;
//     new_coefs.b = current_coefs.b - learning_rate * grad.b;

//     return new_coefs;
// }


// /*
//  *
//  * main function
//  * 
//  */

// void main(){
// 	// This is the driver function
// 	int i, n = 1;
// 	struct ab true_coefs, current_coefs;

// 	true_coefs = (struct ab) {.a = TRUE_A, .b = TRUE_B}; // This is the result
//     current_coefs = (struct ab) {.a = 1, .b = 0};

// 	for (i = 0; i < STEPS; i++) {
//         current_coefs = do_step(data, n, current_coefs, LEARNING_RATE);
//     }

// }

static void init() {
    DEBUG_PRINT("Hello Crazyflie 2.1 deck world!\n");
      xTaskCreate(ekfOptiEngineTask, "ekfOptiEngineTask",
        3* configMINIMAL_STACK_SIZE /* Stack size in terms of WORDS (usually 4 bytes) */,
        NULL, /*priority*/2, NULL);
}

static bool test() {
  DEBUG_PRINT("ekfOptiEngine test passed!\n");
    return true;
}

const DeckDriver ekfOptiEngine = {
    .vid = 0,
    .pid = 0,
    .name = "ekfOptiEngine",

    .usedGpio = 0,  // FIXME: set the used pins

    .init = init,
    .test = test,
};

DECK_DRIVER(ekfOptiEngine);