#!/usr/bin/python3

import collections
import functools

from lib.ArgParser import get_config, get_model_from_config
from lib.Controller import simulate, ReplayController
from lib.Cost import DeviationCost, NegativeCost
from lib.Env import CentralHeatingHistoryEnv

# ARGS
config = get_config()
config['model_past_fields'] = ['temp_in', 'temp_out', 'mode']
config['model_future_fields'] = ['temperature', 'wind_speed', 'radiation', 'humid']

plot = collections.OrderedDict()

for model_type in ('neural', 'linear'):
	config['model_type'] = model_type

	# MODEL
	model = get_model_from_config(config)
	print(repr(model.model))

	# PLOT
	config['plot_fields'] = ['time', 'temp_in', 'temp_in_calc']
	config['past_fields'] = ['mode']
	config['past_values'] = 1
	config['future_values'] = 0

	# ENV
	env = CentralHeatingHistoryEnv()
	env.cost_class = DeviationCost
	#env.cost_class = functools.partial(NegativeCost, 19.5, 22)
	env.model = model
	env.config = config
	env.reset()

	# MAIN
	controller = ReplayController()
	simulate(env, controller, render=False)

	plot['time'] = env.window.plot['time']
	plot['temp_in'] = env.window.plot['temp_in']
	plot['temp_in_' + model_type] = env.window.plot['temp_in_calc']

print(plot.keys())
env.window.plot = plot

env.render()