import numpy as np
import numpy.typing as npt
from shapely import Point, creation
from navsim.planning.simulation.planner.pdm_planner.observation.pdm_occupancy_map import (
    PDMDrivableMap,
)
from navsim.planning.simulation.planner.pdm_planner.utils.pdm_path import (
    PDMPath,
)

#1. 计算车辆在行驶过程中与车道中心线的平均距离
#2. 计算车辆在一段时间或距离内保持在车道内的时间占总行驶时间的百分比
#3. 计算车辆在行驶过程中其位置偏离车道中心线的标准差（稳定性）
# lane keeping metric
# def calc_lk(trajectories,
#             num_proposals,
#             drivable_area_map: PDMDrivableMap) -> npt.NDArray:
#     """
#     vocab_size = 4096 or 8192
#     trajectories: [
#     PDM-Closed Trajectory + vocab_size trajs,
#     1+40 (current pose + 4 secs * 10Hz poses),
#     11: StateIndex, navsim/planning/simulation/planner/pdm_planner/utils/pdm_enums.py]
#     num_proposals: PDM-Closed Trajectory + vocab_size trajs
#     """
#     dummy_scores = np.ones_like(num_proposals)
#     return dummy_scores


def calc_lk(ego_coords,
            centerline,
            BBCoordsIndex,
            num_proposals,
            multi_metrics):
    """
    Calculate the average distance from the vehicle to the centerline.
    """

    progress_in_meter = np.zeros(num_proposals, dtype=np.float64)
    for proposal_idx in range(num_proposals):
        count = 0
        total_distance = 0.0
        for time_idx in range(ego_coords.shape[1]):
            vehicle_position = Point(*ego_coords[proposal_idx, time_idx, BBCoordsIndex.CENTER])

            # Interpolate the centerline position
            distance_along_centerline = centerline.project(vehicle_position)
            centerline_position = centerline.interpolate([distance_along_centerline], as_array=True)

            # Calculate the distance to the centerline
            distance = vehicle_position.distance(Point(*centerline_position))
            total_distance += distance
            count += 1
        progress_in_meter[proposal_idx] = total_distance / count
    # progress_in_meter = progress_in_meter[1:]
    #要不要这些metric去计算？
    # multiplicate_metric_scores = multi_metrics.prod(axis=0)
    # raw_progress = progress_in_meter * multiplicate_metric_scores
    # N = raw_progress.shape[0]
    # pdm_progress = np.repeat(raw_progress[0], N)[..., None]
    # combined_progress = np.concatenate([raw_progress[..., None], pdm_progress], axis=1)
    # max_raw_progress = np.max(
    #     combined_progress,
    #     axis=1
    # )
    # combined_progress
    #tl是越小越好，所以tl越小，分数越大
    # three cases:
    # 1. bigger than t ---------- normalize（x）best->(1-0/max) pdm->(1-0.2/max) ours->(1-0.38/max) 这里的max如何选取？
    # 2. smaller than t & score!=0-------- 1 停车（特殊判断，放最后）max<t
    # 3. bigger than t' & score==0 -------- 0 太大了，分数为0  max> (其实不用管）

    normalized_progress = np.ones_like(progress_in_meter)
    #设置一个阈值，如果偏移量小于这个阈值，就设置为满分
    # good_score = multiplicate_metric_scores != 0.0
    smaller_than_dis_mask = progress_in_meter <= 0.5  # (4096,1)
    bigger_than_dis_mask = np.logical_not(smaller_than_dis_mask)
    normalized_progress[bigger_than_dis_mask] = 0.5 / progress_in_meter[bigger_than_dis_mask]
    normalized_progress[smaller_than_dis_mask] = 1.0

    # normalized_progress = progress_in_meter / max_raw_progress
    # average_distance = total_distance / count
    return normalized_progress