#include <Servo.h>
#include <EEPROM.h>
#include <math.h>

// ----------------- Pinos -----------------
const uint8_t PIN_SERVO_BASE   = 9;
const uint8_t PIN_SERVO_OMBRO  = 10;
const uint8_t PIN_SERVO_BRACO  = 11;
const uint8_t PIN_SERVO_GARRA  = 6;

const uint8_t PIN_J1_X  = A0;
const uint8_t PIN_J1_Y  = A1;
const uint8_t PIN_J1_SW = 2;
const uint8_t PIN_J2_X  = A2;
const uint8_t PIN_J2_Y  = A3;
const uint8_t PIN_J2_SW = 3;

// ----------------- Servos -----------------
Servo servoBase, servoOmbro, servoBraco, servoGarra;

// ----------------- Limites mecânicos -----------------
float baseMin  = 0.0f,   baseMax  = 180.0f, baseInit  = 90.0f;
float ombroMin = 10.0f,  ombroMax = 170.0f, ombroInit = 90.0f;
float bracoMin = 15.0f,  bracoMax = 165.0f, bracoInit = 90.0f;
float garraMin = 10.0f,  garraMax = 120.0f, garraInit = 60.0f;

// Soft-limits (graus antes do limite)
const float SOFT_ZONE_BASE  = 12.0f;
const float SOFT_ZONE_OMBRO = 12.0f;
const float SOFT_ZONE_BRACO = 12.0f;
const float SOFT_ZONE_GARRA = 10.0f;

// Sentido (+1 normal, -1 invertido)
int8_t INVERT_BASE  = +1;
int8_t INVERT_OMBRO = -1;
int8_t INVERT_BRACO = -1;
int8_t INVERT_GARRA = +1;

// ----------------- Controle base -----------------
const uint16_t CONTROL_PERIOD_MS = 10; // 100 Hz
const int      DEADZONE = 50;

// Velocidade máx no MANUAL (referência de v)
float MAX_RATE_BASE  = 180.0f;
float MAX_RATE_OMBRO = 150.0f;
float MAX_RATE_BRACO = 150.0f;
float MAX_RATE_GARRA = 200.0f;

// Expo do joystick
float EXPO_BASE  = 1.5f;
float EXPO_OMBRO = 1.5f;
float EXPO_BRACO = 1.5f;
float EXPO_GARRA = 1.2f;

// Soft-start
const float START_RAMP_DEG_PER_S = 90.0f;

// ----------------- Estados gerais -----------------
float angBase, angOmbro, angBraco, angGarra;
bool  preciseMode = false;
bool  invertGarra = false;
bool  estop = false;

bool  lastJ2 = HIGH;
uint32_t lastToggleMs = 0;
const uint16_t TOGGLE_DEBOUNCE_MS = 200;

uint32_t estopPressStart = 0;
bool j1Prev = HIGH, j2Prev = HIGH;

// EEPROM centros
struct CalData { uint32_t magic; int cx1, cy1, cx2, cy2; } cal = {0,512,512,512,512};
const uint32_t CAL_MAGIC = 0xA55AF11A;

// ----------------- Macro -----------------
struct Pose { float base, ombro, braco, garra; };
const uint8_t MAX_POSES = 16;
Pose poses[MAX_POSES];
uint8_t poseCount = 0;

bool recordMode = false;
bool playbackMode = false;
uint8_t playIndex = 0;
const uint16_t LONG_PRESS_MS = 2000;
uint32_t j1PressStart = 0, j2PressStart = 0;

// ----------------- Playback: AxisJL (posição jerk-limited) -----------------
struct AxisJL {
  float p, v, a;
  float Vmax, Amax, Jmax;
  float target;
  float posEps=0.4f, velEps=0.5f, accEps=2.0f;

  void setLimits(float vMax, float aMax, float jMax){ Vmax=vMax; Amax=aMax; Jmax=jMax; }
  void setStateFromAngle(float angle){ p=angle; v=0; a=0; }
  void setTarget(float tgt){ target=tgt; }
  bool reached() const { return (fabs(p-target)<=posEps)&& (fabs(v)<=velEps)&& (fabs(a)<=accEps); }

