#include <EEPROM.h>

// Parameter
const int stateCount = 256; // Zustände entsprechen den Helligkeitsstufen von 0 bis 255
const int actionCount = 2; // Aktionen: 0 = Dunkler, 1 = Heller
const float learningRate = 0.01; //Gibt die Lernrate an
const float discountFactor = 0.9; //Gibt an wie stark zukünftige Belohnungen gewichtet werden
const float forgettingFactor = 0.95; // Vergessensfaktor (zwischen 0 und 1)
const int reward = 1;  // Belohnungswert  
const int penalty = -1; //Bestrafungswert
const int noiseLevelAction = 5; // Maximales Rauschen für Aktionen
const int noiseLevelReward = 1; // Maximales Rauschen für Belohnungen

// Variablen
int currentState = 1;
int action;
int newState;
int rewardValue;
int prevBrightness = 0;
unsigned long iterationCount = 0; // Zähler für Iterationen

// Pin für die LED
const int ledPin = 9;

void setup() {
  // Initialisiere LED Pin als Ausgang
  pinMode(ledPin, OUTPUT);
  analogWrite(ledPin, currentState);

  // Initialisiere Zufallszahlengenerator
  randomSeed(analogRead(0));

  // Initialisiere serielle Kommunikation
  Serial.begin(9600);
  while (!Serial) {
    ; // Warte auf serielle Verbindung
  }
  Serial.println("Bestärkendes Lernen mit Q-Learning auf Arduino gestartet");
}

void loop() {
  // Wenn der Zustand 255 erreicht wurde, beende die Optimierung
  if (currentState == 255) {
    Serial.print("Optimierung beendet nach ");
    Serial.print(iterationCount);
    Serial.println(" Iterationen.");
    while (true); // Stoppe die Schleife
  }

  // Wähle Aktion (e-greedy)
  action = chooseAction(currentState);

  // Führe Aktion aus und beobachte neuen Zustand
  newState = takeAction(currentState, action);

  // Erhalte Belohnung oder Bestrafung
  rewardValue = getReward(newState);

  // Update Q-Werte
  updateQ(currentState, action, rewardValue, newState);

  // Setze neuen Zustand als aktuellen Zustand
  currentState = newState;

  // Ausgabe der aktuellen Werte zur Überwachung
  Serial.print("Aktueller Zustand: ");
  Serial.print(currentState);
  Serial.print(", Aktion: ");
  Serial.print(action == 0 ? "Dunkler" : "Heller");
  Serial.print(", Belohnung: ");
  Serial.print(rewardValue);
  Serial.print(", Neuer Zustand: ");
  Serial.print(newState);
  Serial.print(", Q-Wert: ");
  Serial.println(getQ(currentState, action));

  // Zähle die Iterationen
  iterationCount++;

  // Warten
  delay(5);
}

int chooseAction(int state) {
 //Wahrscheinlichkeit für zufällige Aktion
  if (random(0, 10) < 1) {
    return random(0, actionCount);
  } 
  else { // Wenn Q(state,0) > Q(state,1) dann wird 0 zurückgegeben sonst 1
    return (getQ(state, 0) > getQ(state, 1)) ? 0 : 1;
  }
}

int takeAction(int state, int action) {
  int noise = random(-noiseLevelAction, noiseLevelAction + 1); // Zufällige Störgröße für die Aktion
  int brightness = state;

  if (action == 0) {
    brightness = max(0, state - 1 + noise);
  } else {
    brightness = min(255, state + 1 + noise);
  }

  // Begrenze den Wert von brightness auf 0 bis 255
  brightness = constrain(brightness, 0, 255);

  analogWrite(ledPin, brightness);
  return brightness;
}

int getReward(int state) {
  // Zufällige Störgröße für die Belohnung
  int rewardNoise = random(-noiseLevelReward, noiseLevelReward + 1);

  // Belohnung wenn Helligkeit zunimmt, Bestrafung wenn abnimmt
  if (state > prevBrightness) {
    prevBrightness = state;
    return reward + rewardNoise;
  } else if (state < prevBrightness) {
    prevBrightness = state;
    return penalty + rewardNoise;
  } else {
    return rewardNoise;
  }
}

float getQ(int state, int action) {
  // Liest Q-Wert aus dem EEPROM
  int address = (state * actionCount + action) * sizeof(byte);
  byte qByte = EEPROM.read(address);
  return (float)qByte / 255.0; // Skaliere von 0-255 auf 0-1
}

void updateQ(int currentstate, int action, int rewardValue, int newState) {
  // Q-Learning-Update
  float oldQ = getQ(currentstate, action);
  float maxQValue = maxQ(newState);
  float newQ = oldQ + learningRate * (rewardValue + discountFactor * maxQValue - oldQ);

  // Anwenden des Vergessensfaktors auf den neuen Q-Wert
  newQ *= forgettingFactor;

  // Speichere neuen Q-Wert im EEPROM, wenn er sich signifikant geändert hat
  if (abs(newQ - oldQ) > 0.01) { // Schwellenwert von 0.01
    int address = (currentstate * actionCount + action) * sizeof(byte);
    byte qByte = (byte)(newQ * 255.0); // Skaliere von 0-1 auf 0-255
    EEPROM.update(address, qByte);
  }
}

float maxQ(int state) {
  // Gibt den maximalen Q-Wert für einen Zustand zurück
  return max(getQ(state, 0), getQ(state, 1));
}


/*
Arbeitsauftrag:
Bitte verändert die folgenden Parameter und lasst das System dann jeweils 3-mal durchlaufen:

learningRate
discountFactor
forgettingFactor
reward
penalty

Notiert bitte auf dem Arbeitsblatt was ihr beobachtet habt.
*/
