#include <Adafruit_Sensor.h>
#include <Wire.h>
#include <Adafruit_MPU6050.h>
#include <WiFi.h>
#include <PubSubClient.h>
#include <ESP32Servo.h>

const char* ssid = "Vienna's Orange iPhone 13";
const char* password = "123456789";
const char* mqtt_server = "broker.hivemq.com";
const char* mqtt_topic = "FYCJHVEU";

WiFiClient espClient;
PubSubClient client(espClient);

float roll = 0;
float pitch = 0;
float yawAngle = 0;
float rollRaw;
float pitchRaw;
float yawRaw;
float gyroZOffset = 0;

Servo rudderServo;
const int SERVO_PIN = 1; // Changed from 0 to avoid boot issues

unsigned long lastPublish = 0;
unsigned long tStart = 0;

// FIR Filter
const float firCoeff[5] = {0.20f, 0.20f, 0.20f, 0.20f, 0.20f};
const int FIR_TAPS = 5;
float bufAx[FIR_TAPS] = {0};
float bufAy[FIR_TAPS] = {0};
float bufAz[FIR_TAPS] = {0};
float bufGz[FIR_TAPS] = {0};
int firIndex = 0;

float applyFIR(float* buf) {
  float out = 0.0f;
  for (int i = 0; i < FIR_TAPS; i++) {
    int idx = (firIndex + i) % FIR_TAPS;
    out += firCoeff[i] * buf[idx];
  }
  return out;
}

Adafruit_MPU6050 mpu;

void calibrateGyro() {
  Serial.println("Calibrating gyro - keep stationary...");
  delay(1000);
  float sumZ = 0;
  int numPoints = 1000;
  for (int i = 0; i < numPoints; i++) {
    sensors_event_t a, g, temp;
    mpu.getEvent(&a, &g, &temp);
    sumZ += g.gyro.z;
    delay(10);
  }
  gyroZOffset = sumZ / numPoints;
  Serial.println("Gyro calibration done!");
}

void connectWiFi() {
  WiFi.begin(ssid, password);
  Serial.print("Connecting to WiFi");
  
  int attempts = 0;
  while (WiFi.status() != WL_CONNECTED && attempts < 20) {
    delay(500);
    Serial.print(".");
    attempts++;
  }

  if (WiFi.status() != WL_CONNECTED) {
    Serial.println("\nWiFi FAILED - check SSID/password");
  } else {
    Serial.println("\nWiFi connected!");
    Serial.print("IP Address: ");
    Serial.println(WiFi.localIP());
  }
}

void reconnect() {
  while (!client.connected()) {
    Serial.print("Connecting to MQTT...");
    String clientId = "ESP32FlightSim-" + String(random(0xffff), HEX);
    if (client.connect(clientId.c_str())) {
      Serial.println("MQTT connected!");
    } else {
      Serial.println("MQTT failed, retrying in 5s");
      delay(5000);
    }
  }
}

void setup() {
  Serial.begin(115200);
  Wire.begin(4, 5); // SDA=GP4, SCL=GP5

  connectWiFi();

  client.setServer(mqtt_server, 1883);

  Wire.begin(4, 5);
  if (!mpu.begin()) {
    Serial.println("MPU6050 not found - check wiring!");
    while (1) delay(10);
  }
  Serial.println("MPU6050 started!");

  mpu.setAccelerometerRange(MPU6050_RANGE_2_G);
  mpu.setFilterBandwidth(MPU6050_BAND_21_HZ);
  mpu.setGyroRange(MPU6050_RANGE_2000_DEG);

  rudderServo.attach(SERVO_PIN);
  rudderServo.write(90);

  calibrateGyro();
  tStart = millis();
}

void loop() {
  if (!client.connected()) { reconnect(); }
  client.loop();

  // Calculate dt first thing
  unsigned long now = millis();
  float dt = (now - tStart) / 1000.0;
  tStart = now;

  sensors_event_t a, g, temp;
  mpu.getEvent(&a, &g, &temp);

  float rawAx = a.acceleration.x / 9.81;
  float rawAy = a.acceleration.y / 9.81;
  float rawAz = a.acceleration.z / 9.81;
  float rawGz = -(g.gyro.z - gyroZOffset);

  // Fill FIR buffers
  bufAx[firIndex] = rawAx;
  bufAy[firIndex] = rawAy;
  bufAz[firIndex] = rawAz;
  bufGz[firIndex] = rawGz;
  firIndex = (firIndex + 1) % FIR_TAPS;

  // Apply FIR filter
  float fAx = applyFIR(bufAx);
  float fAy = applyFIR(bufAy);
  float fAz = applyFIR(bufAz);
  float fGz = applyFIR(bufGz);

  // Pitch and roll from filtered accelerometer
  pitch = atan2(fAy, sqrt(fAz * fAz + fAx * fAx)) * 180.0 / PI;
  roll  = atan2(fAx, sqrt(fAz * fAz + fAy * fAy)) * 180.0 / PI;

  // Yaw from gyro integration using correct dt
  yawAngle = yawAngle + dt * fGz * 180.0 / PI;

  // Map yaw 0-360 to servo 0-180 as per lab spec
  float yawMod = fmod(yawAngle, 360.0);
  if (yawMod < 0) yawMod += 360.0;
  int servoAngle = constrain((int)(yawMod / 2.0), 0, 180);
  rudderServo.write(servoAngle);

  // Publish at 20Hz
  if (millis() - lastPublish >= 50) {
    lastPublish = millis();
    char payload[100];
    snprintf(payload, sizeof(payload),
             "{\"pitch\": %.2f, \"roll\": %.2f, \"yaw\": %.2f}",
             pitch, roll, yawAngle);
    client.publish(mqtt_topic, payload);
    Serial.println(payload);
  }

  // Serial plotter output
  Serial.print("Pitch:"); Serial.print(pitch); Serial.print(',');
  Serial.print("Roll:");  Serial.print(roll);  Serial.print(',');
  Serial.print("Yaw:");   Serial.print(yawAngle); Serial.print(',');
  Serial.print("UL:"); Serial.print(45); Serial.print(',');
  Serial.print("LL:"); Serial.println(-45);

  delay(1);
}