from commonroad.common.file_reader import CommonRoadFileReader

"""
    Provides utility functions necessary to compute the CMs of scenarios.
"""


def find_egos_from_problem_sets(scenario_path: str) -> list[int]:
    """
    Finds the identifiers of the ego vehicles which were used as the basis for the planning problems in a set.
    :param scenario_path: Path to the scenario file to be read.
    :return: List of ego vehicle IDs extracted from the planning problem set.
    """
    scenario, planning_problem_set = CommonRoadFileReader(scenario_path).open()
    # Planing problem ID seems to be generated by adding 90000 to the ID of the ego vehicle it is based on.
    return [problem.planning_problem_id - 90000 for problem in planning_problem_set.planning_problem_dict.values()]


def compute_center_lanelet(ego_trajectory, scenario_with_ego):
    """
    Computes the lanelet the ego vehicle is on at each time step.
    :param ego_trajectory: List of states representing the ego vehicle's trajectory at each time step.
    :param scenario_with_ego: Scenario object includes the trajectory of the ego vehicle.
    :return: Dictionary mapping each time step to the corresponding lanelet for the ego vehicle.
    """
    # Calculate lanelets of ego vehicle at all timesteps
    ego_center_lanelet_dict = dict(
        zip(
            [state.time_step for state in ego_trajectory],
            scenario_with_ego.lanelet_network.find_lanelet_by_position([state.position for state in ego_trajectory]),
        )
    )
    return ego_center_lanelet_dict
