#define F_CPU 14745600
#define BAUD_PRESCALE 23

#define DAY_IN_SECONDS 86400L
#define I_LAYER_COUNT 17
#define J_LAYER_COUNT 16
#define WEIGHT_RESOLUTION 32.0F
#define BASE_ETA 0.01 * WEIGHT_RESOLUTION
#define FORECAST_PERIOD 1800
#define TRAINING_PERIODS 3

#include <stdio.h>
#include <stdbool.h>
#include <stdlib.h>
#include <math.h>
#include <avr/io.h>
#include <avr/interrupt.h>
#include <avr/wdt.h>
#include <util/delay.h>
#include "wti_serial.h"
#include "wti_ringbuffer.h"
#include "wti_switch.h"

typedef enum {SUCCESS, CRC_MISMATCH} network_restore_status;

void initializeHardware();
void initializeNeuralNetwork();
int8_t randInt8();
float sig(float);
void calculateSensorValues(float *sensors, int32_t theTime, int32_t lastCoffeeTime);
void calculateHiddenLayerValues(float *hiddenNodes, float *sensors);
void calculateOutput(float *output, float *hiddenNodes);
float forecast(int32_t theTime, int32_t lastCoffeeTime);
void train(float target, int32_t theTime, int32_t lastCoffeeTime, float eta);
void randAdd(int8_t *integer, float flt);
void lightLights(float forecast);
void eeWrite(uint16_t*, uint8_t);
uint8_t eeRead(uint16_t*);
void backupNetwork();
network_restore_status restoreNetwork();

static FILE serialStdout = FDEV_SETUP_STREAM(serialWrite, NULL, _FDEV_SETUP_WRITE);

volatile int32_t timerCounter;
volatile int16_t forecastTimer = 10;
volatile int16_t trainingTimer = 60;
volatile int16_t backupTimer = 7237;

// Network Data
int8_t ijMatrix[I_LAYER_COUNT][J_LAYER_COUNT];
int8_t jkMatrix[J_LAYER_COUNT];
int8_t jBias[J_LAYER_COUNT];
int8_t kBias;

int main(void) {
  // Training Data
  int32_t tLastCoffeeTimes[TRAINING_PERIODS] = {0};
  int32_t tTheTimes[TRAINING_PERIODS] = {0};
  int32_t lastCoffeeTime = 0;
  float tTargets[TRAINING_PERIODS] = {0.1};
  float eta;

  initializeHardware();

  _delay_ms(1000);
  printf("7");
  _delay_ms(1000);
  printf("6");
  _delay_ms(1000);
  printf("5");
  _delay_ms(1000);
  printf("4");
  _delay_ms(1000);
  printf("3");
  _delay_ms(1000);
  printf("2");
  _delay_ms(1000);
  printf("1");
  _delay_ms(1000);
  printf("0\n");
  if (restoreNetwork() != SUCCESS) {
    _delay_ms(3000);
    initializeNeuralNetwork();
    backupNetwork();
  }

  while(1) {
    wdt_reset();
    if ((PINC & (1<<PINC0)) == 0) {
      for (uint8_t k = 0; k < TRAINING_PERIODS; k++) {
        tTargets[k] = 0.9;
      }

      lastCoffeeTime = timerCounter;
      printf("Button Press.\n");
    }

    if (trainingTimer == 0) {
      trainingTimer = FORECAST_PERIOD / TRAINING_PERIODS;

      eta = BASE_ETA + 0.4 * WEIGHT_RESOLUTION * exp(-(double)timerCounter/(DAY_IN_SECONDS * 2));
      train(tTargets[0], tTheTimes[0], tLastCoffeeTimes[0], eta);

      for (uint8_t k = 1; k < TRAINING_PERIODS; k++) {
        tTargets[k - 1] = tTargets[k];
        tTheTimes[k - 1] = tTheTimes[k];
        tLastCoffeeTimes[k - 1] = tLastCoffeeTimes[k];
      }

      tTargets[TRAINING_PERIODS - 1] = 0.1;
      tTheTimes[TRAINING_PERIODS - 1] = timerCounter;
      tLastCoffeeTimes[TRAINING_PERIODS - 1] = lastCoffeeTime;

      printf("\nijMatrix:\n");
      for (uint8_t i = 0; i < I_LAYER_COUNT; i++) {
        for (uint8_t j = 0; j < J_LAYER_COUNT; j++) {
          printf("%i, ", ijMatrix[i][j]);
        }
        printf("\n");
      }

      printf("\njBiases:\n");
      for (uint8_t j = 0; j < J_LAYER_COUNT; j++) {
        printf("%i, ", jBias[j]);
      }
      printf("\n");

      printf("\njkMatrix:\n");
      for (uint8_t j = 0; j < J_LAYER_COUNT; j++) {
        printf("%i, ", jkMatrix[j]);
      }
      printf("\n\nkBias: %i\n\n", kBias);
    }

    if (forecastTimer == 0) {
      float currentForecast = forecast(timerCounter, lastCoffeeTime);
      forecastTimer = 30;
      printf("Current forecast: %f at %li, %li\n", currentForecast, (timerCounter % 86400), (lastCoffeeTime % 86400));
      lightLights(currentForecast);
    }

    if (backupTimer == 0) {
      backupTimer = 7200;
      backupNetwork();
    }

    _delay_ms(100);
  }
}

