#include <map> // Bluetooth
#include <BluetoothSerial.h> // Bluetooth
#include <M5Core2.h>
#include "icons.cpp" // icons
#include <Preferences.h>  // Library for writing in the memory

// MAC address of EEG headset device
String MAC_ADDRESS = "98:07:2d:7f:ea:15";

// store screen mode and brightness in the memory of the Core
Preferences preferences;
const char* key1 = "screenMode";
const char* key2 = "brightness";

// UART communication variables
#define RX_PIN 19
#define TX_PIN 27

// GUI variables
int brightnessValues[4] = {2500, 2750, 2900, 3150};
int brightnessValue;
ButtonColors noDraw = { NODRAW, NODRAW, NODRAW };
Button screenModeButton(0, 190, 50, 50, false, "", noDraw);
Button brightnessButton(65, 190, 50, 50, false, "", noDraw);
Button restartButton(200, 191, 45, 50, false, "", noDraw);
Button powerOffButton(260, 186, 55, 55, false, "", noDraw);
boolean buttonPressed = false;
boolean connectingBluetooth = true;
enum screenMode {DARK, LIGHT};
screenMode COLOR_MODE;
boolean setupGUI = false;

uint16_t COLORS[2] = {WHITE, BLACK};

// Bluetooth Serial, variables
#define BT_DISCOVER_TIME  10000
#define BAUDRATE 115200
BluetoothSerial SerialBT;
esp_spp_sec_t sec_mask=ESP_SPP_SEC_NONE;
esp_spp_role_t role=ESP_SPP_ROLE_SLAVE;
boolean isConnected = false;
boolean wasConnected = false;

// checksum variables
byte  generatedChecksum = 0;
byte  checksum = 0; 
int   payloadLength = 0;
byte  payloadData[64] = {0};
boolean bigPacket = false;

// EEG variables
byte  poorQuality = 0;
byte  attention = 0;
byte  meditation = 0;
short raw;
uint32_t eegPower[8];

// battery level variables
float batVoltage;
float batPercentage;

// Multi threading
TaskHandle_t TaskHandle_1;

////////////////////////////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////////////////////////

//////////////////////////
// Microprocessor Setup //
//////////////////////////

void setup(){
  Serial.begin(BAUDRATE);
  Serial2.begin(9600, SERIAL_8N1, RX_PIN, TX_PIN); // UART communication
  preferences.begin("Key");
  if (preferences.getUInt(key1) == 0) {
    COLOR_MODE = DARK;
  } else {
    COLOR_MODE = LIGHT;
  }  
  M5.begin();
  M5.Axp.SetLcdVoltage(3150); // Set initial brightness
  
  xTaskCreate(task1, "task1", 4096, NULL, 1, &TaskHandle_1);

  M5.Axp.SetLcdVoltage(brightnessValues[preferences.getUInt(key2)]); // Set initial brightness

  screenModeButton.addHandler(screenModeHandler, E_TOUCH);
  brightnessButton.addHandler(brightnessHandler, E_TOUCH);
  restartButton.addHandler(restartHandler, E_TOUCH);
  powerOffButton.addHandler(powerOffHandler, E_TOUCH);
}

////////////////////////////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////////////////////////

///////////////
// Main loop //
///////////////

