#include <Arduino.h>
#include <AccelStepper.h>
#include <math.h>
#include <ctype.h>
#include <string.h>

#define EN1   4
#define STEP1 5
#define DIR1  6

#define EN2   7
#define STEP2 8
#define DIR2  9

#define BUTTON_PIN 12

AccelStepper m1(AccelStepper::DRIVER, STEP1, DIR1);
AccelStepper m2(AccelStepper::DRIVER, STEP2, DIR2);

// ----------------- USER SETTINGS -----------------
float TOTAL_TIME_S = 62.0f;   // total run time target
float TURN_TIME_S  = 1.2f;   // each 90-deg turn must take this long (seconds)

// new wheels: 3 7/8 inch diameter
const float WHEEL_DIAMETER_CM = 3.875f * 2.54f;

// steps per full wheel revolution (already includes microstepping if you set it that way)
const int   STEPS_PER_REV = 3200;

// wheel path radii from robot turn center (mm)
const float WHEEL_OFFSET_M1_MM = 90.66f;  // left wheel center offset
const float WHEEL_OFFSET_M2_MM = 90.66f;  // right wheel center offset

// 90-degree arc length for each wheel (cm): s = (pi/2)*r
const float TURN_ARC_CM_M1 = (PI * 0.5f) * (WHEEL_OFFSET_M1_MM / 10.0f);
const float TURN_ARC_CM_M2 = (PI * 0.5f) * (WHEEL_OFFSET_M2_MM / 10.0f);

// accelstepper tuning caps
float STRAIGHT_ACCEL_STEPSPS2 = 8000.0f;
float TURN_ACCEL_STEPSPS2     = 5000.0f;

float STRAIGHT_VMAX_MAX = 4500.0f;
float TURN_VMAX_MAX     = 2200.0f;
float VMIN_STEPSPS      = 80.0f;

bool ENABLE_ACTIVE_LOW = true;

// Script syntax:
// s50;  -> forward 50 cm
// b20;  -> backward 20 cm
// l;    -> 90° left
// r;    -> 90° right
// a;    -> 180° turn-around (two rights)
const char SCRIPT[] = R"(
s178.5;
a;
s100;
l;
s50;
r;
s50;
b96.5;
l;
s50;
l;
s50;
l;
s50;
b46.5;
l;
s150;
b50;
r;
s100;
l;
s50;
r;
s50;
r;
s50;
a;
s50;
r;
s50;
b48.5;
r;
s56.7;
)";
// --------------------------------------------------

// segment type MUST be defined before any function uses it
enum SegType : uint8_t { SEG_FWD = 0, SEG_TURN_L = 1, SEG_TURN_R = 2 };

struct Segment {
  SegType type;
  float   cm;   // only meaningful for SEG_FWD
  long    s1;   // steps for motor 1 (signed)
  long    s2;   // steps for motor 2 (signed)
};

const int MAX_SEGS = 150;
Segment segs[MAX_SEGS];
int segCount = 0;
int segIdx   = 0;

bool started = false;
unsigned long runStartMs = 0;

// ----------------- Math helpers -----------------
static inline int cmToSteps(float cm) {
  return (int)lroundf((cm / (PI * WHEEL_DIAMETER_CM)) * (float)STEPS_PER_REV);
}

// solve vmax for trapezoid/triangle time T with accel a and distance d_steps (absolute)
static bool solve_vmax(long d_steps, float T, float a, float &vmax_out) {
  if (T <= 0.0f || a <= 0.0f) return false;

  double d = (double)((d_steps >= 0) ? d_steps : -d_steps);
  double A = 1.0 / (double)a;
  double B = -(double)T;
  double C = d;

  double disc = B*B - 4.0*A*C;
  if (disc < 0.0) return false;

  double sqrt_disc = sqrt(disc);
  double v1 = (-B - sqrt_disc) / (2.0*A);
  double v2 = (-B + sqrt_disc) / (2.0*A);

  double v = v1;
  if (v <= 0.0) v = v2;
  if (v <= 0.0) return false;

  vmax_out = (float)v;
  return true;
}

// ----------------- Parser helpers -----------------
static inline void skipWS(const char* &p) {
  while (*p && (isspace((unsigned char)*p) || *p == '\r' || *p == '\n')) p++;
}

static bool parseFloat(const char* &p, float &out) {
  skipWS(p);
  char* endp = nullptr;
  out = strtof(p, &endp);
  if (endp == p) return false;
  p = endp;
  return true;
}

