/*
 * AI Variable Voltage Controller
 * PID + Neural Network Implementation
 * Author: Global Byte Shop Thailand
 */

#include <WiFi.h>
#include <Wire.h>
#include <Adafruit_SSD1306.h>
#include <ArduinoJson.h>
#include <math.h>

// Display configuration
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Pin definitions
#define VOLTAGE_SENSOR_PIN A0
#define CURRENT_SENSOR_PIN A1
#define PWM_OUTPUT_PIN 9
#define ENCODER_CLK_PIN 2
#define ENCODER_DT_PIN 3
#define ENCODER_SW_PIN 4

// AI Controller parameters
struct PIDController {
  float kp = 2.0;
  float ki = 0.5;
  float kd = 0.1;
  float integral = 0.0;
  float previous_error = 0.0;
  float output = 0.0;
};

struct NeuralNetwork {
  // Network architecture: 3 inputs, 2 hidden layers (10, 5), 1 output
  float weights_ih[3][10];  // Input to hidden layer 1
  float weights_h1h2[10][5]; // Hidden layer 1 to 2
  float weights_ho[5][1];   // Hidden layer 2 to output
  float bias_h1[10];
  float bias_h2[5];
  float bias_o[1];
  float learning_rate = 0.001;
};

// Global variables
PIDController pid;
NeuralNetwork nn;
float target_voltage = 12.0;
float current_voltage = 0.0;
float current_current = 0.0;
float pwm_output = 0.0;
unsigned long last_update = 0;
unsigned long last_learning = 0;

// Data logging
struct DataPoint {
  unsigned long timestamp;
  float voltage;
  float current;
  float target;
  float error;
  float pwm;
};

DataPoint data_buffer[100];
int data_index = 0;

void setup() {
  Serial.begin(115200);
  Serial.println("🚀 AI Variable Voltage Controller Starting...");
  
  // Initialize pins
  pinMode(PWM_OUTPUT_PIN, OUTPUT);
  pinMode(ENCODER_CLK_PIN, INPUT_PULLUP);
  pinMode(ENCODER_DT_PIN, INPUT_PULLUP);
  pinMode(ENCODER_SW_PIN, INPUT_PULLUP);
  
  // Initialize display
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("❌ SSD1306 allocation failed");
    for(;;);
  }
  
  // Initialize neural network
  initializeNeuralNetwork();
  
  // Initialize WiFi for data logging
  WiFi.begin("YourWiFi", "YourPassword");
  
  Serial.println("✅ AI Controller initialized successfully!");
  
  // Show startup screen
  showStartupScreen();
  delay(2000);
}

void loop() {
  unsigned long current_time = millis();
  
  // Read sensors (10ms cycle)
  if (current_time - last_update >= 10) {
    readSensors();
    
    // AI control processing (50ms cycle)
    if ((current_time - last_update) % 50 == 0) {
      float error = target_voltage - current_voltage;
      
      // PID calculation
      float pid_output = calculatePID(error);
      
      // Neural network prediction
      float nn_output = neuralNetworkPredict(current_voltage, current_current, target_voltage);
      
      // Fusion of PID and NN outputs
      pwm_output = fusionControl(pid_output, nn_output);
      
      // Apply PWM output
      analogWrite(PWM_OUTPUT_PIN, constrain(pwm_output * 255, 0, 255));
      
      // Learning phase (1000ms cycle)
      if (current_time - last_learning >= 1000) {
        neuralNetworkLearn(error);
        last_learning = current_time;
      }
      
      // Log data
      logData(current_time, current_voltage, current_current, target_voltage, error, pwm_output);
    }
    
    last_update = current_time;
  }
  
  // Handle user input
  handleEncoder();
  
  // Update display (100ms cycle)
  if (current_time % 100 == 0) {
    updateDisplay();
  }
  
  // Send data via WiFi (5000ms cycle)
  if (current_time % 5000 == 0 && WiFi.status() == WL_CONNECTED) {
    sendDataToCloud();
  }
}

