#include <OneWire.h>
#include <DallasTemperature.h>
#include <Adafruit_HX711.h>
#include <ArduinoJson.h>
#include <avr/wdt.h>

const uint8_t hx711_data = A3;
const uint8_t hx711_clk = A2;

Adafruit_HX711 hx711(hx711_data, hx711_clk);

const uint8_t pump_relay = A1;
const uint8_t solenoid_relay = A0;

const uint8_t pump_tacho = 2;
const uint8_t flow_tacho = 3;
const uint8_t fans_tacho = 8;
const uint8_t temp_1wire = 4;

const uint8_t pump_pwm = 5;
const uint8_t fans_pwm = 10;

volatile uint8_t default_pump_speed = 255;
volatile uint8_t default_fan_speed = 250;

volatile uint8_t pump_pulse_count = 0;
volatile uint8_t flow_pulse_count = 0;
volatile uint8_t fans_pulse_count = 0;

DeviceAddress air_therm = { 0x28, 0xCC, 0xAE, 0x46, 0xD4, 0x1B, 0x7D, 0xC9 };
DeviceAddress water_therm = { 0x28, 0xD8, 0x79, 0x46, 0xD4, 0x76, 0x2C, 0x17 };
const uint8_t temp_precision = 12;

OneWire oneWire(temp_1wire);
DallasTemperature sensors(&oneWire);

//we will use this function to periodically self-reset to avoid timer overflows
void(* resetFunc) (void) = 0;

void setup() {
  Serial.begin(115200);

  sensors.begin();
  sensors.setResolution(water_therm, temp_precision);
  sensors.setResolution(air_therm, temp_precision);
  sensors.setWaitForConversion(false);

  hx711.begin();

  //set pwm timer to highest frequency
  TCCR1B = (TCCR1B & 0b11111000) | 0x01;

  pinMode(pump_pwm, OUTPUT);
  pinMode(fans_pwm, OUTPUT);
  analogWrite(pump_pwm, default_pump_speed);
  analogWrite(fans_pwm, default_fan_speed);

  pinMode(pump_tacho, INPUT);
  attachInterrupt(digitalPinToInterrupt(pump_tacho), count_pump_pulse, FALLING);

  pinMode(flow_tacho, INPUT);
  attachInterrupt(digitalPinToInterrupt(flow_tacho), count_flow_pulse, FALLING);

  pinMode(fans_tacho, INPUT);


  pinMode(pump_relay, OUTPUT);
  pinMode(solenoid_relay, OUTPUT);

  wdt_enable(WDTO_2S);

  Serial.print("Hello!");
  Serial.println();
}

//value limits correspond to +1 mbar and -700mbar. the sensor reports ~30k at gain 64 per 100mbar. all an approximation, don't need super high accuracy.
float pressure_raw_to_mbar(int32_t pressure_raw) {
  return (pressure_raw - 390000) * (1700.0 / (5600000 - 390000)) - 700;
}

void count_pump_pulse() {
  pump_pulse_count++;
}

void count_flow_pulse() {
  flow_pulse_count++;
}

void count_fans_pulse(uint32_t* last_fans_tacho_pulse, uint32_t* cur_loop_timestamp) {
  static bool last_state = HIGH;
  static bool cur_state;

  cur_state = digitalRead(fans_tacho);

  if (cur_state == LOW && last_state == HIGH) {
    fans_pulse_count++;
    *last_fans_tacho_pulse = *cur_loop_timestamp;
    while (digitalRead(fans_tacho) == LOW);
    last_state = HIGH;
  }
  
  last_state = cur_state;
}

float pump_ema(uint8_t new_reading, uint32_t* last_timestamp, uint32_t* cur_timestamp) {
  static const float alpha = 0.8;
  static float ema = new_reading;

  ema = alpha * (new_reading  * (60000 / (*cur_timestamp - *last_timestamp))) + (1 - alpha) * ema;

  return ema;
}

//ema calculation to get flowrate in liters per hour with an aquacomputer flow monitor (approximate)
float flow_ema(uint8_t new_reading, uint32_t* last_timestamp, uint32_t* cur_timestamp) {
  static const float alpha = 0.8;
  static float ema;

  ema = alpha * (new_reading  * (3600000 / (*cur_timestamp - *last_timestamp)) / 170) + (1 - alpha) * ema;

  return ema;
}

