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


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


class TempController(AttributeChange):
	def __init__(self, dt, params):
		AttributeChange.__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_model = Kalman(dt, params['Kalman'])
		self.kalman_plant = Kalman(dt, params['Kalman'])
		self.y = -1
		self.state = States.IDLE
		self.use_kalman = True
		self.kalman_model.initial((self.theta_ist_set, 0))
		self.kalman_plant.initial((self.theta_ist_set, 0))
		self.pid_hold.set_params(params['Hold'])
		self.pid_rate.set_params(params['Heat'])
		self.model = Pot(dt, params['Model'])
		self.theta_ist_delay_model = Delay(dt, params['Model']['Td'])
		self.dtheta_ist_delay_model = Delay(dt, params['Model']['Td'])

	def set_theta_ist(self, value):
		self.theta_ist_set = value

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

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

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

	def process(self):
		# Process Kalman of Model
		Z_model = self.kalman_model.process_measurement((self.model.get_temperature_intermediate(), 0), 0.0)
		xp_model = self.kalman_model.process(Z_model)
		theta_ist_model = xp_model[0, 0]
		heatrate_ist_model = xp_model[1, 0] * 60
		self.theta_ist_delay_model.put(theta_ist_model)
		self.dtheta_ist_delay_model.put(heatrate_ist_model)

		# Process Kalman of Plant
		Z_plant = self.kalman_plant.process_measurement((self.theta_ist_set, 0), 0.0)
		xp_plant = self.kalman_plant.process(Z_plant)
		theta_ist_plant = xp_plant[0, 0]
		heatrate_ist_plant = xp_plant[1, 0] * 60

		# 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()

		self.theta_ist = theta_ist_plant
		self.heatrate_ist = heatrate_ist_plant

		if 0:
			theta_err = self.theta_soll_set - (theta_ist_plant - self.theta_ist_delay_model.get() + theta_ist_model)
		else:
			theta_err = self.theta_soll_set - theta_ist_plant

		heatrate_err = self.heatrate_soll - (heatrate_ist_plant - self.dtheta_ist_delay_model.get() + heatrate_ist_model)

		# Process state
		state_next = self.state
		diff = self.theta_soll_set - self.theta_ist

		THRESH_HEAT = 1.0
		THRESH_IDLE = 1.0
		THRESH_IDLE_HOLD = 0.1
		THRESH_HEAT_HOLD = 2.0

		if self.state == States.IDLE:
			if diff > THRESH_HEAT:
				state_next = States.HEAT
				self.pid_rate.reset()
			elif diff > -THRESH_IDLE_HOLD:
				state_next = States.HOLD
				self.pid_rate.reset()
		elif self.state == States.HOLD:
			if diff > THRESH_HEAT:
				state_next = States.HEAT
				self.pid_rate.reset()
			elif diff < -THRESH_IDLE_HOLD:
				state_next = States.IDLE
		elif self.state == States.HEAT:
			if diff < -THRESH_IDLE:
				state_next = States.IDLE
			elif diff < THRESH_HEAT_HOLD:
				state_next = States.HOLD
				self.pid_hold.reset()

		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()
		self.model.set_power(max(0, 250 + 3500 * self.y))
		self.model.process()

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

	def get_power(self):
		return self.y


if __name__ == '__main__':
	dt = 1.0

	tc_params = {
		"Model": {
			"theta"		: 20,
			"C"			: 4190,
			"M"			: 20,
			"L"			: 0.05,
			"Td"		: 30,
			"kn"		: 0.2
		},
		"Kalman": {
			"var_P"	: 1.0,
			"var_Q"	: 0.0001,
			"var_R"	: 1.0
		},
		"Hold": {
			"kp": 1.0,
			"ki": 0.0,
			"kd": 0.0,
			"kt": 0.0
		},
		"Heat": {
			"kp": 0.2,
			"ki": 0.02,
			"kd": 0.0,
			"kt": 1.5
		}
	}

	pot_params = {
		"dt"		: 1.0,
		"theta"		: 20,
		"C"			: 4190,
		"M"			: 20 + 4,
		"L"			: 0.07,
		"Td"		: 30 + 5,
		"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()
			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_set)
			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")