"""Module containing code that loads parameters from user input and runs the simulations
"""
from time import time
import signal
import importlib
import pickle
import json
import copy
import os
from typing import Tuple
from tqdm import tqdm
import numpy as np
from shapely.geometry import shape, LineString
from .fire_simulator.visualizer import Visualizer
from .fire_simulator.logger import Logger
from .fire_simulator.fire import FireSim
from .fire_simulator.wind import Wind
from .utilities.file_io import SimFolderSelector, LoaderWindow
from .utilities.fire_util import UtilFuncs
from .base_classes.control_base import ControlClass

def initialize(params: dict) -> Tuple[FireSim, bool, Visualizer]:
    """Initializes a fireSim object and a Visualizer(if selected) based on user inputs

    :param params: dictionary containing all user's inputs and preferenes, generated by
                   SimFolderSelector class.
    :type params: dict
    :return: tuple containing a FireSim object, a bool variable specifying if real-time 
             visualization was selected, and a Visualizer object. If real-time visualization not
             selected Visualizer will be None.
    :rtype: Tuple[FireSim, bool, Visualizer]
    """
    # Folder Setup
    input_folder = params['input']
    input_foldername = os.path.basename(input_folder)

    # Real-time viz
    viz_on = params['viz_on']

    # Map Setup
    cell_size = params['cell_size'] # meters
    time_step = params['t_step'] # seconds
    sim_duration = params['sim_time']

    # Load wind forecast and initialize wind vector
    wind_forecast = []

    if not params['zero_wind']:

        wind_file = params['wind']

        with open(wind_file, 'r') as file:
            wind_data = json.load(file)

        forecast_time_step = wind_data['time_step_min']

        for entry in wind_data['data']:
            wind_forecast.append((entry['speed_m_s'], entry['direction']))

    else:
        forecast_time_step = 1e10
        wind_forecast = [(0, 0)]

    wind_forecast = np.array(wind_forecast)

    # Initialize wind vector
    wind_vec = Wind(wind_forecast, forecast_time_step)

    # Read JSON map file
    with open(input_folder + '/' + input_foldername + '.json', 'r') as f:
        data = json.load(f)

    # Load data files
    topography_map = np.load(data['topography']['file'])
    fuel_map = np.load(data['fuel']['file'])
    ignition_dicts = data['initial_ignition']
    initial_ignition = [shape(d) for d in ignition_dicts]
    fire_breaks = [{'geometry': LineString(fb['geometry']['coordinates']),
                    'fuel_value': fb['fuel_value']} for fb in data['fire_breaks']]

    if 'roads' in data:
        with open(data['roads']['file'], 'rb') as f:
            roads = pickle.load(f)
    else:
        roads = None

    # Save resolution, width, and height into variables
    topography_resolution = data['topography']['resolution']
    fuel_resolution = data['fuel']['resolution']
    width_m = data['topography']['width_m']
    height_m = data['topography']['height_m']

    # Calculate simSize
    num_cols = int(np.floor(width_m/(np.sqrt(3)*cell_size)))
    num_rows = int(np.floor(height_m/(1.5*cell_size)))

    sim_size = (num_rows, num_cols)
    print("simSize: " + str(sim_size))

    # Initialize the simulator
    fire = FireSim(fuel_map, fuel_resolution, topography_map, topography_resolution, wind_vec,
                   roads, fire_breaks, time_step = time_step, initial_ignition = initial_ignition,
                   size = (width_m, height_m), cell_size=cell_size, duration_s = sim_duration)

    if viz_on:
        viz = Visualizer(fire)
    else:
        viz = None

    return fire, viz_on, viz

