#include <Arduino.h>
#include <Wire.h>
#include <SPI.h>
#include <Adafruit_GFX.h>
#include <Adafruit_ILI9341.h>
#include <WiFi.h>
#include <PubSubClient.h>
#include <math.h>


const char* WIFI_SSID     = "test";
const char* WIFI_PASSWORD = "test12345";

const char* MQTT_SERVER = "test.mosquitto.org";
const uint16_t MQTT_PORT = 1883;
const char* MQTT_PUB_TOPIC = "wokwi/esp32/mpu6050/data";
const char* MQTT_SUB_TOPIC = "wokwi/esp32/mpu6050/cmd";

const unsigned long UPDATE_INTERVAL_MS = 1000; // dashboard & publish interval

// ---------------- PINS ----------------
#define SDA_PIN 21
#define SCL_PIN 22

#define TFT_CS   5
#define TFT_DC   16
#define TFT_RST  17

#define LED_PIN  2

// MPU6050 I2C address (AD0 = GND)
#define MPU_ADDR 0x68

// ---------------- DISPLAY ----------------
Adafruit_ILI9341 tft = Adafruit_ILI9341(TFT_CS, TFT_DC, TFT_RST);

// ---------------- MQTT / WiFi ----------------
WiFiClient wifiClient;
PubSubClient mqttClient(wifiClient);

// ---------------- MPU6050 calibration / state ----------------
// MPU registers used
#define MPU_PWR_MGMT_1   0x6B
#define MPU_SMPLRT_DIV   0x19
#define MPU_CONFIG       0x1A
#define MPU_GYRO_CONFIG  0x1B
#define MPU_ACCEL_CONFIG 0x1C
#define MPU_INT_ENABLE   0x38
#define MPU_ACCEL_XOUT_H 0x3B

// scale factors for sensor ranges chosen (±2g / ±250deg/s)
const float ACCEL_SENS = 16384.0f; // LSB/g for ±2g
const float GYRO_SENS  = 131.0f;   // LSB/°/s for ±250°/s

// complementary filter parameter
const float ALPHA = 0.98f; // weight for gyro integration

// orientation state
float roll = 0.0f;   // degrees
float pitch = 0.0f;  // degrees

// time tracking for gyro integration
unsigned long lastMicros = 0;

// last publish time
unsigned long lastPublishMs = 0;

// ---------------- Helper: I2C low-level ----------------
void i2cWriteByte(uint8_t addr, uint8_t reg, uint8_t data) {
  Wire.beginTransmission(addr);
  Wire.write(reg);
  Wire.write(data);
  Wire.endTransmission();
}

uint8_t i2cReadByte(uint8_t addr, uint8_t reg) {
  Wire.beginTransmission(addr);
  Wire.write(reg);
  Wire.endTransmission();
  Wire.requestFrom(addr, (uint8_t)1);
  if (Wire.available()) return Wire.read();
  return 0;
}

void i2cReadBytes(uint8_t addr, uint8_t reg, uint8_t* buf, uint8_t len) {
  Wire.beginTransmission(addr);
  Wire.write(reg);
  Wire.endTransmission();
  Wire.requestFrom(addr, len);
  uint8_t i = 0;
  while (Wire.available() && i < len) {
    buf[i++] = Wire.read();
  }
}

// ---------------- MPU6050 functions ----------------
bool mpuCheckIdentity() {
  // WHO_AM_I register is 0x75, should return 0x68 for MPU6050
  uint8_t who = i2cReadByte(MPU_ADDR, 0x75);
  return (who == 0x68);
}

void mpuInit() {
  // Wake up device - clear sleep bit
  i2cWriteByte(MPU_ADDR, MPU_PWR_MGMT_1, 0x00);
  delay(100);

  // Sample rate divider: sample_rate = gyro_output_rate / (1 + SMPLRT_DIV)
  // gyro_output_rate = 8kHz when DLPF disabled, but we'll set DLPF to 42Hz below
  i2cWriteByte(MPU_ADDR, MPU_SMPLRT_DIV, 0x07); // ~1kHz/(1+7)=125Hz sample rate

  // Configure DLPF (low pass filter) - register CONFIG (0x1A)
  // DLPF_CFG = 3 -> ~44Hz bandwidth for accel and gyro
  i2cWriteByte(MPU_ADDR, MPU_CONFIG, 0x03);

  // Gyro config: ±250 °/s (0 << 3)
  i2cWriteByte(MPU_ADDR, MPU_GYRO_CONFIG, 0x00);

  // Accel config: ±2g (0 << 3)
  i2cWriteByte(MPU_ADDR, MPU_ACCEL_CONFIG, 0x00);

  // Enable interrupts (not strictly required here)
  i2cWriteByte(MPU_ADDR, MPU_INT_ENABLE, 0x01);
}

