/**
 *    ||          ____  _ __
 * +------+      / __ )(_) /_______________ _____  ___
 * | 0xBC |     / __  / / __/ ___/ ___/ __ `/_  / / _ \
 * +------+    / /_/ / / /_/ /__/ /  / /_/ / / /_/  __/
 *  ||  ||    /_____/_/\__/\___/_/   \__,_/ /___/\___/
 *
 * Crazyflie control firmware
 *
 * Copyright (C) 2018 Bitcraze AB
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 *
 * lighthouse.c: lighthouse tracking system receiver
 */

#include <math.h>
#include <stdbool.h>
#include <stdint.h>

#include "system.h"
#include "deck.h"
#include "log.h"

#include "config.h"
#include "FreeRTOS.h"
#include "task.h"

#include "svd3.h"
#include "math3d.h"

#define DEBUG_MODULE "LH"
#include "debug.h"
#include "uart1.h"
#include "lh_bootloader.h"

#include "pulse_processor.h"
#include "lighthouse.h"

#include "estimator.h"

#ifdef LH_FLASH_DECK
#include "lh_flasher.h"
#endif

// #ifndef DISABLE_LIGHTHOUSE_DRIVER
//   #define DISABLE_LIGHTHOUSE_DRIVER 1
// #endif

baseStationGeometry_t lighthouseBaseStationsGeometry[]  = {
{.origin = {-0.685524, 1.671130, -0.077716, }, .mat = {{0.269844, 0.615614, -0.740408, }, {-0.113244, 0.783886, 0.610491, }, {0.956222, -0.080891, 0.281241, }, }},
{.origin = {0.504731, 1.683439, 0.613580, }, .mat = {{0.999760, -0.021903, 0.000000, }, {0.016498, 0.753046, 0.657761, }, {-0.014407, -0.657603, 0.753227, }, }},
};

// Uncomment if you want to force the Crazyflie to reflash the deck at each startup
// #define FORCE_FLASH true

#if DISABLE_LIGHTHOUSE_DRIVER == 0

#ifndef FORCE_FLASH
#define FORCE_FLASH false
#endif

#define STR2(x) #x
#define STR(x) STR2(x)

#define INCBIN(name, file) \
    __asm__(".section .rodata\n" \
            ".global incbin_" STR(name) "_start\n" \
            ".align 4\n" \
            "incbin_" STR(name) "_start:\n" \
            ".incbin \"" file "\"\n" \
            \
            ".global incbin_" STR(name) "_end\n" \
            ".align 1\n" \
            "incbin_" STR(name) "_end:\n" \
            ".byte 0\n" \
            ".align 4\n" \
            STR(name) "Size:\n" \
            ".int incbin_" STR(name) "_end - incbin_" STR(name) "_start\n" \
    ); \
    extern const __attribute__((aligned(4))) void* incbin_ ## name ## _start; \
    extern const void* incbin_ ## name ## _end; \
    extern const int name ## Size; \
    static const __attribute__((used)) unsigned char* name = (unsigned char*) & incbin_ ## name ## _start; \

INCBIN(bitstream, "lighthouse.bin");

static void checkVersionAndBoot();

static bool isInit = false;

static pulseProcessorResult_t angles[PULSE_PROCESSOR_N_SENSORS];

// Stats
static int serialFrameCount = 0;
static int frameCount = 0;
static int cycleCount = 0;
static int positionCount = 0;

static float serialFrameRate = 0.0;
static float frameRate = 0.0;
static float cycleRate = 0.0;
static float positionRate = 0.0;

static uint16_t pulseWidth[PULSE_PROCESSOR_N_SENSORS];

static uint32_t latestStatsTimeMs = 0;

typedef union frame_u {
  struct {
    uint32_t timestamp:29;
    uint32_t sensor:3;
    uint16_t width;
    uint8_t sync;
  } __attribute__((packed));
  char data[7];
} __attribute__((packed)) frame_t;

static bool getFrame(frame_t *frame)
{
  int syncCounter = 0;
  for(int i=0; i<7; i++) {
    uart1Getchar(&frame->data[i]);
    if (frame->data[i] != 0) {
      syncCounter += 1;
    }
  }
  return (frame->sync == 0 || (syncCounter==7));
}

