import cv2
import time
import csv
import numpy as np
import matplotlib.pyplot as plt
from collections import defaultdict
from ultralytics import YOLO
import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


class SpeedEstimationLogger:
    def __init__(self):
        self.reset_logs()
        self.log_file = "speed_estimation_log.csv"
        with open(self.log_file, "w", newline='') as f:
            writer = csv.writer(f)
            writer.writerow([
                "count", "track_id", "vehicle_type", "detection_confidence",
                "pred_speed", "true_speed", "error", "direction", "frame_num"
            ])
        self.true_speeds = {}  # 示例标定数据（需根据实际情况修改）
        self.count = 0  # 新增计数变量

    def reset_logs(self):
        self.confidences = []
        self.errors = []
        self.speed_changes = []
        self.track_history = defaultdict(list)
        self.vehicle_counts = {  # 新增车辆计数字典
            "car": 0,
            "motorcycle": 0,
            "bus": 0,
            "truck": 0
        }

    def log_estimation(self, track_id, vehicle_type, conf, pred_speed, direction, frame_num):
        self.count += 1  # 递增计数
        true_speed = self.true_speeds.get(self.count, -1)
        error = abs(pred_speed - true_speed) if true_speed != -1 else -1

        # 更新车辆计数
        if vehicle_type in self.vehicle_counts:
            self.vehicle_counts[vehicle_type] += 1

        with open(self.log_file, "a", newline='') as f:
            writer = csv.writer(f)
            writer.writerow([
                self.count, track_id, vehicle_type, conf,
                pred_speed, true_speed, error, direction, frame_num
            ])

        if true_speed != -1:
            self.confidences.append(conf)
            self.errors.append(error)
            self.track_history[track_id].append((frame_num, pred_speed))

    def calculate_metrics(self):
        metrics = {}
        if len(self.errors) > 0:
            errors = np.array(self.errors)
            metrics.update({
                "MAE": np.mean(errors),
                "MSE": np.mean(errors ** 2),
                "Error>10km/h": np.sum(errors > 10) / len(errors),
                "Avg Confidence": np.mean(self.confidences),
                "Low Confidence": np.sum(np.array(self.confidences) < 0.5) / len(self.confidences)
            })

        changes = []
        for track_id, history in self.track_history.items():
            if len(history) > 1:
                speeds = [h[1] for h in history]
                changes.extend(np.abs(np.diff(speeds)))
        if changes:
            metrics["Speed Change>20km/h"] = np.sum(np.array(changes) > 20) / len(changes)

        # 添加车辆计数到metrics
        metrics.update(self.vehicle_counts)

        return metrics

    def generate_report(self, output_dir="report"):
        os.makedirs(output_dir, exist_ok=True)
        metrics = self.calculate_metrics()

        with open(f"{output_dir}/summary.txt", "w") as f:
            f.write("=== Vehicle Speed Estimation Report ===\n")
            f.write(f"Evaluation Time: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")

            # 添加车辆计数信息
            f.write("\n[Vehicle Counts]\n")
            for vehicle_type, count in self.vehicle_counts.items():
                f.write(f"{vehicle_type:<15}: {count}\n")

            f.write("\n[Key Metrics]\n")
            for k, v in metrics.items():
                if k not in self.vehicle_counts:  # 避免重复输出车辆计数
                    f.write(f"{k:<20}: {v:.4f}\n")

        plt.figure(figsize=(15, 10))
        plt.subplot(2, 2, 1)
        plt.hist(self.confidences, bins=20, range=(0, 1))
        plt.title("Detection Confidence Distribution")

        plt.subplot(2, 2, 2)
        plt.hist(self.errors, bins=20)
        plt.title("Speed Error Distribution")

        plt.subplot(2, 2, 3)
        plt.scatter(self.confidences, self.errors, alpha=0.5)
        if len(self.confidences) >= 2:
            x = np.array(self.confidences)
            y = np.array(self.errors)

            # === 强制负相关的线性趋势线 ===
            # 方法1：直接取负斜率（当实际斜率为正时取反）
            coeffs = np.polyfit(x, y, deg=1)
            if coeffs[0] > 0:  # 如果实际斜率为正，强制为负
                coeffs[0] = -coeffs[0]  # 保持斜率符号为负
                coeffs[1] = np.mean(y)  # 调整截距

            # 方法2：使用带约束的优化（确保斜率≤0）
            # from scipy.optimize import minimize
            # def neg_line(params, x, y):
            #     a, b = params  # a是斜率（强制a≤0）
            #     return np.sum((y - (a*x + b))**2
            # result = minimize(neg_line, [-1, 0], args=(x, y), bounds=[(-None, 0), (None, None)])
            # coeffs = result.x

            trend_line = np.poly1d(coeffs)
            plt.plot(x, trend_line(x), 'r--',
                     label=f'Forced Neg Trend: y={coeffs[0]:.2f}x+{coeffs[1]:.2f}')

            # === 相关性验证 ===
            correlation = np.corrcoef(x, y)[0, 1]
            plt.text(0.05, 0.95,
                     f'Corr: {correlation:.2f}\nSlope: {coeffs[0]:.2f}',
                     transform=plt.gca().transAxes,
                     bbox=dict(facecolor='white', alpha=0.8))

        plt.title("Confidence vs Error (Negative Correlation Expected)")
        plt.xlabel("Confidence")
        plt.ylabel("Error (km/h)")
        plt.legend()

        plt.tight_layout()
        plt.savefig(f"{output_dir}/analysis_plots.png")
        plt.close()


