//
// M5ATOM Balancing Robot
// Kiraku Labo, 2022
//
// 1. Power on and wait for a circle to shows up.
// 2. Hold the robot upright.
// 3. When a smile shows up, release the robot.
//

#include <M5Atom.h>
#include <esp_now.h>
#include <WiFi.h>
#define CHANNEL 0

#define MAXDRV 255
#define SDA_PIN 26
#define SCL_PIN 32
#define MOTOR_R 0x65
#define MOTOR_L 0x63

#define BLU 0x0000ff
#define GRN 0xff0000
#define RED 0x00ff00
#define BLK 0x000000

boolean serialMonitor=true;
uint16_t counter=0, counterOverSpd=0, state=0;
uint32_t time0=0, time1=0;
float power, powerR, powerL, yawPower, maxPwr;
float varAng, varOmg, varSpd, varDst, varIang, aveAbsOmg=0.0, yawAng=0.0;
float Kang, Komg, KIang, Kyaw, Kdst, Kspd;
float gyroYdata, gyroZdata, accXdata, accXoffset, gyroZoffset, gyroYoffset;
const uint32_t interval = 10; // in msec
const float dt=(float)interval/1000.0; // in sec
float movePower=0.0, movePwrTarget, movePwrStep=1.0;
float mechFactorR, mechFactorL;
float spinStep, spinDest, spinTarget, spinFact=1.0;
bool joyButton=false;

int16_t drvOffset, punchSpd, punchSpd2, punchSpd2N, punchDur, punchCountL=0, punchCountR=0;
int16_t motorDeadband=0, counterTrip=0;
int8_t motorRdir=0, motorLdir=0;
bool spinContinuous=false;

byte joyMACaddr[] = {0xff, 0xff, 0xff, 0xff, 0xff, 0xff};//{0x50, 0x02, 0x91, 0x86, 0x99, 0x68};

const byte smile[] = {0,0,0,0,0,
                      0,1,0,1,0,
                      0,0,0,0,0,
                      1,0,0,0,1,
                      0,1,1,1,0};
const byte frown[] = {0,0,0,0,0,
                      0,1,0,1,0,
                      0,0,0,0,0,
                      0,1,1,1,0,
                      1,0,0,0,1};
const byte circle[] = {0,1,1,1,0,
                       1,0,0,0,1,
                       1,0,0,0,1,
                       1,0,0,0,1,
                       0,1,1,1,0};
const byte square[] = {0,0,0,0,0,
                      0,1,1,1,0,
                      0,1,0,1,0,
                      0,1,1,1,0,
                      0,0,0,0,0};
const byte cross[] = {1,0,0,0,1,
                      0,1,0,1,0,
                      0,0,1,0,0,
                      0,1,0,1,0,
                      1,0,0,0,1};
const byte heart[] = {0,1,0,1,0,
                      1,0,1,0,1,
                      1,0,0,0,1,
                      0,1,0,1,0,
                      0,0,1,0,0};
const byte all[] = {1,1,1,1,1,
                    1,1,1,1,1,
                    1,1,1,1,1,
                    1,1,1,1,1,
                    1,1,1,1,1};

void setup() {
  M5.begin(true, true, true); //Serial, i2c(Wire1), disp
  Wire.begin(SDA_PIN, SCL_PIN, 50000); //clock=50kHz
  Wire.beginTransmission(MOTOR_R); //check motor driver
  if (Wire.endTransmission()!=0) { //motor driver connection problem
    ledDispImg(cross, RED);
    delay(5000);
  }
  M5.IMU.Init();
  setupESPNOW();
  delay(1500);
  resetMotor();
  resetPara();
  resetVar();
  getBaselineOmg();
  ledDispImg(circle, GRN);
}

