/*
 * ============================================================
 *  ARMx — EMG Prosthetic Hand Controller Firmware
 *  Teensy 4.1
 * ============================================================
 *
 *  Reads 8 EMG channels + MPU6050 IMU at 1 kHz.
 *  Streams filtered, normalized data to laptop via USB Serial.
 *  Receives gesture commands and drives servo motors.
 *
 *  Serial Protocol (OUT):
 *    EMG:0.42,0.87,0.13,0.65;IMU:0.12,-0.34,9.78,1.2,-0.5,0.8\n
 *
 *  Serial Protocol (IN):
 *    OPEN / FIST / PINCH / THUMB / INDEX / REST
 *    CAL_START / CAL_STOP / GET_CAL
 *    SET_RATE:500
 *    SERVO:0,90
 *    PING / STATUS / STREAM_ON / STREAM_OFF
 *
 *  Hardware:
 *    EMG Channels:  A0–A7 (pins 14–21)
 *    MPU6050:       SDA=25, SCL=24  (Wire2 — avoids A4/A5 conflict)
 *    Servos:        Pins 2,3,4,5,6 (thumb, index, middle, ring, pinky)
 *
 *  Version: 2.0.0
 *  Date:    2026-05-01
 * ============================================================
 */

#include <Wire.h>
#include <Servo.h>

// ─── CONFIGURATION ──────────────────────────────────────────
#define FIRMWARE_VERSION   "2.0.0"
#define NUM_EMG_CHANNELS   8        // Change to 8 for full array
#define SAMPLE_RATE_HZ     1000     // ADC sampling rate
#define SERIAL_BAUD        115200
#define ADC_RESOLUTION     12       // Teensy 4.1 supports 10/12-bit
#define ADC_MAX_VALUE      4095.0f  // 2^12 - 1

// Moving average filter window
#define MA_WINDOW_SIZE     8

// IIR Bandpass filter
#define USE_IIR_FILTER     false

// Servo configuration
#define NUM_SERVOS         5
const int SERVO_PINS[NUM_SERVOS] = {2, 3, 4, 5, 6};

// Servo angle limits
#define SERVO_MIN_ANGLE    0
#define SERVO_MAX_ANGLE    180
#define SERVO_SPEED_LIMIT  5    // Max degrees per update cycle
#define SERVO_UPDATE_MS    20   // Servo update interval (50 Hz)

// EMG channel pins
const int EMG_PINS[8] = {A0, A1, A2, A3, A4, A5, A6, A7};

// MPU6050 I2C address
#define MPU6050_ADDR       0x68

// I2C bus selection:
//   Wire  -> SDA=18, SCL=19  (conflicts with A4/A5 = CH5/CH6)
//   Wire2 -> SDA=25, SCL=24  (safe for all 8 EMG channels)
#define IMU_WIRE           Wire2

// Heartbeat interval
#define HEARTBEAT_MS       5000

// ─── DATA STRUCTURES ───────────────────────────────────────

struct MAFilter {
    float buffer[MA_WINDOW_SIZE];
    int   index;
    float sum;

    void reset() {
        for (int i = 0; i < MA_WINDOW_SIZE; i++) buffer[i] = 0.0f;
        index = 0;
        sum = 0.0f;
    }

    float update(float value) {
        sum -= buffer[index];
        buffer[index] = value;
        sum += value;
        index = (index + 1) % MA_WINDOW_SIZE;
        return sum / MA_WINDOW_SIZE;
    }
};

struct BiquadFilter {
    float b0, b1, b2, a1, a2;
    float x1, x2, y1, y2;

    void reset() { x1 = x2 = y1 = y2 = 0.0f; }

    float update(float x0) {
        float y0 = b0 * x0 + b1 * x1 + b2 * x2 - a1 * y1 - a2 * y2;
        x2 = x1; x1 = x0;
        y2 = y1; y1 = y0;
        return y0;
    }
};

struct CalibrationData {
    float minVal;
    float maxVal;
    float baseline;
    bool  active;

    void reset() {
        minVal = 1.0f;
        maxVal = 0.0f;
        baseline = 0.0f;
        active = false;
    }

    float normalize(float raw) {
        if (!active || (maxVal - minVal) < 0.001f) return raw;
        float norm = (raw - minVal) / (maxVal - minVal);
        return constrain(norm, 0.0f, 1.0f);
    }
};

struct ServoState {
    Servo   servo;
    int     currentAngle;
    int     targetAngle;
    int     pin;

