/*
  Copyright 2025
              Thibaut Ferrand / Ulysse Cura
                                           */

#include <Adafruit_SSD1306.h>
#include <CodeCell.h>
#include <ESP32Servo.h>
#include <cmath>

#define PIN_TIRETTE 5
#define PIN_BUTTON_COLOR 6
#define PIN_MOTOR1 1
#define PIN_MOTOR2 2
#define PIN_SERVO 7

#define DANCING_ACTION_DELTA_ANGLE 2

#define GAIN_KD 50
#define ANGULAR_SPEED 10 // °/s

using std::abs;

enum class Axes {
  X,
  Y,
  Z
};

class Motor {
 public:
  int init(int pin, int channel)
  {
    m_pin = pin;
    m_channel = channel;

    pinMode(m_pin, OUTPUT);

    ledcAttachChannel(m_pin, 5000, 12, m_channel);
    ledcWriteChannel(m_channel, 0);
    //ledcSetup(m_channel, 5000, 12); // Channel, Frequency, Resolution
    //ledcAttachPin(m_pin, m_channel); // Attach pin to channel

    return 0;
  }

  void setSpeed(int speed)
  {
    if(speed > 4095) speed = 4095;

    ledcWriteChannel(m_channel, speed);
  }

 private:
  int m_pin;
  int m_channel;
};

class IO {
 public:
  IO(CodeCell *code_cell) : m_code_cell(code_cell)
  {}

  int init()
  {
    int nb_errors {0};

    pinMode(PIN_TIRETTE, INPUT_PULLUP);
    pinMode(PIN_BUTTON_COLOR, INPUT_PULLUP);

    m_code_cell->Motion_RotationRead(m_init_x, m_init_y, m_init_z);

    m_initMotors();
    m_initServo();

    nb_errors += m_initScreen();

    return nb_errors;
  }

  int update()
  {
    m_is_tirette_pulled = (digitalRead(PIN_TIRETTE) == LOW);
    m_is_color_blue = (digitalRead(PIN_BUTTON_COLOR) == HIGH);

    if(m_is_motor_control_activated)
    {
      m_updateMotorsControl();
    }

    if(m_is_dancing)
    {
      m_updateDancingAction();
    }

    m_updateScreen();

    return 0;
  }

  bool isTirettePulled()
  {
    return m_is_tirette_pulled;
  }

  bool isSelectedColorBlue()
  {
    return m_is_color_blue;
  }

  void motorControlOn()
  {
    m_is_motor_control_activated = true;
  }

  void motorControlOff()
  {
    m_is_motor_control_activated = false;

    m_motors[0].setSpeed(0);
    m_motors[1].setSpeed(0);
  }

  void setDirWithAngularSpeed(float dir, float angular_speed = ANGULAR_SPEED)
  {
    static unsigned long prev_time {millis()};
    float err_dir = dir - m_dir;
    float angular_displacement = angular_speed * static_cast<float>(millis() / 1000 - prev_time / 1000);

    if(abs(err_dir) < angular_displacement)
    {
      m_dir = dir;
    }
    else
    {
      m_dir += angular_displacement * (err_dir < 0) ? -1 : 1;
    }

    prev_time = millis();
  }

  void setDir(float dir)
  {
    m_dir = dir;
  }

  void setSpeed(float speed)
  {
    m_speed = speed;
  }
  
  float getAngle(Axes axis)
  {
    float x, y, z;

    m_code_cell->Motion_RotationRead(x, y, z);

    switch(axis)
    {
     case Axes::X:
      return x - m_init_x;

     case Axes::Y:
      return y - m_init_y;

     default:
      return z - m_init_z;
    }
  }

  void startDancingAction()
  {
    m_is_dancing = true;
  }

  Adafruit_SSD1306 *getScreen()
  {
    return &m_screen;
  }

  Motor *getMotor(int motor)
  {
    return &(m_motors[motor]);
  }

 private:
  // Init Motors
  void m_initMotors()
  {
    m_motors[0].init(PIN_MOTOR1, 0);
    m_motors[1].init(PIN_MOTOR2, 1);
  }

  // Init Servo
  void m_initServo()
  {
    m_servo.setPeriodHertz(50);
    m_servo.attach(PIN_SERVO, 500, 2400);

    m_servo.write(87);
  }

  // Init Screen
  int m_initScreen()
  {
    if(!m_screen.begin(SSD1306_SWITCHCAPVCC, 0x3C))
    {
      Serial.println("SSD1306 allocation failed.");
      return -1;
    }

    m_screen.fillScreen(SSD1306_BLACK);

    return 0;
  }

  void m_updateDancingAction()
  {
    static unsigned long prev_time {millis()};
    static unsigned long delta_time {0};
    static int actual_angle {87};
    static int8_t actual_dir {-1};

    delta_time += millis() - prev_time;

    if(delta_time >= 10)
    {
      delta_time = 0;

      actual_angle += actual_dir * DANCING_ACTION_DELTA_ANGLE;
    }

    if(actual_angle <= 43)
    {
      actual_dir = 1;
    }
    else if(actual_angle >= 130)
    {
      actual_dir = -1;
    }

    m_servo.write(actual_angle);

    prev_time = millis();
  }

  void m_updateMotorsControl()
  {
    float actual_angle = getAngle(Axes::Z);
    float error = m_dir - actual_angle;

    float correction = error * GAIN_KD;

    int m1_speed = static_cast<int>(m_speed + correction);
    int m2_speed = static_cast<int>(m_speed - correction);

    if(m1_speed < 0)
    {
      m1_speed = 0;
      //m2_speed = static_cast<int>(m_speed - 2 * correction);
    }
    else if(m2_speed < 0)
    {
      m2_speed = 0;
      //m1_speed = static_cast<int>(m_speed + 2 * correction);
    }

    m_motors[0].setSpeed(m1_speed);
    m_motors[1].setSpeed(m2_speed);
  }

  void m_updateScreen()
  {
    m_screen.display();
  }

  bool m_is_motor_control_activated {false};
  bool m_is_tirette_pulled {false};
  bool m_is_color_blue {false};
  bool m_is_dancing {false};

  Motor m_motors[2];
  Servo m_servo;

  Adafruit_SSD1306 m_screen {128, 64, &Wire, -1};

  CodeCell *m_code_cell;

  float m_init_x, m_init_y, m_init_z;
  float m_dir {0};
  float m_speed {0};
};