static void resetStats() {
  serialFrameCount = 0;
  frameCount = 0;
  cycleCount = 0;
  positionCount = 0;
}

static void calculateStats(uint32_t nowMs) {
  double time = (nowMs - latestStatsTimeMs) / 1000.0;
  serialFrameRate = serialFrameCount / time;
  frameRate = frameCount / time;
  cycleRate = cycleCount / time;
  positionRate = positionCount / time;

  resetStats();
}

static vec3d position;
static positionMeasurement_t ext_pos;
static poseMeasurement_t ext_pose;
struct quat ori;
static float deltaLog;

static void estimatePosition(pulseProcessorResult_t angles[]) {
  memset(&ext_pos, 0, sizeof(ext_pos));
  int sensorsUsed = 0;
  float delta;

  // Prepare IR positions for the orientation calculation
  float sensor_z_pos = 0.0f;
  float x_distance = 0.015f;
  float y_distance = 0.0075f;
  float X[3][4] = { // Origin points
    // sensor 0 (left bottom),sensor 1 (right bottom),sensor 2 (left up),sensor 3 (right up)
    {-x_distance,-x_distance,x_distance,x_distance},
    {y_distance,-y_distance,y_distance,-y_distance},
    {sensor_z_pos,sensor_z_pos,sensor_z_pos,sensor_z_pos}
  }; 
  float q[3][4]; // Current points
  float Y[3][4]; // Current points centered

  // Average over all sensors with valid data
  for (size_t sensor = 0; sensor < PULSE_PROCESSOR_N_SENSORS; sensor++) {
      if (angles[sensor].validCount == 4) {
        lighthouseGeometryGetPosition(lighthouseBaseStationsGeometry, (void*)angles[sensor].correctedAngles, position, &delta);

        deltaLog = delta;

        ext_pos.x -= position[2];
        ext_pos.y -= position[0];
        ext_pos.z += position[1];
        sensorsUsed++;

        // Save the point for orientation estimation
        q[0][sensor] = position[2];
        q[1][sensor] = position[0];
        q[2][sensor] = position[1];

        positionCount++;
      }
  }

  ext_pos.x /= sensorsUsed;
  ext_pos.y /= sensorsUsed;
  ext_pos.z /= sensorsUsed;

  DEBUG_PRINT("sensorsUsed: %d\n", sensorsUsed);

  /* Orientation estimation */
  // Implements the optimization explained in URL
  // The objective is to estimate the transformation T between the 4 IR points of the lighthouse
  // deck in the origin and the same 4 points rotated and translated in the space.

  if(sensorsUsed == PULSE_PROCESSOR_N_SENSORS){ // IMPROVE: 3 sensor are enought
    // The points of the origin (p) are always the same origin points
    // Current positions points (q) must be centered in the origin. To do that
    // the centroid of the point cloud (ext_pos) is substracted
    for (size_t sensor = 0; sensor < PULSE_PROCESSOR_N_SENSORS; sensor++) {
      Y[0][sensor] = q[0][sensor] - ext_pos.x;
      Y[1][sensor] = q[1][sensor] - ext_pos.y;
      Y[2][sensor] = q[2][sensor] - ext_pos.z;
    }

    // S matrix: S = X*Y.transpose()
    float   S11, S12, S13,
            S21, S22, S23,
            S31, S32, S33;

    S11 = X[0][0]*Y[0][0]+X[0][1]*Y[0][1]+X[0][2]*Y[0][2]+X[0][3]*Y[0][3];
    S12 = X[0][0]*Y[1][0]+X[0][1]*Y[1][1]+X[0][2]*Y[1][2]+X[0][3]*Y[1][3];
    S13 = X[0][0]*Y[2][0]+X[0][1]*Y[2][1]+X[0][2]*Y[2][2]+X[0][3]*Y[2][3];

    S21 = X[1][0]*Y[0][0]+X[1][1]*Y[0][1]+X[1][2]*Y[0][2]+X[1][3]*Y[0][3];
    S22 = X[1][0]*Y[1][0]+X[1][1]*Y[1][1]+X[1][2]*Y[1][2]+X[1][3]*Y[1][3];
    S23 = X[1][0]*Y[2][0]+X[1][1]*Y[2][1]+X[1][2]*Y[2][2]+X[1][3]*Y[2][3];

    S31 = X[2][0]*Y[0][0]+X[2][1]*Y[0][1]+X[2][2]*Y[0][2]+X[2][3]*Y[0][3];
    S32 = X[2][0]*Y[1][0]+X[2][1]*Y[1][1]+X[2][2]*Y[1][2]+X[2][3]*Y[1][3];
    S33 = X[2][0]*Y[2][0]+X[2][1]*Y[2][1]+X[2][2]*Y[2][2]+X[2][3]*Y[2][3];

    // SVD of S

    float   u11, u12, u13, 
            u21, u22, u23, 
            u31, u32, u33;

    float   k11, k12, k13, 
            k21, k22, k23, 
            k31, k32, k33;

    float   v11, v12, v13, 
            v21, v22, v23, 
            v31, v32, v33;

    svd(S11, S12, S13, S21, S22, S23, S31, S32, S33,
        &u11, &u12, &u13, &u21, &u22, &u23, &u31, &u32, &u33,
        &k11, &k12, &k13, &k21, &k22, &k23, &k31, &k32, &k33,
        &v11, &v12, &v13, &v21, &v22, &v23, &v31, &v32, &v33);

    // I suppose there are not reflections
    // Estimate R = V*U.transpose()
    float   R11, R12, R13, 
            R21, R22, R23, 
            R31, R32, R33;
    multAtB( v11,  v21,  v31,  v12,  v22,  v32,  v13,  v23,  v33,  // V transpose
            u11,  u21,  u31,  u12,  u22,  u32,  u13,  u23,  u33,   // U transpose
            &R11, &R12, &R13, &R21, &R22, &R23, &R31, &R32, &R33);

    // R its converted to mat33 of crazyflie and then to quaternion
    struct mat33 R;
    R.m[0][0] = R11;  R.m[0][1] = R12;  R.m[0][2] = R13;
    R.m[1][0] = R21;  R.m[1][1] = R22;  R.m[1][2] = R23;
    R.m[2][0] = R31;  R.m[2][1] = R32;  R.m[2][2] = R33;
    ori = mat2quat(R);

    // Pose struct is filled
    ext_pose.stdDevPos = 0.01;
    ext_pose.stdDevQuat = 0.001;
    /*ext_pose.x = ext_pos.x;
    ext_pose.y = ext_pos.y;
    ext_pose.z = ext_pos.z;
    ext_pose.quat.x = ori.x;
    ext_pose.quat.y = ori.y;
    ext_pose.quat.z = ori.z;
    ext_pose.quat.w = ori.w;*/
    ext_pose.x = 1.0;
    ext_pose.y = 2.0;
    ext_pose.z = 3.0;
    ext_pose.quat.x = 0;
    ext_pose.quat.y = 0;
    ext_pose.quat.z = 0;
    ext_pose.quat.w = 1;

    DEBUG_PRINT("   x: %f,   y: %f,   z: %f\n\n", (double)ext_pose.x,(double)ext_pose.y,(double)ext_pose.z);
    DEBUG_PRINT("   %f   %f   %f\n", (double)R11,(double)R12,(double)R13);
    DEBUG_PRINT("   %f   %f   %f\n", (double)R21,(double)R22,(double)R23);
    DEBUG_PRINT("   %f   %f   %f\n\n", (double)R31,(double)R32,(double)R33);

    // Make sure we feed sane data into the estimator
    if (!isfinite(ext_pos.pos[0]) || !isfinite(ext_pos.pos[1]) || !isfinite(ext_pos.pos[2])) {
      return;
    }
    estimatorEnqueuePose(&ext_pose);

    /*ext_pos.stdDev = 0.01;
    estimatorEnqueuePosition(&ext_pos);*/
  }else{  // With less points, only feed the position
    // Make sure we feed sane data into the estimator
    if (!isfinite(ext_pos.pos[0]) || !isfinite(ext_pos.pos[1]) || !isfinite(ext_pos.pos[2])) {
      return;
    }
    ext_pos.stdDev = 0.01;
    estimatorEnqueuePosition(&ext_pos);
    DEBUG_PRINT("ONLY POSITION!\n");
  }
}