    void init(int p) {
        pin = p;
        servo.attach(pin);
        currentAngle = 90;
        targetAngle = 90;
        servo.write(currentAngle);
    }

    void setTarget(int angle) {
        targetAngle = constrain(angle, SERVO_MIN_ANGLE, SERVO_MAX_ANGLE);
    }

    void update() {
        if (currentAngle == targetAngle) return;
        int diff = targetAngle - currentAngle;
        int step = constrain(diff, -SERVO_SPEED_LIMIT, SERVO_SPEED_LIMIT);
        currentAngle += step;
        servo.write(currentAngle);
    }
};

struct IMUData {
    float ax, ay, az;  // Accelerometer (g)
    float gx, gy, gz;  // Gyroscope (deg/s)
};

// ─── GLOBAL STATE ──────────────────────────────────────────

MAFilter         maFilters[8];
BiquadFilter     bpFilters[8];
CalibrationData  calibration[8];
ServoState       servos[NUM_SERVOS];
IMUData          imu;

float            emgRaw[8];
float            emgFiltered[8];
float            emgNormalized[8];

volatile bool    sampleReady = false;
IntervalTimer    sampleTimer;

bool             calibrating = false;
bool             streaming = true;
unsigned long    lastServoUpdate = 0;
unsigned long    lastSerialSend = 0;
unsigned long    lastHeartbeat = 0;
int              serialSendRateHz = 100;

// Gesture-to-servo angle presets: [thumb, index, middle, ring, pinky]
const int GESTURE_OPEN[5]  = {180, 180, 180, 180, 180};
const int GESTURE_FIST[5]  = {  0,   0,   0,   0,   0};
const int GESTURE_PINCH[5] = {  0,   0, 180, 180, 180};
const int GESTURE_THUMB[5] = {  0, 180, 180, 180, 180};
const int GESTURE_INDEX[5] = {180,   0, 180, 180, 180};
const int GESTURE_REST[5]  = { 90,  90,  90,  90,  90};

// ─── INITIALIZATION ────────────────────────────────────────

void initMPU6050() {
    IMU_WIRE.begin();
    IMU_WIRE.setClock(400000);

    IMU_WIRE.beginTransmission(MPU6050_ADDR);
    IMU_WIRE.write(0x6B);
    IMU_WIRE.write(0x00);
    IMU_WIRE.endTransmission(true);

    // Accelerometer: ±4g
    IMU_WIRE.beginTransmission(MPU6050_ADDR);
    IMU_WIRE.write(0x1C);
    IMU_WIRE.write(0x08);
    IMU_WIRE.endTransmission(true);

    // Gyroscope: ±500 deg/s
    IMU_WIRE.beginTransmission(MPU6050_ADDR);
    IMU_WIRE.write(0x1B);
    IMU_WIRE.write(0x08);
    IMU_WIRE.endTransmission(true);

    // DLPF: 44 Hz bandwidth
    IMU_WIRE.beginTransmission(MPU6050_ADDR);
    IMU_WIRE.write(0x1A);
    IMU_WIRE.write(0x03);
    IMU_WIRE.endTransmission(true);
}

void readMPU6050() {
    IMU_WIRE.beginTransmission(MPU6050_ADDR);
    IMU_WIRE.write(0x3B);
    IMU_WIRE.endTransmission(false);
    IMU_WIRE.requestFrom(MPU6050_ADDR, 14, true);

    int16_t rawAx = (IMU_WIRE.read() << 8) | IMU_WIRE.read();
    int16_t rawAy = (IMU_WIRE.read() << 8) | IMU_WIRE.read();
    int16_t rawAz = (IMU_WIRE.read() << 8) | IMU_WIRE.read();
    IMU_WIRE.read(); IMU_WIRE.read();  // Skip temperature
    int16_t rawGx = (IMU_WIRE.read() << 8) | IMU_WIRE.read();
    int16_t rawGy = (IMU_WIRE.read() << 8) | IMU_WIRE.read();
    int16_t rawGz = (IMU_WIRE.read() << 8) | IMU_WIRE.read();

    imu.ax = rawAx / 8192.0f;
    imu.ay = rawAy / 8192.0f;
    imu.az = rawAz / 8192.0f;
    imu.gx = rawGx / 65.5f;
    imu.gy = rawGy / 65.5f;
    imu.gz = rawGz / 65.5f;
}

