import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import matplotlib.patheffects as path_effects
from mpl_toolkits.mplot3d import Axes3D


# ==========================================
# 核心理论模块：通信迟滞容忍的动力学风险场发生器
# (Latency-Tolerant Kinematic Risk Field Generator)
# ==========================================
def generate_kinematic_risk_field(X, Y, obj):
    """
    根据车辆的物理状态和通信迟滞，生成高斯风险势能场
    """
    x0, y0 = obj['pos']
    vx, vy = obj['vel']
    L, W = obj['size']
    dt = obj['latency']  # V2X 通信延迟 (秒)

    # 1. 运动学位置补偿 (由于延迟，目标在真实物理世界已经往前移动了)
    x_real = x0 + vx * dt
    y_real = y0 + vy * dt

    # 2. 时空耦合：协方差膨胀 (速度越快、延迟越高，沿运动方向的不确定性风险拖尾越长)
    speed = np.hypot(vx, vy)
    # 基础物理边界 (方差，代表车辆本身的尺寸)
    sigma_x = L / 2.0
    sigma_y = W / 2.0

    if speed > 0.1:
        # ⚠️ 顶刊核心公式：沿着运动方向剧烈拉伸协方差！
        stretch_factor = 1.0 + 0.8 * speed * dt
        sigma_x_dilated = sigma_x * stretch_factor
        theta = np.arctan2(vy, vx)
    else:
        sigma_x_dilated = sigma_x
        theta = 0.0

    # 3. 旋转协方差矩阵，对齐到运动方向
    cos_t, sin_t = np.cos(theta), np.sin(theta)
    R = np.array([[cos_t, -sin_t],
                  [sin_t, cos_t]])
    S = np.array([[sigma_x_dilated ** 2, 0],
                  [0, sigma_y ** 2]])
    Cov = R @ S @ R.T  # 膨胀后的 2D 协方差矩阵
    Cov_inv = np.linalg.inv(Cov)

    # 4. 计算二维高斯势能曲面 (Mahalanobis Distance)
    dx = X - x_real
    dy = Y - y_real
    # 矢量化二次型计算
    E = np.exp(-0.5 * (Cov_inv[0, 0] * dx ** 2 + 2 * Cov_inv[0, 1] * dx * dy + Cov_inv[1, 1] * dy ** 2))
    return E * obj['risk_weight']


# ==========================================
# 场景构建：BEV 上帝视角物理空间初始化
# ==========================================
# 设定一个 50米 x 50米 的路口物理网格 (高分辨率)
x_grid = np.linspace(0, 50, 300)
y_grid = np.linspace(0, 50, 300)
X, Y = np.meshgrid(x_grid, y_grid)

# 定义场景中的交通参与者 (完美复现你的任务书痛点)
objects = [
    {
        'name': '高速来车 (带极端V2X延迟)',
        'pos': (10, 25), 'vel': (18, 0),  # 速度 18m/s (约65km/h)向右
        'size': (4.8, 2.0), 'latency': 0.4,  # 恐怖的 400ms 网络延迟！
        'risk_weight': 1.0
    },
    {
        'name': '非标事件: 散落物/掉落轮胎',
        'pos': (35, 12), 'vel': (0, 0),  # 静止
        'size': (1.5, 1.5), 'latency': 0.0,
        'risk_weight': 0.8
    },
    {
        'name': '鬼探头行人 (突然窜出)',
        'pos': (28, 42), 'vel': (0, -4),  # 速度 4m/s 向下横穿
        'size': (0.8, 0.8), 'latency': 0.1,
        'risk_weight': 0.9
    }
]

# ==========================================
# 计算全局连续风险势能场 (Superposition of Risk Fields)
# ==========================================
Total_Risk_Field = np.zeros_like(X)
for obj in objects:
    E_obj = generate_kinematic_risk_field(X, Y, obj)
    Total_Risk_Field = np.maximum(Total_Risk_Field, E_obj)  # 多风险源叠加取极值