void initializeHardware() {
  cli();
  stdout = &serialStdout;

  //Initialize USART, its buffer and TWI
  usartInit(BAUD_PRESCALE);
  rbInit(&serialTxBuffer);

  // turn off Timer/Counter2, Timer/Counter0, Serial Peripheral Interface, ADC
  PRR = (1<<PRTIM2) | (1<<PRTIM0) | (1<<PRSPI) | (1<<PRADC);

  MCUSR &= ~(1<<WDRF); //Disable System Reset Flag in MCU Status Register
  wdt_reset(); //C.f. doc note on page 53.
  //"The Watchdog Timer should be reset before any change of the WDP bits"
  WDTCSR = (1<<WDCE) | (1<<WDE); //Write logic one to WDCE and WDE bits of WDTCSR simultaneously
  //to begin wd config procedure
  WDTCSR = (1<<WDP3) | (1<<WDP0);  //  Set WD prescaler for 8.0 seconds

  OCR1A  = 0x3840; // Compare match at 1.00 seconds
  TCCR1B = (1<<CS12) | (1<<CS10);
  TIMSK1 = (1<<OCIE1A) | (1<<TOIE1);

  PORTC = (1<<PORTC0) | (1<<PORTC1) | (1<<PORTC2) | (1<<PORTC3);
  DDRC = (1<<PORTC1) | (1<<PORTC2) | (1<<PORTC3);

  // Set up EEPROM
  EECR = 0x00; // Mode = atomic, Interrupt = disabled
  sei();
}

void initializeNeuralNetwork() {
  srandom(252);

  for (int j = 0; j < J_LAYER_COUNT; j++) {
    for (int i = 0; i < I_LAYER_COUNT; i++) {
      ijMatrix[i][j] = randInt8() / 4;
    }
  }

  for (int j = 0; j < J_LAYER_COUNT; j++) {
    jkMatrix[j] = randInt8() / 4;
    jBias[j] = randInt8() / 4;

    if(jkMatrix[j] == 0) {
      jkMatrix[j] = 4;
    }
  }

  kBias = randInt8() / 4;

  for (int32_t tTime = 0; tTime < 86400 * 6; tTime += 60 * 5) {
    train(0.3, tTime, tTime - 3600, 0.1 * WEIGHT_RESOLUTION);
  }
}

int8_t randInt8() {
  return (rand() % 128) - 64;
}

float readClockSensorSin(int32_t theTime, int32_t period) {
  return sin(2 * M_PI * (double)(theTime % period) / period);
}

float readClockSensorCos(int32_t theTime, int32_t period) {
  return cos(2 * M_PI * (double)(theTime % period) / period);
}

float sig(float x) {
  return 1 / (1 + exp(-x));
}

