import pandas as pd
from config import ROAD_FUNCTION_LEVEL_MAP

class SpeedingViolation(TrafficViolation):
    """超速违规类"""

    def __init__(self, df_data: pd.DataFrame):
        super().__init__()  # 调用父类构造函数
        self.traffic_violations_type = "超速违规类"
        self.data = df_data
        
        # 初始化违规统计
        self.violation_counts = {
            'over_speed_50': 0,
            'over_speed_20_50': 0,
            'over_speed_20': 0,
            'under_speed_set': 0
        }
        
        # 处理数据
        self.process_violations()

    def process_violations(self):
        """处理数据帧，检查是否存在超速违规"""
        # 定义速度限制
        self.data['speed_limit_max'] = self.data['road_speed_max']
        self.data['speed_limit_min'] = self.data['road_speed_min']

        # 提取有效道路类型（高速和城市快速路）
        valid_road_types = [ROAD_FUNCTION_LEVEL_MAP[1], ROAD_FUNCTION_LEVEL_MAP[2]]
        
        # 违规判定
        conditions = [
            (self.data['road_type'].isin(valid_road_types)) & 
            (self.data['v'] > self.data['speed_limit_max'] * 1.5),  # 超速50%
            (self.data['road_type'].isin(valid_road_types)) & 
            (self.data['v'] > self.data['speed_limit_max'] * 1.2) &
            (self.data['v'] <= self.data['speed_limit_max'] * 1.5),  # 超速20%-50%
            (self.data['road_type'].isin(valid_road_types)) & 
            (self.data['v'] > self.data['speed_limit_max']) &
            (self.data['v'] <= self.data['speed_limit_max'] * 1.2),  # 超速20%
            (self.data['road_type'].isin(valid_road_types)) & 
            (self.data['v'] < self.data['speed_limit_min'])  # 低于最低时速
        ]

        violation_types = ['heavy_speeding', 'moderate_speeding', 'minor_speeding', 'below_minimum_speed']
        
        # 设置违规类型
        for condition, violation_type in zip(conditions, violation_types):
            self.data.loc[condition, 'violation_type'] = violation_type

        # 统计违规情况
        self.violation_counts['over_speed_50'] = self.data['violation_type'].value_counts().get('heavy_speeding', 0)
        self.violation_counts['over_speed_20_50'] = self.data['violation_type'].value_counts().get('moderate_speeding', 0)
        self.violation_counts['over_speed_20'] = self.data['violation_type'].value_counts().get('minor_speeding', 0)
        self.violation_counts['under_speed_set'] = (self.data['v'] < self.data['speed_limit_min']).sum()

    def get_violation_info(self) -> str:
        return (
            f"{super().get_violation_info()} - 超速统计: "
            f"超过50%: {self.violation_counts['over_speed_50']}, "
            f"超过20%-50%: {self.violation_counts['over_speed_20_50']}, "
            f"超过20%以内: {self.violation_counts['over_speed_20']}, "
            f"低于最低时速: {self.violation_counts['under_speed_set']}"
        )