from __future__ import absolute_import
from __future__ import print_function

import os
import datetime
from shutil import copyfile

from STL_simulation import Simulation
from generator import TrafficGenerator
from visualization import Visualization
from utils import import_STL_configuration, set_sumo, set_STL_path

if __name__ == "__main__":

    config = import_STL_configuration(config_file='STL_settings.ini')
    sumo_cmd = set_sumo(config['gui'], config['sumocfg_file_name'], config['max_steps'])
    path = set_STL_path(config['models_path_name'])

    TrafficGen = TrafficGenerator(
        config['max_steps'], 
        config['n_cars_generated']
    )

    Visualization = Visualization(
        path, 
        dpi=96
    )
        
    Simulation = Simulation(
        TrafficGen,
        sumo_cmd,
        config['max_steps'],
        config['green_duration'],
        config['yellow_duration']
    )

    episode = 0
    timestamp_start = datetime.datetime.now()

    while episode < config['total_episodes']:
        print('\n----- Episode', str(episode+1), 'of', str(config['total_episodes']))
        simulation_time = Simulation.run(episode)  # run the simulation
        print('Simulation time:', simulation_time, 's')
        episode += 1

    print("\n----- Start time:", timestamp_start)
    print("----- End time:", datetime.datetime.now())
    print("----- Session info saved at:", path)

    copyfile(src='STL_settings.ini', dst=os.path.join(path, 'STL_settings.ini'))

    Visualization.save_data_and_plot(data=Simulation.cumulative_wait_store, filename='delay', xlabel='Episode', ylabel='Cumulative delay (s)')
    Visualization.save_data_and_plot(data=Simulation.avg_queue_length_store, filename='queue', xlabel='Episode', ylabel='Average queue length (vehicles)')