#include <SPI.h>
#include <TFT_eSPI.h>
#include <WiFi.h>
#include <PubSubClient.h>
#include <DHT.h>
#include <ArduinoJson.h>
#include <freertos/FreeRTOS.h>
#include <freertos/task.h>

// ----------------- CHANGE ME -----------------------------
// Patient Information
const char* patientName = "DTJ HU";   // Your Name
const char* patientID = "patient037";   // Your ID
// ---------------------------------------------------------

// Initialize TFT display
TFT_eSPI tft = TFT_eSPI();

// Configuration
bool wifiRequired = true; // Set to true if you want to run only when WiFi is connected

TaskHandle_t mqttTaskHandle = NULL;
TaskHandle_t mqttPublishTaskHandle = NULL;

// Pin Definitions
#define HEART_RATE_PIN 36    // VP pin
#define BP_SYSTOLIC_PIN 34   
#define BP_DIASTOLIC_PIN 32  
#define OXYGEN_LEVEL_PIN 33  
#define GLUCOSE_PIN 35
#define TEMPERATURE_PIN 14   
#define DHTTYPE DHT22

// WiFi and MQTT configuration
const char* ssid = "Wokwi-GUEST";
const char* password = "";
const char* mqtt_server = "13.48.43.195";
//const char* mqtt_server = "broker.hivemq.com";
//const char* mqtt_server = "test.mosquitto.org";
const int mqtt_port = 1883;
String mqtt_topic_base = "medical/";
String mqtt_topic = "";

// Normal ranges for vital signs
#define HR_MIN 75
#define HR_MAX 100
#define BP_SYS_MIN 100
#define BP_SYS_MAX 140
#define BP_DIA_MIN 60
#define BP_DIA_MAX 90
#define SPO2_MIN 94
#define SPO2_MAX 100
#define TEMP_MIN_C 35.0
#define TEMP_MAX_C 40.0
#define TEMP_MIN_F 95.0  // 35.0°C in Fahrenheit
#define TEMP_MAX_F 104.0 // 40.0°C in Fahrenheit
#define GLUCOSE_MIN 70
#define GLUCOSE_MAX 140

// Fluctuation ranges
#define HR_FLUCTUATION 1      // ±1 BPM
#define BP_SYS_FLUCTUATION 1  // ±2 mmHg
#define BP_DIA_FLUCTUATION 1  // ±2 mmHg
#define SPO2_FLUCTUATION 0.3  // ±0.3%
#define TEMP_FLUCTUATION 0.036 // ±0.036°F (equivalent to ±0.02°C)
#define GLUCOSE_FLUCTUATION 1.5 // ±1.5 mg/dL

// Color definitions
#define BACKGROUND TFT_BLACK
#define GRID_COLOR TFT_DARKGREY
#define ECG_COLOR TFT_GREEN
#define TEXT_COLOR TFT_WHITE
#define HEART_COLOR TFT_RED
#define OXYGEN_COLOR TFT_CYAN
#define TEMP_COLOR TFT_YELLOW
#define BP_COLOR TFT_MAGENTA
#define GLUCOSE_COLOR TFT_ORANGE

// Initialize components
DHT dht(TEMPERATURE_PIN, DHTTYPE);
WiFiClient espClient;
PubSubClient client(espClient);

// Variables to store sensor values
float heartRate = 75.0;
int systolicBP = 120;
int diastolicBP = 80;
float oxygenLevel = 98.0;
float bodyTemperature = 37.0;
float glucoseLevel = 100.0;

// ECG simulation variables
const int ECG_BUFFER_SIZE = 240;
int ecgBuffer[ECG_BUFFER_SIZE];
int ecgIndex = 0;
unsigned long lastEcgUpdate = 0;
const int ECG_UPDATE_INTERVAL = 15; // ms

int baseline = 150;       

// Timers
unsigned long lastVitalUpdateTime = 0;
const long vitalUpdateInterval = 1000;  // Update vital signs every 2 seconds
unsigned long lastMqttAttempt = 0;
const long mqttRetryInterval = 2000; // Retry MQTT connection every 2 seconds