static void lighthouseTask(void *param)
{
  bool synchronized = false;
  int syncCounter = 0;
  char c;
  static frame_t frame;
  static pulseProcessor_t ppState = {};

  int basestation;
  int axis;

  systemWaitStart();

#ifdef LH_FLASH_DECK
  // Flash deck bootloader using SPI (factory and recovery flashing)
  lhflashInit();
  lhflashFlashBootloader();
#endif

  // Boot the deck firmware
  checkVersionAndBoot();

  while(1) {
    // Synchronize
    syncCounter = 0;
    while (!synchronized) {
      
      uart1Getchar(&c);
      if (c != 0) {
        syncCounter += 1;
      } else {
        syncCounter = 0;
      }
      synchronized = syncCounter == 7;
    }

    DEBUG_PRINT("Synchronized!\n");

    // Receive data until being desynchronized
    synchronized = getFrame(&frame);
    while(synchronized) {
      if (frame.sync != 0) {
        synchronized = getFrame(&frame);
        memset(pulseWidth, 0, sizeof(pulseWidth[0])*PULSE_PROCESSOR_N_SENSORS);
        continue;
      }

      serialFrameCount++;

      pulseWidth[frame.sensor] = frame.width;

      if (pulseProcessorProcessPulse(&ppState, frame.sensor, frame.timestamp, frame.width, angles, &basestation, &axis)) {
        frameCount++;
        if (basestation == 1 && axis == 1) {
          cycleCount++;

          pulseProcessorApplyCalibration(&ppState, angles);

          estimatePosition(angles);
          for (size_t sensor = 0; sensor < PULSE_PROCESSOR_N_SENSORS; sensor++) {
            angles[sensor].validCount = 0;
          }
        }
      }

      uint32_t nowMs = T2M(xTaskGetTickCount());
      if ((nowMs - latestStatsTimeMs) > 1000) {
        calculateStats(nowMs);
        latestStatsTimeMs = nowMs;
      }

      synchronized = getFrame(&frame);
      if (frame.sync != 0) {
        synchronized = getFrame(&frame);
        continue;
      }
    }
  }
}