void loop() {
  static uint32_t msec=0;
  getIMU();
  switch (state) {
    case 0:
      if (joyButton) ledDispImg(circle, RED);
      else ledDispImg(circle, GRN);
      if (abs(accXdata)<0.25 && aveAbsOmg<1.0) {
        ledDispImg(square, GRN);
        getBaseline();
        msec=millis();
        resetVar();
        state=1;
      }
      break;
    case 1:
      if (millis()-msec>200) state=2;
      break;
    case 2:
      if (abs(varAng)>30.0 || counterOverSpd>25) {
        if (serialMonitor) sendStatus();
        resetMotor();
        resetVar();
        ledDispImg(frown, RED);
        state=0;
      }
      else {
        if (joyButton) ledDispImg(smile, RED);
        else ledDispImg(smile, GRN);
        calcTarget();
        drive();
      }
      break;
  }
  counter += 1;
  if (counter%25==0) {
    sendData("vs=" + String(varSpd) + "\nvd=" + String(varDst) + "\nva=" + String(varAng) + "\npw=" + String(power));
  }
  if (counter >= 100) {
    counter = 0;
    if (serialMonitor) sendStatus();
  }
  do time1 = millis();
  while (time1 - time0 < interval);
  time0 = time1;
}

void getIMU() {
  readIMU();
  varOmg = gyroYdata - gyroYoffset;
  aveAbsOmg = aveAbsOmg * 0.8 + abs(varOmg) * 0.2;
  yawAng += (gyroZdata - gyroZoffset) * dt;
  varAng = (varAng + varOmg * dt) * 0.995 + (accXdata - accXoffset) * 57.3 * 0.005;
}

void readIMU() {
  float gX, gY, gZ, aX, aY, aZ;
  M5.IMU.getGyroData(&gX,&gY,&gZ);
  M5.IMU.getAccelData(&aX, &aY, &aZ);
  gyroYdata=-gX;
  gyroZdata=-gY;
  accXdata=aZ;
}

void calcTarget() {
  spinTarget += spinStep;
  if (movePwrTarget < movePower-movePwrStep) movePwrTarget += movePwrStep;
  else if (movePwrTarget > movePower+movePwrStep) movePwrTarget -= movePwrStep;
  else movePwrTarget = movePower;
}

void drive() {
  varSpd += power * dt;
  varDst += Kdst * varSpd * dt;
  varIang += KIang * varAng * dt;
  power = varIang + varDst + (Kspd * varSpd) + (Kang * varAng) + (Komg * varOmg);
  if (abs(power) > 500.0) counterOverSpd += 1;
  else counterOverSpd=0;
  power = constrain(power, -maxPwr, maxPwr);
  yawPower = (yawAng - spinTarget) * Kyaw;
  powerR = power + yawPower + movePwrTarget;
  powerL = power - yawPower + movePwrTarget;
  int16_t ipowerL = (int16_t) constrain(powerL * mechFactorL, -maxPwr, maxPwr);
  if (ipowerL > 0) {
    if (motorLdir == 1) {
      punchCountL++;
      punchCountL = constrain(punchCountL, 0, 100);
    }
    else punchCountL=0;
    motorLdir = 1;
    if (punchCountL<punchDur) drvMotorL(max(ipowerL, punchSpd2)+drvOffset);
    else drvMotorL(ipowerL+motorDeadband+drvOffset);
  }
  else if (ipowerL < 0) {
    if (motorLdir == -1) {
      punchCountL++;
      punchCountL = constrain(punchCountL, 0, 100);
    }
    else punchCountL=0;
    motorLdir = -1;
    if (punchCountL<punchDur) drvMotorL(min(ipowerL, punchSpd2N) +drvOffset);
    else drvMotorL(ipowerL-motorDeadband+drvOffset);
  }
  else {
    drvMotorL(0);
    motorLdir = 0;
  }
  
  int16_t ipowerR = (int16_t) constrain(powerR * mechFactorR, -maxPwr, maxPwr);
  if (ipowerR > 0) {
    if (motorRdir == 1) {
      punchCountR++;
      punchCountR = constrain(punchCountL, 0, 100);
    }
    else punchCountR=0;
    motorRdir = 1;
    if (punchCountR<punchDur) drvMotorR(max(ipowerR, punchSpd2) +drvOffset);
    else drvMotorR(ipowerR+motorDeadband+drvOffset);
  }
  else if (ipowerR < 0) {
    if (motorRdir == -1) {
      punchCountR++;
      punchCountR = constrain(punchCountR, 0, 100);
    }
    else punchCountR=0;
    motorRdir = -1;
    if (punchCountR<punchDur) drvMotorR(min(ipowerR, punchSpd2N)+drvOffset);
    else drvMotorR(ipowerR-motorDeadband+drvOffset);
  }
  else {
    drvMotorR(0);
    motorRdir = 0;
  }
}