void loop() {
  if (!isConnected) {
    drawBluetoothConnecting();
    connectBluetooth();
  }

  if (!setupGUI && isConnected) {
    setupGUI = true;
    drawGui();
  }

  M5.update();
  drawBattery();
  
  if (!SerialBT.isClosed() && SerialBT.connected()) {
    if ((readOneByte() == 170) && readOneByte() == 170) {
      payloadLength = readOneByte();
      if (payloadLength > 169)  // Payload length can not be greater than 169
      return;
      
      generatedChecksum = 0;

      for (int i = 0; i < payloadLength; i++) {  
        payloadData[i] = readOneByte();  // Read payload into memory
        generatedChecksum += payloadData[i];
      }

      checksum = readOneByte();  // Read checksum byte from stream      
      generatedChecksum = 255 - generatedChecksum;

      if(checksum == generatedChecksum) {
        for(int i = 0; i < payloadLength; i++) {  // Parse the payload
          switch (payloadData[i]) {
          case 2:
            i++;            
            poorQuality = payloadData[i];
            bigPacket = true;            
            break;
          case 4:
            i++;
            attention = payloadData[i];                        
            break;
          case 5:
            i++;
            meditation = payloadData[i];
            break;
          case 0x80:
            i++;
            raw = (payloadData[i++] << 8) | payloadData[i++];
            break;
          case 0x83:
            for (int j = 0; j < 8; j++) {
              uint8_t a,b,c;
              a = payloadData[++i];
              b = payloadData[++i];
              c = payloadData[++i];
              eegPower[j] = ((uint32_t)a << 16) | ((uint32_t)b << 8) | (uint32_t)c;
            }
            break;
          default:
            break;
          }
        }

        if(bigPacket) {
            updatePoorQuality();
            updateMeditation();
            updateAttention();

          // UART OUTPUT
          Serial2.printf("%3d", attention);
          Serial2.printf("%3d", meditation);
          Serial2.printf("%3d", poorQuality);
          Serial2.printf("%+8d", raw);

          // SERIAL OUTPUT
          Serial.print("\n");
          Serial.print("Attention: ");
          Serial.print(attention);
          Serial.print("\n");
          Serial.print("Poor quality: ");
          Serial.print(poorQuality);
          Serial.print("\n");
          Serial.print("Meditation: ");
          Serial.print(meditation);
          Serial.print("\n");
          Serial.print("Raw: ");
          Serial.print(raw);
          Serial.print("\n");            
          Serial.print("EEG POWER: ");
          Serial.print(String(eegPower[0]) + ", " + String(eegPower[1]) + ", " + String(eegPower[2]) + ", " + eegPower[3] + ", " + eegPower[4] + ", " + eegPower[5] + ", " + eegPower[6]  + ", " +  eegPower[7]);
          Serial.print("\n");
        }
        bigPacket = false;
      }
    }
  } else {
    drawBluetoothFailed();
    delay(5000);
    if (wasConnected) {
      M5.shutdown(1);
    }
  }
}

//////////////////////////////////////////
// Task 1 for blinking while connecting //
//////////////////////////////////////////

void task1(void *pvParameters) {
  while (1) {
    delay(500);
    if (connectingBluetooth) {
      bluetoothBlinking();
    }
  }
}

////////////////////////////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////////////////////////

//////////////////////////
// Bluetooth connection //
//////////////////////////

void connectBluetooth() {
  connectingBluetooth = true;
  if(! SerialBT.begin("ESP32test", true) ) {
    Serial.println("========== serialBT failed!");
  }

  Serial.println("Starting discoverAsync...");
  BTScanResults* btDeviceList = SerialBT.getScanResults();
  if (SerialBT.discoverAsync([](BTAdvertisedDevice* pDevice) {
      Serial.printf(">>>>>>>>>>>Found a new device asynchronously: %s\n", pDevice->toString().c_str());
    } )
    ) {
    delay(BT_DISCOVER_TIME);
    Serial.print("Stopping discoverAsync... ");
    SerialBT.discoverAsyncStop();
    Serial.println("discoverAsync stopped");
    delay(5000);
    if(btDeviceList->getCount() > 0) {
      BTAddress addr;
      int channel=0;
      Serial.println("Found devices:");
      for (int i=0; i < btDeviceList->getCount(); i++) {
        BTAdvertisedDevice *device=btDeviceList->getDevice(i);
        if (String(device->getAddress().toString().c_str()) == MAC_ADDRESS) {
          Serial.printf(" ----- %s  %s %d\n", device->getAddress().toString().c_str(), device->getName().c_str(), device->getRSSI());
          std::map<int,std::string> channels=SerialBT.getChannels(device->getAddress());
          Serial.printf("scanned for services, found %d\n", channels.size());
          for(auto const &entry : channels) {
            Serial.printf("     channel %d (%s)\n", entry.first, entry.second.c_str());
          }
          if(channels.size() > 0) {
            addr = device->getAddress();
            channel=channels.begin()->first;
          }
        }
      }
      if(addr) {
        Serial.printf("connecting to %s - %d\n", addr.toString().c_str(), channel);
        SerialBT.connect(addr, channel, sec_mask, role);
        isConnected = true;
        wasConnected = true;
      }
    } else {
      Serial.println("Didn't find any devices");
    }
  } else {
    Serial.println("Error on discoverAsync f.e. not workin after a \"connect\"");
  }
  connectingBluetooth = false;
}