void readSensors() {
  // Read voltage sensor (0-25V range)
  int voltage_raw = analogRead(VOLTAGE_SENSOR_PIN);
  current_voltage = (voltage_raw / 1023.0) * 25.0;
  
  // Apply moving average filter
  static float voltage_buffer[10] = {0};
  static int voltage_index = 0;
  voltage_buffer[voltage_index] = current_voltage;
  voltage_index = (voltage_index + 1) % 10;
  
  float voltage_sum = 0;
  for (int i = 0; i < 10; i++) {
    voltage_sum += voltage_buffer[i];
  }
  current_voltage = voltage_sum / 10.0;
  
  // Read current sensor (ACS712 30A)
  int current_raw = analogRead(CURRENT_SENSOR_PIN);
  current_current = ((current_raw / 1023.0) * 5.0 - 2.5) / 0.066; // 66mV/A sensitivity
  
  // Apply calibration
  current_voltage = current_voltage * 1.02 - 0.05; // Calibration factors
  current_current = abs(current_current); // Absolute value for current
}

float calculatePID(float error) {
  // Proportional term
  float proportional = pid.kp * error;
  
  // Integral term
  pid.integral += error * 0.01; // 10ms sampling time
  pid.integral = constrain(pid.integral, -10, 10); // Anti-windup
  float integral = pid.ki * pid.integral;
  
  // Derivative term
  float derivative = pid.kd * (error - pid.previous_error) / 0.01;
  pid.previous_error = error;
  
  // PID output
  pid.output = proportional + integral + derivative;
  return constrain(pid.output, 0, 1);
}

void initializeNeuralNetwork() {
  // Initialize weights with random values
  for (int i = 0; i < 3; i++) {
    for (int j = 0; j < 10; j++) {
      nn.weights_ih[i][j] = random(-100, 100) / 100.0;
    }
  }
  
  for (int i = 0; i < 10; i++) {
    for (int j = 0; j < 5; j++) {
      nn.weights_h1h2[i][j] = random(-100, 100) / 100.0;
    }
  }
  
  for (int i = 0; i < 5; i++) {
    nn.weights_ho[i][0] = random(-100, 100) / 100.0;
  }
  
  // Initialize biases
  for (int i = 0; i < 10; i++) nn.bias_h1[i] = 0.0;
  for (int i = 0; i < 5; i++) nn.bias_h2[i] = 0.0;
  nn.bias_o[0] = 0.0;
  
  Serial.println("✅ Neural Network initialized");
}

float sigmoid(float x) {
  return 1.0 / (1.0 + exp(-x));
}

float neuralNetworkPredict(float voltage, float current, float target) {
  // Normalize inputs
  float inputs[3] = {
    voltage / 25.0,        // Normalize to [0,1]
    current / 30.0,        // Normalize to [0,1]
    target / 25.0          // Normalize to [0,1]
  };
  
  // Hidden layer 1
  float hidden1[10];
  for (int i = 0; i < 10; i++) {
    float sum = nn.bias_h1[i];
    for (int j = 0; j < 3; j++) {
      sum += inputs[j] * nn.weights_ih[j][i];
    }
    hidden1[i] = sigmoid(sum);
  }
  
  // Hidden layer 2
  float hidden2[5];
  for (int i = 0; i < 5; i++) {
    float sum = nn.bias_h2[i];
    for (int j = 0; j < 10; j++) {
      sum += hidden1[j] * nn.weights_h1h2[j][i];
    }
    hidden2[i] = sigmoid(sum);
  }
  
  // Output layer
  float output = nn.bias_o[0];
  for (int i = 0; i < 5; i++) {
    output += hidden2[i] * nn.weights_ho[i][0];
  }
  
  return sigmoid(output);
}

void neuralNetworkLearn(float error) {
  // Simple gradient descent learning
  // This is a simplified version - in practice, you'd implement full backpropagation
  
  float learning_signal = error * nn.learning_rate;
  
  // Adjust output layer weights
  for (int i = 0; i < 5; i++) {
    nn.weights_ho[i][0] += learning_signal * 0.1;
  }
  
  // Adjust learning rate based on performance
  static float previous_error = 0;
  if (abs(error) < abs(previous_error)) {
    nn.learning_rate *= 1.01; // Increase learning rate if improving
  } else {
    nn.learning_rate *= 0.99; // Decrease if getting worse
  }
  nn.learning_rate = constrain(nn.learning_rate, 0.0001, 0.01);
  previous_error = error;
}