void initBandpassFilters() {
    for (int i = 0; i < 8; i++) {
        bpFilters[i].b0 =  0.5625f;
        bpFilters[i].b1 =  0.0f;
        bpFilters[i].b2 = -0.5625f;
        bpFilters[i].a1 = -0.1250f;
        bpFilters[i].a2 = -0.1250f;
        bpFilters[i].reset();
    }
}

// ─── TIMER ISR ─────────────────────────────────────────────

void sampleISR() {
    sampleReady = true;
}

// ─── SETUP ─────────────────────────────────────────────────

void setup() {
    Serial.begin(SERIAL_BAUD);
    analogReadResolution(ADC_RESOLUTION);

    for (int i = 0; i < 8; i++) {
        maFilters[i].reset();
        calibration[i].reset();
    }
    initBandpassFilters();

    for (int i = 0; i < NUM_SERVOS; i++) {
        servos[i].init(SERVO_PINS[i]);
    }

    initMPU6050();

    sampleTimer.begin(sampleISR, 1000000 / SAMPLE_RATE_HZ);

    delay(500);
    Serial.print("ARMX:READY:V");
    Serial.println(FIRMWARE_VERSION);
    Serial.print("CHANNELS:");
    Serial.println(NUM_EMG_CHANNELS);
    Serial.print("SAMPLE_RATE:");
    Serial.println(SAMPLE_RATE_HZ);
}

// ─── MAIN LOOP ─────────────────────────────────────────────

void loop() {
    // ── 1. Sample EMG channels when timer fires ──
    if (sampleReady) {
        sampleReady = false;

        for (int ch = 0; ch < NUM_EMG_CHANNELS; ch++) {
            int rawADC = analogRead(EMG_PINS[ch]);
            float rawNorm = rawADC / ADC_MAX_VALUE;
            emgRaw[ch] = rawNorm;

            float filtered = maFilters[ch].update(rawNorm);

            if (USE_IIR_FILTER) {
                filtered = bpFilters[ch].update(filtered);
                filtered = abs(filtered);
            }

            emgFiltered[ch] = filtered;

            if (calibrating) {
                if (filtered < calibration[ch].minVal) calibration[ch].minVal = filtered;
                if (filtered > calibration[ch].maxVal) calibration[ch].maxVal = filtered;
            }

            emgNormalized[ch] = calibration[ch].normalize(filtered);
        }

        // Read IMU at 100 Hz (every 10th sample)
        static int imuCounter = 0;
        if (++imuCounter >= 10) {
            imuCounter = 0;
            readMPU6050();
        }
    }

    // ── 2. Send data to laptop at decimated rate ──
    unsigned long now = millis();
    unsigned long sendInterval = 1000 / serialSendRateHz;

    if (streaming && (now - lastSerialSend >= sendInterval)) {
        lastSerialSend = now;
        sendSerialData();
    }

    // ── 3. Heartbeat ──
    if (now - lastHeartbeat >= HEARTBEAT_MS) {
        lastHeartbeat = now;
        Serial.println("HB");
    }

    // ── 4. Receive commands from laptop ──
    if (Serial.available()) {
        String cmd = Serial.readStringUntil('\n');
        cmd.trim();
        processCommand(cmd);
    }

    // ── 5. Update servos with smoothing ──
    if (now - lastServoUpdate >= SERVO_UPDATE_MS) {
        lastServoUpdate = now;
        for (int i = 0; i < NUM_SERVOS; i++) {
            servos[i].update();
        }
    }
}

// ─── SERIAL OUTPUT ─────────────────────────────────────────

void sendSerialData() {
    Serial.print("EMG:");
    for (int i = 0; i < NUM_EMG_CHANNELS; i++) {
        Serial.print(emgNormalized[i], 4);
        if (i < NUM_EMG_CHANNELS - 1) Serial.print(",");
    }

    Serial.print(";IMU:");
    Serial.print(imu.ax, 3); Serial.print(",");
    Serial.print(imu.ay, 3); Serial.print(",");
    Serial.print(imu.az, 3); Serial.print(",");
    Serial.print(imu.gx, 2); Serial.print(",");
    Serial.print(imu.gy, 2); Serial.print(",");
    Serial.print(imu.gz, 2);

    Serial.println();
}

// ─── COMMAND PROCESSING ────────────────────────────────────