// Flag to track display initialization
bool displayInitialized = false;
bool wifiConnected = false;

void setup() {
  Serial.begin(115200);
  delay(1000);  // Give serial a moment to start
  Serial.println("\n\n----- Patient IoT Monitor Starting -----");
  
  // Initialize TFT
  tft.init();
  displayInitialized = true;
  Serial.println("TFT initialized");

  if (displayInitialized) {
    tft.setRotation(1); // Landscape
    tft.fillScreen(BACKGROUND);
    
    // Draw welcome screen
    tft.setTextColor(TEXT_COLOR, BACKGROUND);
    tft.setTextSize(2);
    tft.setCursor(60, 80);
    tft.println("Medical IoT Monitor");
    tft.setCursor(60, 120);
    tft.println("Initializing...");
  }
  
  // Initialize DHT sensor
  dht.begin();
  Serial.println("DHT sensor initialized");
  
  // Initialize analog pins
  pinMode(HEART_RATE_PIN, INPUT);
  pinMode(BP_SYSTOLIC_PIN, INPUT);
  pinMode(BP_DIASTOLIC_PIN, INPUT);
  pinMode(OXYGEN_LEVEL_PIN, INPUT);
  pinMode(GLUCOSE_PIN, INPUT);
  Serial.println("Analog pins configured");
  
  // Setup initial ECG buffer
  for (int i = 0; i < ECG_BUFFER_SIZE; i++) {
    ecgBuffer[i] = baseline;
  }
  
  // Set up MQTT topic using patientID
  mqtt_topic = mqtt_topic_base + patientID + "/vitals";
  Serial.print("MQTT topic: ");
  Serial.println(mqtt_topic);
  
  // Connect to WiFi
  setupWifi();
  
  // Connect to MQTT if WiFi is connected
  if (wifiConnected) {
    client.setServer(mqtt_server, mqtt_port);
     // Create MQTT reconnection task
  xTaskCreatePinnedToCore(
    mqttReconnectTask,   // Task function
    "MQTTReconnectTask", // Name of task
    4096,               // Stack size (bytes)
    NULL,               // Parameter to pass
    1,                  // Task priority
    &mqttTaskHandle,    // Task handle
    0                   // Run on core 0
  );
  xTaskCreatePinnedToCore(
  mqttPublishTask,     // Task function
  "MQTTPublishTask",   // Name of task
  4096,               // Stack size (bytes)
  NULL,               // Parameter to pass
  1,                  // Task priority
  &mqttPublishTaskHandle, // Task handle
  0                   // Run on core 0
);
  }
  
  // Draw initial screen
  if (displayInitialized) {
    drawScreenLayout();
  }
  
  Serial.println("Setup complete!");
}

void mqttReconnectTask(void * parameter) {
  for(;;) {
    if (wifiConnected && !client.connected()) {
      Serial.print("Attempting MQTT connection from task...");
      String clientId = "ESP32Client-";
      clientId += String(random(0xffff), HEX);
      
      if (client.connect(clientId.c_str())) {
        Serial.println("connected");
      } else {
        Serial.print("failed, rc=");
        Serial.print(client.state());
        Serial.println(" will retry");
      }
    }
    // Wait before next attempt
    vTaskDelay(5000 / portTICK_PERIOD_MS);
  }
}

void mqttPublishTask(void * parameter) {
  for(;;) {
    ulTaskNotifyTake(pdTRUE, portMAX_DELAY); // Wait for notification
    if (wifiConnected && client.connected()) {
      publishData();
    }
  }
}