def run_sim(fire: FireSim, viz: Visualizer, progress_bar: tqdm, loader_window: LoaderWindow,
            viz_on: bool, user_code_path: str, user_code_class: str, i:int, num_runs:int):
    """Function that runs the fire simulation and user code. Runs until the sim duration is reached
    or the fire is extinguished.

    :param fire: FireSim object to be run
    :type fire: FireSim
    :param viz: Visualizer object used for real-time visualization
    :type viz: Visualizer
    :param progress_bar: tqdm progress bar for command line display
    :type progress_bar: tqdm
    :param loader_window: LoaderWindow object used for generating a tk window that displays 
                          progress
    :type loader_window: LoaderWindow
    :param viz_on: boolean variable specifying whether the sim should be displayed in real-time
                   or not
    :type viz_on: bool
    :param user_code_path: Path to the user's selected module to run in parallel with sim
    :type user_code_path: str
    :param user_code_class: Name of the user's selected class to run in parallel with sim
    :type user_code_class: str
    :param i: Current number of simulation running
    :type i: int
    :param num_runs: Number of total simulations to be run
    :type num_runs: int
    """
    start = time()

    time_since_last_update = 20
    last_update = 0

    if user_code_path == "":
        user = None
    else:
        # Strip the extension and path
        user_module_name = os.path.splitext(os.path.basename(user_code_path))[0]

        # Dynamically import the module
        user_module = importlib.import_module(user_module_name)

        # Access the class within the module
        UserCodeClass = getattr(user_module, user_code_class)

        # Check if the user's class is a subclass of the abstract base class
        if not issubclass(UserCodeClass, ControlClass):
            raise TypeError(f"{user_code_class} is not a subclass of ControlClass")

        user = UserCodeClass(fire)  # Construct the user object

    fire.logger.start_status_log()

    while not fire.finished:

        fire.iterate()

        if user:
            user.process_state(fire)

        progress_bar.update()
        loader_window.increment_progress()

        # ----------- Perform real-time plotting ------------- #
        if viz_on:
            time_since_last_update = fire.curr_time_s - last_update

            if time_since_last_update >= fire.display_frequency:
                viz.update_grid(fire)
                last_update = fire.curr_time_s

    iters_skipped = ((fire.sim_duration/fire.time_step) + 1) - fire.iters
    progress_bar.update(iters_skipped)
    loader_window.increment_progress(increment_value=iters_skipped)
    end = time()
    elapsed = end - start

    msg = (f"Simulation complete. Sim time: {UtilFuncs.get_time_str(fire.curr_time_s)}, "
        f"took {UtilFuncs.get_time_str(elapsed, True)} seconds to compute.")

    fire.logger.log_message(msg)

    loader_window.set_text(f"Saving log file for simulation ({i+1}/{num_runs})")
    fire.logger.finish(fire)

def handle_interrupt(logger, fire):
    """Function to handle user interrupts while a sim is running.
    """
    print("\nUser interrupt detected. Finalizing logging...")
    logger.finish(fire, True)

def sim_loop(params: dict):
    """Main function that gets user input from GUI and runs the specified simulation(s)
    """
    num_runs = params['num_runs']
    duration = params['sim_time']
    time_step = params['t_step']

    # Get user code info
    user_code_path = params['user_path']
    user_code_class = params['class_name']

    # Create loader window object
    loader_window = LoaderWindow('Fire sim', max_value=num_runs*(duration/time_step)+1)
    loader_window.set_text("Initializing simulations...")

    fire, viz_on, viz = initialize(params)

    data_logger = Logger(params['log'])
    data_logger.store_metadata(params, fire)
    data_logger.store_init_fire_obj(fire)

    print("Running simulations... ")
    print("")

    # Create command line progress bar
    progress_bar = tqdm(total=(num_runs*(fire.sim_duration/fire.time_step)),
                        desc='Total progress', position=1, leave=True)

    curr_fire = None

    for i in range(num_runs):
        if num_runs > 1:
            curr_fire = copy.deepcopy(fire)

        else:
            curr_fire = fire

        # Register the signal handler
        signal.signal(signal.SIGINT, lambda signum, frame: handle_interrupt(data_logger, curr_fire))

        curr_fire.logger = data_logger
        loader_window.set_text(f"Running simulation ({i+1}/{num_runs})")
        run_sim(curr_fire, viz, progress_bar, loader_window, viz_on, user_code_path,
                user_code_class, i, num_runs)

        if viz_on:
            viz.reset_figure(done = i==num_runs-1)

    progress_bar.close()
    loader_window.set_text(f"Finished running {num_runs} simulations(s)")
    loader_window.close()

    print("")
    print(f"Finished running {num_runs} simulation(s)")

def main():
    folder_selector = SimFolderSelector(sim_loop)
    folder_selector.run()

if __name__ == "__main__":
    main()
