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


class TempController(APid):
	def __init__(self, dt, params, model_params):
		APid.__init__(self)
		self.pid_hold = Pid(dt)
		self.pid_rate = Pid(dt)
		self.theta_ist_set = 0
		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.model_params = model_params
		self.kalman_model = Kalman(dt, params['Kalman'])
		self.kalman_model_delay = Kalman(dt, params['Kalman'])
		self.kalman_plant = Kalman(dt, params['Kalman'])
		self.y = -1
		self.state = States.INIT
		self.use_kalman = True
		self.pid_hold.set_params(params['Hold'])
		self.pid_rate.set_params(params['Heat'])
		self.model = Pot(dt, model_params)
		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_model.initial((value, 0))
			self.kalman_model_delay.initial((value, 0))
			self.kalman_plant.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 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

		# Process Kalman of Model
		k_model = self.kalman_model.process_measurement((self.model.get_temperature_intermediate(), 0), 0.0)
		xp_model = self.kalman_model.process(k_model)
		theta_ist_model = xp_model[0, 0]
		heatrate_ist_model = xp_model[1, 0] * 60

		# Process Kalman of delayed Model
		k_model_delay = self.kalman_model_delay.process_measurement((self.model.get_temperature(), 0), 0.0)
		xp_model_delay = self.kalman_model_delay.process(k_model_delay)
		theta_ist_model_delay = xp_model_delay[0, 0]
		dtheta_ist_model_delay = xp_model_delay[1, 0] * 60

		self.theta_ist_plant = theta_ist_plant
		self.dtheta_ist_plant = heatrate_ist_plant
		self.theta_ist_model = theta_ist_model
		self.dtheta_ist_model = heatrate_ist_model
		self.theta_ist_model_delay = theta_ist_model_delay
		self.dtheta_ist_model_delay = dtheta_ist_model_delay
		self.theta_ist = theta_ist_plant
		self.heatrate_ist = heatrate_ist_plant

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

#		print ("Model    : T_ist={:2.2f}, dT_ist={:2.2f}".format(theta_ist_model, heatrate_ist_model))
#		print ("Model*z-1: T_ist={:2.2f}, dT_ist={:2.2f}".format(theta_ist_model_delay, dtheta_ist_model_delay))
#		print ("Plant    : T_ist={:2.2f}, dT_ist={:2.2f}".format(theta_ist_plant, heatrate_ist_plant))
		if 0:
			theta_err = self.theta_soll_set - (theta_ist_plant - theta_ist_model_delay + theta_ist_model)
		else:
			theta_err = self.theta_soll_set - theta_ist_plant

		heatrate_err = self.heatrate_soll - (heatrate_ist_plant - dtheta_ist_model_delay + heatrate_ist_model)
		diff = self.theta_soll_set - self.theta_ist

		self.process_fsm(diff)
		self.process_pid(theta_err, heatrate_err)

	def process_fsm(self, diff):
		# Process state
		state_next = self.state
		if self.state == States.INIT:
			if not self.is_startup:
				state_next = States.IDLE
		elif self.state == States.IDLE:
			if diff >= THRESH_IDLE_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_HOLD_HEAT:
				state_next = States.HEAT
				self.pid_rate.reset()
			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()

		if state_next != self.state:
			self.state = state_next
			print("New state = {}".format(state_next))
			if state_next == States.HEAT:
				self.model.initial(self.theta_ist)
				self.kalman_model.initial((self.theta_ist, 0))
				self.kalman_model_delay.initial((self.theta_ist, 0))

	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)

		if self.state == States.IDLE:
			self.y = 0
		else:
			self.y = self.pid_rate.get_y()

		self.model.set_power(max(0, 3500 * self.y))
		self.model.process()

	def get_power(self):
		return self.y


if __name__ == '__main__':
	dt = 1.0

	temp_soll = 20
	ctrl = TempController(dt, Test.tc_ctrl_params, Test.tc_model_params)
	plant = Pot(dt, Test.tc_pot_params)
	_y = np.empty(0)
	_fb = np.empty(0)
	_t = np.empty(0)
	_temp_soll = np.empty(0)
	_temp_ist_kalman = np.empty(0)
	_heatrate_ist_kalman = np.empty(0)
	_temp_ist_kalman_plant = np.empty(0)
	_heatrate_ist_kalman_plant = np.empty(0)
	_temp_ist_kalman_model = np.empty(0)
	_heatrate_ist_kalman_model = 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
			plant.process()
			temp_ist = plant.get_temperature()
			ctrl.set_theta_ist(temp_ist)
			ctrl.process()

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

			_y = np.append(_y, y)
			_fb = np.append(_fb, fb)
			_t = np.append(_t, t)
			_temp_soll = np.append(_temp_soll, temp_soll)
			_temp_ist_kalman_plant = np.append(_temp_ist_kalman_plant, ctrl.theta_ist_plant)
			_heatrate_ist_kalman_plant = np.append(_heatrate_ist_kalman_plant, max(-1, min(3, ctrl.dtheta_ist_plant)))
			_temp_ist_kalman_model = np.append(_temp_ist_kalman_model, ctrl.theta_ist_model_delay)
			_heatrate_ist_kalman_model = np.append(_heatrate_ist_kalman_model, max(-1, min(3, ctrl.dtheta_ist_model_delay)))
			t += 1

	figure(1)
	subplot(3, 1, 1)
	plot(_t, _temp_ist_kalman_plant, _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_plant, '-b', linewidth=1)
	legend(["heatrate"])
	grid(True)

	figure(2)
	subplot(2, 1, 1)
	plot(_t, _temp_ist_kalman_plant, '-b', _t, _temp_ist_kalman_model, 'r-', linewidth=1)
	legend(["plant", "model"])
	grid(True)
	subplot(2, 1, 2)
	plot(_t, _heatrate_ist_kalman_plant, '-b', _t, _heatrate_ist_kalman_model, '-r', linewidth=1)
	legend(["plant", "model"])
	grid(True)
	show()

	print("End of program")