static void checkVersionAndBoot()
{

  uint8_t bootloaderVersion = 0;
  lhblGetVersion(&bootloaderVersion);
  DEBUG_PRINT("Lighthouse bootloader version: %d\n", bootloaderVersion);

  // Wakeup mem
  lhblFlashWakeup();
  vTaskDelay(M2T(1));

  // Checking the bitstreams are identical
  // Also decoding bitstream version for console
  static char deckBitstream[65];
  lhblFlashRead(LH_FW_ADDR, 64, (uint8_t*)deckBitstream);
  deckBitstream[64] = 0;
  int deckVersion = strtol(&deckBitstream[2], NULL, 10);
  int embeddedVersion = strtol((char*)&bitstream[2], NULL, 10);

  bool identical = true;
  for (int i=0; i<=bitstreamSize; i+=64) {
    int length = ((i+64)<bitstreamSize)?64:bitstreamSize-i;
    lhblFlashRead(LH_FW_ADDR + i, length, (uint8_t*)deckBitstream);
    if (memcmp(deckBitstream, &bitstream[i], length)) {
      DEBUG_PRINT("Fail comparing firmware\n");
      identical = false;
      break;
    }
  }

  if (identical == false || FORCE_FLASH) {
    DEBUG_PRINT("Deck has version %d and we embeed version %d\n", deckVersion, embeddedVersion);
    DEBUG_PRINT("Updating deck with embedded version!\n");

    // Erase LH deck FW
    lhblFlashEraseFirmware();

    // Flash LH deck FW
    if (lhblFlashWriteFW((uint8_t*)bitstream, bitstreamSize)) {
      DEBUG_PRINT("FW updated [OK]\n");
    } else {
      DEBUG_PRINT("FW updated [FAILED]\n");
    }
  }

  // Launch LH deck FW
  DEBUG_PRINT("Firmware version %d verified, booting deck!\n", deckVersion);
  lhblBootToFW();
}