# ==========================================
# 惊艳大招：计算势能梯度（多车协同调速排斥力场） F = -∇E
# ==========================================
dEy, dEx = np.gradient(Total_Risk_Field)
Force_X = -dEx
Force_Y = -dEy

# ==========================================
# 顶刊级数据可视化渲染 (Matplotlib 画图)
# ==========================================
fig = plt.figure(figsize=(16, 7), facecolor='#111111')  # 暗黑极客底色
fig.suptitle("End-to-End Continuous Risk Potential Field for V2X Cooperative Driving",
             fontsize=18, fontweight='bold', color='white', y=0.98)

# --- 子图 1: 3D 连续风险势能面 ---
ax1 = fig.add_subplot(1, 2, 1, projection='3d')
ax1.set_facecolor('#111111')
surf = ax1.plot_surface(X, Y, Total_Risk_Field, cmap=cm.inferno, alpha=0.9, rstride=3, cstride=3, linewidth=0,
                        antialiased=True)

ax1.set_title("(a) 3D Spatiotemporal Risk Surface\n【任务2】时空演化连续风险曲面", fontsize=14, fontweight='bold',
              color='white', pad=15)
ax1.set_xlabel("BEV X-Coordinate (m)", color='white')
ax1.set_ylabel("BEV Y-Coordinate (m)", color='white')
ax1.set_zlabel("Risk Potential Energy ($E_{risk}$)", color='white')
ax1.tick_params(colors='white')
ax1.xaxis.pane.fill = False
ax1.yaxis.pane.fill = False
ax1.zaxis.pane.fill = False
ax1.view_init(elev=40, azim=-45)  # 绝佳的观察视角

# --- 子图 2: 2D 梯度力场与规控闭环 ---
ax2 = fig.add_subplot(1, 2, 2)
ax2.set_facecolor('#111111')
# 画势能等高线
contour = ax2.contourf(X, Y, Total_Risk_Field, levels=30, cmap=cm.inferno, alpha=0.8)

# 画排斥力场 (Quiver 矢量箭头)
step = 10  # 箭头采样稀疏度
Q = ax2.quiver(X[::step, ::step], Y[::step, ::step], Force_X[::step, ::step], Force_Y[::step, ::step],
               color='cyan', scale=1.5, width=0.003, alpha=0.9)


# 标注解释
def add_label(x, y, text):
    ax2.text(x, y, text, color='white', ha='center', fontsize=10,
             path_effects=[path_effects.withStroke(linewidth=2, foreground='k')])


add_label(35, 15, "Static Debris\n(Isotropic Field)")
add_label(20, 25, "Latency-Dilated Comet Tail\n($\Delta t = 400ms$)")
add_label(28, 38, "Crossing Pedestrian")

ax2.set_title("(b) Gradient-Driven Repulsive Force Field ($-\\nabla E$)\n【任务3】风险梯度驱动的多车协同调速诱导力场",
              fontsize=14, fontweight='bold', color='white', pad=15)
ax2.set_xlabel("BEV X-Coordinate (m)", color='white')
ax2.set_ylabel("BEV Y-Coordinate (m)", color='white')
ax2.tick_params(colors='white')
ax2.set_xlim(0, 50)
ax2.set_ylim(0, 50)
ax2.set_aspect('equal')
ax2.grid(color='white', linestyle='--', linewidth=0.3, alpha=0.2)

# 添加 Colorbar
cbar = fig.colorbar(surf, ax=[ax1, ax2], shrink=0.5, aspect=15, pad=0.05)
cbar.set_label('Collision Probability / Risk Intensity', color='white')
cbar.ax.yaxis.set_tick_params(color='white')
plt.setp(plt.getp(cbar.ax.axes, 'yticklabels'), color='white')

plt.tight_layout(pad=3.0)
plt.show()
# 如果你想保存高清原图，取消下一行的注释
# plt.savefig("v2x_risk_field.png", dpi=300, facecolor=fig.get_facecolor(), edgecolor='none')