#include <ESP8266WiFi.h>
#include <ModbusIP_ESP8266.h>

#include "DHT.h"
#include "config.h"
#include "rgb_lcd.h"
#include "LittleFS.h"

#define TCP_PORT 15000
// After how much time without connection we will 
// reboot (in millis)
#define NON_CONNECTED_REBOOT_TIMEOUT 60000 

#define SENSOR_UPDATE_INTERVAL 500
#define SCREEN_UPDATE_INTERVAL 200
#define STATE_SAVE_INTERVAL 5000
// Interval between allowed switch-states of the pump
#define PUMP_SWITCH_INTERVAL 5000
// Interval considering "rather long" to work for pump without pressure increase
#define PUMP_UNPRODUCTIVE_TIMEOUT 60000
// amount of pressure what we consider "productive increase" (in 0.01 of bar)
#define PRESSURE_PRODUCTIVE_INCREASE 50
// In fact, pump needs to be able to increase pressure by PRESSURE_PRODUCTIVE_INCREASE
// in PUMP_UNPRODUCTIVE_TIMEOUT time of working

// Interval between allowed switch-states of the flush valve
#define FLUSH_SWITCH_INTERVAL 1000

#define MODE_AUTO_ON 0
#define MODE_AUTO_OFF 1
#define MODE_MANUAL 2

#define REGISTER_CURRENT_TEMP 0
#define REGISTER_CURRENT_HUM 1
#define REGISTER_CURRENT_PRESS 2
#define REGISTER_PUMP_STATE 3
#define REGISTER_WORK_MODE 4
#define REGISTER_P_ON 5
#define REGISTER_P_OFF 6
#define REGISTER_T_ON 7
#define REGISTER_T_OFF 8
#define REGISTER_HEAT_STATE 10
#define REGISTER_PUMP_ERROR 11
#define REGISTER_FLUSH_STATE 12
#define REGISTER_FORCE_PUMP 13
#define REGISTER_FORCE_HEAT 14
#define REGISTER_FORCE_FLUSH 15

#define DHTPIN D3
#define DHTTYPE AM2301

#define PUMP_PIN D4
#define HEATER_PIN D0
#define FLUSH_PIN D6

class ModbusIPNoDelay : public ModbusIP
{
  public:
    void serverNoDelay(uint16_t port)
    {
      server(port);
      tcpserver->setNoDelay(true);
    }
};

ModbusIPNoDelay mb;
DHT dht(DHTPIN, DHTTYPE);
rgb_lcd lcd;

char id[64];

unsigned long last_update;

int p_on = 150;
int p_off = 250;
int cur_work_mode = MODE_AUTO_ON;
int cur_pump_state = 0;
int cur_p = 0;
int cur_t = 0;
int cur_hum = 0;
int t_on = 200; // 2 degrees
int t_off = 1500; // 15 degrees
int cur_heat_state = 0;
int cur_flush_state = 0;
int force_pump = 0;
int force_heat = 0;
int force_flush = 0;

bool error_pump = false;
bool error_heat = false;

int ReadInt(File &f)
{
  String s = f.readStringUntil('\n');
	return atoi(s.c_str());	
}

void LoadState()
{
	File f = LittleFS.open("/state.txt", "r");
	if (!f) return;
  t_on = ReadInt(f);
  t_off = ReadInt(f);
  p_on = ReadInt(f);
  p_off = ReadInt(f);
  cur_work_mode = ReadInt(f);
	f.close();
}

void SaveState()
{
	File f = LittleFS.open("/state.txt", "w");
	if (!f) return;
	f.println(t_on);
  f.println(t_off);
  f.println(p_on);
  f.println(p_off);
  f.println(cur_work_mode);
	f.close();
}