bool mpuReadRaw(int16_t &ax, int16_t &ay, int16_t &az, int16_t &gx, int16_t &gy, int16_t &gz) {
  uint8_t buf[14];
  i2cReadBytes(MPU_ADDR, MPU_ACCEL_XOUT_H, buf, 14);
  // combine bytes
  ax = (int16_t)((buf[0] << 8) | buf[1]);
  ay = (int16_t)((buf[2] << 8) | buf[3]);
  az = (int16_t)((buf[4] << 8) | buf[5]);
  // temp skipped: buf[6], buf[7]
  gx = (int16_t)((buf[8] << 8) | buf[9]);
  gy = (int16_t)((buf[10] << 8) | buf[11]);
  gz = (int16_t)((buf[12] << 8) | buf[13]);
  return true;
}

// ---------------- Display helpers ----------------
void drawHeader() {
  tft.fillScreen(ILI9341_BLACK);
  tft.setRotation(1);
  tft.setTextSize(2);
  tft.setTextColor(ILI9341_WHITE);
  tft.setCursor(8, 8);
  tft.print("MPU6050 Dashboard");

  tft.setTextSize(1);
  tft.setCursor(8, 34);
  tft.print("ESP32 | ILI9341 | MPU6050");
}

void drawSensorValues(float axg, float ayg, float azg, float gxds, float gyds, float gzds, float r, float p, bool mqttConnected) {
  // Values area
  tft.setTextSize(2);
  tft.setTextColor(ILI9341_YELLOW);
  tft.setCursor(8, 60);
  tft.printf("Roll:  %6.2f", r);

  tft.setCursor(8, 88);
  tft.printf("Pitch: %6.2f", p);

  tft.setTextSize(1);
  tft.setTextColor(ILI9341_CYAN);
  tft.setCursor(8, 120);
  tft.printf("Ax: %6.3fg  Ay: %6.3fg", axg, ayg);

  tft.setCursor(8, 138);
  tft.printf("Az: %6.3fg");

  // print az (we place value manually)
  tft.setCursor(88, 138);
  tft.printf("%6.3fg", azg);

  tft.setCursor(8, 156);
  tft.printf("Gx: %6.2f dps", gxds);

  tft.setCursor(8, 174);
  tft.printf("Gy: %6.2f dps", gyds);

  tft.setCursor(8, 192);
  tft.printf("Gz: %6.2f dps", gzds);

  // MQTT status
  tft.setTextSize(1);
  tft.setCursor(8, 216);
  tft.setTextColor(mqttConnected ? ILI9341_GREEN : ILI9341_RED);
  tft.print("MQTT: ");
  tft.setTextColor(ILI9341_WHITE);
  tft.print(mqttConnected ? "Connected" : "Disconnected");
}

// ---------------- MQTT callback ----------------
void mqttCallback(char* topic, byte* payload, unsigned int length) {
  // Copy to string
  String msg;
  for (unsigned int i = 0; i < length; i++) msg += (char)payload[i];
  msg.trim();

  Serial.print("MQTT recv: ");
  Serial.println(msg);

  if (msg.equalsIgnoreCase("LED ON")) {
    digitalWrite(LED_PIN, HIGH);
  } else if (msg.equalsIgnoreCase("LED OFF")) {
    digitalWrite(LED_PIN, LOW);
  } else if (msg.equalsIgnoreCase("RESET")) {
    // Reset display content to header
    drawHeader();
  }
}

// ---------------- WiFi & MQTT ----------------
void connectWiFi() {
  Serial.print("Connecting WiFi ");
  Serial.print(WIFI_SSID);
  WiFi.begin(WIFI_SSID, WIFI_PASSWORD);
  int tries = 0;
  while (WiFi.status() != WL_CONNECTED && tries < 40) {
    delay(250);
    Serial.print(".");
    tries++;
  }
  Serial.println();
  if (WiFi.status() == WL_CONNECTED) {
    Serial.print("WiFi connected, IP: ");
    Serial.println(WiFi.localIP());
  } else {
    Serial.println("WiFi connection failed");
  }
}

