import numpy as np
from matplotlib.pyplot import plot, step, figure, subplot, grid, show, legend, yticks
from components.sud import Sud, SudState
from components.pid.temp_controller_smith import TempController
from components.plant.pot import Pot
from components.actor.stirrersim import StirrerSim
from components.actor.heater_sim import HeaterSim

# Mirrors tasks/sud.py's SudTask, but driven synchronously for a demo run.
TEMP_REACHED_TOLERANCE = 0.2

# A real user would click "Confirm" in the GUI; auto-confirm after this many
# simulated seconds so the demo can run to completion unattended.
WAIT_USER_AUTO_CONFIRM_S = 30.0


if __name__ == '__main__':
	dt = 1.0
	theta_amb = 20

	ctrl_params = {
		"Hold": {
			"kp": 0.4,
			"ki": 0.0,
			"kd": 0.0,
			"kt": 0.0
		},
		"Heat": {
			"kp": 0.08,
			"ki": 0.008,
			"kd": 0.0,
			"kt": 1.5
		}
	}

	plant_params = {
		"C"			: 4190,
		"M"			: 20,
		"L"			: 0.2,
		"Td"		: 30
	}

	sud = Sud("sude/sud_0010.json")
	ctrl = TempController(dt, ctrl_params, plant_params, theta_amb)
	plant = Pot(dt, plant_params, theta_amb)
	stirrer = StirrerSim(dt)
	heater = HeaterSim({})

	# Mirror brewpi.py's data flow: the controller's internal Smith-predictor
	# model sees the (continuous) requested power, the real plant sees the
	# heater's actual (discretized) output power.
	heater.set_on_changed('power_set', ctrl.set_model_power)
	heater.set_on_changed('power_eff', plant.set_power)

	def on_rast_changed(rast):
		if rast is not None:
			ctrl.set_theta_soll(rast['temp'])
			ctrl.set_heatrate_soll(rast['heatRate'])

	def on_state_changed(state):
		if state == SudState.RAMPING:
			stirrer.set_duty_cycle(1.0)
			stirrer.set_speed(sud.stirr_speed_heat)
		elif state == SudState.HOLDING:
			stirrer.set_cycle_time(sud.stirr_cycle_time)
			stirrer.set_duty_cycle(sud.stirr_duty_rast)
			stirrer.set_speed(sud.stirr_speed_rast)
		elif state in (SudState.WAIT_USER, SudState.DONE):
			stirrer.set_duty_cycle(1.0)
			stirrer.set_speed(0)

	sud.set_on_changed('rast', on_rast_changed)
	sud.set_on_changed('state', on_state_changed)

	sud.start()

	_t = []
	_temp_ist = []
	_temp_soll = []
	_heatrate_ist = []
	_heatrate_soll = []
	_sud_state = []
	_rast_index = []
	_stirrer_speed = []
	_stirrer_duty = []
	_heater_power_set = []
	_heater_power_eff = []

	wait_user_elapsed = 0.0
	t = 0.0
	steps = 0
	max_steps = 30000
	while sud.state != SudState.DONE and steps < max_steps:
		plant.process()
		ctrl.set_theta_ist(plant.get_temperature())
		ctrl.process()

		y = ctrl.get_power()
		heater.set_power(max(0, 3500*y))
		heater.process()

		stirrer.process()

		if sud.state == SudState.RAMPING:
			if abs(ctrl.get_theta_ist() - ctrl.get_theta_soll_set()) < TEMP_REACHED_TOLERANCE:
				sud.temp_reached()

		if sud.state == SudState.WAIT_USER:
			wait_user_elapsed += dt
			if wait_user_elapsed >= WAIT_USER_AUTO_CONFIRM_S:
				sud.confirm()
				wait_user_elapsed = 0.0
		else:
			wait_user_elapsed = 0.0

		sud.tick(dt)

		_t.append(t)
		_temp_ist.append(ctrl.get_theta_ist())
		_temp_soll.append(ctrl.get_theta_soll_set())
		_heatrate_ist.append(ctrl.get_heatrate_ist())
		_heatrate_soll.append(ctrl.get_heatrate_soll())
		_sud_state.append(sud.state.value)
		_rast_index.append(sud.index)
		_stirrer_speed.append(stirrer.speed * stirrer.isOn)
		_stirrer_duty.append(stirrer.dutyCycle * 100)
		_heater_power_set.append(heater.power_set)
		_heater_power_eff.append(heater.power_eff)

		t += dt
		steps += 1

	if sud.state != SudState.DONE:
		print("Sud did not finish within {} steps (stuck in {})".format(max_steps, sud.state))
	print("Sud finished after {:.0f} s ({:.1f} min), final state = {}".format(t, t/60.0, sud.state))

	_t_min = np.array(_t) / 60.0

	figure(1)
	subplot(3, 1, 1)
	plot(_t_min, _temp_ist, '-b', _t_min, _temp_soll, '-r', linewidth=1)
	legend(["theta_ist", "theta_soll"])
	grid(True)
	subplot(3, 1, 2)
	plot(_t_min, _heatrate_ist, '-b', _t_min, _heatrate_soll, '-r', linewidth=1)
	legend(["heatrate_ist", "heatrate_soll"])
	grid(True)
	subplot(3, 1, 3)
	plot(_t_min, _heater_power_set, '-r', _t_min, _heater_power_eff, '-b', linewidth=1)
	legend(["heater power_set", "heater power_eff"])
	grid(True)

	figure(2)
	subplot(3, 1, 1)
	step(_t_min, _sud_state, where='post', color='m')
	yticks([0, 1, 2, 3, 4], ["IDLE", "RAMPING", "HOLDING", "WAIT_USER", "DONE"])
	legend(["Sud state"])
	grid(True)
	subplot(3, 1, 2)
	step(_t_min, _rast_index, where='post', color='k')
	legend(["Rast index"])
	grid(True)
	subplot(3, 1, 3)
	plot(_t_min, _stirrer_speed, '-g', _t_min, _stirrer_duty, '-c', linewidth=1)
	legend(["stirrer speed (effective)", "stirrer duty cycle [%]"])
	grid(True)

	show()

	print("End of program")