#!/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

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

	# 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_params = config['Plant']
	pot = Pot(DT, pot_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'], config['Model'], 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 (optional)
	sud_path = config.get('sud')
	if sud_path:
		sud = Sud(sud_path)
		taskmgr.add(SudTask(sud, tc, stirrer, 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
	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)

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