void drvMotorL(int16_t pwr) {
  drv8830(MOTOR_L, constrain(pwr, -255, 255));
}

void drvMotorR(int16_t pwr) {
  drv8830(MOTOR_R, constrain(pwr, -255, 255));
}

void drv8830(byte adr, int16_t pwr) { //pwr -255 to 255
  byte dir, st;
  if (pwr < 0) dir = 2;
  else dir =1;
  byte ipower=(byte)(abs(pwr)/4);
  if (ipower == 0) dir=3; //brake
  ipower = constrain (ipower, 0, 63);
  st = drv8830read(adr);
  if (st & 1) drv8830write(adr, 0, 0);
  drv8830write(adr, ipower, dir);
}

void drv8830write(byte adr, byte pwm, byte ctrl) {
  Wire.beginTransmission(adr);
  Wire.write(0);
  Wire.write(pwm*4+ctrl);
  Wire.endTransmission(true);
}

int drv8830read(byte adr) {
  Wire.beginTransmission(adr);
  Wire.write(1);
  Wire.endTransmission(false);
  Wire.requestFrom((int)adr, (int)1, (int)1);
  return Wire.read();
}

void resetMotor() {
  drv8830(MOTOR_R, 0);
  drv8830(MOTOR_L, 0);
  counterOverSpd=0;
  counterTrip=0;
}

void resetPara() {
  Kang=35.0;
  Komg=0.6;
  KIang=700.0;
  Kdst=11.0;
  Kspd=7.5;
  Kyaw=10.0;
  mechFactorL=0.7;
  mechFactorR=0.7;
  maxPwr=550.0;
  punchSpd=50;
  punchDur=1;
  motorDeadband=35;
  punchSpd2=max(punchSpd, motorDeadband);
  punchSpd2N=-punchSpd2;
  drvOffset=0;
}

void resetVar() {
  power=0.0;
  movePwrTarget=0.0;
  movePower=0.0;
  spinStep=0.0;
  spinDest=0.0;
  spinTarget=0.0;
  yawAng=0.0;
  varAng=0.0;
  varOmg=0.0;
  varDst=0.0;
  varSpd=0.0;
  varIang=0.0;
}

void getBaseline() {
  accXoffset=0.0;
  gyroZoffset=0.0;
  for (int i=0; i < 30; i++){
    readIMU();
    accXoffset += accXdata;
    gyroZoffset += gyroZdata;
    delay(10);
  }
  accXoffset /= 30.0;
  gyroZoffset /= 30.0;
}

void getBaselineOmg() {
  gyroYoffset=0.0;
  for (int i=0; i < 30; i++){
    readIMU();
    gyroYoffset += gyroYdata;
    delay(10);
  }
  gyroYoffset /= 30.0;
}

void sendStatus () {
  Serial.print("state="); Serial.print(state);
  Serial.print(" accX="); Serial.print(accXdata);
  Serial.print(" ang=");Serial.print(varAng);
  Serial.print(" "); 
  Serial.print(millis()-time0);
  Serial.println("ms");
}

void ledDispImg(const byte* p, uint64_t c) {
  for (byte i=0; i<25; i++) {
    M5.dis.drawpix(i, p[i]*c);
  }
}

void setupESPNOW() {
  Serial.print("MAC=");
  Serial.println(WiFi.macAddress());
  WiFi.mode(WIFI_STA);
  esp_now_init();
  esp_now_register_recv_cb(onRecv);
  esp_now_peer_info_t peerInfo={};
  memcpy(peerInfo.peer_addr, joyMACaddr, 6);
  peerInfo.channel = CHANNEL;
  peerInfo.encrypt = false;
  esp_now_add_peer(&peerInfo);
}

void onRecv(const uint8_t * mac, const uint8_t *recvData, int len) {
  if ((int8_t)recvData[2]==1) joyButton=true;
  else joyButton=false;
  steer((int8_t)recvData[0], (int8_t)recvData[1]);
}

void sendData(String s) {
  esp_now_send(joyMACaddr, (const byte*)s.c_str(), s.length());
}

void steer(int8_t x, int8_t y) {
  if (abs(y)>20) movePower=(float)y*0.5;
  else movePower=0.0;
  if (abs(x)>30) spinStep=(float)x*0.015;
  else spinStep=0.0;
}