/**************************************************************************
 This is an example for our Monochrome OLEDs based on SSD1306 drivers

 Pick one up today in the adafruit shop!
 ------> http://www.adafruit.com/category/63_98

 This example is for a 128x64 pixel display using I2C to communicate
 3 pins are required to interface (two I2C and one reset).

 Adafruit invests time and resources providing this open
 source code, please support Adafruit and open-source
 hardware by purchasing products from Adafruit!

 Written by Limor Fried/Ladyada for Adafruit Industries,
 with contributions from the open source community.
 BSD license, check license.txt for more information
 All text above, and the splash screen below must be
 included in any redistribution.
 **************************************************************************/

#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <SPI.h>
#include <Wire.h>
#include <EEPROM.h>

#include "RP2040_PWM.h"

#define SCREEN_WIDTH 128 // OLED display width, in pixels
#define SCREEN_HEIGHT 64 // OLED display height, in pixels

#define OLED_RESET -1
#define SCREEN_ADDRESS 0x3C
#define SCL 21
#define SDA 20

#define STAT_LED_A 22
#define STAT_LED_B 23

#define STAT_LED_C 17

#define ROT_A 8
#define ROT_B 9

#define WAVE_X1 5
#define WAVE_Y1 25
#define WAVE_X2 123
#define WAVE_Y2 48

#define IN1 26
#define IN2 22
#define ENA 27

#define MOT_ENABLE 23
#define BTN 7

#define EEPROM_WRITE_OFFSET 0x0

bool stat_led_a = 0;
bool stat_led_b = 0;
bool stat_led_c = 0;

float pwm_freq = 20000.0;
int8_t pwr = 0;
uint8_t pwr_norm = 0;
float dc_norm = 0;
bool fwd = true;

struct TuningParam {
  float duty_cycle;
  float freq;

  bool operator==(TuningParam that) {
    return this->duty_cycle == that.duty_cycle && this->freq == that.freq;
  }

  bool operator!=(TuningParam that) {
    return !(*this == that);
  }
};

enum DisplayState { RUN, MAP_WAIT, MAP_WRITE };
DisplayState state = RUN;

// display
uint32_t refresh_timer_ms = 0;
const uint32_t refresh_timeout_ms = 33;

// rotary encoder polling
const uint8_t rot_a_int = digitalPinToInterrupt(ROT_A);
const uint8_t rot_b_int = digitalPinToInterrupt(ROT_B);
int8_t renc_change = 1;
uint32_t poll_timer_ms = 0;
const uint32_t poll_timeout_ms = 2;
bool poll_wait = false;
bool poll_wait_start = false;

// tuning map
TuningParam tuning_map[100];
bool tuning_map_set = false;

// serial read timeout
uint32_t serial_timer_ms = 0;
const uint32_t serial_timeout_ms = 2000;

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
RP2040_PWM* PWM_Instance;

void setup() {
  Wire.setSCL(SCL);
  Wire.setSDA(SDA);

  pinMode(STAT_LED_A, OUTPUT);
  pinMode(STAT_LED_B, OUTPUT);
  pinMode(STAT_LED_C, OUTPUT);
  
  pinMode(ROT_A, INPUT_PULLUP);
  pinMode(ROT_B, INPUT_PULLUP);
  
  pinMode(MOT_ENABLE, INPUT_PULLUP);
  pinMode(BTN, INPUT_PULLUP);
  
  pinMode(IN1, OUTPUT);
  pinMode(IN2, OUTPUT);

  delay(500);

  if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    digitalWrite(STAT_LED_C, HIGH);
  }

  display.clearDisplay();
  display.setTextColor(SSD1306_WHITE);
  display.setTextSize(2);
  display.write("initializing...");
  display.display();

  delay(500);
  Serial.begin();
  delay(500);
  
  if (!digitalRead(BTN)) {
    if (!Serial) {
      display.setTextSize(1);
      display.setCursor(5, 5);
      display.print("Waiting for serial connection...");
      display.display();
    }
    
    bool pix_on = 1;

    while (!Serial) {
      delay(100);
      stat_led_c = !(stat_led_c);
      digitalWrite(STAT_LED_C, stat_led_c);
    }
    
    display.display();
  }
  
  Serial.println("start");

  arm_rotary_encoder();
  
  EEPROM.begin(1000);

  tuning_map_set = load_tuning_map();

  PWM_Instance = new RP2040_PWM(ENA, pwm_freq, dc_norm);
}