/////////////////////////////////////////
// Read one byte from Serial Bluetooth //
/////////////////////////////////////////

byte readOneByte() {
  int ByteRead;
  if(! SerialBT.isClosed() && SerialBT.connected()) {
    ByteRead = SerialBT.read();
    
    return ByteRead;
  } else {
    Serial.println("not connected");
  }
}

///////////////////
// GUI functions //
///////////////////

void drawGui() {
  M5.Lcd.clear();
  M5.Lcd.fillScreen((COLORS[COLOR_MODE] - 1) % 2);
  M5.Lcd.setTextColor(COLORS[COLOR_MODE]);

  // HEADER
  updatePoorQuality();

  if (COLOR_MODE == DARK) {
    M5.lcd.drawBitmap(240, 5, 30, 50, (uint16_t *)bluetoothDark);
  } else {
    M5.lcd.drawBitmap(240, 5, 30, 50, (uint16_t *)bluetoothLight);
  }

  drawBattery();

  M5.Lcd.drawLine(0, 60, 320, 60, COLORS[COLOR_MODE]);

  // BODY
  M5.Lcd.drawString("Attention", 20, 80, 4);
  M5.Lcd.drawString("Meditation", 160, 80, 4);

  updateMeditation();
  updateAttention();

  // BOTTOM
  M5.Lcd.drawLine(0, 180, 320, 180, COLORS[COLOR_MODE]);
  screenModeButton.draw();
  if (COLOR_MODE == DARK) {
    M5.lcd.drawBitmap(0, 191, 50, 50, (uint16_t *)modeSwitchDark);
    M5.lcd.drawBitmap(65, 191, 50, 50, (uint16_t *)brightness_white);
    M5.lcd.drawBitmap(200, 191, 45, 50, (uint16_t *)restart_white);
    M5.lcd.drawBitmap(260, 186, 55, 55, (uint16_t *)power_off_white);
  } else {
    M5.lcd.drawBitmap(0, 191, 50, 50, (uint16_t *)modeSwitchLight);
    M5.lcd.drawBitmap(65, 191, 50, 50, (uint16_t *)brightness);
    M5.lcd.drawBitmap(200, 191, 45, 50, (uint16_t *)restart);
    M5.lcd.drawBitmap(260, 186, 55, 55, (uint16_t *)power_off);
  }
}

void drawBluetoothConnecting() {
  M5.Lcd.clear();
  M5.Lcd.fillScreen((COLORS[COLOR_MODE] - 1) % 2);
  M5.Lcd.setTextColor(COLORS[COLOR_MODE]);

  M5.Lcd.drawLine(0, 60, 320, 60, COLORS[COLOR_MODE]);

  M5.lcd.fillRect(0, 65, 320, 110, (COLORS[COLOR_MODE] - 1) % 2);
  M5.Lcd.drawString("connecting Bluetooth", 35, 110, 4);

  M5.Lcd.drawLine(0, 180, 320, 180, COLORS[COLOR_MODE]);

  drawBattery();
}

void drawBluetoothFailed() {
  M5.lcd.fillRect(0, 65, 320, 110, (COLORS[COLOR_MODE] - 1) % 2);
  if (wasConnected) {
    M5.lcd.fillRect(0, 0, 150, 60, (COLORS[COLOR_MODE] - 1) % 2);
    M5.lcd.fillRect(0, 185, 320, 240, (COLORS[COLOR_MODE] - 1) % 2);
    M5.lcd.fillRect(0, 185, 320, 240, (COLORS[COLOR_MODE] - 1) % 2);
    M5.Lcd.drawString("Bluetooth: disconnected", 20, 110, 4);
    for (int i = 0; i < 4; i++) {  // thicker line
      M5.lcd.drawLine(240 + i, 55, 270 + i, 5, RED);
    }
  } else {
    M5.Lcd.drawString("connecting failed", 50, 110, 4);
  }
}

