#!/usr/bin/env python3
import asyncio
import json
from utils import ChangedFloat
from ws.message import MessageDispatcher
from ws.server.ws_server_multi_user import WsServerMultiUser
from components.sensor import TempSensorFactory
from components.pid import PidFactory
from components.plant import Pot
from components.actor import HeaterFactory, StirrerFactory
from components.sud import Sud
from tasks import TaskManager, TempSensorTask, HeaterTask, PotTask, TcTask, StirrerTask, TracerTask, SudTask
from tracer import Tracer

import argparse as ap

# Default lumped thermal params, shared by both the real Pot plant and the
# temperature controller's internal Smith-predictor model.
DEFAULT_PLANT_PARAMS = {
	"C": 4190,
	"M": 25,
	"L": 0.2,
	"Td": 30
}

if __name__ == '__main__':
	parser = ap.ArgumentParser()
	parser.add_argument("-d", "--debug", action="store_true")
	parser.add_argument("-c", "--config", default="config.json")
	parser.add_argument("-m", "--model", default="model.json")
	parser.add_argument("-s", "--sim", action="store_true")

	args = parser.parse_args()
	config = json.load(open(args.config))
	dispatcher = MessageDispatcher()
	server = WsServerMultiUser(listener=dispatcher)
	taskmgr = TaskManager()

	DT = config['Controller']['dt']
	DT_TASK = 1.0 / config['Controller']['sim_warp_factor']
	DT_TASK_TRACER = DT_TASK / DT
	theta_amb = config['ambient_temperature']
	plant_sim = "sim" in config['Controller']['plant_name']

	# Sensor
	sensor = TempSensorFactory.create(config['Controller']['sensor_name'], temp_offset=-0.15, variance=0.01)
	sensor_task = TempSensorTask(sensor, DT_TASK, dispatcher.msgio_get("Sensor"))
	taskmgr.add(sensor_task)

	# Plant
	pot = Pot(DT, DEFAULT_PLANT_PARAMS, theta_amb)
	taskmgr.add(PotTask(pot, DT_TASK, dispatcher.msgio_get("Pot")))

	# Heater
	heater = HeaterFactory.create(config['Controller']['heater_name'], config['Heater'])
	heater.set_on_changed("power_eff", ChangedFloat(pot.set_power, prec=0).set)
	heater_task = HeaterTask(heater, DT_TASK, dispatcher.msgio_get("Heater"))
	taskmgr.add(heater_task)

	# Temperature Controller
	tc = PidFactory.create(config['Controller']['pid_type'], DT, config['TempCtrl'], DEFAULT_PLANT_PARAMS, theta_amb=theta_amb)
	tc_task = TcTask(tc, DT_TASK, dispatcher.msgio_get("TempCtrl"))
	taskmgr.add(tc_task)

	tc_trace_vars = []

	trace_tc = Tracer(tc, tc_trace_vars, name='tc_trace')

	# Stirrer
	stirrer = StirrerFactory.create(config['Controller']['stirrer_name'], DT, config['Stirrer'])
	stirrer_task = StirrerTask(stirrer, DT_TASK, dispatcher.msgio_get("Stirrer"))
	taskmgr.add(stirrer_task)

	# Mash schedule - starts out empty; a client loads one of sude/*.json's
	# several schedules onto it later (see components/sud.py's Sud).
	sud = Sud()
	taskmgr.add(SudTask(sud, tc, stirrer, pot, DT, DT_TASK, dispatcher.msgio_get("Sud")))

	# Tracer
	taskmgr.add(TracerTask(sensor, heater, tc, trace_tc, DT_TASK_TRACER, dispatcher.msgio_get("Tracer")))

	# Assign data flow
	# Assign tc control value to heater
	tc.set_on_changed("y", heater_task.actor)

	# Assign temp. sensor readings to tc
	sensor_task.set_on_changed("temp", ChangedFloat(tc.set_theta_ist, prec=2).set)

	# Assign heater power set to tc model - "Normal" has no internal model
	if hasattr(tc, "set_model_power"):
		heater.set_on_changed("power_set", tc.set_model_power)

	# For simulation
	if "sim" in config['Controller']['sensor_name']:
		pot.set_on_changed("temp", ChangedFloat(sensor.set_fake_temp, prec=3).set)
		sensor.set_fake_temp(pot.temp)

	# Static, global info that doesn't belong to any one task - sent once;
	# a client connecting later still gets it via the dispatcher's replay
	# of accumulated global_state on subscribe. AmbientTemp can also be
	# changed live by a client (e.g. the GUI's ambient temperature field).
	msg_system = dispatcher.msgio_get("System")

	async def send_system_info():
		await msg_system.send({'AmbientTemp': theta_amb, 'WarpFactor': DT / DT_TASK, 'PlantSim': plant_sim})
	asyncio.ensure_future(send_system_info())

	async def on_system_recv(data):
		global theta_amb
		if 'AmbientTemp' in data:
			theta_amb = data['AmbientTemp']
			pot.set_ambient_temperature(theta_amb)
			if hasattr(tc, 'set_ambient_temperature'):
				tc.set_ambient_temperature(theta_amb)
			await msg_system.send({'AmbientTemp': theta_amb})
	msg_system.set_recv_handler(on_system_recv)

	# Message dispatcher
	h_dispatcher = taskmgr.start()
	h_server = server.listen("0.0.0.0", 8765)
	asyncio.gather(h_dispatcher, h_server)
	server.run_forever()