#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SH110X.h>
#include <Adafruit_QMC5883P.h>
#include <EEPROM.h>
#include <NMEAGPS.h>

// ----------------- Hardware -----------------
static const uint8_t OLED_ADDR = 0x3C;
static const uint8_t MAG_ADDR  = 0x2C;
static const uint8_t BTN_PIN   = 4;   // D4 -> GND, INPUT_PULLUP

Adafruit_SH1106G display(128, 64, &Wire, -1);
Adafruit_QMC5883P mag;

// GPS (NeoGPS)
NMEAGPS gps;
gps_fix fix;
#define gpsPort Serial1   // Pro Micro hardware UART on D0/D1

// ----------------- Settings -----------------
static const uint32_t SCREEN_TIMEOUT_MS = 300000;
static const uint32_t DEBOUNCE_MS       = 35;
static const uint32_t LONGPRESS_MS      = 1200;  // start/stop calibration
static const uint32_t RESETCAL_MS       = 4000;  // clear EEPROM cal
static const uint32_t CAL_DURATION_MS   = 12000; // auto-finish

static const float DECLINATION_DEG = 0.0f;

// You mounted sensor 90°; you said North showed East => -90 fixes it
static const float SENSOR_TO_SCREEN_OFFSET_DEG = -90.0f;

static const float SMOOTH_ALPHA = 0.18f;

// ----------------- UI -----------------
enum ScreenMode : uint8_t { MODE_DATA = 0, MODE_COMPASS = 1 };
ScreenMode mode = MODE_DATA;

bool screenOn = true;
uint32_t lastActivityMs = 0;

// button state
uint8_t lastBtnRead = HIGH, stableBtn = HIGH;
uint32_t lastDebounceMs = 0, pressStartMs = 0;
bool longPressFired = false;

// heading smoothing
float headingSmoothed = NAN;

// ----------------- Calibration (min/max during calibration) -----------------
bool calibrating = false;
uint32_t calStartMs = 0;
float minX, maxX, minY, maxY;

// ----------------- Stored calibration (fixed-point) -----------------
// We store offsets and scales as int16 to save flash/RAM.
// offX/offY in milli-gauss (mG), scale in Q12 fixed point (4096 = 1.0)
int16_t offX_mG = 0, offY_mG = 0;
int16_t scX_q12 = 4096, scY_q12 = 4096;

// EEPROM struct
struct CalStore16 {
  uint32_t magic;
  int16_t offX_mG, offY_mG;
  int16_t scX_q12, scY_q12;
};
static const uint32_t CAL_MAGIC = 0x514D4332UL; // "QMC2"
static const int EEPROM_ADDR = 0;

// ----------------- Helpers -----------------
static float wrap360(float deg) {
  while (deg < 0) deg += 360.0f;
  while (deg >= 360.0f) deg -= 360.0f;
  return deg;
}

static float smoothAngle(float prev, float cur) {
  if (isnan(prev)) return cur;
  float d = cur - prev;
  if (d > 180) d -= 360;
  if (d < -180) d += 360;
  return wrap360(prev + SMOOTH_ALPHA * d);
}

static float headingFromMag(float mx, float my) {
  // NOTE: you previously needed E/W fix; if still swapped, use atan2(my, -mx)
  return atan2(my, -mx) * 180.0f / PI;
}

static const char* cardinal(uint16_t deg) {
  if (deg < 23 || deg >= 338) return "N";
  if (deg < 68)  return "NE";
  if (deg < 113) return "E";
  if (deg < 158) return "SE";
  if (deg < 203) return "S";
  if (deg < 248) return "SW";
  if (deg < 293) return "W";
  return "NW";
}

static void p2(uint8_t v) { if (v < 10) display.print('0'); display.print(v); }
static void p4(uint16_t v) {
  if (v < 1000) display.print('0');
  if (v < 100)  display.print('0');
  if (v < 10)   display.print('0');
  display.print(v);
}

static void screenOff() {
  display.clearDisplay();
  display.display();
  display.oled_command(0xAE);
  screenOn = false;
}
static void screenOnFn() {
  display.oled_command(0xAF);
  screenOn = true;
  lastActivityMs = millis();
}

// ---- EEPROM cal ----
static void loadCal() {
  CalStore16 s;
  EEPROM.get(EEPROM_ADDR, s);
  if (s.magic == CAL_MAGIC &&
      s.scX_q12 >= 512 && s.scX_q12 <= 16384 &&
      s.scY_q12 >= 512 && s.scY_q12 <= 16384) {
    offX_mG = s.offX_mG; offY_mG = s.offY_mG;
    scX_q12 = s.scX_q12; scY_q12 = s.scY_q12;
  } else {
    offX_mG = offY_mG = 0;
    scX_q12 = scY_q12 = 4096;
  }
}