int fan_PID(float* air_temp, float* water_temp, uint32_t* cur_loop_timestamp) {
  static const float kp = 120.0;
  static const float ki = 0.16;
  static const float kd = 4.0;

  static float integral = 0;
  static float derivative = 0;
  static float last_error = 0;
  static float error;
  static float delta;
  static float last_time = *cur_loop_timestamp;
  static const float min_water_temp = 5.0;
  static const float max_water_temp = 40.0;
  static const uint8_t min_fan_speed = 90;
  static const uint8_t max_fan_speed = 255;

  static const uint8_t temp_target_offset = 4;
  static const uint8_t fan_offset = 10;

  static int16_t fan_speed;

  if (*water_temp >= max_water_temp) {
    return max_fan_speed;
  }

  else if (*water_temp <= min_water_temp) {
    return 0;
  }

  else {
    error = *water_temp - min(*air_temp + temp_target_offset, max_water_temp);
    delta = *cur_loop_timestamp - last_time;

    //mitigate unlimited integral windup
    if (fan_speed == max_fan_speed) {
      integral += error * delta;
    }

    if (*air_temp + temp_target_offset > *water_temp - 1) {
      integral = 0;
    }

    derivative = (error - last_error) / delta;
    fan_speed = round(constrain(min_fan_speed + fan_offset + (kp * error + ki * integral + kd * derivative), min_fan_speed, max_fan_speed));
    last_error = error;
    return fan_speed;
  }
}

void loop() {
  static uint32_t cur_loop_timestamp;
  static uint32_t serial_last_sent;
  static uint32_t last_temp_request_timestamp;
  static uint32_t low_pressure_confirmation_timestamp;
  static uint32_t last_fans_tacho_pulse;
  static uint32_t last_pressure_check;
  static float new_water_temp;
  static float new_air_temp;
  static float water_temp;
  static float air_temp;
  static bool sucking = false;
  static bool checking_low_pressure = false;
  static int16_t low_pressure_threshold = -500;
  static int16_t high_pressure_threshold = -600;
  static float loop_pressure = pressure_raw_to_mbar(hx711.readChannelRaw(CHAN_A_GAIN_64));
  static float new_loop_pressure = loop_pressure;
  static uint8_t fan_speed;

  static StaticJsonDocument<64> output_stats;
  static JsonObject temps = output_stats.createNestedObject("temps");
  static JsonObject tacho = output_stats.createNestedObject("tacho");

  cur_loop_timestamp = millis();

  //reset the system when approaching timer overflow
  if (cur_loop_timestamp >= 4000000000) {
    resetFunc();
  }

  if (cur_loop_timestamp - last_temp_request_timestamp >= 1000) {
    new_water_temp = sensors.getTempC(water_therm);
    new_air_temp = sensors.getTempC(air_therm);

    if (new_water_temp != -127) {
      water_temp = new_water_temp;
    }

    if (new_air_temp != -127) {
      air_temp = new_air_temp;
    }

    sensors.requestTemperatures();

    fan_speed = fan_PID(&air_temp, &water_temp, &cur_loop_timestamp);
    analogWrite(fans_pwm, fan_speed);

    output_stats["mbar"] = loop_pressure;

    temps["air"] = air_temp;
    temps["water"] = water_temp;

    tacho["fan"] = fan_speed;
    tacho["pump"] = pump_ema(pump_pulse_count, &last_temp_request_timestamp, &cur_loop_timestamp);
    tacho["flow"] = flow_ema(flow_pulse_count, &last_temp_request_timestamp, &cur_loop_timestamp);

    serializeJson(output_stats, Serial);
    Serial.println();

    fans_pulse_count = 0;
    pump_pulse_count = 0;
    flow_pulse_count = 0;

    last_temp_request_timestamp = cur_loop_timestamp;
    wdt_reset();
  }

  if (cur_loop_timestamp - last_pressure_check >= 100) {
    loop_pressure = pressure_raw_to_mbar(hx711.readChannelRaw(CHAN_A_GAIN_64)); 

    if (sucking == false) {
      if (loop_pressure > low_pressure_threshold) {
        if (checking_low_pressure == false) {
          checking_low_pressure = true;
          low_pressure_confirmation_timestamp = cur_loop_timestamp;
        }
        if (cur_loop_timestamp - low_pressure_confirmation_timestamp >= 1000) {
          digitalWrite(pump_relay, HIGH);
          digitalWrite(solenoid_relay, HIGH);
          sucking = true;
          checking_low_pressure = false;
        }
      }

      else if (cur_loop_timestamp - low_pressure_confirmation_timestamp >= 1000) {
        checking_low_pressure = false;
      }
    }

    else {
      if (loop_pressure < high_pressure_threshold) {
        digitalWrite(pump_relay, LOW);
        digitalWrite(solenoid_relay, LOW);
        sucking = false;
      }
    }
    last_pressure_check = cur_loop_timestamp;
  }
}