#include <Encoder.h>
#define ENCODER_OPTIMIZE_INTERRUPTS

// Uncomment the following line to enable debug output
#define DEBUG

// Define encoder pins using const int
const uint8_t ENCODER_PIN_A = 2;
const uint8_t ENCODER_PIN_B = 3;

Encoder myEnc(ENCODER_PIN_A, ENCODER_PIN_B);

const uint8_t RELAY_PIN = 8;  // Pin connected to the relay

// Encoder parameters
const int32_t PULSES_PER_REVOLUTION = 20; // Adjust based on your encoder
const int32_t COUNTS_PER_PULSE = 4;        // For quadrature encoders
const int32_t COUNTS_PER_REVOLUTION = PULSES_PER_REVOLUTION * COUNTS_PER_PULSE;

// Constants
const unsigned long RELAY_ACTIVE_TIME = 15000UL;          // Relay activation time in milliseconds
const unsigned long INACTIVITY_TIMEOUT = 5000UL;          // Inactivity timeout before movement starts
const unsigned long MOVEMENT_INACTIVITY_TIMEOUT = 4000UL; // Inactivity timeout during movement
const uint8_t TOLERANCE_PERCENT = 20;                     // Tolerance percentage for movements
const unsigned long LOOP_DELAY = 5UL;                     // Loop delay in milliseconds
const unsigned long MOVEMENT_PAUSE_TIME = 650UL;          // Time to consider wheel has stopped moving
const int32_t MOVEMENT_THRESHOLD = 1;                     // Threshold for significant movement

// Smoothing parameters
const uint8_t SMOOTHING_FACTOR_NUMERATOR = 1;  // Numerator of smoothing factor (alpha)
const uint8_t SMOOTHING_FACTOR_DENOMINATOR = 4; // Denominator of smoothing factor

const uint8_t NUM_MOVEMENTS = 3; // Number of movements in the sequence

// Movement degrees (expected degrees for each movement)
const int16_t MOVE_DEGREES[NUM_MOVEMENTS] = {
    90,    // First move: 90 degrees clockwise
    -180,  // Second move: 180 degrees counter-clockwise
    360    // Third move: 360 degrees clockwise
};

// Movement parameters
struct Movement {
    int32_t expectedCounts;  // Expected encoder counts for the movement
    int32_t tolerance;       // Tolerance in counts
    int32_t lowerBound;      // Lower bound of acceptable counts
    int32_t upperBound;      // Upper bound of acceptable counts
};

Movement movements[NUM_MOVEMENTS]; // Array to hold movement parameters

// State machine states
enum State {
    STATE_WAITING,
    STATE_MOVING,
    STATE_RELAY_ON,
};

State state = STATE_WAITING;

int32_t lastPosition = 0;                // Position at the last state change
unsigned long relayActivatedTime = 0UL;  // Time when relay was activated
unsigned long lastActionTime = 0UL;      // Time when the last action occurred
uint8_t movementIndex = 0;               // Index of the current movement in the sequence

int32_t previousPosition = 0;            // Previous encoder position
unsigned long lastEncoderChangeTime = 0UL; // Time when encoder last changed

bool movementStarted = false;            // Flag to indicate if movement has started

// Smoothing variables
int32_t smoothedPosition = 0;
bool isFirstSample = true; // Flag to handle the first sample differently

void setup() {
#ifdef DEBUG
    Serial.begin(115200);
#endif
    pinMode(RELAY_PIN, OUTPUT);
    digitalWrite(RELAY_PIN, LOW);   // Ensure relay is off at startup

    initializeMovements();          // Initialize the movements array

    int32_t initialPosition = myEnc.read();
    smoothedPosition = initialPosition;
    previousPosition = initialPosition;
    lastPosition = initialPosition;
    isFirstSample = false;

    lastEncoderChangeTime = millis(); // Initialize lastEncoderChangeTime
    lastActionTime = millis();        // Initialize lastActionTime
}