void setupWifi() {
  if (displayInitialized) {
    tft.setTextColor(TEXT_COLOR, BACKGROUND);
    tft.setCursor(60, 160);
    tft.print("Connecting to WiFi...");
  }

  Serial.print("Connecting to WiFi...");

  // Ensure clean WiFi state
  WiFi.disconnect(true); // Disconnect any previous connections
  WiFi.mode(WIFI_STA);   // Set to Station mode
  delay(1000);           // Give time for mode to stabilize

  // Begin WiFi connection
  WiFi.begin(ssid, password);

  // Wait for connection with a reasonable timeout
  int attempts = 0;
  while (WiFi.status() != WL_CONNECTED && attempts < 20) {
    delay(500);
    Serial.print(".");
    attempts++;
  }

  if (WiFi.status() == WL_CONNECTED) {
    Serial.println("\nWiFi connected");
    Serial.print("IP address: ");
    Serial.println(WiFi.localIP());
    wifiConnected = true;

    if (displayInitialized) {
      tft.fillRect(60, 160, 240, 30, BACKGROUND);
      tft.setCursor(60, 160);
      tft.print("WiFi connected!");
    }
  } else {
    Serial.println("\nWiFi connection FAILED");
    wifiConnected = false;

    if (displayInitialized) {
      tft.fillRect(60, 160, 240, 30, BACKGROUND);
      tft.setCursor(60, 160);
      tft.print("WiFi failed - using simulation");
    }

    if (wifiRequired) {
      if (displayInitialized) {
        tft.fillRect(60, 190, 240, 30, BACKGROUND);
        tft.setCursor(60, 190);
        tft.setTextColor(HEART_COLOR, BACKGROUND);
        tft.print("WiFi required! System halted.");
      }
      Serial.println("WiFi connection required but failed.");
    }
  }
}

void drawScreenLayout() {
  if (!displayInitialized) return;
  
  // Clear screen
  tft.fillScreen(BACKGROUND);
  
  // Draw title with patient name
  tft.setTextColor(TEXT_COLOR, BACKGROUND);
  tft.setTextSize(2);
  tft.setCursor(5, 5);
  tft.print(patientName);
  tft.setCursor(200, 5);
  tft.print("Monitor");
  
  // Draw grid for ECG
  tft.drawRect(0, 40, 320, 100, GRID_COLOR);
  for (int x = 0; x < 320; x += 20) {
    for (int y = 40; y < 140; y += 20) {
      tft.drawPixel(x, y, GRID_COLOR);
    }
  }
  
  // Draw vital signs areas
  tft.drawRect(0, 150, 105, 45, HEART_COLOR);
  tft.drawRect(0, 195, 105, 45, BP_COLOR);
  tft.drawRect(106, 150, 105, 45, OXYGEN_COLOR);
  tft.drawRect(106, 195, 105, 45, GLUCOSE_COLOR);
  tft.drawRect(212, 150, 108, 90, TEMP_COLOR);
  
  // Labels for vital signs
  tft.setTextSize(1);
  
  // Heart rate
  tft.setTextColor(HEART_COLOR);
  tft.setCursor(5, 155);
  tft.print("HEART RATE");
  
  // Blood Pressure
  tft.setTextColor(BP_COLOR);
  tft.setCursor(5, 200);
  tft.print("BLOOD PRESS");
  
  // Oxygen
  tft.setTextColor(OXYGEN_COLOR);
  tft.setCursor(111, 155);
  tft.print("SpO2");
  
  // Glucose
  tft.setTextColor(GLUCOSE_COLOR);
  tft.setCursor(111, 200);
  tft.print("GLUCOSE");
  
  // Temperature
  tft.setTextColor(TEMP_COLOR);
  tft.setCursor(217, 155);
  tft.print("TEMPERATURE");
}