void calculateSensorValues(float *sensors, int32_t theTime, int32_t lastCoffeeTime) {
  sensors[0] = readClockSensorSin(theTime, DAY_IN_SECONDS);
  sensors[1] = readClockSensorCos(theTime, DAY_IN_SECONDS);

  sensors[2] = readClockSensorSin(theTime, DAY_IN_SECONDS/2);
  sensors[3] = readClockSensorCos(theTime, DAY_IN_SECONDS/2);

  sensors[4] = readClockSensorSin(theTime, DAY_IN_SECONDS/3);
  sensors[5] = readClockSensorCos(theTime, DAY_IN_SECONDS/3);

  sensors[6] = readClockSensorSin(theTime, DAY_IN_SECONDS/4);
  sensors[7] = readClockSensorCos(theTime, DAY_IN_SECONDS/4);

  sensors[8] = readClockSensorSin(theTime, DAY_IN_SECONDS/6);
  sensors[9] = readClockSensorCos(theTime, DAY_IN_SECONDS/6);

  sensors[10] = readClockSensorSin(theTime, DAY_IN_SECONDS/8);
  sensors[11] = readClockSensorCos(theTime, DAY_IN_SECONDS/8);

  sensors[12] = readClockSensorSin(theTime, DAY_IN_SECONDS/12);
  sensors[13] = readClockSensorCos(theTime, DAY_IN_SECONDS/12);

  sensors[14] = exp(-(double)(theTime - lastCoffeeTime) / ((double)60L * 60 * 2));
  sensors[15] = exp(-(double)(theTime - lastCoffeeTime) / ((double)60L * 60 * 8));
  sensors[16] = exp(-(double)(theTime - lastCoffeeTime) / ((double)60L * 60 * 24));
}

void calculateHiddenLayerValues(float *hiddenNodes, float *sensors) {
  for (uint8_t j = 0; j < J_LAYER_COUNT; j++) {
    hiddenNodes[j] = (float)jBias[j] / WEIGHT_RESOLUTION;

    for (uint8_t i = 0; i < I_LAYER_COUNT; i++) {
      hiddenNodes[j] += sensors[i] * (float)ijMatrix[i][j] / WEIGHT_RESOLUTION;
    }

    hiddenNodes[j] = sig(hiddenNodes[j]);
  }
}

void calculateOutput(float *output, float *hiddenNodes) {
  *output = (float)kBias / WEIGHT_RESOLUTION;
  for (uint8_t j = 0; j < J_LAYER_COUNT; j++) {
    *output += hiddenNodes[j] * (float)jkMatrix[j] / WEIGHT_RESOLUTION;
  }
  *output = sig(*output);
}

float forecast(int32_t theTime, int32_t lastCoffeeTime) {
  float sensors[I_LAYER_COUNT];
  float hiddenNodes[J_LAYER_COUNT];
  float output;

  calculateSensorValues(sensors, theTime, lastCoffeeTime);
  calculateHiddenLayerValues(hiddenNodes, sensors);
  calculateOutput(&output, hiddenNodes);

  return output;
}

void train (float target, int32_t theTime, int32_t lastCoffeeTime, float eta) {
  float sensors[I_LAYER_COUNT];
  float hiddenNodes[J_LAYER_COUNT];
  float deltaJs[J_LAYER_COUNT];
  float output, deltaK;

  calculateSensorValues(sensors, theTime, lastCoffeeTime);
  calculateHiddenLayerValues(hiddenNodes, sensors);
  calculateOutput(&output, hiddenNodes);

  printf("\nTraining:\n");
  printf("Target: %f, Output: %f\n", target, output);

  deltaK = output * (1 - output) * (output - target);

  for (uint8_t j = 0; j < J_LAYER_COUNT; j++) {
    deltaJs[j] = hiddenNodes[j] * (1 - hiddenNodes[j]) * deltaK * jkMatrix[j];
  }

  for (uint8_t j = 0; j < J_LAYER_COUNT; j++) {
    for (uint8_t i = 0; i < I_LAYER_COUNT; i++) {
      randAdd(&ijMatrix[i][j], -eta * deltaJs[j] * sensors[i]);
    }
  }

  for (uint8_t j = 0; j < J_LAYER_COUNT; j++) {
    randAdd(&jBias[j], -eta * deltaJs[j]);
  }

  for (uint8_t j = 0; j < J_LAYER_COUNT; j++) {
    randAdd(&jkMatrix[j], -eta * deltaK * hiddenNodes[j]);
  }

  randAdd(&kBias, -eta * deltaK);

  calculateSensorValues(sensors, theTime, lastCoffeeTime);
  calculateHiddenLayerValues(hiddenNodes, sensors);
  calculateOutput(&output, hiddenNodes);
  printf("\nUpdated Output: %f\n", output);
}

void randAdd(int8_t *integer, float flt) {
  int8_t sign;
  int16_t result = *integer;

  if (flt < 0) {
    flt = -flt;
    sign = -1;
  } else {
    sign = 1;
  }

  while (flt > 1) {
    flt--;
    result += sign;
    printf("|");
  }

  if (flt * 1000 >= rand() % 1000) {
    printf(".");
    result += sign;
  }

  if (result > 127) {
    result = 127;
  } else if (result < -128) {
    result = -128;
  }

  *integer = (int8_t)result;
}