  float stoppingDistance() const {
    float dv=fabs(v), da=fabs(a);
    return (dv*dv)/(2.0f*Amax) + (da*da)/(2.0f*Jmax) + 0.1f;
  }
  void step(float dt){
    float d = target - p;
    float dir = (d>=0)?1.0f:-1.0f;
    bool needBrake = (fabs(d) <= stoppingDistance());

    float j = 0.0f;
    if (needBrake){
      if (v*dir > velEps || a*dir > accEps){
        if (a*dir > 0.0f) j = -Jmax*dir;
        else              j = (-Jmax)*((v*dir>0.0f)?dir:-dir);
      } else {
        j = (a>0)?-Jmax:((a<0)?+Jmax:0);
      }
    } else {
      if (v*dir < Vmax){
        if (a*dir < Amax) j = +Jmax*dir;
        else              j = 0.0f;
      } else {
        if (fabs(a) > accEps) j = (a>0)?-Jmax:+Jmax;
        else                  j = 0.0f;
      }
    }
    a += j*dt; if (a> Amax) a= Amax; if (a<-Amax) a=-Amax;
    v += a*dt; if (v> Vmax) v= Vmax; if (v<-Vmax) v=-Vmax;
    p += v*dt;

    if ( ((target-p)*dir<0.0f) && (fabs(v)<velEps) && (fabs(a)<accEps) ){ p=target; v=0; a=0; }
  }
};

// ----------------- Manual: VelJL (velocidade jerk-limited) -----------------
struct VelJL {
  float v=0.0f, a=0.0f;
  float Vmax=180.0f;
  float Amax=300.0f;
  float Jmax=2500.0f;
  float velEps=0.2f, accEps=2.0f;

  void setLimits(float vMax, float aMax, float jMax){
    Vmax=vMax; Amax=aMax; Jmax=jMax;
  }
  float step(float v_cmd, float dt){
    if (v_cmd >  Vmax) v_cmd =  Vmax;
    if (v_cmd < -Vmax) v_cmd = -Vmax;

    float dv = v_cmd - v;
    float dir = (dv>=0)?1.0f:-1.0f;
    float j = 0.0f;

    if (fabs(dv) > velEps){
      if ( (a*dir) < Amax ) j = +Jmax*dir;
      else                  j = 0.0f;
      if ( (v_cmd==0.0f && fabs(v)<velEps) || (dv*dir<0.0f && fabs(dv)<2*velEps) ){
        j = (a>0)?-Jmax:((a<0)?+Jmax:0.0f);
      }
    } else {
      if (fabs(a) > accEps) j = (a>0)?-Jmax:+Jmax;
      else                  j = 0.0f;
    }

    a += j*dt; if (a> Amax) a= Amax; if (a<-Amax) a=-Amax;
    v += a*dt; if (v> Vmax) v= Vmax; if (v<-Vmax) v=-Vmax;
    return v;
  }
};

AxisJL jlBase, jlOmbro, jlBraco, jlGarra;
VelJL  vmBase, vmOmbro, vmBraco, vmGarra;