void setup() 
{
	if (!LittleFS.begin())
	{
		Serial.println("Unable to mount the filesystem");
		ESP.restart();
	}
	LoadState();
  pinMode(FLUSH_PIN, OUTPUT);
  pinMode(PUMP_PIN, OUTPUT);
	pinMode(HEATER_PIN, OUTPUT);
	pinMode(DHTPIN, INPUT);
	pinMode(A0, INPUT);

	lcd.begin(16, 2);
	lcd.setRGB(0, 255, 0);
	Serial.begin(115200);
	Serial.println("Initialized...");

	WiFi.mode(WIFI_STA);
	WiFi.begin(WIFI_NAME, WIFI_PASS);
	Serial.print("Connecting");
	while (WiFi.status() != WL_CONNECTED)
	{
		delay(500);
		Serial.print(".");
	}
	Serial.println();

	Serial.print("Connected, IP address: ");
	Serial.println(WiFi.localIP());
	uint8_t macAddr[6];
	WiFi.macAddress(macAddr);
	sprintf(id, "%02x:%02x:%02x:%02x:%02x:%02x", macAddr[0], macAddr[1], macAddr[2], macAddr[3], macAddr[4], macAddr[5]);
	Serial.printf("Connected, mac address: %s\n", id);
	// DHT sensor init
	dht.begin();
	// Modbus init	
	mb.serverNoDelay(15000);
	mb.addHreg(REGISTER_CURRENT_TEMP, cur_t);
	mb.addHreg(REGISTER_CURRENT_HUM, cur_hum);
	mb.addHreg(REGISTER_CURRENT_PRESS, cur_p);
	mb.addHreg(REGISTER_PUMP_STATE, cur_pump_state);
	mb.addHreg(REGISTER_WORK_MODE, cur_work_mode);
	mb.addHreg(REGISTER_P_ON, p_on);
	mb.addHreg(REGISTER_P_OFF, p_off);
	mb.addHreg(REGISTER_T_ON, t_on);
	mb.addHreg(REGISTER_T_OFF, t_off);
	mb.addHreg(REGISTER_HEAT_STATE, cur_heat_state);
	mb.addHreg(REGISTER_PUMP_ERROR, error_pump);
	mb.addHreg(REGISTER_FLUSH_STATE, cur_flush_state);
  mb.addHreg(REGISTER_FORCE_PUMP, force_pump);
  mb.addHreg(REGISTER_FORCE_HEAT, force_heat);
  mb.addHreg(REGISTER_FORCE_FLUSH, force_flush);
	// 
	Serial.print("Heater pin ");
	Serial.println(HEATER_PIN);
	Serial.print("Pump pin ");
	Serial.println(PUMP_PIN);
}