void readSensors() {
  // Skip if WiFi is required but not connected
  if (wifiRequired && !wifiConnected) {
    return;
  }
  
  // Add time-based fluctuation factor (0.0 to 1.0) based on seconds
  float timeFactor = (millis() % 10000) / 10000.0;
  float fluctuationFactor = sin(timeFactor * 2 * PI);
  
  // Read heart rate with natural fluctuation over time
  int rawHeartRate = analogRead(HEART_RATE_PIN);
  heartRate = map(rawHeartRate, 0, 4095, HR_MIN, HR_MAX);
  heartRate += fluctuationFactor * HR_FLUCTUATION;
  heartRate += random(-1, 1); // ±0.5 BPM
  heartRate = constrain(heartRate, HR_MIN, HR_MAX);
  
  // Read blood pressure
  int rawSystolic = analogRead(BP_SYSTOLIC_PIN);
  int rawDiastolic = analogRead(BP_DIASTOLIC_PIN);
  systolicBP = map(rawSystolic, 0, 4095, BP_SYS_MIN, BP_SYS_MAX);
  diastolicBP = map(rawDiastolic, 0, 4095, BP_DIA_MIN, BP_DIA_MAX);
  systolicBP += fluctuationFactor * BP_SYS_FLUCTUATION;
  diastolicBP += fluctuationFactor * BP_DIA_FLUCTUATION;
  if (diastolicBP >= systolicBP - 30) {
    diastolicBP = systolicBP - 30 - random(0, 10);
  }
  
  // Read oxygen level
  int rawOxygen = analogRead(OXYGEN_LEVEL_PIN);
  oxygenLevel = map(rawOxygen, 0, 4095, SPO2_MIN, SPO2_MAX);
  oxygenLevel += fluctuationFactor * SPO2_FLUCTUATION;
  oxygenLevel += random(-2, 3) * 0.1; // ±0.2%
  oxygenLevel = constrain(oxygenLevel, SPO2_MIN, SPO2_MAX);
  
 // Read temperature in Celsius and convert to Fahrenheit
  float ambientTemp = dht.readTemperature();
  if (!isnan(ambientTemp)) {
    // Map sensor range (-40°C to 80°C) to human body range (35.0°C to 40.0°C)
    float tempInCelsius = map(ambientTemp * 100, -4000, 8000, TEMP_MIN_C * 100, TEMP_MAX_C * 100) / 100.0;
    // Convert to Fahrenheit
    bodyTemperature = (tempInCelsius * 9.0 / 5.0) + 32.0;
    // Apply fluctuations in Fahrenheit
    bodyTemperature += fluctuationFactor * TEMP_FLUCTUATION;
    bodyTemperature += random(-2, 3) * 0.018; // ±0.018°F (equivalent to ±0.01°C)
  } else {
    Serial.println("Failed to read from DHT sensor!");
    // Simulate temperature in human body range
    float tempInCelsius = TEMP_MIN_C + (TEMP_MAX_C - TEMP_MIN_C) * 0.5; // Center at 37.5°C
    bodyTemperature = (tempInCelsius * 9.0 / 5.0) + 32.0;
    bodyTemperature += fluctuationFactor * TEMP_FLUCTUATION;
    bodyTemperature += random(-2, 3) * 0.018;
  }
  bodyTemperature = constrain(bodyTemperature, TEMP_MIN_F, TEMP_MAX_F);
  
  // Read glucose level
  int rawGlucose = analogRead(GLUCOSE_PIN);
  glucoseLevel = map(rawGlucose, 0, 4095, GLUCOSE_MIN, GLUCOSE_MAX);
  glucoseLevel += fluctuationFactor * GLUCOSE_FLUCTUATION;
  glucoseLevel = constrain(glucoseLevel, GLUCOSE_MIN, GLUCOSE_MAX);

}

bool publishData() {
  if (!wifiConnected) {
    Serial.println("Not connected to WiFi, skipping publish");
    return false;
  }

  if (!client.connected()) {
    Serial.println("Not connected to MQTT, skipping publish");
    return false;
  }
  
  StaticJsonDocument<256> doc;
  
  doc["patient_id"] = patientID;
  doc["patient_name"] = patientName;
  doc["heart_rate"] = round(heartRate * 10) / 10.0;
  doc["blood_pressure_systolic"] = systolicBP;
  doc["blood_pressure_diastolic"] = diastolicBP;
  doc["oxygen"] = round(oxygenLevel * 10) / 10.0;
  doc["temperature"] = round(bodyTemperature * 100) / 100.0;
  doc["glucose"] = round(glucoseLevel);
  
  char jsonBuffer[256];
  serializeJson(doc, jsonBuffer);
  
  Serial.print("Publishing to MQTT topic: ");
  Serial.println(mqtt_topic);
  Serial.print("Payload: ");
  Serial.println(jsonBuffer);
  
  if (client.publish(mqtt_topic.c_str(), jsonBuffer)) {
    Serial.println("Published to MQTT successfully");
    return true;
  } else {
    Serial.println("Failed to publish to MQTT");
    return false;
  }
}