void mqttReconnect() {
  if (mqttClient.connected()) return;
  Serial.print("Connecting MQTT...");
  String clientId = "ESP32_MPU_" + String(random(0xffff), HEX);
  if (mqttClient.connect(clientId.c_str())) {
    Serial.println("connected");
    mqttClient.subscribe(MQTT_SUB_TOPIC);
  } else {
    Serial.print("failed, rc=");
    Serial.println(mqttClient.state());
  }
}

// ---------------- Setup & Loop ----------------
void setup() {
  Serial.begin(115200);
  delay(100);

  // pins
  pinMode(LED_PIN, OUTPUT);
  digitalWrite(LED_PIN, LOW);

  // I2C
  Wire.begin(SDA_PIN, SCL_PIN);
  delay(10);

  // Display init
  tft.begin();
  tft.setRotation(1);
  drawHeader();

  // MPU init
  if (!mpuCheckIdentity()) {
    Serial.println("MPU6050 not found at 0x68. Check wiring or address.");
    tft.setCursor(8, 60);
    tft.setTextSize(2);
    tft.setTextColor(ILI9341_RED);
    tft.print("MPU NOT FOUND!");
    // still continue; reading will likely give garbage
  } else {
    Serial.println("MPU6050 found, initializing...");
    mpuInit();
    Serial.println("MPU initialized.");
  }

  // WiFi + MQTT
  connectWiFi();
  mqttClient.setServer(MQTT_SERVER, MQTT_PORT);
  mqttClient.setCallback(mqttCallback);

  // init timing
  lastMicros = micros();
  lastPublishMs = millis() - UPDATE_INTERVAL_MS; // force immediate publish
}

void loop() {
  // Ensure MQTT connected when WiFi available
  if (WiFi.status() == WL_CONNECTED) {
    mqttReconnect();
    mqttClient.loop();
  }

  unsigned long nowMs = millis();
  if (nowMs - lastPublishMs >= UPDATE_INTERVAL_MS) {
    lastPublishMs = nowMs;

    // read raw sensors
    int16_t rawAx, rawAy, rawAz, rawGx, rawGy, rawGz;
    mpuReadRaw(rawAx, rawAy, rawAz, rawGx, rawGy, rawGz);

    // convert to physical values
    float axg = (float)rawAx / ACCEL_SENS; // in g
    float ayg = (float)rawAy / ACCEL_SENS;
    float azg = (float)rawAz / ACCEL_SENS;

    float gx_dps = (float)rawGx / GYRO_SENS; // deg/s
    float gy_dps = (float)rawGy / GYRO_SENS;
    float gz_dps = (float)rawGz / GYRO_SENS;

    // compute dt (seconds) using micros to integrate gyro
    unsigned long curMicros = micros();
    float dt = (curMicros - lastMicros) / 1000000.0f;
    if (dt <= 0 || dt > 0.5f) dt = 0.01f; // sanity (avoid huge dt if first run)
    lastMicros = curMicros;

    // accel-based angles (degrees)
    float accRoll  = atan2(ayg, azg) * 180.0f / PI;
    float accPitch = atan2(-axg, sqrt(ayg * ayg + azg * azg)) * 180.0f / PI;

    // integrate gyro rates (deg/s)
    // apply complementary filter
    roll  = ALPHA * (roll + gx_dps * dt) + (1.0f - ALPHA) * accRoll;
    pitch = ALPHA * (pitch + gy_dps * dt) + (1.0f - ALPHA) * accPitch;

    // Draw on display
    bool mqttConnected = mqttClient.connected();
    drawHeader();
    drawSensorValues(axg, ayg, azg, gx_dps, gy_dps, gz_dps, roll, pitch, mqttConnected);

    // Build JSON payload
    char payload[256];
    snprintf(payload, sizeof(payload),
             "{\"ax_g\":%.3f, \"ay_g\":%.3f, \"az_g\":%.3f, \"gx_dps\":%.2f, \"gy_dps\":%.2f, \"gz_dps\":%.2f, \"roll\":%.2f, \"pitch\":%.2f, \"ip\":\"%s\"}",
             axg, ayg, azg, gx_dps, gy_dps, gz_dps, roll, pitch,
             WiFi.status() == WL_CONNECTED ? WiFi.localIP().toString().c_str() : "0.0.0.0");

    Serial.print("Payload: ");
    Serial.println(payload);

    // publish via MQTT if connected
    if (mqttConnected) {
      bool ok = mqttClient.publish(MQTT_PUB_TOPIC, payload);
      Serial.print("Published? ");
      Serial.println(ok ? "yes" : "no");
    } else {
      Serial.println("MQTT not connected, skipping publish.");
    }
  }

  // small background delay to yield
  delay(10);
}