static bool eatChar(const char* &p, char c) {
  skipWS(p);
  if (*p != c) return false;
  p++;
  return true;
}

static void pushSeg(SegType type, float cm, long s1, long s2) {
  if (segCount >= MAX_SEGS) return;
  segs[segCount].type = type;
  segs[segCount].cm   = cm;
  segs[segCount].s1   = s1;
  segs[segCount].s2   = s2;
  segCount++;
}

static bool parseScript() {
  segCount = 0;
  const char* p = SCRIPT;

  while (1) {
    skipWS(p);
    if (!*p) break;

    char c = (char)tolower((unsigned char)*p);

    if (c == 's') {
      p++;
      float cm = 0.0f;
      if (!parseFloat(p, cm)) return false;
      if (!eatChar(p, ';')) return false;

      long s = (long)cmToSteps(cm);
      if (segCount > 0 && segs[segCount-1].type == SEG_FWD && ((segs[segCount-1].cm > 0) == (cm > 0))) {
        segs[segCount-1].cm += cm;
        segs[segCount-1].s1 += s;
        segs[segCount-1].s2 -= s;
      } else {
        pushSeg(SEG_FWD, cm, s, -s);
      }
      continue;
    }

    if (c == 'b') {
      p++;
      float cm = 0.0f;
      if (!parseFloat(p, cm)) return false;
      if (!eatChar(p, ';')) return false;

      float cmNeg = -cm;
      long s = (long)cmToSteps(cmNeg);
      if (segCount > 0 && segs[segCount-1].type == SEG_FWD && ((segs[segCount-1].cm > 0) == (cmNeg > 0))) {
        segs[segCount-1].cm += cmNeg;
        segs[segCount-1].s1 += s;
        segs[segCount-1].s2 -= s;
      } else {
        pushSeg(SEG_FWD, cmNeg, s, -s);
      }
      continue;
    }
    
    if (c == 'l') {
      p++;
      if (!eatChar(p, ';')) return false;

      long s1 = (long)cmToSteps(TURN_ARC_CM_M1);
      long s2 = (long)cmToSteps(TURN_ARC_CM_M2);
      pushSeg(SEG_TURN_L, 0.0f, -s1, -s2);
      continue;
    }

    if (c == 'r') {
      p++;
      if (!eatChar(p, ';')) return false;

      long s1 = (long)cmToSteps(TURN_ARC_CM_M1);
      long s2 = (long)cmToSteps(TURN_ARC_CM_M2);
      pushSeg(SEG_TURN_R, 0.0f,  s1,  s2);
      continue;
    }

    if (c == 'a') {
      p++;
      if (!eatChar(p, ';')) return false;

      long s1 = (long)cmToSteps(TURN_ARC_CM_M1);
      long s2 = (long)cmToSteps(TURN_ARC_CM_M2);
      pushSeg(SEG_TURN_R, 0.0f, s1, s2);
      pushSeg(SEG_TURN_R, 0.0f, s1, s2);
      continue;
    }

    return false;
  }

  return segCount > 0;
}

// ----------------- Time-based speed logic -----------------
static int turnsRemainingFrom(int fromIdx) {
  int n = 0;
  for (int i = fromIdx; i < segCount; i++) if (segs[i].type != SEG_FWD) n++;
  return n;
}

static float straightStepsRemainingFrom(int fromIdx) {
  double sum = 0.0;
  for (int i = fromIdx; i < segCount; i++) {
    if (segs[i].type == SEG_FWD) sum += (double)labs(segs[i].s1);
  }
  return (float)sum;
}