static void saveCal() {
  CalStore16 s;
  s.magic = CAL_MAGIC;
  s.offX_mG = offX_mG; s.offY_mG = offY_mG;
  s.scX_q12 = scX_q12; s.scY_q12 = scY_q12;
  EEPROM.put(EEPROM_ADDR, s);
}

static void clearCal() {
  offX_mG = offY_mG = 0;
  scX_q12 = scY_q12 = 4096;
  CalStore16 s = {0, 0, 0, 4096, 4096};
  EEPROM.put(EEPROM_ADDR, s);
  headingSmoothed = NAN;
}

// Apply fixed-point cal to gauss values
static void applyCal(float &mx, float &my) {
  // Convert offsets from mG to Gauss: 1000 mG = 1 Gauss
  float ox = (float)offX_mG / 1000.0f;
  float oy = (float)offY_mG / 1000.0f;

  mx -= ox;
  my -= oy;

  mx = mx * ((float)scX_q12 / 4096.0f);
  my = my * ((float)scY_q12 / 4096.0f);
}

// ---- Calibration flow ----
static void startCal() {
  calibrating = true;
  calStartMs = millis();
  minX = minY =  1e9f;
  maxX = maxY = -1e9f;
}

static void finishCal() {
  calibrating = false;

  // offsets (gauss)
  float offX = (maxX + minX) * 0.5f;
  float offY = (maxY + minY) * 0.5f;

  // ranges
  float rx = (maxX - minX) * 0.5f;
  float ry = (maxY - minY) * 0.5f;
  if (rx < 1e-6f) rx = 1.0f;
  if (ry < 1e-6f) ry = 1.0f;
  float avg = (rx + ry) * 0.5f;

  float scX = avg / rx;
  float scY = avg / ry;

  // store as fixed-point
  offX_mG = (int16_t) (offX * 1000.0f);
  offY_mG = (int16_t) (offY * 1000.0f);

  int32_t sx = (int32_t)(scX * 4096.0f);
  int32_t sy = (int32_t)(scY * 4096.0f);
  if (sx < 512) sx = 512; if (sx > 16384) sx = 16384;
  if (sy < 512) sy = 512; if (sy > 16384) sy = 16384;
  scX_q12 = (int16_t)sx;
  scY_q12 = (int16_t)sy;

  saveCal();
  headingSmoothed = NAN;
}

// ---- Button ----
static void handleButton() {
  uint32_t now = millis();
  uint8_t r = digitalRead(BTN_PIN);

  if (r != lastBtnRead) { lastBtnRead = r; lastDebounceMs = now; }
  if ((uint32_t)(now - lastDebounceMs) < DEBOUNCE_MS) return;

  if (r != stableBtn) {
    stableBtn = r;

    if (stableBtn == LOW) {
      pressStartMs = now;
      longPressFired = false;
      lastActivityMs = now;

      if (!screenOn) { screenOnFn(); return; } // wake-only
    } else {
      if (!longPressFired && screenOn && !calibrating) {
        mode = (mode == MODE_DATA) ? MODE_COMPASS : MODE_DATA;
      }
    }
  }

  if (stableBtn == LOW && !longPressFired) {
    uint32_t held = now - pressStartMs;

    if (held >= RESETCAL_MS) {
      longPressFired = true;
      calibrating = false;
      clearCal();
      lastActivityMs = now;
    } else if (held >= LONGPRESS_MS) {
      longPressFired = true;
      if (!calibrating) startCal(); else finishCal();
      lastActivityMs = now;
    }
  }
}

// ---- Compass dial ----
static void drawCompass(uint16_t headingDeg) {
  const int cx = 64, cy = 32, r = 28;

  display.drawCircle(cx, cy, r, 1);
  display.drawCircle(cx, cy, r - 1, 1);

  for (int d = 0; d < 360; d += 10) {
    float rad = (d - 90) * PI / 180.0f;
    int len = (d % 30 == 0) ? 6 : 3;

    int x1 = cx + (int)((r - 1) * cos(rad));
    int y1 = cy + (int)((r - 1) * sin(rad));
    int x2 = cx + (int)((r - 1 - len) * cos(rad));
    int y2 = cy + (int)((r - 1 - len) * sin(rad));
    display.drawLine(x1, y1, x2, y2, 1);
  }

  display.setCursor(cx - 3, cy - r - 8); display.print('N');
  display.setCursor(cx + r + 2, cy - 3); display.print('E');
  display.setCursor(cx - 3, cy + r + 2); display.print('S');
  display.setCursor(cx - r - 8, cy - 3); display.print('W');

  float a = ((float)headingDeg - 90.0f) * PI / 180.0f;
  int nx = cx + (int)((r - 8) * cos(a));
  int ny = cy + (int)((r - 8) * sin(a));

  display.drawLine(cx, cy, nx, ny, 1);
  display.drawLine(cx + 1, cy, nx, ny, 1);
  display.drawLine(cx, cy + 1, nx, ny, 1);
  display.fillCircle(cx, cy, 2, 1);

  display.setCursor(0, 0);
  display.print(headingDeg);
  display.print(' ');
  display.print(cardinal(headingDeg));
}