static void lighthouseInit(DeckInfo *info)
{
  if (isInit) return;

  uart1Init(230400);
  lhblInit(I2C1_DEV);
  
  xTaskCreate(lighthouseTask, "LH",
              2*configMINIMAL_STACK_SIZE, NULL, /*priority*/1, NULL);

  isInit = true;
}


static const DeckDriver lighthouse_deck = {
  .vid = 0xBC,
  .pid = 0x10,
  .name = "bcLighthouse4",

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

  .init = lighthouseInit,
};


DECK_DRIVER(lighthouse_deck);

LOG_GROUP_START(lighthouse)
LOG_ADD(LOG_FLOAT, rawAngle0x, &angles[0].angles[0][0])
LOG_ADD(LOG_FLOAT, rawAngle0y, &angles[0].angles[0][1])
LOG_ADD(LOG_FLOAT, rawAngle1x, &angles[0].angles[1][0])
LOG_ADD(LOG_FLOAT, rawAngle1y, &angles[0].angles[1][1])
LOG_ADD(LOG_FLOAT, angle0x, &angles[0].correctedAngles[0][0])
LOG_ADD(LOG_FLOAT, angle0y, &angles[0].correctedAngles[0][1])
LOG_ADD(LOG_FLOAT, angle1x, &angles[0].correctedAngles[1][0])
LOG_ADD(LOG_FLOAT, angle1y, &angles[0].correctedAngles[1][1])

LOG_ADD(LOG_FLOAT, angle0x_1, &angles[1].correctedAngles[0][0])
LOG_ADD(LOG_FLOAT, angle0y_1, &angles[1].correctedAngles[0][1])
LOG_ADD(LOG_FLOAT, angle1x_1, &angles[1].correctedAngles[1][0])
LOG_ADD(LOG_FLOAT, angle1y_1, &angles[1].correctedAngles[1][1])

LOG_ADD(LOG_FLOAT, angle0x_2, &angles[2].correctedAngles[0][0])
LOG_ADD(LOG_FLOAT, angle0y_2, &angles[2].correctedAngles[0][1])
LOG_ADD(LOG_FLOAT, angle1x_2, &angles[2].correctedAngles[1][0])
LOG_ADD(LOG_FLOAT, angle1y_2, &angles[2].correctedAngles[1][1])

LOG_ADD(LOG_FLOAT, angle0x_3, &angles[3].correctedAngles[0][0])
LOG_ADD(LOG_FLOAT, angle0y_3, &angles[3].correctedAngles[0][1])
LOG_ADD(LOG_FLOAT, angle1x_3, &angles[3].correctedAngles[1][0])
LOG_ADD(LOG_FLOAT, angle1y_3, &angles[3].correctedAngles[1][1])

LOG_ADD(LOG_FLOAT, x, &position[0])
LOG_ADD(LOG_FLOAT, y, &position[1])
LOG_ADD(LOG_FLOAT, z, &position[2])

LOG_ADD(LOG_FLOAT, delta, &deltaLog)

LOG_ADD(LOG_FLOAT, serRt, &serialFrameRate)
LOG_ADD(LOG_FLOAT, frmRt, &frameRate)
LOG_ADD(LOG_FLOAT, cycleRt, &cycleRate)
LOG_ADD(LOG_FLOAT, posRt, &positionRate)

LOG_ADD(LOG_UINT16, width0, &pulseWidth[0])
#if PULSE_PROCESSOR_N_SENSORS > 1
LOG_ADD(LOG_UINT16, width1, &pulseWidth[1])
#endif
#if PULSE_PROCESSOR_N_SENSORS > 2
LOG_ADD(LOG_UINT16, width2, &pulseWidth[2])
#endif
#if PULSE_PROCESSOR_N_SENSORS > 3
LOG_ADD(LOG_UINT16, width3, &pulseWidth[3])
#endif
LOG_GROUP_STOP(lighthouse)


#endif // DISABLE_LIGHTHOUSE_DRIVER