void loop() {
    int32_t currentRawPosition = myEnc.read();

    // Apply exponential moving average
    smoothedPosition = (SMOOTHING_FACTOR_NUMERATOR * currentRawPosition +
                       (SMOOTHING_FACTOR_DENOMINATOR - SMOOTHING_FACTOR_NUMERATOR) * smoothedPosition) /
                       SMOOTHING_FACTOR_DENOMINATOR;

    int32_t currentPosition = smoothedPosition;

    unsigned long currentTime = millis();

    // Update lastEncoderChangeTime if encoder position has changed significantly
    if (abs(currentPosition - previousPosition) >= MOVEMENT_THRESHOLD) {
        updateLastEncoderChangeTime();
        updateLastActionTime();
        previousPosition = currentPosition;
    }

    switch (state) {
        case STATE_WAITING:
            // Check for significant movement
            if (abs(currentPosition - lastPosition) >= MOVEMENT_THRESHOLD) {
#ifdef DEBUG
                Serial.println(F("Movement detected, starting sequence."));
                printMovementDetails(movementIndex);
#endif
                state = STATE_MOVING;
                lastPosition = currentPosition;
                updateLastActionTime();
                movementIndex = 0; // Start from the first movement
                updateLastEncoderChangeTime();
                movementStarted = false; // Reset movementStarted flag
            }
            break;

        case STATE_MOVING:
            // Check for inactivity during movement
            if (timeoutOccurred(currentTime, lastEncoderChangeTime, MOVEMENT_INACTIVITY_TIMEOUT)) {
                if (movementStarted && timeoutOccurred(currentTime, lastActionTime, MOVEMENT_PAUSE_TIME)) {
#ifdef DEBUG
                    Serial.println(F("No movement detected during movement, resetting to WAITING state."));
#endif
                    resetToWaitingState();
                    break;
                }
            }

            if (!movementStarted) {
                // Waiting for wheel to start moving for current movement
                if (abs(currentPosition - lastPosition) >= MOVEMENT_THRESHOLD) {
                    movementStarted = true;
                    updateLastEncoderChangeTime();
                    lastPosition = currentPosition;
#ifdef DEBUG
                    Serial.print(F("Started movement "));
                    Serial.println(movementIndex + 1);
#endif
                } else if (timeoutOccurred(currentTime, lastActionTime, INACTIVITY_TIMEOUT)) {
#ifdef DEBUG
                    Serial.println(F("No movement detected for inactivity timeout, resetting to WAITING state."));
#endif
                    resetToWaitingState();
                    break;
                }
            } else {
                // Movement has started
                // Check if the wheel has stopped moving
                if (timeoutOccurred(currentTime, lastEncoderChangeTime, MOVEMENT_PAUSE_TIME)) {
                    // Wheel has stopped moving
                    int32_t movementDifference = currentPosition - lastPosition;
                    int32_t lowerBound = movements[movementIndex].lowerBound;
                    int32_t upperBound = movements[movementIndex].upperBound;

#ifdef DEBUG
                    int16_t movementDegrees = countsToDegrees(movementDifference);
                    Serial.print(F("Movement difference for movement "));
                    Serial.print(movementIndex + 1);
                    Serial.print(F(": "));
                    Serial.print(movementDegrees);
                    Serial.println(F(" degrees"));
#endif

                    // Now check if the movement is within the acceptable range
                    if (movementDifference >= lowerBound && movementDifference <= upperBound) {
                        // Movement is within acceptable range
#ifdef DEBUG
                        Serial.print(F("Correct movement "));
                        Serial.print(movementIndex + 1);
                        Serial.println(F(" detected."));
                        if (movementIndex + 1 < NUM_MOVEMENTS) {
                            Serial.print(F("Moving to movement "));
                            Serial.println(movementIndex + 2);
                            printMovementDetails(movementIndex + 1);
                        }
#endif
                        // Move to the next movement
                        movementIndex++;
                        lastPosition = currentPosition;
                        updateLastActionTime();
                        updateLastEncoderChangeTime();
                        movementStarted = false;

                        if (movementIndex >= NUM_MOVEMENTS) {
                            state = STATE_RELAY_ON;
                            relayActivatedTime = currentTime;
                            digitalWrite(RELAY_PIN, HIGH);
#ifdef DEBUG
                            Serial.println(F("Relay activated."));
#endif
                        }
                    } else {
                        // Movement not within acceptable range (including overshoot)
                        if ((movements[movementIndex].expectedCounts >= 0 && movementDifference > upperBound) ||
                            (movements[movementIndex].expectedCounts < 0 && movementDifference < lowerBound)) {
#ifdef DEBUG
                            Serial.print(F("Overshoot on movement "));
                            Serial.print(movementIndex + 1);
                            Serial.println(F(", resetting to WAITING state."));
#endif
                        } else {
#ifdef DEBUG
                            Serial.print(F("Incorrect movement "));
                            Serial.print(movementIndex + 1);
                            Serial.println(F(", resetting to WAITING state."));
#endif
                        }
                        resetToWaitingState();
                    }
                }
            }
            break;

        case STATE_RELAY_ON:
            if (timeoutOccurred(currentTime, relayActivatedTime, RELAY_ACTIVE_TIME)) {
                digitalWrite(RELAY_PIN, LOW);
#ifdef DEBUG
                Serial.println(F("Relay deactivated."));
#endif
                resetToWaitingState();
            }
            break;
    }

    // Optional: delay can be adjusted or removed based on performance
    delay(LOOP_DELAY);
}

