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

# 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.02,
			"kd": 0.0,
			"kt": 1.5
		},
		"Cool": {
			"kp": 0.08,
			"ki": 0.02,
			"kd": 0.0,
			"kt": 1.5
		}
	}

	sud = Sud()
	with open("sude/sud_0010.json") as f:
		sud.load(json.load(f))

	first_step = sud.schedule[0]
	plant_params = {
		**sud.derive_plant_params(first_step['grain_mass'], first_step['water_mass']),
		"L"			: 0.2,
		"Td"		: 60
	}

	ctrl = TempController(dt, plant_params, theta_amb)
	ctrl.set_params(ctrl_params)
	plant = Pot(dt)
	plant.set_plant_params(plant_params)
	plant.set_ambient_temperature(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 apply_stirrer(phase):
		stirrer_cfg = phase.get('stirrer', {})
		speed = stirrer_cfg.get('speed', 0)
		interval_time = stirrer_cfg.get('interval_time', 0)
		on_ratio = stirrer_cfg.get('on_ratio', 1.0)
		if interval_time > 0:
			stirrer.set_cycle_time(interval_time)
			stirrer.set_duty_cycle(on_ratio)
		else:
			stirrer.set_cycle_time(1.0)
			stirrer.set_duty_cycle(1.0 if speed > 0 else 0.0)
		stirrer.set_speed(speed)

	def apply_plant_params(step):
		params = sud.derive_plant_params(step['grain_mass'], step['water_mass'])
		plant.set_thermal_params(params['M'], params['C'])
		ctrl.set_model_params(params['M'], params['C'])

	def on_step_changed(step):
		# A step may carry both 'ramp' and 'hold' (ramp to temp, then hold);
		# sud.state tells us which phase is currently active.
		if step is not None:
			ramping = sud.state == SudState.RAMPING
			phase = step['ramp'] if ramping else step.get('hold', {})

			print(f"Step {step['number']}: {step['descr']}")
			apply_plant_params(step)
			if ramping and step['temperature'] is not None:
				ctrl.set_theta_soll(step['temperature'])
				ctrl.set_heatrate_soll(step['ramp']['rate'])
			apply_stirrer(phase)

	def on_state_changed(state):
		if state == SudState.WAIT_USER and sud.user_message:
			print("Sud: {}".format(sud.user_message))
		if state in (SudState.DONE, SudState.IDLE):
			stirrer.set_duty_cycle(1.0)
			stirrer.set_speed(0)
			ctrl.set_enabled(False)
		else:
			ctrl.set_enabled(True)

	sud.set_on_changed('step', on_step_changed)
	sud.set_on_changed('state', on_state_changed)

	sud.start()

	_t = []
	_temp_ist = []
	_temp_soll = []
	_heatrate_ist = []
	_heatrate_soll = []
	_sud_state = []
	_step_index = []
	_stirrer_speed = []
	_heater_power_set = []
	_heater_power_eff = []
	_plant_M = []
	_plant_C = []

	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 ctrl.is_holding():
				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)
		_step_index.append(sud.index)
		_stirrer_speed.append(stirrer.speed * stirrer.isOn)
		_heater_power_set.append(heater.power_set)
		_heater_power_eff.append(heater.power_eff)
		_plant_M.append(plant.M)
		_plant_C.append(plant.C)

		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_arr = np.array(_t)

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

	figure(2)
	subplot(3, 1, 1)
	step(_t_arr, _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_arr, _step_index, where='post', color='k')
	legend(["Step index"])
	grid(True)
	subplot(3, 1, 3)
	plot(_t_arr, _stirrer_speed, '-g', linewidth=1)
	legend(["stirrer speed (effective)"])
	grid(True)

	figure(3)
	subplot(2, 1, 1)
	plot(_t_arr, _plant_M, '-b', linewidth=1)
	legend(["plant M (kg)"])
	grid(True)
	subplot(2, 1, 2)
	plot(_t_arr, _plant_C, '-b', linewidth=1)
	legend(["plant C (J/(kg*K))"])
	grid(True)

	show()

	print("End of program")