zhaoyuan/modules/metric/comfort.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
##################################################################
#
# Copyright (c) 2023 CICV, Inc. All Rights Reserved
#
##################################################################
"""
@Authors:           zhanghaiwen(zhanghaiwen@china-icv.cn), yangzihao(yangzihao@china-icv.cn)
@Data:              2023/06/25
@Last Modified:     2023/06/25
@Summary:           Comfort metrics
"""

import sys
import math
import pandas as pd
import numpy as np
import scipy.signal
from pathlib import Path 
from typing import Dict, List, Any, Optional, Callable, Union, Tuple

from modules.lib.score import Score
from modules.lib.common import get_interpolation, get_frame_with_time
from modules.lib import data_process

from modules.lib.log_manager import LogManager

COMFORT_INFO = [
    "simTime",
    "simFrame",
    "speedX",
    "speedY",
    "accelX",
    "accelY",
    "curvHor",
    "lightMask",
    "v",
    "lat_acc",
    "lon_acc",
    "time_diff",
    "lon_acc_diff",
    "lon_acc_roc",
    "speedH",
    "accelH",
]
# ----------------------
# 独立指标计算函数
# ----------------------
def weaving(data_processed) -> dict:
    """计算蛇行指标"""
    comfort = ComfortCalculator(data_processed)
    zigzag_count = comfort.calculate_zigzag_count()
    return {"weaving": float(zigzag_count)}

def shake(data_processed) -> dict:
    """计算晃动指标"""
    comfort = ComfortCalculator(data_processed)
    shake_count = comfort.calculate_shake_count()
    return {"shake": float(shake_count)}

def cadence(data_processed) -> dict:
    """计算顿挫指标"""
    comfort = ComfortCalculator(data_processed)
    cadence_count = comfort.calculate_cadence_count()
    return {"cadence": float(cadence_count)}

def slamBrake(data_processed) -> dict:
    """计算急刹车指标"""
    comfort = ComfortCalculator(data_processed)
    slam_brake_count = comfort.calculate_slam_brake_count()
    return {"slamBrake": float(slam_brake_count)}

def slamAccelerate(data_processed) -> dict:
    """计算急加速指标"""
    comfort = ComfortCalculator(data_processed)
    slam_accel_count = comfort.calculate_slam_accel_count()
    return {"slamAccelerate": float(slam_accel_count)}


# 装饰器保持不变
def peak_valley_decorator(method):
    def wrapper(self, *args, **kwargs):
        peak_valley = self._peak_valley_determination(self.df)
        pv_list = self.df.loc[peak_valley, ['simTime', 'speedH']].values.tolist()
        if len(pv_list) != 0:
            flag = True
            p_last = pv_list[0]

            for i in range(1, len(pv_list)):
                p_curr = pv_list[i]

                if self._peak_valley_judgment(p_last, p_curr):
                    # method(self, p_curr, p_last)
                    method(self, p_curr, p_last, flag, *args, **kwargs)
                else:
                    p_last = p_curr

            return method
        else:
            flag = False
            p_curr = [0, 0]
            p_last = [0, 0]
            method(self, p_curr, p_last, flag, *args, **kwargs)
            return method

    return wrapper


class ComfortRegistry:
    """舒适性指标注册器"""
    
    def __init__(self, data_processed):
        self.logger = LogManager().get_logger()  # 获取全局日志实例
        self.data = data_processed
        self.comfort_config = data_processed.comfort_config["comfort"]
        self.metrics = self._extract_metrics(self.comfort_config)
        self._registry = self._build_registry()
    
    def _extract_metrics(self, config_node: dict) -> list:
        """DFS遍历提取指标"""
        metrics = []
        def _recurse(node):
            if isinstance(node, dict):
                if 'name' in node and not any(isinstance(v, dict) for v in node.values()):
                    metrics.append(node['name'])
                for v in node.values():
                    _recurse(v)
        _recurse(config_node)
        self.logger.info(f'评比的舒适性指标列表:{metrics}')
        return metrics
    
    def _build_registry(self) -> dict:
        """自动注册指标函数"""
        registry = {}
        for metric_name in self.metrics:
            try:
                registry[metric_name] = globals()[metric_name]
            except KeyError:
                self.logger.error(f"未实现指标函数: {metric_name}")
        return registry
    
    def batch_execute(self) -> dict:
        """批量执行指标计算"""
        results = {}
        for name, func in self._registry.items():
            try:
                result = func(self.data)
                results.update(result)
            except Exception as e:
                self.logger.error(f"{name} 执行失败: {str(e)}", exc_info=True)
                results[name] = None
        self.logger.info(f'舒适性指标计算结果:{results}')
        return results


