#实战项目三：自动驾驶感知

import torch
import torch.nn as nn

class AutonomousDrivingModel(nn.Module):
    def __init__(self):
        super().__init__()
        
        # 共享编码器
        self.backbone = nn.Sequential(
            nn.Conv2d(3, 64, 7, stride=2, padding=3),
            nn.ReLU(),
            nn.MaxPool2d(3, stride=2),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 256, 3, padding=1),
            nn.ReLU(),
        )
        
        # 任务一：车道线检测
        self.lane_detection = nn.Sequential(
            nn.Conv2d(256, 128, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 1, 1),
            nn.Sigmoid(),
        )
        
        # 任务二：车辆检测
        self.vehicle_detection = nn.Sequential(
            nn.Conv2d(256, 128, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 5, 1),  # 4 个坐标 + 1 个置信度
        )
        
        # 任务三：距离估算
        self.distance_estimation = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 1),
        )
    
    def forward(self, x):
        features = self.backbone(x)
        
        lane = self.lane_detection(features)
        vehicles = self.vehicle_detection(features)
        distance = self.distance_estimation(features)
        
        return lane, vehicles, distance

model = AutonomousDrivingModel()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 损失函数
lane_loss_fn = nn.BCELoss()
vehicle_loss_fn = nn.SmoothL1Loss()
distance_loss_fn = nn.MSELoss()

for epoch in range(100):
    for images, lane_labels, vehicle_labels, distance_labels in train_loader:
        # 前向传播
        lane_pred, vehicle_pred, distance_pred = model(images)
        
        # 计算损失
        lane_loss = lane_loss_fn(lane_pred, lane_labels)
        vehicle_loss = vehicle_loss_fn(vehicle_pred, vehicle_labels)
        distance_loss = distance_loss_fn(distance_pred, distance_labels)
        
        # 加权组合
        total_loss = 0.5 * lane_loss + 0.3 * vehicle_loss + 0.2 * distance_loss
        
        # 反向传播
        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()
    
    print(f"Epoch {epoch+1}, Loss: {total_loss.item():.4f}")

class AutonomousDrivingSystem:
    def __init__(self, model_path):
        self.model = AutonomousDrivingModel()
        self.model.load_state_dict(torch.load(model_path))
        self.model.eval()
    
    def process_frame(self, frame):
        # 预处理
        image = torch.from_numpy(frame).permute(2, 0, 1).float() / 255.0
        image = image.unsqueeze(0)
        
        # 推理
        with torch.no_grad():
            lane, vehicles, distance = self.model(image)
        
        # 后处理
        lane_map = lane.squeeze().numpy()
        vehicle_boxes = vehicles.squeeze().numpy()
        distance_value = distance.item()
        
        # 决策
        decision = self.make_decision(lane_map, vehicle_boxes, distance_value)
        
        return decision
    
    def make_decision(self, lane_map, vehicle_boxes, distance):
        # 车道线中心
        lane_center = np.argmax(np.sum(lane_map, axis=0))
        
        # 转向决策
        if lane_center < 320:
            steering = -0.5  # 左转
        elif lane_center > 320:
            steering = 0.5   # 右转
        else:
            steering = 0.0   # 直行
        
        # 速度决策
        if distance < 10:
            throttle = -1.0  # 制动
        elif distance < 20:
            throttle = 0.0   # 保持
        else:
            throttle = 0.5   # 加速
        
        return {"steering": steering, "throttle": throttle}