from components.plant.pot import Pot
from matplotlib.pyplot import plot, figure, subplot, grid, show, legend
from components.pid import Pid, Kalman
from components import APid
import numpy as np
from enum import Enum


class States(Enum):
	IDLE = 0,
	HEAT = 1,
	HOLD = 2


class TempController(APid):
	def __init__(self, dt, params):
		APid.__init__(self)
		self.pid_hold = Pid(dt)
		self.pid_rate = Pid(dt)
		self.theta_ist_set = 20
		self.theta_soll_set = 0
		self.heatrate_ist_set = 0
		self.heatrate_soll_set = 1.0
		self.heatrate_soll = 1.0
		self.theta_ist = 0
		self.heatrate_ist = 0
		self.params = params
		self.kalman = Kalman(dt, params['Kalman'])
		self.y = -1
		self.state = States.IDLE
		self.use_kalman = True
		self.pid_hold.set_params(params['Hold'])
		self.pid_rate.set_params(params['Heat'])
		self.is_startup = True

	def set_theta_ist(self, value):
		self.theta_ist_set = value
		if self.is_startup:
			self.is_startup = False
			self.kalman.initial((value, 0))

	def get_theta_ist(self):
		return self.theta_ist

	def set_heatrate_ist(self, value):
		self.heatrate_ist_set = value

	def get_heatrate_ist(self):
		return self.heatrate_ist

	def set_theta_soll(self, value):
		self.theta_soll_set = value

	def get_theta_soll(self):
		return self.theta_soll

	def get_theta_soll_set(self):
		return self.theta_soll_set

	def set_heatrate_soll(self, value):
		self.heatrate_soll_set = value

	def get_heatrate_soll(self):
		return self.heatrate_soll

	def get_heatrate_soll_set(self):
		return self.heatrate_soll_set

	def process(self):
		# Process Kalman
		if self.use_kalman:
			Z = self.kalman.process_measurement((self.theta_ist_set, 0), 0.0)
			xp = self.kalman.process(Z)
			self.theta_ist = xp[0, 0]
			self.heatrate_ist = xp[1, 0] * 60
		else:
			self.theta_ist = self.theta_ist_set
			self.heatrate_ist = self.heatrate_ist_set

		# Compensate for max heat rate to reduce overshoot
		if self.heatrate_soll_set > 0:
			self.pid_hold.scale(1.0/self.heatrate_soll_set)

		self.heatrate_soll = self.heatrate_soll_set * self.pid_hold.get_y()
		theta_err = self.theta_soll_set - self.theta_ist
		heatrate_err = self.heatrate_soll - self.heatrate_ist

		diff = self.theta_soll_set - self.theta_ist
		self.process_fsm(diff)
		self.process_pid(theta_err, heatrate_err)

	def process_fsm(self, diff):
		THRESH_HOLD_IDLE = 0.1
		THRESH_HOLD_HEAT = 2.0
		THRESH_IDLE_HEAT = 1.0
		THRESH_IDLE_HOLD = 0.1
		THRESH_HEAT_HOLD = 2.0
		THRESH_HEAT_IDLE = 2.0

		# Process state
		state_next = self.state
		if self.state == States.IDLE:
			self.pid_rate.reset()
			if diff > THRESH_IDLE_HEAT:
				state_next = States.HEAT
			elif diff > -THRESH_IDLE_HOLD:
				state_next = States.HOLD
		elif self.state == States.HOLD:
			self.pid_rate.reset()
			if diff > THRESH_HOLD_HEAT:
				state_next = States.HEAT
			elif diff < -THRESH_HOLD_IDLE:
				state_next = States.IDLE
		elif self.state == States.HEAT:
			if diff < -THRESH_HEAT_IDLE:
				state_next = States.IDLE
			elif diff < THRESH_HEAT_HOLD:
				state_next = States.HOLD
				self.pid_hold.reset()
				self.pid_rate.reset()

		if state_next != self.state:
			self.state = state_next
			print("New state = {}".format(state_next))

	def process_pid(self, theta_err, heatrate_err):
		self.pid_hold.process(theta_err, -self.theta_ist)
		self.pid_rate.process(heatrate_err, -self.heatrate_ist)

		self.y = self.pid_rate.get_y()

	def get_power(self):
		return self.y


if __name__ == '__main__':
	dt = 1.0

	tc_params = {
		"Kalman": {
			"var_P"	: 1.0,
			"var_Q"	: 0.0001,
			"var_R"	: 1.0
		},
		"Hold": {
			"kp": 0.5,
			"ki": 0.0,
			"kd": 0.0,
			"kt": 0.0
		},
		"Heat": {
			"kp": 0.1,
			"ki": 0.01,
			"kd": 0.0,
			"kt": 1.5
		}
	}

	pot_params = {
		"theta"		: 20,
		"C"			: 4190,
		"M"			: 20,
		"L"			: 0.05,
		"Td"		: 80,
		"kn"		: 0.2
	}

	temp_ist = 0
	temp_soll = 20
	ctrl = TempController(dt, tc_params)
	plant = Pot(dt, pot_params)
	_temp_ist = np.empty(0)
	_temp_soll = np.empty(0)
	_y = np.empty(0)
	_fb = np.empty(0)
	_t = np.empty(0)
	_heatrate_ist_kalman = np.empty(0)
	_temp_ist_kalman = np.empty(0)
	a = 0.5
	fb = 0
	rho = 0.02
	temps = [{'Temp': 20, 'Duration': 1000}, {'Temp': 40, 'Duration': 1000}, {'Temp': 50, 'Duration': 1000}, {'Temp': 60, 'Duration': 1000}, {'Temp': 70, 'Duration': 1000}, {'Temp': 80, 'Duration': 1000}, {'Temp': 78, 'Duration': 1000}]

	t = 0
	for temp in temps:
		temp_soll = temp['Temp']
		hold_counter = temp['Duration']
		hold = False
		ctrl.set_theta_soll(temp_soll)
		ctrl.set_heatrate_soll(1.0)

		while True:
			if hold:
				if hold_counter == 0:
					break
				hold_counter -= 1
			ctrl.process()
			plant.process()
			temp_ist = plant.get_temperature() + 0.0 * np.random.randn()
			ctrl.set_theta_ist(temp_ist)

			y = 3500*ctrl.get_power()
			power = max(0, 250+y)
			plant.set_power(power)
			fb = plant.get_power()
			if abs(temp_ist - temp_soll) < 0.1:
				hold = True

			temp_ist -= rho
			_temp_ist = np.append(_temp_ist, temp_ist)
			_temp_soll = np.append(_temp_soll, temp_soll)
			_y = np.append(_y, y)
			_fb = np.append(_fb, fb)
			_t = np.append(_t, t)
			_heatrate_ist_kalman = np.append(_heatrate_ist_kalman, max(-1, min(3, ctrl.heatrate_ist)))
			_temp_ist_kalman = np.append(_temp_ist, ctrl.theta_ist)
			t += 1

	figure(1)
	subplot(3, 1, 1)
	plot(_t, _temp_ist, _t, _temp_soll, 'r-', linewidth=1)
	legend(["ist", "soll"])
	grid(True)
	subplot(3, 1, 2)
	plot(_t, _y, '-b', _t, _fb, '-r', linewidth=1)
	legend(["y", "pot"])
	grid(True)
	subplot(3, 1, 3)
	plot(_t, _heatrate_ist_kalman, '-b', linewidth=1)
	legend(["heatrate"])
	grid(True)
	show()

	print("End of program")