void processCommand(String cmd) {
    // ── Heartbeat ──
    if (cmd == "PING") {
        Serial.println("PONG");
        return;
    }

    // ── Gesture commands ──
    if (cmd == "OPEN") {
        applyGesture(GESTURE_OPEN);
        Serial.println("ACK:OPEN");
    }
    else if (cmd == "FIST") {
        applyGesture(GESTURE_FIST);
        Serial.println("ACK:FIST");
    }
    else if (cmd == "PINCH") {
        applyGesture(GESTURE_PINCH);
        Serial.println("ACK:PINCH");
    }
    else if (cmd == "THUMB") {
        applyGesture(GESTURE_THUMB);
        Serial.println("ACK:THUMB");
    }
    else if (cmd == "INDEX") {
        applyGesture(GESTURE_INDEX);
        Serial.println("ACK:INDEX");
    }
    else if (cmd == "REST") {
        applyGesture(GESTURE_REST);
        Serial.println("ACK:REST");
    }
    // ── Direct servo control ──
    else if (cmd.startsWith("SERVO:")) {
        int commaIdx = cmd.indexOf(',', 6);
        if (commaIdx > 6) {
            int servoIdx = cmd.substring(6, commaIdx).toInt();
            int angle = cmd.substring(commaIdx + 1).toInt();
            if (servoIdx >= 0 && servoIdx < NUM_SERVOS) {
                servos[servoIdx].setTarget(angle);
                Serial.print("ACK:SERVO:");
                Serial.print(servoIdx);
                Serial.print(",");
                Serial.println(angle);
            }
        }
    }
    // ── Calibration ──
    else if (cmd == "CAL_START") {
        calibrating = true;
        for (int i = 0; i < NUM_EMG_CHANNELS; i++) {
            calibration[i].reset();
        }
        Serial.println("ACK:CAL_START");
    }
    else if (cmd == "CAL_STOP") {
        calibrating = false;
        for (int i = 0; i < NUM_EMG_CHANNELS; i++) {
            calibration[i].active = true;
            calibration[i].baseline = (calibration[i].minVal + calibration[i].maxVal) / 2.0f;
        }
        Serial.print("CAL_RESULT:");
        for (int i = 0; i < NUM_EMG_CHANNELS; i++) {
            Serial.print(calibration[i].minVal, 4);
            Serial.print(",");
            Serial.print(calibration[i].maxVal, 4);
            if (i < NUM_EMG_CHANNELS - 1) Serial.print(";");
        }
        Serial.println();
        Serial.println("ACK:CAL_STOP");
    }
    else if (cmd == "GET_CAL") {
        Serial.print("CAL_STATE:");
        Serial.print(calibrating ? "active" : "inactive");
        Serial.print(";DATA:");
        for (int i = 0; i < NUM_EMG_CHANNELS; i++) {
            Serial.print(calibration[i].active ? "1" : "0");
            Serial.print(",");
            Serial.print(calibration[i].minVal, 4);
            Serial.print(",");
            Serial.print(calibration[i].maxVal, 4);
            if (i < NUM_EMG_CHANNELS - 1) Serial.print(";");
        }
        Serial.println();
    }
    // ── Configuration ──
    else if (cmd.startsWith("SET_RATE:")) {
        int rate = cmd.substring(9).toInt();
        if (rate >= 10 && rate <= 1000) {
            serialSendRateHz = rate;
            Serial.print("ACK:SET_RATE:");
            Serial.println(rate);
        }
    }
    else if (cmd == "STREAM_ON") {
        streaming = true;
        Serial.println("ACK:STREAM_ON");
    }
    else if (cmd == "STREAM_OFF") {
        streaming = false;
        Serial.println("ACK:STREAM_OFF");
    }
    else if (cmd == "STATUS") {
        Serial.print("STATUS:fw=");
        Serial.print(FIRMWARE_VERSION);
        Serial.print(",channels=");
        Serial.print(NUM_EMG_CHANNELS);
        Serial.print(",rate=");
        Serial.print(SAMPLE_RATE_HZ);
        Serial.print(",send_rate=");
        Serial.print(serialSendRateHz);
        Serial.print(",streaming=");
        Serial.print(streaming ? "on" : "off");
        Serial.print(",calibrated=");
        Serial.print(calibration[0].active ? "yes" : "no");
        Serial.print(",calibrating=");
        Serial.println(calibrating ? "yes" : "no");
    }
    else {
        Serial.print("ERR:UNKNOWN_CMD:");
        Serial.println(cmd);
    }
}

void applyGesture(const int angles[5]) {
    for (int i = 0; i < NUM_SERVOS; i++) {
        servos[i].setTarget(angles[i]);
    }
}