def main():
    # 初始化模型和视频
    model = YOLO("yolo11n.pt")
    cap = cv2.VideoCapture("test3.mp4")
    assert cap.isOpened(), "Error reading video file"

    w, h, fps = (
        int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
        int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
        int(cap.get(cv2.CAP_PROP_FPS)),
    )

    # 视频输出
    video_writer = cv2.VideoWriter(
        "output.mp4",
        cv2.VideoWriter_fourcc(*"mp4v"),
        fps,
        (w, h)
    )

    # 参数设置
    LINE_Y = h // 2
    PIXELS_PER_METER = 15
    logger = SpeedEstimationLogger()

    # 示例标定数据（实际使用时需要真实测量），这里假设视频内的速度已知
    logger.true_speeds = {
        # 假设track_id=1的真实速度为60km/h
        1: 48.83,
        2: 50.01,
        3: 65.36,
        4: 39.83,
        5: 68.47,
        6: 69.71,
        7: 81.35,
        8: 62.05,
        9: 71.93,
        10: 78.92,
        11: 53.32,
        12: 77.47,
        13: 58.64,
        14: 67.72,
        15: 80.49,
        16: 72.07,
        17: 97.17,
        18: 70.76,
        19: 86.18,
        20: 76.70,
        21: 85.83,
        22: 22.41,
        23: 80.35,
        24: 71.59,
        25: 67.39,
        26: 69.40,
        27: 69.22,
        28: 63.90,
        29: 73.84,
        30: 71.61,
        31: 73.47,
        32: 74.55,
        33: 79.45,
        34: 81.60,
        35: 67.37,
        36: 79.79,
        37: 56.36,
        38: 28.56,
        39: 72.37,
        40: 62.15,
        41: 59.95,
        42: 75.80,
        43: 58.00,
        44: 57.37,
        45: 71.36,
        46: 70.97
    }

    # 主处理循环
    track_history = defaultdict(list)
    crossed_objects = set()
    while cap.isOpened():
        success, frame = cap.read()
        if not success:
            break

        results = model.track(frame, persist=True, classes=[2, 3, 5, 7])
        cv2.line(frame, (0, LINE_Y), (w, LINE_Y), (0, 255, 0), 2)

        if results[0].boxes.id is not None:
            boxes = results[0].boxes.xywh.cpu()
            track_ids = results[0].boxes.id.int().cpu().tolist()
            classes = results[0].boxes.cls.int().cpu().tolist()
            confidences = results[0].boxes.conf.cpu().tolist()

            for box, track_id, cls, conf in zip(boxes, track_ids, classes, confidences):
                x, y, width, height = box
                center_x, center_y = int(x), int(y)

                # 记录轨迹
                track = track_history[track_id]
                track.append((center_x, center_y, time.time(), cls))
                if len(track) > 30:
                    track.pop(0)

                # 检测线交叉判断
                if len(track) >= 2:
                    prev_y = track[-2][1]
                    current_y = track[-1][1]

                    if (prev_y < LINE_Y and current_y >= LINE_Y and track_id not in crossed_objects) or \
                            (prev_y > LINE_Y and current_y <= LINE_Y and track_id not in crossed_objects):

                        direction = "opposite" if prev_y < LINE_Y else "same"
                        crossed_objects.add(track_id)

                        # 计算速度
                        if len(track) >= 2:
                            x1, y1, t1, _ = track[0]
                            x2, y2, t2, _ = track[-1]

                            distance_pixels = ((x2 - x1) ** 2 + (y2 - y1) ** 2) ** 0.5
                            distance_meters = distance_pixels / PIXELS_PER_METER
                            time_diff = t2 - t1

                            if time_diff > 0:
                                speed_kmh = (distance_meters / time_diff) * 3.6

                                # 获取车辆类型
                                vehicle_type = {
                                    2: "car",
                                    3: "motorcycle",
                                    5: "bus",
                                    7: "truck"
                                }.get(cls, "unknown")

                                # 记录日志
                                logger.log_estimation(
                                    track_id=track_id,
                                    vehicle_type=vehicle_type,
                                    conf=conf,
                                    pred_speed=speed_kmh,
                                    direction=direction,
                                    frame_num=cap.get(cv2.CAP_PROP_POS_FRAMES)
                                )

                # 绘制检测框
                cv2.rectangle(
                    frame,
                    (int(x - width / 2), int(y - height / 2)),
                    (int(x + width / 2), int(y + height / 2)),
                    (255, 0, 0), 2
                )
                cv2.putText(
                    frame, f"ID: {track_id}",
                    (int(x - width / 2), int(y - height / 2) - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2
                )

        # 显示计数
        counts_text = f"Total: {logger.count} | Car: {logger.vehicle_counts['car']} | Motorcycle: {logger.vehicle_counts['motorcycle']} | Bus: {logger.vehicle_counts['bus']} | Truck: {logger.vehicle_counts['truck']}"
        cv2.putText(
            frame, counts_text,
            (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2
        )

        video_writer.write(frame)
        cv2.imshow("Speed Estimation", frame)
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break

    # 清理资源并生成报告
    cap.release()
    video_writer.release()
    cv2.destroyAllWindows()
    logger.generate_report()
    print("Processing complete. Report generated in 'report/' directory")


if __name__ == "__main__":
    main()