import numpy as np
import pybullet as p
import time
import pickle
import os
from gym_pybullet_drones.envs.VelocityAviary import VelocityAviary
from gym_pybullet_drones.utils.enums import Physics
from utils import step_compute

def run_realtime_simulation():
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
    components_path = os.path.join(BASE_DIR, "W_drone_components.pkl")
    if not os.path.exists(components_path):
        print(f"{components_path} not found!")
        return
        
    with open(components_path, "rb") as f:
        components = pickle.load(f)
    W, Win, Wout, mask, alpha, Nx = components['W'], components['Win'], components['Wout'], components['mask'], components['alpha'], components['W'].shape[0]

    # Realistic physics engine features: Ground Effect, Drag, and Downwash
    env = VelocityAviary(
        gui=True, 
        initial_xyzs=np.array([[0, 0, .5]]), 
        initial_rpys=np.array([[0, 0, 0]]),
        physics=Physics.PYB_GND_DRAG_DW
    )
    
    # Close PyBullet GUI side panels and enable shadows for realistic visuals
    p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=env.CLIENT)
    p.configureDebugVisualizer(p.COV_ENABLE_SHADOWS, 1, physicsClientId=env.CLIENT)

    obs, info = env.reset()
    state = obs[0] # VelocityAviary state is size 20
    x = np.random.random((Nx, 1))

    # --- VIDEO RECORDING SETTINGS ---
    # 1. Third Person camera tracking the drone
    log_id_tp = p.startStateLogging(p.STATE_LOGGING_VIDEO_MP4, "drone_third_person.mp4")
    
    # 2. First Person camera (ego-view)
    log_id_fp = p.startStateLogging(p.STATE_LOGGING_VIDEO_MP4, "drone_first_person.mp4")

    print("Simulation starting, recording video...")
    
    try:
        for i in range(43200): # Match simulation length
            # Blind Flight: Hide global X and Y positions from the network
            state[0] = 0
            state[1] = 0
            
            x, y = step_compute(x, W, Win, Wout, state, alpha, mask)
            
            # Velocity references are directly generated by ESN (Autonomous flight)
            obs, reward, terminated, truncated, info = env.step(np.array([y.flatten()]))
            state = obs[0]

            # --- SPECTATOR CAMERA (FREECAM) CONTROL ---
            cam_info = p.getDebugVisualizerCamera(physicsClientId=env.CLIENT)
            cam_yaw = cam_info[8]
            cam_pitch = cam_info[9]
            cam_dist = cam_info[10]
            cam_target = list(cam_info[11])
            
            forward = np.array(cam_info[5]) # Camera forward direction
            right = np.array(cam_info[6])   # Camera right direction
            
            # Reset Z axis for flat movement in X/Y plane
            forward[2] = 0
            if np.linalg.norm(forward) > 0: forward = forward / np.linalg.norm(forward)
            right[2] = 0
            if np.linalg.norm(right) > 0: right = right / np.linalg.norm(right)

            keys = p.getKeyboardEvents()
            cam_speed = 0.05
            moved = False
            
            # Arrow Keys (Forward-Backward, Left-Right based on camera direction)
            if p.B3G_UP_ARROW in keys and keys[p.B3G_UP_ARROW] & p.KEY_IS_DOWN:
                cam_target[0] += forward[0] * cam_speed
                cam_target[1] += forward[1] * cam_speed
                moved = True
            if p.B3G_DOWN_ARROW in keys and keys[p.B3G_DOWN_ARROW] & p.KEY_IS_DOWN:
                cam_target[0] -= forward[0] * cam_speed
                cam_target[1] -= forward[1] * cam_speed
                moved = True
            if p.B3G_RIGHT_ARROW in keys and keys[p.B3G_RIGHT_ARROW] & p.KEY_IS_DOWN:
                cam_target[0] += right[0] * cam_speed
                cam_target[1] += right[1] * cam_speed
                moved = True
            if p.B3G_LEFT_ARROW in keys and keys[p.B3G_LEFT_ARROW] & p.KEY_IS_DOWN:
                cam_target[0] -= right[0] * cam_speed
                cam_target[1] -= right[1] * cam_speed
                moved = True
                
            # Move Up (Space) / Move Down (Shift)
            if 32 in keys and keys[32] & p.KEY_IS_DOWN:
                cam_target[2] += cam_speed
                moved = True
            if p.B3G_SHIFT in keys and keys[p.B3G_SHIFT] & p.KEY_IS_DOWN:
                cam_target[2] -= cam_speed
                moved = True
                
            if moved:
                p.resetDebugVisualizerCamera(
                    cameraDistance=cam_dist, 
                    cameraYaw=cam_yaw, 
                    cameraPitch=cam_pitch, 
                    cameraTargetPosition=cam_target,
                    physicsClientId=env.CLIENT
                )
            
            time.sleep(env.CTRL_TIMESTEP)
            
    except KeyboardInterrupt:
        print("Stopped by user.")
    finally:
        # --- CLOSE VIDEO RECORDING ---
        p.stopStateLogging(log_id_tp)
        p.stopStateLogging(log_id_fp)
        env.close()
        print("Videos saved: drone_third_person.mp4 and drone_first_person.mp4")

if __name__ == "__main__":
    run_realtime_simulation()