void loop() {
  uint32_t now = millis();
  
  if (poll_wait_start && !poll_wait) {
    poll_timer_ms = now;
    poll_wait = true;
    poll_wait_start = false;
  }
  
  uint32_t last_poll_delta = now - poll_timer_ms;

  if (poll_wait && last_poll_delta >= poll_timeout_ms) {
    // a sum of 2 means the encoder has settled
    uint8_t pin_sum = digitalRead(ROT_A) + digitalRead(ROT_B);

    if (pin_sum == 2) {
      uint8_t vel_mult = 1;

      if (last_poll_delta < poll_timeout_ms * 2) vel_mult = 2;
      pwr += renc_change * vel_mult;

      // nested because we don't need to check this until ready for next poll
      poll_wait = false;

      arm_rotary_encoder();
    }
  }

  if (Serial.available()) {
    char cmd = Serial.read();
    
    if (state == MAP_WAIT) {
      switch (cmd) {
        case 'R': {
          clear_buffer();

          state = RUN;
          arm_rotary_encoder();

          Serial.println("returning");
          display.clearDisplay();
          break;
        }
        case 'C': {
          tuning_map_set = false;
          Serial.println("disabled tuning map");
          break;
        }
        case 'E': {
          tuning_map_set = true;
          Serial.println("enabled tuning map");
          break;
        }
        case 'S': {
          Serial.println("saving");
          save_tuning_map();
          Serial.println("saved");
          break;
        }
        case 'L': {
          Serial.println("loading");
          if (load_tuning_map()) {
            Serial.println("loaded");
          } else {
            Serial.println("no tuning map present");
          }
          break;
        }
        case 'D': {
          size_t idx = Serial.parseInt();
          if (idx > 99) break;
          tuning_map[idx].duty_cycle = Serial.parseFloat();
          tuning_map[idx].freq = Serial.parseFloat();

          display.drawPixel(12 + idx, 30, SSD1306_WHITE);
          display.display();
          Serial.println("ok");
          break;
        }
        default: {
          clear_buffer();
          Serial.println("invalid command");
          break;
        }
      }
    } else {
      switch (cmd) {
        case 'M': {
          state = MAP_WAIT;
          
          disarm_rotary_encoder();

          display.clearDisplay();
          display.setTextSize(1);
          display.setCursor(5, 5);
          display.print("receiving tuning map...");
          display.display();
          
          Serial.println("ready");

          break;
        }
        default: {
          clear_buffer();
          Serial.println("invalid command");
          break;
        }
      }
    }
  }
  
  switch (state) {
    case RUN:
      stat_led_c = !(stat_led_c);
      digitalWrite(STAT_LED_C, stat_led_c);

      if (now - refresh_timer_ms >= refresh_timeout_ms) {
        display.setTextSize(2);

        if (!tuning_map_set) {
          PWM_Instance->setPWM(ENA, pwm_freq, 0);
          display.clearDisplay();
          display.setTextSize(1);
          display.setCursor(5, 5);
          display.print("critical: no tuning map!");
          display.display();
          return;
        }
    
        fwd = pwr >= 0;
        pwr_norm = abs(pwr);
        if (pwr_norm == 0) {
          dc_norm = 0;
          pwm_freq = 0;
        } else {
          dc_norm = tuning_map[pwr_norm].duty_cycle;
          pwm_freq = tuning_map[pwr_norm].freq;
        }

        bool pwm_on = !digitalRead(MOT_ENABLE) && pwr_norm != 0; // since it's input_pullup

        display.clearDisplay();

        if (digitalRead(ROT_A)) display.drawPixel(0, 0, SSD1306_WHITE);

        if (digitalRead(ROT_B)) display.drawPixel(2, 0, SSD1306_WHITE);

        if (digitalRead(BTN)) display.drawPixel(4, 0, SSD1306_WHITE);
        
        // draw text
        display.setCursor(2, 2);
        display.print("pwr");
        
        if (dc_norm > 0 && pwm_on) {
          // draw number
          display.setCursor(48, 2);
          
          display.print(pwr_norm);
          display.print('%');
          
          // draw fwd/rev
          display.setCursor(90, 2);

          if (fwd)
            display.print("FWD");
          else
            display.print("REV");
        } else {
          display.setCursor(72, 2);
          display.print("OFF");
        }

        // draw wave
        uint16_t peak_edge_x = WAVE_X1 + dc_norm;

        display.drawLine(WAVE_X1, WAVE_Y1, peak_edge_x, WAVE_Y1, SSD1306_WHITE);
        display.drawLine(peak_edge_x, WAVE_Y1, peak_edge_x, WAVE_Y2, SSD1306_WHITE);
        display.drawLine(peak_edge_x, WAVE_Y2, WAVE_X2, WAVE_Y2, SSD1306_WHITE);

        // draw freq
        display.setTextSize(1);
        display.setCursor(5, 50);
        display.print(pwm_freq);
        display.print("Hz");

        // draw duty cycle
        display.setTextSize(1);
        display.setCursor(75, 50);
        display.print(dc_norm);
        display.print("%");

        display.display();
        
        // drive motor
        if (pwm_on) {
          if (fwd) {
            digitalWrite(IN1, HIGH);
            digitalWrite(IN2, LOW);
          } else {
            digitalWrite(IN1, LOW);
            digitalWrite(IN2, HIGH);
          }
          PWM_Instance->setPWM(ENA, pwm_freq, dc_norm);
        } else {
          digitalWrite(IN1, LOW);
          digitalWrite(IN2, LOW);
          PWM_Instance->setPWM(ENA, pwm_freq, 0);
        }

        refresh_timer_ms = now;
      }
      
      break;
    case MAP_WAIT:
      break;
  }

}
// cw interrupt
void cw_b_isr() {
  disarm_rotary_encoder();
  attachInterrupt(rot_a_int, cw_a_isr, FALLING);
}

