# 기본 패키지 설치
pip install stable-baselines3[extra]
pip install gymnasium
pip install carla  # CARLA Python API

# CARLA 서버는 별도로 다운로드 필요
# https://carla.org/에서 CARLA 0.9.15 다운로드

import gymnasium as gym
import numpy as np
import carla

class CarlaEnv(gym.Env):
    """CARLA 기반 자율주행 강화학습 환경"""
    
    def __init__(self):
        super(CarlaEnv, self).__init__()
        
        # CARLA 서버 연결
        self.client = carla.Client('localhost', 2000)
        self.client.set_timeout(10.0)
        self.world = self.client.get_world()
        
        # 행동 공간 정의: [조향각(-1~1), 스로틀(0~1), 브레이크(0~1)]
        self.action_space = gym.spaces.Box(
            low=np.array([-1.0, 0.0, 0.0]),
            high=np.array([1.0, 1.0, 1.0]),
            dtype=np.float32
        )
        
        # 관측 공간 정의: 카메라 이미지 (84x84x3)
        self.observation_space = gym.spaces.Box(
            low=0, high=255,
            shape=(84, 84, 3),
            dtype=np.uint8
        )
        
        self.vehicle = None
        self.camera = None
        self.collision_sensor = None
        self.current_obs = None
        self.collision_occurred = False
        
    def reset(self, seed=None, options=None):
        """에피소드 초기화"""
        self._cleanup()
        
        # 차량 스폰
        blueprint_library = self.world.get_blueprint_library()
        vehicle_bp = blueprint_library.find('vehicle.tesla.model3')
        spawn_points = self.world.get_map().get_spawn_points()
        spawn_point = np.random.choice(spawn_points)
        
        self.vehicle = self.world.spawn_actor(vehicle_bp, spawn_point)
        
        # 카메라 센서 부착
        camera_bp = blueprint_library.find('sensor.camera.rgb')
        camera_bp.set_attribute('image_size_x', '84')
        camera_bp.set_attribute('image_size_y', '84')
        camera_transform = carla.Transform(carla.Location(x=2.5, z=0.7))
        self.camera = self.world.spawn_actor(
            camera_bp, camera_transform, attach_to=self.vehicle
        )
        self.camera.listen(self._process_image)
        
        # 충돌 센서 부착
        collision_bp = blueprint_library.find('sensor.other.collision')
        self.collision_sensor = self.world.spawn_actor(
            collision_bp, carla.Transform(), attach_to=self.vehicle
        )
        self.collision_sensor.listen(self._on_collision)
        
        self.collision_occurred = False
        
        # 초기 관측값 반환
        obs = self._get_obs()
        return obs, {}
    
    def step(self, action):
        """한 스텝 실행"""
        steer, throttle, brake = action
        
        # CARLA 차량 제어
        control = carla.VehicleControl(
            steer=float(steer),
            throttle=float(throttle),
            brake=float(brake)
        )
        self.vehicle.apply_control(control)
        
        # 시뮬레이터 한 틱 진행
        self.world.tick()
        
        # 관측값, 보상, 종료 여부 계산
        obs = self._get_obs()
        reward = self._compute_reward()
        terminated = self.collision_occurred
        truncated = False
        
        return obs, reward, terminated, truncated, {}
    
    def _compute_reward(self):
        """보상 함수 설계 - 핵심 부분!"""
        reward = 0.0
        
        # 충돌 시 큰 패널티
        if self.collision_occurred:
            return -100.0
        
        # 속도 보상 (적절한 속도 유지)
        velocity = self.vehicle.get_velocity()
        speed = np.sqrt(velocity.x**2 + velocity.y**2)
        target_speed = 30 / 3.6  # 30 km/h를 m/s로
        speed_reward = 1.0 - abs(speed - target_speed) / target_speed
        reward += speed_reward * 2.0
        
        # 차선 유지 보상 (waypoint 기반)
        waypoint = self.world.get_map().get_waypoint(
            self.vehicle.get_location()
        )
        lane_offset = abs(
            self.vehicle.get_location().distance(waypoint.transform.location)
        )
        lane_reward = max(0, 1.0 - lane_offset / 2.0)
        reward += lane_reward
        
        return reward
    
    def _process_image(self, image):
        """카메라 이미지 처리"""
        array = np.frombuffer(image.raw_data, dtype=np.uint8)
        array = array.reshape((image.height, image.width, 4))
        self.current_obs = array[:, :, :3]  # RGB만 사용
    
    def _on_collision(self, event):
        self.collision_occurred = True
    
    def _get_obs(self):
        if self.current_obs is None:
            return np.zeros((84, 84, 3), dtype=np.uint8)
        return self.current_obs
    
    def _cleanup(self):
        """액터 정리"""
        if self.camera:
            self.camera.destroy()
        if self.collision_sensor:
            self.collision_sensor.destroy()
        if self.vehicle:
            self.vehicle.destroy()

from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.callbacks import EvalCallback

# 환경 생성
env = DummyVecEnv([lambda: CarlaEnv()])

# PPO 에이전트 생성
model = PPO(
    "CnnPolicy",          # CNN 기반 정책 (이미지 입력)
    env,
    learning_rate=3e-4,
    n_steps=2048,
    batch_size=64,
    n_epochs=10,
    gamma=0.99,           # 할인율
    gae_lambda=0.95,      # GAE 람다
    clip_range=0.2,       # PPO 클리핑 범위
    verbose=1,
    tensorboard_log="./tb_logs/"
)

# 평가 콜백 설정
eval_env = DummyVecEnv([lambda: CarlaEnv()])
eval_callback = EvalCallback(
    eval_env,
    best_model_save_path="./best_model/",
    eval_freq=10000,
    n_eval_episodes=5
)

# 학습 시작! (1,000,000 스텝)
model.learn(
    total_timesteps=1_000_000,
    callback=eval_callback
)

# 모델 저장
model.save("carla_ppo_autonomous")