float fusionControl(float pid_output, float nn_output) {
  // Adaptive fusion based on system performance
  static float pid_weight = 0.7;
  static float nn_weight = 0.3;
  
  // Adjust weights based on recent performance
  static float recent_errors[10] = {0};
  static int error_index = 0;
  
  float current_error = abs(target_voltage - current_voltage);
  recent_errors[error_index] = current_error;
  error_index = (error_index + 1) % 10;
  
  float avg_error = 0;
  for (int i = 0; i < 10; i++) {
    avg_error += recent_errors[i];
  }
  avg_error /= 10.0;
  
  // If error is decreasing, trust NN more
  if (current_error < avg_error) {
    nn_weight = min(nn_weight + 0.01, 0.8);
    pid_weight = 1.0 - nn_weight;
  } else {
    pid_weight = min(pid_weight + 0.01, 0.8);
    nn_weight = 1.0 - pid_weight;
  }
  
  return (pid_output * pid_weight) + (nn_output * nn_weight);
}

void handleEncoder() {
  static int last_clk = HIGH;
  static unsigned long last_button_press = 0;
  
  int clk = digitalRead(ENCODER_CLK_PIN);
  int dt = digitalRead(ENCODER_DT_PIN);
  
  if (clk != last_clk && clk == LOW) {
    if (dt == HIGH) {
      target_voltage += 0.1;
    } else {
      target_voltage -= 0.1;
    }
    target_voltage = constrain(target_voltage, 0, 20);
  }
  last_clk = clk;
  
  // Button press for menu
  if (digitalRead(ENCODER_SW_PIN) == LOW && millis() - last_button_press > 500) {
    // Toggle display mode or enter settings
    last_button_press = millis();
  }
}

void updateDisplay() {
  display.clearDisplay();
  
  // Title
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.println("AI Voltage Controller");
  
  // Current values
  display.setCursor(0, 12);
  display.print("Target: ");
  display.print(target_voltage, 2);
  display.println("V");
  
  display.setCursor(0, 22);
  display.print("Output: ");
  display.print(current_voltage, 2);
  display.println("V");
  
  display.setCursor(0, 32);
  display.print("Current: ");
  display.print(current_current, 2);
  display.println("A");
  
  display.setCursor(0, 42);
  display.print("PWM: ");
  display.print(pwm_output * 100, 1);
  display.println("%");
  
  // Error and learning status
  float error = target_voltage - current_voltage;
  display.setCursor(0, 52);
  display.print("Error: ");
  display.print(error, 3);
  display.print("V LR:");
  display.print(nn.learning_rate, 4);
  
  display.display();
}

void logData(unsigned long timestamp, float voltage, float current, float target, float error, float pwm) {
  data_buffer[data_index] = {timestamp, voltage, current, target, error, pwm};
  data_index = (data_index + 1) % 100;
}

void sendDataToCloud() {
  if (WiFi.status() != WL_CONNECTED) return;
  
  // Create JSON payload
  StaticJsonDocument<200> doc;
  doc["timestamp"] = millis();
  doc["voltage"] = current_voltage;
  doc["current"] = current_current;
  doc["target"] = target_voltage;
  doc["error"] = target_voltage - current_voltage;
  doc["pwm"] = pwm_output;
  doc["learning_rate"] = nn.learning_rate;
  
  String jsonString;
  serializeJson(doc, jsonString);
  
  // Send to cloud service (implement your preferred method)
  Serial.println("📡 Data sent: " + jsonString);
}

void showStartupScreen() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(10, 10);
  display.println("AI VVC");
  
  display.setTextSize(1);
  display.setCursor(0, 35);
  display.println("Variable Voltage");
  display.setCursor(0, 45);
  display.println("Controller v2.0");
  display.setCursor(0, 55);
  display.println("Global Byte Shop");
  
  display.display();
}

// Utility functions
float getSystemEfficiency() {
  float power_in = 24.0 * current_current; // Assuming 24V input
  float power_out = current_voltage * current_current;
  return (power_out / power_in) * 100.0;
}

void printSystemStats() {
  Serial.println("=== System Statistics ===");
  Serial.print("Target Voltage: "); Serial.println(target_voltage);
  Serial.print("Output Voltage: "); Serial.println(current_voltage);
  Serial.print("Output Current: "); Serial.println(current_current);
  Serial.print("PWM Output: "); Serial.println(pwm_output * 100);
  Serial.print("System Efficiency: "); Serial.println(getSystemEfficiency());
  Serial.print("Learning Rate: "); Serial.println(nn.learning_rate);
  Serial.println("========================");
}