// recompute straight vmax so straights finish in remaining straight-time
static float computeStraightVmaxNow(int idx) {
  float elapsed = (millis() - runStartMs) / 1000.0f;
  float targetLeft = TOTAL_TIME_S - elapsed;
  if (targetLeft < 0.05f) targetLeft = 0.05f;

  int tRem = turnsRemainingFrom(idx);
  float timeTurnsLeft = (float)tRem * TURN_TIME_S;
  float timeStraightsLeft = targetLeft - timeTurnsLeft;
  if (timeStraightsLeft < 0.05f) timeStraightsLeft = 0.05f;

  float stepsStraightsLeft = straightStepsRemainingFrom(idx);
  if (stepsStraightsLeft < 1.0f) stepsStraightsLeft = 1.0f;

  float vavg = stepsStraightsLeft / timeStraightsLeft;

  float stepsThis = (float)labs(segs[idx].s1);
  if (stepsThis < 1.0f) stepsThis = 1.0f;

  float Ti = stepsThis / vavg;
  if (Ti < 0.02f) Ti = 0.02f;

  float vmax = 0.0f;
  bool ok = solve_vmax((long)stepsThis, Ti, STRAIGHT_ACCEL_STEPSPS2, vmax);
  if (!ok) vmax = stepsThis / Ti;

  if (vmax < VMIN_STEPSPS) vmax = VMIN_STEPSPS;
  if (vmax > STRAIGHT_VMAX_MAX) vmax = STRAIGHT_VMAX_MAX;
  return vmax;
}

// compute per-wheel vmax for turns so the turn completes in TURN_TIME_S
static void computeTurnVmax(float &v1, float &v2, long steps1, long steps2) {
  long d1 = (long)labs(steps1);
  long d2 = (long)labs(steps2);
  if (d1 < 1) d1 = 1;
  if (d2 < 1) d2 = 1;

  float tv1 = 0.0f, tv2 = 0.0f;
  bool ok1 = solve_vmax(d1, TURN_TIME_S, TURN_ACCEL_STEPSPS2, tv1);
  bool ok2 = solve_vmax(d2, TURN_TIME_S, TURN_ACCEL_STEPSPS2, tv2);

  if (!ok1) tv1 = (float)d1 / TURN_TIME_S;
  if (!ok2) tv2 = (float)d2 / TURN_TIME_S;

  if (tv1 < VMIN_STEPSPS) tv1 = VMIN_STEPSPS;
  if (tv2 < VMIN_STEPSPS) tv2 = VMIN_STEPSPS;
  if (tv1 > TURN_VMAX_MAX) tv1 = TURN_VMAX_MAX;
  if (tv2 > TURN_VMAX_MAX) tv2 = TURN_VMAX_MAX;

  v1 = tv1;
  v2 = tv2;
}

// ----------------- Button + segment start -----------------
static bool debouncePressed() {
  static int lastRead = HIGH;
  static unsigned long t0 = 0;
  static int stable = HIGH;

  int r = digitalRead(BUTTON_PIN);
  if (r != lastRead) { lastRead = r; t0 = millis(); }
  if ((millis() - t0) > 25) stable = r;

  static int lastStable = HIGH;
  bool pressed = (lastStable == HIGH && stable == LOW);
  lastStable = stable;
  return pressed;
}

static void startSegment(int idx) {
  if (segs[idx].type == SEG_FWD) {
    float vmax = computeStraightVmaxNow(idx);
    m1.setMaxSpeed(vmax);
    m2.setMaxSpeed(vmax);
    m1.setAcceleration(STRAIGHT_ACCEL_STEPSPS2);
    m2.setAcceleration(STRAIGHT_ACCEL_STEPSPS2);
  } else {
    float v1 = 0.0f, v2 = 0.0f;
    computeTurnVmax(v1, v2, segs[idx].s1, segs[idx].s2);
    m1.setMaxSpeed(v1);
    m2.setMaxSpeed(v2);
    m1.setAcceleration(TURN_ACCEL_STEPSPS2);
    m2.setAcceleration(TURN_ACCEL_STEPSPS2);
  }

  m1.move(segs[idx].s1);
  m2.move(segs[idx].s2);
}

void setup() {
  pinMode(EN1, OUTPUT);
  pinMode(EN2, OUTPUT);
  digitalWrite(EN1, ENABLE_ACTIVE_LOW ? LOW : HIGH);
  digitalWrite(EN2, ENABLE_ACTIVE_LOW ? LOW : HIGH);

  pinMode(BUTTON_PIN, INPUT_PULLUP);

  m1.setCurrentPosition(0);
  m2.setCurrentPosition(0);
}

void loop() {
  if (!started && debouncePressed()) {
    if (parseScript()) {
      started = true;
      segIdx = 0;
      runStartMs = millis();
      startSegment(segIdx);
    }
  }

  m1.run();
  m2.run();

  if (started && m1.distanceToGo() == 0 && m2.distanceToGo() == 0) {
    segIdx++;
    if (segIdx >= segCount) {
      started = false;
    } else {
      startSegment(segIdx);
    }
  }
}