void UpdateDHTData()
{
	float h = dht.readHumidity();
	float t = dht.readTemperature();
	cur_hum = (int)(h * 100);
	cur_t = (int)(t * 100);
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Workings with pressure sensor
///////////////////////////////////////////////////////////////////////////////////////////////////

const int MEASURES_AMOUNT = 10;
int measures[MEASURES_AMOUNT];
int measure_pos = 0;

void UpdatePressure()
{
	int p_int = analogRead(A0);
  // we have 10-bit ADC, so get value between 0 and 1 and then multiply by max voltage
	// (it is 3.3)
	// we divide sensor output by two (in order not to burn our ADC), se here we multiply it back
	float p_volt = (p_int / 1024.0f) * 3.3f * 2.f;
	// sensor outputs 0.5V --> 0bar and 4.5V --> 10.2bar (150psi)
	// so to get value in bar, we subtract 0.5V, then divide by 4.5 - 0.5 = 4V 
	// and then multiply by sensor max value (10.2 bar)
	// finally multiply by 100 as we have int registers
	if (p_volt >= 0.5f && p_volt <= 4.5f)
	{ 
		measures[measure_pos] = (int)((((p_volt - 0.5f) / 4.f) * 10.2) * 100);
		measure_pos++;
	}
	// We will use average of MEASURES_AMOUNT last sensor readings
	// excluding minimal and maximal value (to get rid of peaks)
	if (measure_pos >= MEASURES_AMOUNT) measure_pos = 0;
	int min_pos = 0;
	int max_pos = 0;
	for (int i = 0; i < MEASURES_AMOUNT; i++)
	{
		if (measures[i] > measures[max_pos]) max_pos = i;
		if (measures[i] < measures[min_pos]) min_pos = i;
	}
	int sum = 0;
	for (int i = 0; i < MEASURES_AMOUNT; i++)
	{
		if (i == min_pos || i == max_pos) continue;
		sum += measures[i];
	}
	cur_p = sum / (MEASURES_AMOUNT - 2);
}

///////////////////////////////////////////////////////////////////////////////////////////////////

unsigned long last_screen_update = 0;

void UpdateDisplay()
{
	unsigned long current = millis();
	if (current < last_screen_update || current - last_screen_update >= SCREEN_UPDATE_INTERVAL)
	{
		lcd.setCursor(0, 0);
		lcd.print("P: ");
		lcd.print((cur_p / 100.0), 2);
		lcd.print(" T: ");
		lcd.print((cur_t / 100.0), 1);
		lcd.setCursor(0, 1);
		lcd.print("H: ");
		lcd.print(cur_hum / 100);
		if (error_pump)
		{
			lcd.print(" PUMP ERR");
		}
		else
		{
			lcd.print(" WiFi: ");
			lcd.print(WiFi.status());
		}
		//
		if (error_pump)
		{
			lcd.setRGB(255, 0, 0);
		}
		else
		{
			if (cur_pump_state)
			{
				lcd.setRGB(0, 0, 255);
			}
			else
			{
				lcd.setRGB(0, 255, 0);
			}
		}
		last_screen_update = current;
	}
}

unsigned long last_pump_switch = 0;
unsigned long last_pump_on = 0;
int last_pump_p = 0;

void UpdatePump()
{
	unsigned long current = millis();
	int new_state = cur_pump_state;
	if (cur_work_mode == MODE_AUTO_ON)
	{
     // turning on pump only if temperature is high enough
    if (cur_t > t_on)
    {
		  if (cur_p < p_on) new_state = 1;
  		else if (cur_p > p_off) new_state = 0;
    }
	  if (error_pump)
  	{
  		new_state = 0;
	  }
	}
	else if (cur_work_mode == MODE_AUTO_OFF)
	{
		new_state = 0;
	}
  else if (cur_work_mode == MODE_MANUAL)
  {
    new_state = force_pump;
  }
  if ((cur_work_mode == MODE_MANUAL || (current < last_pump_switch || current - last_pump_switch > PUMP_SWITCH_INTERVAL))
			&& new_state != cur_pump_state )
  {
    Serial.print("CHANGE PUMP STATE ");
    Serial.print(cur_p);
    Serial.print(" ");
    Serial.print(p_on);
    Serial.print(" ");
    Serial.print(p_off);
    cur_pump_state = new_state;
    last_pump_switch = current;
    if (new_state == 1)
    {
      last_pump_on = current;
      last_pump_p = cur_p;
    }
  }
	digitalWrite(PUMP_PIN, cur_pump_state);
  mb.Hreg(REGISTER_PUMP_STATE, cur_pump_state);
	// check for unproductive (in auto modes)
	if (cur_work_mode != MODE_MANUAL && cur_pump_state == 1)
	{
		if (current < last_pump_on || current - last_pump_on >= PUMP_UNPRODUCTIVE_TIMEOUT)
		{
			if (cur_p - last_pump_p < PRESSURE_PRODUCTIVE_INCREASE)
			{
				error_pump = true;
				cur_pump_state = 0;
			}
			else
			{
				// the pump is working, will perform new check in PUMP_UNPRODUCTIVE_TIMEOUT
				last_pump_p = cur_p;
				last_pump_on = current;
			}
		}
	}
}

void UpdateHeater()
{
	int new_state = cur_heat_state;
	if (cur_work_mode == MODE_AUTO_ON)
	{
		if (cur_t < t_on) new_state = 1;
		else if (cur_t > t_off) new_state = 0;
	}
  else if (cur_work_mode == MODE_AUTO_OFF)
  {
    new_state = 0;
  }
  else if (cur_work_mode == MODE_MANUAL)
  {
    new_state = force_heat;
  }
  if (new_state != cur_heat_state)
  {
    cur_heat_state = new_state;
  }
	digitalWrite(HEATER_PIN, cur_heat_state);
  mb.Hreg(REGISTER_HEAT_STATE, cur_heat_state);
}

unsigned long last_flush_switch = 0;

void UpdateFlush()
{
	unsigned long current = millis();
  int new_state = cur_flush_state;
  if (cur_work_mode == MODE_AUTO_ON)
  {
    // closing flush valve only when temperature is high enough
    if (cur_t > t_on)
    {
      new_state = 0; // 0 means closed
    }
    else
    {
      new_state = 1; 
    }
  }
  else if (cur_work_mode == MODE_AUTO_OFF)
  {
    new_state = 1;
  }
  else if (cur_work_mode == MODE_MANUAL)
  {
    new_state = force_flush;
  }
  if ((cur_work_mode == MODE_MANUAL || current < last_flush_switch || current - last_flush_switch > FLUSH_SWITCH_INTERVAL)
			&& new_state != cur_flush_state )
  {
    Serial.print("CHANGE FLUSH STATE ");
    Serial.print(new_state);
    cur_flush_state = new_state;
    last_flush_switch = current;
  }
	digitalWrite(FLUSH_PIN, cur_flush_state);
    mb.Hreg(REGISTER_FLUSH_STATE, cur_flush_state);
}

unsigned long last_connected = 0;
unsigned long last_state_save = 0;

void loop() 
{
	unsigned long current = millis();
	int u = 0;
	if (current < last_update || current - last_update >= SENSOR_UPDATE_INTERVAL)
	{
		UpdateDHTData();
		last_update = current;
		u = 1;
	}
	UpdatePressure();
	mb.task();
	p_on = mb.Hreg(REGISTER_P_ON);
	p_off = mb.Hreg(REGISTER_P_OFF);
	cur_work_mode = mb.Hreg(REGISTER_WORK_MODE);
	t_on = mb.Hreg(REGISTER_T_ON);
	t_off = mb.Hreg(REGISTER_T_OFF);
	error_pump= mb.Hreg(REGISTER_PUMP_ERROR);
  force_pump = mb.Hreg(REGISTER_FORCE_PUMP);
  force_heat = mb.Hreg(REGISTER_FORCE_HEAT);
  force_flush = mb.Hreg(REGISTER_FORCE_FLUSH);
	UpdatePump();
	UpdateHeater();
  UpdateFlush();
	if (u) 
	{
    Serial.print(current);
		Serial.print(" current pressure: ");
		Serial.print(cur_p);
		Serial.print(" t ");
		Serial.print(cur_t);
		Serial.print(" h ");
		Serial.print(cur_hum);
		Serial.print(" pump ");
		Serial.print(cur_pump_state);
		Serial.print(" heat ");
		Serial.print(cur_heat_state);
    Serial.print(" flush ");
		Serial.print(cur_flush_state);
    Serial.print(" mode ");
		Serial.println(cur_work_mode);
	}
	mb.Hreg(REGISTER_P_ON, p_on);
	mb.Hreg(REGISTER_P_OFF, p_off);
	mb.Hreg(REGISTER_WORK_MODE, cur_work_mode);
	mb.Hreg(REGISTER_T_ON, t_on);
	mb.Hreg(REGISTER_T_OFF, t_off);
	mb.Hreg(REGISTER_CURRENT_TEMP, cur_t);
	mb.Hreg(REGISTER_CURRENT_HUM, cur_hum);
	mb.Hreg(REGISTER_CURRENT_PRESS, cur_p);
	mb.Hreg(REGISTER_PUMP_ERROR, error_pump);
  mb.Hreg(REGISTER_FORCE_PUMP, force_pump);
  mb.Hreg(REGISTER_FORCE_HEAT, force_heat);
  mb.Hreg(REGISTER_FORCE_FLUSH, force_flush);
	UpdateDisplay();
	if (WiFi.status() != WL_CONNECTED)
	{
    Serial.println("RECONNECTING");
    Serial.println(WiFi.status());
		WiFi.reconnect();
	}
	else
	{
		last_connected = current;
	}
  if (current < last_state_save || current - last_state_save > STATE_SAVE_INTERVAL)
	{
    SaveState();
    last_state_save = current;
  }
	if (current < last_connected || current - last_connected > NON_CONNECTED_REBOOT_TIMEOUT)
	{
    SaveState();
		ESP.restart();
	}
}