void cw_a_isr() {
  detachInterrupt(rot_a_int);
  if (pwr < 100) renc_change = 1;
  poll_wait_start = true;
}

// ccw interrupt
void ccw_a_isr() {
  disarm_rotary_encoder();
  attachInterrupt(rot_b_int, ccw_b_isr, FALLING);
}

void ccw_b_isr() {
  detachInterrupt(rot_b_int);
  if (pwr > -100) renc_change = -1;
  poll_wait_start = true;
}

void clear_buffer() {
  while (Serial.available()) {
    Serial.read();
  }
}

void save_tuning_map() {
  bool tuning_map_present;
  EEPROM.get(EEPROM_WRITE_OFFSET, tuning_map_present);
  
  if (!tuning_map_present) {
    EEPROM.put(EEPROM_WRITE_OFFSET, true);
  }

  size_t offset = EEPROM_WRITE_OFFSET + sizeof(bool);
  for (size_t i = 0; i < 100; i++) {
    TuningParam param = tuning_map[i];
    TuningParam existing_param;
    EEPROM.get(offset, existing_param);
    
    if (existing_param != param) {
      EEPROM.put(offset, tuning_map[i]);
    }

    offset += sizeof(TuningParam);
  }

  EEPROM.commit();
}

int load_tuning_map() {
  bool tuning_map_present;
  EEPROM.get(EEPROM_WRITE_OFFSET, tuning_map_present);
  
  if (tuning_map_present) {
    size_t offset = EEPROM_WRITE_OFFSET + sizeof(bool);
    for (size_t i = 0; i < 100; i++) {
      TuningParam param;
      EEPROM.get(offset, param);
      tuning_map[i] = param;
      offset += sizeof(TuningParam);
    }
    return 1;
  } else {
    return 0;
  }
}

void arm_rotary_encoder() {
  attachInterrupt(rot_b_int, cw_b_isr, FALLING);
  attachInterrupt(rot_a_int, ccw_a_isr, FALLING);
}

void disarm_rotary_encoder() {
  detachInterrupt(rot_a_int);
  detachInterrupt(rot_b_int);
}