void lightLights(float forecast) {
  PORTC &= ~((1<<PORTC1) | (1<<PORTC2) | (1<<PORTC3));

  if (forecast >= 0.05 && 0.15 > forecast) {
    PORTC |= (1 << PORTC1);
  }

  if (forecast >= 0.15 && 0.25 > forecast) {
    PORTC |= (1 << PORTC2);
  }

  if (forecast >= 0.25) {
    PORTC |= (1 << PORTC3);
  }
}

ISR(TIMER1_COMPA_vect) {
  TCNT1 = 0x0000;
  timerCounter++;

  if (trainingTimer > 0) {
    trainingTimer--;
  }

  if (forecastTimer > 0) {
    forecastTimer--;
  }

  if (backupTimer > 0) {
    backupTimer--;
  }
}

void eeWrite(uint16_t *uiAddress, uint8_t ucData) {
  cli();
  /* Wait for completion of previous write */
  while(EECR & (1<<EEPE)) ;

  /* Set up address and Data Registers */
  EEAR = *uiAddress;
  EEDR = ucData;
  /* Write logical one to EEMPE */
  EECR |= (1<<EEMPE);
  /* Start eeprom write by setting EEPE */
  EECR |= (1<<EEPE);
  (*uiAddress)++;
  sei();
}

uint8_t eeRead(uint16_t *uiAddress) {
  /* Wait for completion of previous write */
  while(EECR & (1<<EEPE)) ;
  /* Set up address register */
  EEAR = *uiAddress;
  /* Start eeprom read by writing EERE */
  EECR |= (1<<EERE);
  (*uiAddress)++;
  /* Return data from Data Register */
  return EEDR;
}

uint16_t crc16(uint16_t crc, uint8_t data) {
  crc = crc ^ (uint16_t)data;

  for (uint8_t i = 8; i > 0; i--) {
    if (crc & 0x0001) {
      crc = (crc >> 1)^0x8408;
    } else {
      crc >>= 1;
    }
  }

  return crc;
}

void backupNetwork() {
  uint16_t eeAddress = 0;
  uint16_t crc = 0xFFFF;

  for (int j = 0; j < J_LAYER_COUNT; j++) {
    for (int i = 0; i < I_LAYER_COUNT; i++) {
      eeWrite(&eeAddress, ijMatrix[i][j]);
      crc = crc16(crc, ijMatrix[i][j]);
    }
  }

  for (int j = 0; j < J_LAYER_COUNT; j++) {
    eeWrite(&eeAddress, jkMatrix[j]);
    crc = crc16(crc, jkMatrix[j]);
  }

  for (int j = 0; j < J_LAYER_COUNT; j++) {
    eeWrite(&eeAddress, jBias[j]);
    crc = crc16(crc, jBias[j]);
  }

  eeWrite(&eeAddress, kBias);
  crc = crc16(crc, kBias);

  printf("New CRC: %x\n", crc);
  eeWrite(&eeAddress, (crc>>8));
  eeWrite(&eeAddress, crc);
}

network_restore_status restoreNetwork() {
  uint16_t eeAddress = 0;
  uint16_t crc = 0xFFFF;
  uint16_t storedCrc = 0;

  for (int j = 0; j < J_LAYER_COUNT; j++) {
    for (int i = 0; i < I_LAYER_COUNT; i++) {
      ijMatrix[i][j] = eeRead(&eeAddress);
      crc = crc16(crc, ijMatrix[i][j]);
    }
  }

  for (int j = 0; j < J_LAYER_COUNT; j++) {
    jkMatrix[j] = eeRead(&eeAddress);
    crc = crc16(crc, jkMatrix[j]);
  }

  for (int j = 0; j < J_LAYER_COUNT; j++) {
    jBias[j] = eeRead(&eeAddress);
    crc = crc16(crc, jBias[j]);
  }

  kBias = eeRead(&eeAddress);
  crc = crc16(crc, kBias);

  storedCrc |= (uint16_t)eeRead(&eeAddress);
  storedCrc <<= 8;
  storedCrc |= (uint16_t)eeRead(&eeAddress);
  if (crc != storedCrc) {
    printf("Restore Failed. %x != %x\n", crc, storedCrc);
    return CRC_MISMATCH;
  } else {
    printf("Network restored successfully.\n");
    return SUCCESS;
  }
}