// ----------------- Utils -----------------
inline float constrainf(float x,float a,float b){ if(x<a)return a; if(x>b)return b; return x; }
inline float expoMap(int deflRaw,float expo){
  float x = deflRaw/511.0f; if(x>1)x=1; if(x<-1)x=-1; float s=(x>=0)?1.0f:-1.0f; return s*powf(fabs(x),expo);
}
int readAxisCentered(uint8_t pin,int center){
  int v=analogRead(pin); int d=v-center; if(abs(d)<DEADZONE) return 0; if(d>0)d-=DEADZONE; else d+=DEADZONE; return d;
}
void attachServos(){
  servoBase.attach(PIN_SERVO_BASE);
  servoOmbro.attach(PIN_SERVO_OMBRO);
  servoBraco.attach(PIN_SERVO_BRACO);
  servoGarra.attach(PIN_SERVO_GARRA);
}
void writeServos(){
  servoBase.write((int)angBase);
  servoOmbro.write((int)angOmbro);
  servoBraco.write((int)angBraco);
  servoGarra.write((int)angGarra);
}
void loadCalFromEEPROM(){
  EEPROM.get(0, cal);
  if (cal.magic != CAL_MAGIC){ cal.magic=CAL_MAGIC; cal.cx1=512; cal.cy1=512; cal.cx2=512; cal.cy2=512; EEPROM.put(0,cal); }
}
void saveCalToEEPROM(){ cal.magic=CAL_MAGIC; EEPROM.put(0,cal); }

float softLimitFactorSigned(float ang,float amin,float amax,float zone,float rate){
  if(rate>0){ float d=amax-ang; return (d<zone)?constrainf(d/zone,0,1):1.0f; }
  if(rate<0){ float d=ang-amin; return (d<zone)?constrainf(d/zone,0,1):1.0f; }
  return 1.0f;
}
float rampToward(float current,float target,float maxRate,float dt){
  float delta=target-current, step=maxRate*dt;
  if(delta> step) return current+step;
  if(delta<-step) return current-step;
  return target;
}

// ----------------- Setup -----------------
void setup(){
  Serial.begin(115200);
  pinMode(PIN_J1_SW, INPUT_PULLUP);
  pinMode(PIN_J2_SW, INPUT_PULLUP);

  loadCalFromEEPROM();

  if (digitalRead(PIN_J1_SW)==LOW){
    delay(300);
    cal.cx1=analogRead(PIN_J1_X); cal.cy1=analogRead(PIN_J1_Y);
    cal.cx2=analogRead(PIN_J2_X); cal.cy2=analogRead(PIN_J2_Y);
    saveCalToEEPROM();
    Serial.println(F("[CAL] Centros salvos."));
  }

  attachServos();

  angBase  = constrainf((float)servoBase.read(),  baseMin,  baseMax);
  angOmbro = constrainf((float)servoOmbro.read(), ombroMin, ombroMax);
  angBraco = constrainf((float)servoBraco.read(), bracoMin, bracoMax);
  angGarra = constrainf((float)servoGarra.read(), garraMin, garraMax);

  jlBase .setStateFromAngle(angBase ); jlBase .setLimits(120,400,2500);
  jlOmbro.setStateFromAngle(angOmbro); jlOmbro.setLimits(100,350,2200);
  jlBraco.setStateFromAngle(angBraco); jlBraco.setLimits(100,350,2200);
  jlGarra.setStateFromAngle(angGarra); jlGarra.setLimits(140,500,3000);

  vmBase .setLimits(MAX_RATE_BASE , 300, 2500);
  vmOmbro.setLimits(MAX_RATE_OMBRO, 280, 2200);
  vmBraco.setLimits(MAX_RATE_BRACO, 280, 2200);
  vmGarra.setLimits(MAX_RATE_GARRA, 350, 3000);
}