// ---- Date/Time + Lat/Lon (NeoGPS) ----
static void printDateTimeUTC() {
  if (fix.valid.date) {
    display.print(F("Date:"));
    p2(fix.dateTime.date); display.print('/');
    p2(fix.dateTime.month); display.print('/');
    p4(fix.dateTime.year);
    display.println();
  } else display.println(F("Date: ---"));

  if (fix.valid.time) {
    display.print(F("Time:"));
    p2(fix.dateTime.hours); display.print(':');
    p2(fix.dateTime.minutes); display.print(':');
    p2(fix.dateTime.seconds);
    display.println(F(" UTC"));
  } else display.println(F("Time: ---"));
}

static void printLatLon() {
  if (!fix.valid.location) {
    display.println(F("Lat: ---"));
    display.println(F("Lon: ---"));
    return;
  }

  auto printDeg1e7 = [&](int32_t v) {
    if (v < 0) { display.print('-'); v = -v; } else display.print('+');
    int32_t deg = v / 10000000L;
    uint32_t frac = (uint32_t)(v % 10000000L);
    display.print(deg);
    display.print('.');
    // print 5 decimals (rounded from 7)
    frac = (frac + 50) / 100;
    if (frac < 10000) display.print('0');
    if (frac < 1000)  display.print('0');
    if (frac < 100)   display.print('0');
    if (frac < 10)    display.print('0');
    display.print(frac);
  };

  display.print(F("Lat:")); printDeg1e7(fix.latitudeL());  display.println();
  display.print(F("Lon:")); printDeg1e7(fix.longitudeL()); display.println();
}

void setup() {
  Wire.begin();
  pinMode(BTN_PIN, INPUT_PULLUP);

  if (!display.begin(OLED_ADDR, true)) while (1) {}
  display.setTextSize(1);
  display.setTextColor(1);
  display.clearDisplay();
  display.setCursor(0, 0);
  display.println(F("OK"));
  display.display();

  if (!mag.begin(MAG_ADDR, &Wire)) {
    display.clearDisplay();
    display.setCursor(0, 0);
    display.println(F("MAG?"));
    display.display();
    while (1) {}
  }
  mag.setMode(QMC5883P_MODE_CONTINUOUS);

  gpsPort.begin(9600);
  loadCal();
  lastActivityMs = millis();
}

void loop() {
  handleButton();

  while (gps.available(gpsPort)) {
    fix = gps.read();
    lastActivityMs = millis();
  }

  uint32_t now = millis();

  if (calibrating && (uint32_t)(now - calStartMs) >= CAL_DURATION_MS) finishCal();
  if (screenOn && !calibrating && (uint32_t)(now - lastActivityMs) > SCREEN_TIMEOUT_MS) screenOff();

  // mag + heading
  float mx, my, mz;
  uint16_t headingInt = 0;
  bool headingValid = false;

  if (mag.isDataReady() && mag.getGaussField(&mx, &my, &mz)) {
    if (calibrating) {
      if (mx < minX) minX = mx; if (mx > maxX) maxX = mx;
      if (my < minY) minY = my; if (my > maxY) maxY = my;
    } else {
      applyCal(mx, my);
      float raw = headingFromMag(mx, my);
      raw = wrap360(raw + DECLINATION_DEG + SENSOR_TO_SCREEN_OFFSET_DEG);
      headingSmoothed = smoothAngle(headingSmoothed, raw);
      if (!isnan(headingSmoothed)) {
        headingInt = (uint16_t)(headingSmoothed + 0.5f);
        if (headingInt >= 360) headingInt -= 360;
        headingValid = true;
      }
    }
  }

  if (!screenOn) { delay(10); return; }

  display.clearDisplay();
  display.setCursor(0, 0);

  if (calibrating) {
    display.println(F("CAL"));
    display.println(F("Rotate"));
    uint16_t left = (uint16_t)((CAL_DURATION_MS - (now - calStartMs) + 999) / 1000);
    display.print(F("t=")); display.print(left); display.println(F("s"));
    display.println(F("Hold=Stop"));
  } else if (mode == MODE_COMPASS) {
    if (headingValid) drawCompass(headingInt);
    else { display.println(F("Comp ---")); display.println(F("No mag")); }
  } else {
    display.print(F("Fix:"));
    display.print(fix.valid.status && (fix.status >= gps_fix::STATUS_STD) ? 'Y' : 'N');
    display.print(F(" Sat:"));
    display.println(fix.valid.satellites ? fix.satellites : 0);

    printLatLon();
    printDateTimeUTC();

    display.print(F("MAG:"));
    if (headingValid) { display.print(headingInt); display.print(' '); display.println(cardinal(headingInt)); }
    else display.println(F("---"));

    display.println(F("Btn=Mode"));
    display.println(F("Hold=Cal"));
  }

  display.display();
  delay(20);
}