void resetToWaitingState() {
    state = STATE_WAITING;
    myEnc.write(0);        // Reset encoder count to zero
    lastPosition = 0;
    movementIndex = 0;
    updateLastEncoderChangeTime();
    updateLastActionTime();
    movementStarted = false;
    isFirstSample = true;  // Reset the smoothing flag
    smoothedPosition = 0;  // Reset smoothed position
    previousPosition = 0;  // Reset previous position
}


void initializeMovements() {
    for (uint8_t i = 0; i < NUM_MOVEMENTS; i++) {
        // Calculate expectedCounts using integer arithmetic
        movements[i].expectedCounts = (int32_t)(((int32_t)MOVE_DEGREES[i] * COUNTS_PER_REVOLUTION + 180) / 360);

        movements[i].tolerance = abs(movements[i].expectedCounts) * TOLERANCE_PERCENT / 100;
        if (movements[i].tolerance == 0) {
            movements[i].tolerance = 1;
        }
        movements[i].lowerBound = movements[i].expectedCounts - movements[i].tolerance;
        movements[i].upperBound = movements[i].expectedCounts + movements[i].tolerance;
    }
}

void printMovementDetails(uint8_t index) {
#ifdef DEBUG
    Serial.print(F("Expected degrees for movement "));
    Serial.print(index + 1);
    Serial.print(F(": "));
    Serial.print(MOVE_DEGREES[index]);
    Serial.print(F(" degrees, Tolerance: ±"));
    Serial.print(TOLERANCE_PERCENT);
    Serial.println(F("%"));
#endif
}

void updateLastActionTime() {
    lastActionTime = millis();
}

void updateLastEncoderChangeTime() {
    lastEncoderChangeTime = millis();
}

inline bool timeoutOccurred(unsigned long current, unsigned long previous, unsigned long interval) {
    return (current - previous) >= interval;
}

// Function to convert counts to degrees using integer arithmetic
int16_t countsToDegrees(int32_t counts) {
    // Multiply counts by 360, add half of COUNTS_PER_REVOLUTION for rounding, then divide by COUNTS_PER_REVOLUTION
    int16_t degrees = (int16_t)((counts * 360 + (COUNTS_PER_REVOLUTION / 2)) / COUNTS_PER_REVOLUTION);
    return degrees;
}