void updateVitalsDisplay() {
  if (!displayInitialized) return;
  if (wifiRequired && !wifiConnected) return;

  // Heart rate
  tft.fillRect(5, 165, 95, 30, BACKGROUND);
  tft.setTextSize(2);
  tft.setTextColor(TFT_RED);
  tft.setCursor(5, 170);
  float displayHeartRate = round(heartRate * 10) / 10.0;
  tft.print(displayHeartRate, 1);
  tft.setTextSize(1);
  tft.print(" BPM");

  // Blood Pressure
  tft.fillRect(5, 210, 95, 30, BACKGROUND);
  tft.setTextSize(2);
  tft.setTextColor(BP_COLOR);
  tft.setCursor(5, 215);
  tft.print(systolicBP);
  tft.print("/");
  tft.print(diastolicBP);

  // Oxygen level
  tft.fillRect(111, 165, 95, 30, BACKGROUND);
  tft.setTextSize(2);
  tft.setTextColor(OXYGEN_COLOR);
  tft.setCursor(111, 170);
  float displayOxygen = round(oxygenLevel * 10) / 10.0;
  tft.print(displayOxygen, 1);
  tft.setTextSize(1);
  tft.print("%");

  // Glucose level
  tft.fillRect(111, 210, 95, 30, BACKGROUND);
  tft.setTextSize(2);
  tft.setTextColor(GLUCOSE_COLOR);
  tft.setCursor(111, 215);
  int displayGlucose = round(glucoseLevel);
  tft.print(displayGlucose);

  // Temperature
  tft.fillRect(217, 165, 98, 65, BACKGROUND);
  tft.setTextSize(3);
  tft.setTextColor(TEMP_COLOR);
  tft.setCursor(217, 180);
  tft.print(bodyTemperature, 1);
  tft.setTextSize(1);
  tft.setCursor(217, 215);
  tft.print("Fahrenheit");
  if (bodyTemperature < 100.0) {
    tft.drawRect(295, 180, 5, 20, TEMP_COLOR);
    tft.drawCircle(297, 205, 5, TEMP_COLOR);
    tft.fillCircle(297, 205, 3, TEMP_COLOR);
    tft.fillRect(295, 200, 5, 5, TEMP_COLOR);
  }
}

void generateEcgPoint() {
  if (wifiRequired && !wifiConnected) {
    return;
  }
  
  static int ecgPhase = 0;
  const int cycleLength = 200;
  static unsigned long lastCycleStart = 0;
  
  int beatInterval = 60000 / heartRate;
  beatInterval += random(-beatInterval/33, beatInterval/33);
  unsigned long currentTime = millis();
  
  if (currentTime - lastCycleStart >= beatInterval) {
    ecgPhase = 0;
    lastCycleStart = currentTime;
  }
  
  int position = ecgPhase;
  int ecgValue = baseline;
  
  const int pDuration = 12;
  const int prInterval = 28;
  const int qDuration = 3;
  const int rDuration = 4;
  const int sDuration = 3;
  const int stSegment = 20;
  const int tDuration = 28;
  
  int pEnd = pDuration;
  int prEnd = pEnd + prInterval;
  int qEnd = prEnd + qDuration;
  int rEnd = qEnd + rDuration;
  int sEnd = rEnd + sDuration;
  int stEnd = sEnd + stSegment;
  int tEnd = stEnd + tDuration;
  
  int pAmp = 40;
  int qAmp = -60;
  int rAmp = 250;
  int sAmp = -80;
  int tAmp = 100;
  
  if (position < pEnd) {
    float t = (position * PI) / pDuration;
    ecgValue = baseline + pAmp * (cos(t - PI) + 1) / 2;
  } else if (position < prEnd) {
    ecgValue = baseline;
  } else if (position < qEnd) {
    float t = (position - prEnd) / (float)qDuration;
    ecgValue = baseline + qAmp * t;
  } else if (position < rEnd) {
    float t = (position - qEnd) / (float)rDuration;
    if (t < 0.5) {
      ecgValue = baseline + qAmp + (rAmp - qAmp) * (t / 0.5);
    } else {
      ecgValue = baseline + rAmp - (rAmp - sAmp) * ((t - 0.5) / 0.5);
    }
  } else if (position < sEnd) {
    float t = (position - rEnd) / (float)sDuration;
    ecgValue = baseline + sAmp * (1 - t);
  } else if (position < stEnd) {
    ecgValue = baseline + 10;
  } else if (position < tEnd) {
    float t = (position - stEnd) * PI / tDuration;
    ecgValue = baseline + tAmp * (cos(t - PI) + 1) / 2;
  } else {
    ecgValue = baseline + random(-2, 3);
  }
  
  ecgValue += random(-5, 6);
  
  ecgBuffer[ecgIndex] = ecgValue;
  ecgIndex = (ecgIndex + 1) % ECG_BUFFER_SIZE;
  
  ecgPhase = (ecgPhase + 1) % cycleLength;
}