class ComfortCalculator:
    """舒适性指标计算类 - 提供核心计算功能"""
    
    def __init__(self, data_processed):
        self.data_processed = data_processed
        self.logger = LogManager().get_logger()
        
        self.data = data_processed.ego_data
        self.ego_df = pd.DataFrame()
        self.discomfort_df = pd.DataFrame(columns=['start_time', 'end_time', 'start_frame', 'end_frame', 'type'])
        
        self.time_list = self.data['simTime'].values.tolist()
        self.frame_list = self.data['simFrame'].values.tolist()
        
        self.zigzag_count = 0
        self.shake_count = 0
        self.cadence_count = 0
        self.slam_brake_count = 0
        self.slam_accel_count = 0
        
        self.zigzag_time_list = []
        self.zigzag_stre_list = []
        self.cur_ego_path_list = []
        self.curvature_list = []
        
        self._initialize_data()
    
    def _initialize_data(self):
        """初始化数据"""
        self.ego_df = self.data[COMFORT_INFO].copy()
        self.df = self.ego_df.reset_index(drop=True)
        self._prepare_comfort_parameters()
    
    def _prepare_comfort_parameters(self):
        """准备舒适性计算所需参数"""
        # 计算加减速阈值
        self.ego_df['ip_acc'] = self.ego_df['v'].apply(get_interpolation, point1=[18, 4], point2=[72, 2])
        self.ego_df['ip_dec'] = self.ego_df['v'].apply(get_interpolation, point1=[18, -5], point2=[72, -3.5])
        self.ego_df['slam_brake'] = (self.ego_df['lon_acc'] - self.ego_df['ip_dec']).apply(
            lambda x: 1 if x < 0 else 0)
        self.ego_df['slam_accel'] = (self.ego_df['lon_acc'] - self.ego_df['ip_acc']).apply(
            lambda x: 1 if x > 0 else 0)
        self.ego_df['cadence'] = self.ego_df.apply(
            lambda row: self._cadence_process_new(row['lon_acc'], row['ip_acc'], row['ip_dec']), axis=1)

        # 计算曲率相关参数
        self.ego_df['cur_ego_path'] = self.ego_df.apply(self._cal_cur_ego_path, axis=1)
        self.ego_df['curvHor'] = self.ego_df['curvHor'].astype('float')
        self.ego_df['cur_diff'] = (self.ego_df['cur_ego_path'] - self.ego_df['curvHor']).abs()
        self.ego_df['R'] = self.ego_df['curvHor'].apply(lambda x: 10000 if x == 0 else 1 / x)
        self.ego_df['R_ego'] = self.ego_df['cur_ego_path'].apply(lambda x: 10000 if x == 0 else 1 / x)
        self.ego_df['R_diff'] = (self.ego_df['R_ego'] - self.ego_df['R']).abs()
        
        self.cur_ego_path_list = self.ego_df['cur_ego_path'].values.tolist()
        self.curvature_list = self.ego_df['curvHor'].values.tolist()
    
    def _cal_cur_ego_path(self, row):
        """计算车辆轨迹曲率"""
        try:
            divide = (row['speedX'] ** 2 + row['speedY'] ** 2) ** (3 / 2)
            if not divide:
                res = None
            else:
                res = (row['speedX'] * row['accelY'] - row['speedY'] * row['accelX']) / divide
        except:
            res = None
        return res
    
    def _peak_valley_determination(self, df):
        """确定角速度的峰谷"""
        peaks, _ = scipy.signal.find_peaks(df['speedH'], height=0.01, distance=1, prominence=0.01)
        valleys, _ = scipy.signal.find_peaks(-df['speedH'], height=0.01, distance=1, prominence=0.01)
        peak_valley = sorted(list(peaks) + list(valleys))
        return peak_valley
    
    def _peak_valley_judgment(self, p_last, p_curr, tw=10000, avg=0.02):
        """判断峰谷是否满足蛇行条件"""
        t_diff = p_curr[0] - p_last[0]
        v_diff = abs(p_curr[1] - p_last[1])
        s = p_curr[1] * p_last[1]

        zigzag_flag = t_diff < tw and v_diff > avg and s < 0
        if zigzag_flag and ([p_last[0], p_curr[0]] not in self.zigzag_time_list):
            self.zigzag_time_list.append([p_last[0], p_curr[0]])
        return zigzag_flag
    
    def _cadence_process_new(self, lon_acc, ip_acc, ip_dec):
        """处理顿挫数据"""
        if abs(lon_acc) < 1 or lon_acc > ip_acc or lon_acc < ip_dec:
            return np.nan
        elif abs(lon_acc) == 0:
            return 0
        elif lon_acc > 0 and lon_acc < ip_acc:
            return 1
        elif lon_acc < 0 and lon_acc > ip_dec:
            return -1
        else:
            return 0
    
    @peak_valley_decorator
    def _zigzag_count_func(self, p_curr, p_last, flag=True):
        """计算蛇行次数"""
        if flag:
            self.zigzag_count += 1
        else:
            self.zigzag_count += 0
    
    @peak_valley_decorator
    def _cal_zigzag_strength(self, p_curr, p_last, flag=True):
        """计算蛇行强度"""
        if flag:
            v_diff = abs(p_curr[1] - p_last[1])
            t_diff = p_curr[0] - p_last[0]
            self.zigzag_stre_list.append(v_diff / t_diff)  # 平均角加速度
        else:
            self.zigzag_stre_list = []
    
    def calculate_zigzag_count(self):
        """计算蛇行指标"""
        self._zigzag_count_func()
        return self.zigzag_count
    
    def calculate_shake_count(self):
        """计算晃动指标"""
        self._shake_detector()
        return self.shake_count
    
    def calculate_cadence_count(self):
        """计算顿挫指标"""
        self._cadence_detector()
        return self.cadence_count
    
    def calculate_slam_brake_count(self):
        """计算急刹车指标"""
        self._slam_brake_detector()
        return self.slam_brake_count
    
    def calculate_slam_accel_count(self):
        """计算急加速指标"""
        self._slam_accel_detector()
        return self.slam_accel_count
    
    def _shake_detector(self, Cr_diff=0.05, T_diff=0.39):
        """晃动检测器"""
        time_list = []
        frame_list = []

        df = self.ego_df.copy()
        df = df[df['cur_diff'] > Cr_diff]
        df['frame_ID_diff'] = df['simFrame'].diff()
        filtered_df = df[df.frame_ID_diff > T_diff]

        row_numbers = filtered_df.index.tolist()
        cut_column = pd.cut(df.index, bins=row_numbers)

        grouped = df.groupby(cut_column)
        dfs = {}
        for name, group in grouped:
            dfs[name] = group.reset_index(drop=True)

        for name, df_group in dfs.items():
            # 直道，未主动换道
            df_group['curvHor'] = df_group['curvHor'].abs()
            df_group_straight = df_group[(df_group.lightMask == 0) & (df_group.curvHor < 0.001)]
            if not df_group_straight.empty:
                time_list.extend(df_group_straight['simTime'].values)
                frame_list.extend(df_group_straight['simFrame'].values)
                self.shake_count = self.shake_count + 1

            # 打转向灯，道路为直道
            df_group_change_lane = df_group[(df_group['lightMask'] != 0) & (df_group['curvHor'] < 0.001)]
            df_group_change_lane_data = df_group_change_lane[df_group_change_lane.cur_diff > Cr_diff + 0.2]
            if not df_group_change_lane_data.empty:
                time_list.extend(df_group_change_lane_data['simTime'].values)
                frame_list.extend(df_group_change_lane_data['simFrame'].values)
                self.shake_count = self.shake_count + 1

            # 转弯，打转向灯
            df_group_turn = df_group[(df_group['lightMask'] != 0) & (df_group['curvHor'].abs() > 0.001)]
            df_group_turn_data = df_group_turn[df_group_turn.cur_diff.abs() > Cr_diff + 0.1]
            if not df_group_turn_data.empty:
                time_list.extend(df_group_turn_data['simTime'].values)
                frame_list.extend(df_group_turn_data['simFrame'].values)
                self.shake_count = self.shake_count + 1

        # 分组处理
        TIME_RANGE = 1
        t_list = time_list
        f_list = frame_list
        group_time = []
        group_frame = []
        sub_group_time = []
        sub_group_frame = []
        
        if len(f_list) > 0:
            for i in range(len(f_list)):
                if not sub_group_time or t_list[i] - t_list[i - 1] <= TIME_RANGE:
                    sub_group_time.append(t_list[i])
                    sub_group_frame.append(f_list[i])
                else:
                    group_time.append(sub_group_time)
                    group_frame.append(sub_group_frame)
                    sub_group_time = [t_list[i]]
                    sub_group_frame = [f_list[i]]

            group_time.append(sub_group_time)
            group_frame.append(sub_group_frame)

        # 输出图表值
        shake_time = [[g[0], g[-1]] for g in group_time]
        shake_frame = [[g[0], g[-1]] for g in group_frame]
        self.shake_count = len(shake_time)

        if shake_time:
            time_df = pd.DataFrame(shake_time, columns=['start_time', 'end_time'])
            frame_df = pd.DataFrame(shake_frame, columns=['start_frame', 'end_frame'])
            discomfort_df = pd.concat([time_df, frame_df], axis=1)
            discomfort_df['type'] = 'shake'
            self.discomfort_df = pd.concat([self.discomfort_df, discomfort_df], ignore_index=True)

        return time_list
    
    def _cadence_detector(self):
        """顿挫检测器"""
        data = self.ego_df[['simTime', 'simFrame', 'lon_acc', 'lon_acc_roc', 'cadence']].copy()
        time_list = data['simTime'].values.tolist()

        data = data[data['cadence'] != np.nan]
        data['cadence_diff'] = data['cadence'].diff()
        data.dropna(subset='cadence_diff', inplace=True)
        data = data[data['cadence_diff'] != 0]

        t_list = data['simTime'].values.tolist()
        f_list = data['simFrame'].values.tolist()

        TIME_RANGE = 1
        group_time = []
        group_frame = []
        sub_group_time = []
        sub_group_frame = []
        for i in range(len(f_list)):
            if not sub_group_time or t_list[i] - t_list[i - 1] <= TIME_RANGE:  # 特征点相邻一秒内的，算作同一组顿挫
                sub_group_time.append(t_list[i])
                sub_group_frame.append(f_list[i])
            else:
                group_time.append(sub_group_time)
                group_frame.append(sub_group_frame)
                sub_group_time = [t_list[i]]
                sub_group_frame = [f_list[i]]

        group_time.append(sub_group_time)
        group_frame.append(sub_group_frame)
        group_time = [g for g in group_time if len(g) >= 1]  # 有一次特征点则算作一次顿挫
        group_frame = [g for g in group_frame if len(g) >= 1]

        # 输出图表值
        cadence_time = [[g[0], g[-1]] for g in group_time]
        cadence_frame = [[g[0], g[-1]] for g in group_frame]

        if cadence_time:
            time_df = pd.DataFrame(cadence_time, columns=['start_time', 'end_time'])
            frame_df = pd.DataFrame(cadence_frame, columns=['start_frame', 'end_frame'])
            discomfort_df = pd.concat([time_df, frame_df], axis=1)
            discomfort_df['type'] = 'cadence'
            self.discomfort_df = pd.concat([self.discomfort_df, discomfort_df], ignore_index=True)

        # 将顿挫组的起始时间为组重新统计时间
        cadence_time_list = [time for pair in cadence_time for time in time_list if pair[0] <= time <= pair[1]]

        stre_list = []
        freq_list = []
        for g in group_time:
            # calculate strength
            g_df = data[data['simTime'].isin(g)]
            strength = g_df['lon_acc'].abs().mean()
            stre_list.append(strength)

            # calculate frequency
            cnt = len(g)
            t_start = g_df['simTime'].iloc[0]
            t_end = g_df['simTime'].iloc[-1]
            t_delta = t_end - t_start
            frequency = cnt / t_delta
            freq_list.append(frequency)

        self.cadence_count = len(freq_list)
        cadence_stre = sum(stre_list) / len(stre_list) if stre_list else 0

        return cadence_time_list
    
    def _slam_brake_detector(self):
        """急刹车检测器"""
        data = self.ego_df[['simTime', 'simFrame', 'lon_acc', 'lon_acc_roc', 'ip_dec', 'slam_brake']].copy()
        res_df = data[data['slam_brake'] == 1]
        t_list = res_df['simTime'].values
        f_list = res_df['simFrame'].values.tolist()

        TIME_RANGE = 1
        group_time = []
        group_frame = []
        sub_group_time = []
        sub_group_frame = []
        for i in range(len(f_list)):
            if not sub_group_time or f_list[i] - f_list[i - 1] <= TIME_RANGE:  # 连续帧的算作同一组急刹
                sub_group_time.append(t_list[i])
                sub_group_frame.append(f_list[i])
            else:
                group_time.append(sub_group_time)
                group_frame.append(sub_group_frame)
                sub_group_time = [t_list[i]]
                sub_group_frame = [f_list[i]]

        group_time.append(sub_group_time)
        group_frame.append(sub_group_frame)
        group_time = [g for g in group_time if len(g) >= 2]  # 达到两帧算作一次急刹
        group_frame = [g for g in group_frame if len(g) >= 2]

        # 输出图表值
        slam_brake_time = [[g[0], g[-1]] for g in group_time]
        slam_brake_frame = [[g[0], g[-1]] for g in group_frame]

        if slam_brake_time:
            time_df = pd.DataFrame(slam_brake_time, columns=['start_time', 'end_time'])
            frame_df = pd.DataFrame(slam_brake_frame, columns=['start_frame', 'end_frame'])
            discomfort_df = pd.concat([time_df, frame_df], axis=1)
            discomfort_df['type'] = 'slam_brake'
            self.discomfort_df = pd.concat([self.discomfort_df, discomfort_df], ignore_index=True)

        time_list = [element for sublist in group_time for element in sublist]
        self.slam_brake_count = len(group_time)
        return time_list
    
    def _slam_accel_detector(self):
        """急加速检测器"""
        data = self.ego_df[['simTime', 'simFrame', 'lon_acc', 'ip_acc', 'slam_accel']].copy()
        res_df = data.loc[data['slam_accel'] == 1]
        t_list = res_df['simTime'].values
        f_list = res_df['simFrame'].values.tolist()

        group_time = []
        group_frame = []
        sub_group_time = []
        sub_group_frame = []
        for i in range(len(f_list)):
            if not group_time or f_list[i] - f_list[i - 1] <= 1:  # 连续帧的算作同一组急加速
                sub_group_time.append(t_list[i])
                sub_group_frame.append(f_list[i])
            else:
                group_time.append(sub_group_time)
                group_frame.append(sub_group_frame)
                sub_group_time = [t_list[i]]
                sub_group_frame = [f_list[i]]

        group_time.append(sub_group_time)
        group_frame.append(sub_group_frame)
        group_time = [g for g in group_time if len(g) >= 2]
        group_frame = [g for g in group_frame if len(g) >= 2]

        # 输出图表值
        slam_accel_time = [[g[0], g[-1]] for g in group_time]
        slam_accel_frame = [[g[0], g[-1]] for g in group_frame]

        if slam_accel_time:
            time_df = pd.DataFrame(slam_accel_time, columns=['start_time', 'end_time'])
            frame_df = pd.DataFrame(slam_accel_frame, columns=['start_frame', 'end_frame'])
            discomfort_df = pd.concat([time_df, frame_df], axis=1)
            discomfort_df['type'] = 'slam_accel'
            self.discomfort_df = pd.concat([self.discomfort_df, discomfort_df], ignore_index=True)

        time_list = [element for sublist in group_time for element in sublist]
        self.slam_accel_count = len(group_time)
        return time_list


class ComfortManager:
    """舒适性指标计算主类"""
    
    def __init__(self, data_processed):
        self.data = data_processed
        self.logger = LogManager().get_logger()
        self.registry = ComfortRegistry(self.data)

    def report_statistic(self):
        """生成舒适性评分报告"""
        comfort_result = self.registry.batch_execute()
        # evaluator = Score(self.data.comfort_config)
        # result = evaluator.evaluate(comfort_result) 
        # return result
        return comfort_result


if __name__ == '__main__':
    case_name = 'ICA'
    mode_label = 'PGVIL'
    
    data = data_process.DataPreprocessing(case_name, mode_label)
    comfort_instance = ComfortManager(data)
    
    try:  
        comfort_result = comfort_instance.report_statistic() 
        result = {'comfort': comfort_result}
        print(result) 
    except Exception as e:  
        print(f"An error occurred in Comfort.report_statistic: {e}")