// ----------------- Loop -----------------
void loop(){
  static uint32_t lastMs=0;
  uint32_t now=millis();
  if (now-lastMs < CONTROL_PERIOD_MS) return;
  float dt=(now-lastMs)/1000.0f; lastMs=now;

  bool j1 = (digitalRead(PIN_J1_SW)==LOW);
  bool j2 = (digitalRead(PIN_J2_SW)==LOW);

  preciseMode = j1 && !playbackMode;

  static uint32_t bothPressStart=0;
  if (j1 && j2){
    if(bothPressStart==0) bothPressStart=now;
    else if ((now-bothPressStart)>=1000 && (j1Prev||j2Prev)){
      estop=!estop;
      Serial.print(F("[E-STOP] ")); Serial.println(estop?F("ATIVADO"):F("DESATIVADO"));
      bothPressStart=now+10000;
    }
  } else bothPressStart=0;
  j1Prev=j1; j2Prev=j2;

  if (!estop){
    static uint32_t j1PressStart=0;
    if (j1){
      if (j1PressStart==0) j1PressStart=now;
      else if (!recordMode && !playbackMode && (now-j1PressStart)>=LONG_PRESS_MS){
        recordMode=true; Serial.println(F("[MACRO] Gravacao INICIADA (J2 clique salva pose)."));
        j1PressStart=now+10000;
      }
    } else {
      if (recordMode && j1PressStart && (now-j1PressStart)>=LONG_PRESS_MS){
        recordMode=false; Serial.print(F("[MACRO] Gravacao ENCERRADA. Poses=")); Serial.println(poseCount);
      }
      j1PressStart=0;
    }

    static uint32_t j2PressStart=0;
    if (j2){
      if (j2PressStart==0) j2PressStart=now;
      else if (!playbackMode && (now-j2PressStart)>=LONG_PRESS_MS && poseCount>0){
        playbackMode=true; recordMode=false; playIndex=0;
        jlBase .setTarget(constrainf(poses[0].base , baseMin , baseMax ));
        jlOmbro.setTarget(constrainf(poses[0].ombro, ombroMin, ombroMax));
        jlBraco.setTarget(constrainf(poses[0].braco, bracoMin, bracoMax));
        jlGarra.setTarget(constrainf(poses[0].garra, garraMin, garraMax));
        Serial.println(F("[MACRO] Reproducao INICIADA (jerk-limited)."));
        j2PressStart=now+10000;
      } else if (playbackMode && (now-j2PressStart)>=LONG_PRESS_MS){
        playbackMode=false; Serial.println(F("[MACRO] Reproducao PARADA."));
        j2PressStart=now+10000;
      }
    } else j2PressStart=0;
  }

  bool j2State=(digitalRead(PIN_J2_SW)==LOW);
  if (j2State && !lastJ2 && (now-lastToggleMs)>TOGGLE_DEBOUNCE_MS){
    if (recordMode){
      if (poseCount<MAX_POSES){
        poses[poseCount++] = { angBase, angOmbro, angBraco, angGarra };
        Serial.print(F("[MACRO] Pose salva (#")); Serial.print(poseCount); Serial.println(F(")."));
      } else Serial.println(F("[MACRO] Limite de poses atingido."));
    } else if (!playbackMode){
      invertGarra=!invertGarra;
      Serial.print(F("[GARRA] invertida=")); Serial.println(invertGarra?F("SIM"):F("NAO"));
    }
    lastToggleMs=now;
  }
  lastJ2=j2State;

  static bool rampDone=false;
  if (!rampDone){
    angBase  = rampToward(angBase,  baseInit,  START_RAMP_DEG_PER_S, dt);
    angOmbro = rampToward(angOmbro, ombroInit, START_RAMP_DEG_PER_S, dt);
    angBraco = rampToward(angBraco, bracoInit, START_RAMP_DEG_PER_S, dt);
    angGarra = rampToward(angGarra, garraInit, START_RAMP_DEG_PER_S, dt);
    if (fabs(angBase-baseInit)<0.5f && fabs(angOmbro-ombroInit)<0.5f &&
        fabs(angBraco-bracoInit)<0.5f && fabs(angGarra-garraInit)<0.5f){
      rampDone=true;
      jlBase .setStateFromAngle(angBase );
      jlOmbro.setStateFromAngle(angOmbro);
      jlBraco.setStateFromAngle(angBraco);
      jlGarra.setStateFromAngle(angGarra);
    }
    writeServos(); return;
  }

  if (estop){ writeServos(); return; }

  if (playbackMode && poseCount>0){
    if (jlBase.reached() && jlOmbro.reached() && jlBraco.reached() && jlGarra.reached()){
      playIndex = (playIndex+1) % poseCount;
      const Pose &t = poses[playIndex];
      jlBase .setTarget(constrainf(t.base , baseMin , baseMax ));
      jlOmbro.setTarget(constrainf(t.ombro, ombroMin, ombroMax));
      jlBraco.setTarget(constrainf(t.braco, bracoMin, bracoMax));
      jlGarra.setTarget(constrainf(t.garra, garraMin, garraMax));
    }
    jlBase.step(dt); jlOmbro.step(dt); jlBraco.step(dt); jlGarra.step(dt);
    angBase  = constrainf(jlBase.p , baseMin , baseMax );
    angOmbro = constrainf(jlOmbro.p, ombroMin, ombroMax);
    angBraco = constrainf(jlBraco.p, bracoMin, bracoMax);
    angGarra = constrainf(jlGarra.p, garraMin, garraMax);
    writeServos(); return;
  }

  int x1 = readAxisCentered(PIN_J1_X, cal.cx1);
  int y1 = readAxisCentered(PIN_J1_Y, cal.cy1);
  int x2 = readAxisCentered(PIN_J2_X, cal.cx2);
  int y2 = readAxisCentered(PIN_J2_Y, cal.cy2);

  float ux1 = expoMap(x1, EXPO_BASE );
  float uy1 = expoMap(y1, EXPO_OMBRO);
  float ux2 = expoMap(x2, EXPO_GARRA);
  float uy2 = expoMap(y2, EXPO_BRACO);

  float precisionK = preciseMode ? 0.5f : 1.0f;

  float vrefBase  = INVERT_BASE  * ux1 * MAX_RATE_BASE  * precisionK;
  float vrefOmbro = INVERT_OMBRO * uy1 * MAX_RATE_OMBRO * precisionK;
  float vrefBraco = INVERT_BRACO * uy2 * MAX_RATE_BRACO * precisionK;
  int   garraDir  = (invertGarra ? -1 : +1);
  float vrefGarra = INVERT_GARRA * garraDir * ux2 * MAX_RATE_GARRA * precisionK;

  float fBase  = softLimitFactorSigned(angBase , baseMin , baseMax , SOFT_ZONE_BASE ,  vrefBase );
  float fOmbro = softLimitFactorSigned(angOmbro, ombroMin, ombroMax, SOFT_ZONE_OMBRO, vrefOmbro);
  float fBraco = softLimitFactorSigned(angBraco, bracoMin, bracoMax, SOFT_ZONE_BRACO, vrefBraco);
  float fGarra = softLimitFactorSigned(angGarra, garraMin, garraMax, SOFT_ZONE_GARRA, vrefGarra);

  vrefBase  *= fBase;
  vrefOmbro *= fOmbro;
  vrefBraco *= fBraco;
  vrefGarra *= fGarra;

  float vBase  = vmBase .step(vrefBase , dt);
  float vOmbro = vmOmbro.step(vrefOmbro, dt);
  float vBraco = vmBraco.step(vrefBraco, dt);
  float vGarra = vmGarra.step(vrefGarra, dt);

  angBase  = constrainf(angBase  + vBase  *dt, baseMin , baseMax );
  angOmbro = constrainf(angOmbro + vOmbro *dt, ombroMin, ombroMax);
  angBraco = constrainf(angBraco + vBraco *dt, bracoMin, bracoMax);
  angGarra = constrainf(angGarra + vGarra *dt, garraMin, garraMax);

  writeServos();

  static uint32_t lastPrint=0;
  if (now-lastPrint>350){
    lastPrint=now;
    Serial.print(F("ManualJL vB=")); Serial.print(vBase,1);
    Serial.print(F(" vO=")); Serial.print(vOmbro,1);
    Serial.print(F(" vR=")); Serial.print(vBraco,1);
    Serial.print(F(" vG=")); Serial.print(vGarra,1);
    Serial.print(F(" | A: ")); Serial.print((int)angBase); Serial.print(' ');
    Serial.print((int)angOmbro); Serial.print(' ');
    Serial.print((int)angBraco); Serial.print(' ');
    Serial.print((int)angGarra); Serial.println();
  }
}