from components.pid import pid
from components.pid import kalman
from components.plant.pot import Pot
from matplotlib.pyplot import plot, figure, subplot, title, xlabel, ylabel, grid, show, legend
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, params, params_kalman):
		AttributeChange.__init__(self)

		self.dt = params['dt']
		self.pid_hold = pid.Pid()
		self.pid_rate = pid.Pid()
		self.theta_ist_set = 0
		self.theta_soll_set = 0
		self.heatrate_ist_set = 0
		self.heatrate_soll_set = 2.0
		self.theta_ist = 0
		self.heatrate_ist = 0
		self.params = params
		self.kalman = kalman.Kalman(params_kalman)
		self.y = 0
		self.state = States.HOLD
		self.use_kalman = True
		self.kalman.initial((20, 0))

	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
		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

		theta_err = self.theta_soll_set - self.theta_ist
		heatrate_err = self.heatrate_soll_set - self.heatrate_ist

		# Process state
		state_next = self.state
		if self.state == States.HOLD:
			if (self.theta_ist + 1.0) < self.theta_soll_set:
				state_next = States.HEAT
				self.pid_rate.reset()
		elif self.state == States.HEAT:
			if (self.theta_ist + 1.0) >= self.theta_soll_set:
				state_next = States.HOLD
				self.pid_hold.reset()

		if self.state != States.IDLE:
			self.pid_hold.process(self.dt, self.params['Hold']['Pid'], theta_err)
			self.pid_rate.process(self.dt, self.params['Heat']['Pid'], heatrate_err)

		if self.state == States.HEAT:
			self.y = self.pid_rate.get_y()
		else:
			self.y = self.pid_hold.get_y()

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

		self.state = state_next

	def get_power(self):
		return self.y


if __name__ == '__main__':
	params = {
		"dt": 1.0,
		"Hold": {
			"Pid": {
				"kp": 0.5,
				"ki": 0.005,
				"kd": 0.0,
				"rho": 1.0
			}
		},
		"Heat": {
			"Pid": {
				"kp": 0.1,
				"ki": 0.01,
				"kd": 0.0,
				"rho": 1.0
			}
		}
	}

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

	# Kalman Filter
	kalman_params = {
		"dt"			: 1.0,
		"var_P"				: 1,
		"var_Q"				: 0.0001,
		"var_R"				: 1
	}
	temp_ist = 0
	temp_soll = 20
	ctrl = TempController(params, kalman_params)
	plant = Pot(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(2.0)

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

			y = 3500*ctrl.get_power()
			power = max(0, 0+y)
			plant.setPower(power)
			fb = plant.getPower()
			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, min(3000, 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")