void bluetoothBlinking() {
  while (connectingBluetooth) {
    if (COLOR_MODE == DARK) {
      M5.lcd.fillRect(240, 5, 35, 51, BLACK);
      delay(500);
      drawBattery();
      M5.lcd.drawBitmap(240, 5, 30, 50, (uint16_t *)bluetoothDark);
      delay(500);
    } else {
      M5.lcd.fillRect(240, 5, 35, 51, WHITE);
      delay(500);
      drawBattery();
      M5.lcd.drawBitmap(240, 5, 30, 50, (uint16_t *)bluetoothLight);
      delay(500);
    }
  }

  if (!isConnected) {
    for (int i = 0; i < 4; i++) {
      M5.lcd.drawLine(240 + i, 55, 270 + i, 5, RED);
    }      
  }  
}

void updatePoorQuality() {
  if (COLOR_MODE == DARK) {
    M5.lcd.fillRect(0, 0, 110, 59, BLACK);
    M5.Lcd.drawString(String(poorQuality), 60, 20, 4);

    if (poorQuality < 50) {
      M5.lcd.drawBitmap(5, 5, 50, 50, (uint16_t *)signal_white_3_3);
    } else if ((50 < poorQuality) && (poorQuality < 150)) {
      M5.lcd.drawBitmap(5, 5, 50, 50, (uint16_t *)signal_white_2_3);
    } else if ((poorQuality == 200) || (poorQuality > 230)) {
      M5.lcd.drawBitmap(5, 5, 50, 50, (uint16_t *)signal_white_0_3);
    } else {
      M5.lcd.drawBitmap(5, 5, 50, 50, (uint16_t *)signal_white_1_3);
    }
  } else {
    M5.lcd.fillRect(0, 0, 110, 59, WHITE);
    M5.Lcd.drawString(String(poorQuality), 60, 20, 4);

    if (poorQuality < 50) {
      M5.lcd.drawBitmap(5, 5, 50, 50, (uint16_t *)signal_3_3);
    } else if ((50 < poorQuality) && (poorQuality < 150)) {
      M5.lcd.drawBitmap(5, 5, 50, 50, (uint16_t *)signal_2_3);
    } else if ((poorQuality == 200) || (poorQuality > 230)) {
      M5.lcd.drawBitmap(5, 5, 50, 50, (uint16_t *)signal_0_3);
    } else {
      M5.lcd.drawBitmap(5, 5, 50, 50, (uint16_t *)signal_1_3);
    }
  }

  if (poorQuality == 200) {        // if headset not connected to person
    for (int i = 0; i < 4; i++) {  // thicker line
      M5.lcd.drawLine(5 + i, 55, 55 + i, 5, RED);
    }
  }
}

void updateAttention() {
  int x = 60 - ((countDigit(attention) - 1) * 5);
  if (COLOR_MODE == DARK) {  
    M5.lcd.fillRect(x - 10, 130, 60, 40, BLACK);
    M5.Lcd.drawString(String(attention), x, 130, 4);
  } else {
    M5.lcd.fillRect(x - 10, 130, 60, 40, WHITE);
    M5.Lcd.drawString(String(attention), x, 130, 4);
  }
}

void updateMeditation() {
  int x = 210 - ((countDigit(meditation) - 1) * 5);
  if (COLOR_MODE == DARK) {  
    M5.lcd.fillRect(x - 10, 130, 60, 40, BLACK);
    M5.Lcd.drawString(String(meditation), x, 130, 4);
  } else {
    M5.lcd.fillRect(x - 10, 130, 60, 40, WHITE);
    M5.Lcd.drawString(String(meditation), x, 130, 4);
  }
}