void drawEcg() {
  if (!displayInitialized) return;
  if (wifiRequired && !wifiConnected) {
    return;
  }
  
  static int lastX = 1;
  static int lastY = 0;
  
  int middleY = 40 + 98/2;
  
  int x = (ecgIndex + ECG_BUFFER_SIZE - 1) % ECG_BUFFER_SIZE;
  int y = middleY - (ecgBuffer[x] - baseline) * 0.4;
  y = constrain(y, 41, 139);
  
  int clearX = (lastX + 1) % 318 + 1;
  tft.fillRect(clearX, 41, 2, 98, BACKGROUND);
  
  for (int y = 41; y < 139; y += 20) {
    if (clearX % 20 == 0) {
      tft.drawPixel(clearX, y, GRID_COLOR);
    }
  }
  
  int prevY = middleY - (ecgBuffer[(x + ECG_BUFFER_SIZE - 1) % ECG_BUFFER_SIZE] - baseline) * 0.4;
  prevY = constrain(prevY, 41, 139);
  tft.drawLine(lastX, prevY, lastX + 1, y, ECG_COLOR);
  
  lastX = (lastX + 1) % 318 + 1;
  lastY = y;
}

bool reconnect() {
  if (!wifiConnected) return false;

  unsigned long currentMillis = millis();
  if (currentMillis - lastMqttAttempt < mqttRetryInterval) {
    return false;
  }

  lastMqttAttempt = currentMillis;

  if (!client.connected()) {
    String clientId = "ESP32Client-";
    clientId += String(random(0xffff), HEX);

    if (client.connect(clientId.c_str())) {
      Serial.println("connected");
      return true;
    } else {
      Serial.print("failed, rc=");
      Serial.print(client.state());
      Serial.println(" will retry");
      return false;
    }
  }
  return true;
}

void loop() {
  unsigned long currentMillis = millis();

  // If WiFi is required but not connected, show error and halt
  if (wifiRequired && !wifiConnected) {
    if (displayInitialized) {
      if (currentMillis % 2000 < 1000) {
        tft.setTextColor(HEART_COLOR, BACKGROUND);
      } else {
        tft.setTextColor(BACKGROUND, BACKGROUND);
      }
      tft.setTextSize(2);
      tft.setCursor(60, 120);
      tft.print("WiFi Required!");
      tft.setCursor(60, 150);
      tft.print("System Halted");
    }
    delay(100);
    return;
  }

  if (wifiConnected && client.connected()) {
  client.loop();
}

  // Update ECG
  if (currentMillis - lastEcgUpdate >= ECG_UPDATE_INTERVAL) {
    lastEcgUpdate = currentMillis;
    generateEcgPoint();
    drawEcg();
  }

  // Read sensors and update display every 2 seconds
  if (currentMillis - lastVitalUpdateTime >= vitalUpdateInterval) {
    lastVitalUpdateTime = currentMillis;
    readSensors();
    updateVitalsDisplay();
    xTaskNotifyGive(mqttPublishTaskHandle); // Notify publish task
  }
}