# Computes criticality metrics for a scenario
import logging
import traceback
from datetime import datetime
from multiprocessing import Pool
from pathlib import Path

from commonroad.common.file_reader import CommonRoadFileReader
from commonroad.scenario.scenario import Scenario
from commonroad_crime.data_structure.base import CriMeBase
from commonroad_crime.data_structure.configuration import CriMeConfiguration
from commonroad_crime.data_structure.crime_interface import CriMeInterface

from commonroad_labeling.criticality.computer.crit_util import compute_center_lanelet, find_egos_from_problem_sets
from commonroad_labeling.criticality.trajectory_inserter.trajectory_inserter import TrajectoryInserter


class CMComputer:
    """
    This class is used to compute the Criticality Metrics from CommonRoad scenario files.
    """

    def __init__(
        self,
        metrics: list[type[CriMeBase]],
        verbose=True,
        crime_verbose=False,
        overwrite=True,
    ):
        """
        Initializes the class with the given parameters.
        :param metrics: A list of the CMs that should be computed.
        :param verbose: A boolean flag indicating whether the process should output detailed information.
            Defaults to True.
        :param crime_verbose: A boolean flag indicating whether CommonRoad-CriMe should output
            detailed information about the computation of the CMs. Defaults to False.
        :param overwrite: A boolean flag specifying whether already existing output files should be overwritten.
            When set to False the CM calculation of a scenario will be skipped,
            if an output file (with the same name as the one that would be newly produced) already exists
            in the output directory. Defaults to True.
        """
        self.verbose = verbose
        self.metrics = metrics
        self.crime_verbose = crime_verbose
        self.overwrite = overwrite

    def compute_metrics(
        self,
        scenario_path: str,
        ego_id: int = None,
        save_plots=False,
        show_plots=False,
        make_gif=False,
        do_log=False,
        output_dir=None,
    ):
        """
        Computes the CMs for a given scenario file. Simulates the trajectory of the ego vehicle if needed.
        :param scenario_path: Path to the scenario file to be read.
        :param ego_id: Identifier for the ego vehicle. If None, a trajectory for the ego vehicle will be simulated
            using Reactive Planner.
        :param save_plots: Boolean flag indicating whether to save plots generated during trajectory computation
            of Reactive Planner.
        :param show_plots: Boolean flag indicating whether to show plots generated during trajectory computation
            of Reactive Planner.
        :param make_gif: Boolean flag indicating whether to create a gif from the plots generated by Reactive Planner.
        :param do_log: Boolean flag indicating whether Reactive Planner should create logs.
        :param output_dir: Directory to which the CM output data should be written to.
        """
        scenario, planning_problem_set = CommonRoadFileReader(scenario_path).open()

        if ego_id is None:
            inserter = TrajectoryInserter(
                save_plots=save_plots,
                show_plots=show_plots,
                do_make_gif=make_gif,
                do_log=do_log,
            )
            scenario, ego_id = inserter.insert_ego_trajectory(scenario, planning_problem_set)

        self.compute_metrics_for_id(scenario, ego_id, scenario_path, output_dir=output_dir)

    def compute_metrics_for_id(
        self,
        scenario_with_ego,
        ego_id,
        scenario_path,
        output_dir=str(Path.cwd().joinpath("output")),
    ):
        """
        Computes the CMs for the vehicle with the given identifier.
        :param scenario_with_ego: The scenario data containing the ego vehicle and its trajectory.
        :param ego_id: The identifier for the ego vehicle within the scenario.
        :param scenario_path: The file path of the scenario to be evaluated.
        :param output_dir: The directory path where output files will be saved, defaults to an
            "output" directory in the current working directory.
        """
        config = self.create_crime_config(scenario_with_ego, ego_id, scenario_path)

        ego_obstacle = scenario_with_ego.obstacle_by_id(ego_id)
        ts_start = ego_obstacle.initial_state.time_step
        ts_end = ego_obstacle.prediction.final_time_step

        all_states = ego_obstacle.prediction.trajectory.state_list
        all_states.insert(0, ego_obstacle.initial_state)

        if ego_obstacle.prediction.center_lanelet_assignment is None:
            ego_obstacle.prediction.center_lanelet_assignment = compute_center_lanelet(all_states, scenario_with_ego)

        crime_interface = CriMeInterface(config)

        if self.verbose:
            print(
                f"{datetime.now().strftime('%H:%M:%S')}: Started computing metrics for scenario {scenario_path},"
                f" ego_id {ego_id}"
            )
        crime_interface.evaluate_scenario(self.metrics, ts_start, ts_end, verbose=self.crime_verbose)

        output_dir = output_dir if output_dir is not None else str(Path.cwd().joinpath("output"))
        crime_interface.save_to_file(output_dir)

        if self.verbose:
            print(
                f"{datetime.now().strftime('%H:%M:%S')}: Finished computing metrics for scenario {scenario_path},"
                f" ego_id {ego_id}"
            )
        # crime_interface.visualize(5)
        # utils_vis.plot_criticality_curve(crime_interface)

    def compute_parallel(
        self,
        scenario_dir: str,
        process_count=2,
        save_plots=False,
        show_plots=False,
        make_gif=False,
        do_log=False,
        output_dir: str = None,
    ):
        """
        Computes the CMs for scenarios in the given directory in parallel.

        WARNING: be careful about increasing the number of processes. Depending on the amount of
        Criticality Metrics computed, each process can take multiple GIGABYTES of RAM.

        :param scenario_dir: Directory containing the scenario files to be processed.
        :param process_count: Number of parallel processes to use for computation. Default is 2.
        :param save_plots: Boolean flag indicating whether to save plots generated during trajectory
            computation of Reactive Planner.
        :param show_plots: Boolean flag indicating whether to show plots generated during trajectory
            computation of Reactive Planner.
        :param make_gif: Boolean flag indicating whether to create a gif from the plots generated by Reactive Planner.
        :param do_log: Boolean flag indicating whether Reactive Planner should create logs.
        :param output_dir: Directory to which the CM output data should be written to.
        """
        dir_path = Path(scenario_dir)
        all_scenarios = [str(x.absolute()) for x in list(dir_path.iterdir())]
        scenario_ego_pairs = []
        for scenario in all_scenarios:
            for ego_id in find_egos_from_problem_sets(scenario):
                if (
                    not self.overwrite
                    and Path(output_dir).joinpath(f"CriMe-{Path(scenario).name[:-4]}_veh_{ego_id}.xml").exists()
                ):
                    continue
                scenario_ego_pairs.append(
                    (
                        scenario,
                        ego_id,
                        save_plots,
                        show_plots,
                        make_gif,
                        do_log,
                        output_dir,
                    )
                )
        print(f"Starting parallel computation of {len(scenario_ego_pairs)} tasks.")
        with Pool(processes=process_count) as pool:
            pool.starmap(self.compute_metrics_catching, scenario_ego_pairs)
            pool.close()
            pool.join()

    def compute_sync(
        self,
        scenario_dir: str,
        save_plots=False,
        show_plots=False,
        make_gif=False,
        do_log=False,
        output_dir: str = None,
    ):
        """
        Computes the CMs for scenarios in the given directory in sequence. Useful for debugging.
        :param scenario_dir: Directory containing the scenario files to be processed.
        :param save_plots: Boolean flag indicating whether to save plots generated during trajectory
            computation of Reactive Planner.
        :param show_plots: Boolean flag indicating whether to show plots generated during trajectory
            computation of Reactive Planner.
        :param make_gif: Boolean flag indicating whether to create a gif from the plots generated by Reactive Planner.
        :param do_log: Boolean flag indicating whether Reactive Planner should create logs.
        :param output_dir: Directory to which the CM output data should be written to.
        """
        dir_path = Path(scenario_dir)
        all_scenarios = [str(x.absolute()) for x in list(dir_path.iterdir())]
        scenario_ego_pairs = []
        for scenario in all_scenarios:
            for ego_id in find_egos_from_problem_sets(scenario):
                if (
                    not self.overwrite
                    and Path(output_dir).joinpath(f"CriMe-{Path(scenario).name[:-4]}_veh_{ego_id}.xml").exists()
                ):
                    continue
                scenario_ego_pairs.append((scenario, ego_id))
        print(f"Starting computation of {len(scenario_ego_pairs)} tasks.")
        for scenario, ego_id in scenario_ego_pairs:
            self.compute_metrics_catching(
                scenario,
                ego_id=ego_id,
                output_dir=output_dir,
                save_plots=save_plots,
                show_plots=show_plots,
                make_gif=make_gif,
                do_log=do_log,
            ),

    def compute_metrics_catching(
        self,
        scenario_path: str,
        ego_id: int = None,
        save_plots=False,
        show_plots=False,
        make_gif=False,
        do_log=False,
        output_dir=None,
    ):
        """
        Computes the CMs for a given scenario file, catching and reporting any potential exceptions. Simulates the
        trajectory of the ego vehicle if needed.
        :param scenario_path: Path to the scenario file to be read.
        :param ego_id: Identifier for the ego vehicle. If None, a trajectory for the ego vehicle will be simulated
            using Reactive Planner.
        :param save_plots: Boolean flag indicating whether to save plots generated during trajectory computation
            of Reactive Planner.
        :param show_plots: Boolean flag indicating whether to show plots generated during trajectory computation
            of Reactive Planner.
        :param make_gif: Boolean flag indicating whether to create a gif from the plots generated by Reactive Planner.
        :param do_log: Boolean flag indicating whether Reactive Planner should create logs.
        :param output_dir: Directory to which the CM output data should be written to.
        """
        try:
            self.compute_metrics(
                scenario_path,
                ego_id,
                save_plots,
                show_plots,
                make_gif,
                do_log,
                output_dir,
            )
        except Exception:
            print(
                f"{datetime.now().strftime('%H:%M:%S')}: Exception when computing metrics for scenario "
                f"{scenario_path}, for ego_id {ego_id}.\n {traceback.format_exc()}"
            )

    def create_crime_config(self, scenario_with_ego: Scenario, ego_id: int, scenario_path: str) -> CriMeConfiguration:
        """
        Creates a CriMeConfiguration object necessary for CriMe to function.
        :param scenario_with_ego: The scenario object containing the trajectory of the ego vehicle.
        :param ego_id: The identifier for the ego vehicle within the scenario.
        :param scenario_path: The path to the scenario file.
        :return: An instance of CriMeConfiguration configured based on the provided parameters.
        """
        config = CriMeConfiguration()
        config.general.name_scenario = str(scenario_with_ego.scenario_id)
        path_split = scenario_path.rsplit("/", 1)
        config.general.path_scenarios = path_split[0] + "/"
        config.scenario = scenario_with_ego
        config.vehicle.ego_id = ego_id
        config.debug.save_plots = False

        if self.crime_verbose:
            config.print_configuration_summary()

        if not self.crime_verbose:
            logging.getLogger().setLevel(logging.CRITICAL)

        return config