void drawBattery() {
  batVoltage = M5.Axp.GetBatVoltage();
  batPercentage = (batVoltage < 3.2) ? 0 : (batVoltage - 3.2) * 100;

  if (COLOR_MODE == DARK) {
    if (M5.Axp.isCharging()) {
      M5.lcd.drawBitmap(280, 5, 32, 50, (uint16_t *) batteryCharging_white);
    } else {
      if (batPercentage > 90) {
        M5.lcd.drawBitmap(280, 5, 32, 50, (uint16_t *) batteryStatus4_white);
      } else if ((90 > batPercentage) && (batPercentage > 60)) {
        M5.lcd.drawBitmap(280, 5, 32, 50, (uint16_t *) batteryStatus3_white);
      } else if ((60 > batPercentage) && (batPercentage > 20)) {
        M5.lcd.drawBitmap(280, 5, 32, 50, (uint16_t *) batteryStatus2_white);
      } else if ((20 > batPercentage) && (batPercentage > 10)) {
        M5.lcd.drawBitmap(280, 5, 32, 50, (uint16_t *) batteryStatus1_white);
      } else {
        M5.lcd.drawBitmap(280, 5, 32, 50, (uint16_t *) batteryStatus0_white);      
      }
    }
  } else {
    if (M5.Axp.isCharging()) {
      M5.lcd.drawBitmap(280, 5, 32, 50, (uint16_t *) batteryCharging);
    } else {
      if (batPercentage > 90) {
        M5.lcd.drawBitmap(280, 5, 32, 50, (uint16_t *) batteryStatus4);
      } else if ((90 > batPercentage) && (batPercentage > 60)) {
        M5.lcd.drawBitmap(280, 5, 32, 50, (uint16_t *) batteryStatus3);
      } else if ((60 > batPercentage) && (batPercentage > 20)) {
        M5.lcd.drawBitmap(280, 5, 32, 50, (uint16_t *) batteryStatus2);
      } else if ((20 > batPercentage) && (batPercentage > 10)) {
        M5.lcd.drawBitmap(280, 5, 32, 50, (uint16_t *) batteryStatus1);
      } else {
        M5.lcd.drawBitmap(280, 5, 32, 50, (uint16_t *) batteryStatus0);      
      }
    }
  }
}

////////////////////////////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////////////////////////

/////////////////////
// Button handlers //
/////////////////////

void screenModeHandler(Event &e) {
  M5.Axp.SetLDOEnable(3, true);
  delay(125);
  M5.Axp.SetLDOEnable(3, false);
  if (COLOR_MODE == DARK) {
    COLOR_MODE = LIGHT;
    preferences.putUInt(key1, 1);
  } else {
    COLOR_MODE = DARK;
    preferences.putUInt(key1, 0);
  }

  buttonPressed = true;
  drawGui();
}

void brightnessHandler(Event &e) {
  M5.Axp.SetLDOEnable(3, true);
  delay(125);
  M5.Axp.SetLDOEnable(3, false);
  switch (preferences.getUInt(key2)) {
    case 0:
      preferences.putUInt(key2, 1);
      break;
    case 1:
      preferences.putUInt(key2, 2);
      break;
    case 2:
      preferences.putUInt(key2, 3);
      break;
    case 3:
      preferences.putUInt(key2, 0);
      break;
  }
  M5.Axp.SetLcdVoltage(brightnessValues[preferences.getUInt(key2)]);
}

void restartHandler(Event& e) {
  M5.Axp.SetLDOEnable(3, true);
  delay(250);
  M5.Axp.SetLDOEnable(3, false);  
  M5.shutdown(1);
}

void powerOffHandler(Event& e) {
  M5.Axp.SetLDOEnable(3, true);
  delay(500);
  M5.Axp.SetLDOEnable(3, false);  
  M5.shutdown();  
}

int countDigit(int n) {
  if (n == 0)
      return 1;
  int count = 0;
  while (n != 0) {
      n = n / 10;
      ++count;
  }
  return count;
}