#!/usr/bin/env python3
"""
Standalone Drone Simulation - Runs without ROS 2.
Includes built-in controller for landing demonstration.

Usage: python standalone_simulation.py [--pattern PATTERN] [--speed SPEED]
"""

import argparse
import base64
import json
import math
import time
from typing import Dict, Any, Optional, Tuple

import numpy as np
import pybullet as p
import pybullet_data


class StandaloneSimulation:
    """Complete simulation with built-in controller - no ROS 2 needed."""

    PHYSICS_TIMESTEP = 1.0 / 240.0
    GRAVITY = -9.81
    CONTROL_INTERVAL = 5  # Apply control every N physics steps
    
    DRONE_MASS = 1.0
    DRONE_SIZE = (0.3, 0.3, 0.1)
    DRONE_START_POS = (0.0, 0.0, 5.0)
    
    THRUST_SCALE = 15.0
    PITCH_TORQUE_SCALE = 2.0
    ROLL_TORQUE_SCALE = 2.0
    YAW_TORQUE_SCALE = 1.0
    HOVER_THRUST = DRONE_MASS * abs(GRAVITY)
    
    ROVER_SIZE = (1.0, 1.0, 0.3)
    ROVER_START_POS = [0.0, 0.0, 0.15]
    
    CAMERA_FOV = 60.0

    def __init__(self, rover_pattern='stationary', rover_speed=0.5, rover_amplitude=2.0):
        print("=" * 60)
        print("Standalone Drone Simulation (No ROS 2 Required)")
        print("=" * 60)
        
        self._running = True
        self._step_count = 0
        self._time = 0.0
        
        # Rover movement settings
        self._rover_pattern = rover_pattern
        self._rover_speed = rover_speed
        self._rover_amplitude = rover_amplitude
        self._rover_pos = list(self.ROVER_START_POS)
        
        # Control command
        self._command = {'thrust': 0.0, 'pitch': 0.0, 'roll': 0.0, 'yaw': 0.0}
        
        self._init_physics()
        self._init_objects()
        
        print(f"  Rover Pattern: {rover_pattern}")
        print(f"  Rover Speed: {rover_speed} m/s")
        print("  Press Ctrl+C to exit")
        print("=" * 60)

    def _init_physics(self) -> None:
        self._physics_client = p.connect(p.GUI)
        p.setGravity(0, 0, self.GRAVITY)
        p.setTimeStep(self.PHYSICS_TIMESTEP)
        p.resetDebugVisualizerCamera(
            cameraDistance=8.0,
            cameraYaw=45,
            cameraPitch=-30,
            cameraTargetPosition=[0, 0, 1]
        )
        p.setAdditionalSearchPath(pybullet_data.getDataPath())

    def _init_objects(self) -> None:
        self._ground_id = p.loadURDF("plane.urdf")
        
        # Create rover (landing pad)
        rover_collision = p.createCollisionShape(
            p.GEOM_BOX,
            halfExtents=[s/2 for s in self.ROVER_SIZE]
        )
        rover_visual = p.createVisualShape(
            p.GEOM_BOX,
            halfExtents=[s/2 for s in self.ROVER_SIZE],
            rgbaColor=[0.2, 0.6, 0.2, 1.0]
        )
        self._rover_id = p.createMultiBody(
            baseMass=0,
            baseCollisionShapeIndex=rover_collision,
            baseVisualShapeIndex=rover_visual,
            basePosition=self.ROVER_START_POS
        )
        
        # Create landing pad marker
        self._create_landing_marker()
        
        # Create drone
        drone_collision = p.createCollisionShape(
            p.GEOM_BOX,
            halfExtents=[s/2 for s in self.DRONE_SIZE]
        )
        drone_visual = p.createVisualShape(
            p.GEOM_BOX,
            halfExtents=[s/2 for s in self.DRONE_SIZE],
            rgbaColor=[0.8, 0.2, 0.2, 1.0]
        )
        self._drone_id = p.createMultiBody(
            baseMass=self.DRONE_MASS,
            baseCollisionShapeIndex=drone_collision,
            baseVisualShapeIndex=drone_visual,
            basePosition=self.DRONE_START_POS
        )
        p.changeDynamics(
            self._drone_id, -1,
            linearDamping=0.1,
            angularDamping=0.5
        )

    def _create_landing_marker(self) -> None:
        marker_height = self.ROVER_START_POS[2] + self.ROVER_SIZE[2] / 2 + 0.01
        h_size = 0.3
        line_color = [1, 1, 1]
        
        p.addUserDebugLine([-h_size, 0, marker_height], [h_size, 0, marker_height],
                          lineColorRGB=line_color, lineWidth=3)
        p.addUserDebugLine([-h_size, -h_size, marker_height], [-h_size, h_size, marker_height],
                          lineColorRGB=line_color, lineWidth=3)
        p.addUserDebugLine([h_size, -h_size, marker_height], [h_size, h_size, marker_height],
                          lineColorRGB=line_color, lineWidth=3)

    def run(self) -> None:
        print("\nSimulation running...")
        
        try:
            while self._running:
                loop_start = time.time()
                
                if not p.isConnected():
                    break
                
                # Update rover position
                self._update_rover()
                
                # Run controller
                if self._step_count % self.CONTROL_INTERVAL == 0:
                    self._run_controller()
                
                # Apply physics
                self._apply_controls()
                p.stepSimulation()
                self._step_count += 1
                self._time += self.PHYSICS_TIMESTEP
                
                # Check landing
                if self._check_landing():
                    print("\n*** LANDING SUCCESSFUL! ***\n")
                    time.sleep(2)
                    break
                
                # Timing
                elapsed = time.time() - loop_start
                sleep_time = self.PHYSICS_TIMESTEP - elapsed
                if sleep_time > 0:
                    time.sleep(sleep_time)
                    
        except KeyboardInterrupt:
            print("\nStopping simulation...")
        finally:
            if p.isConnected():
                p.disconnect()

    def _update_rover(self) -> None:
        """Move the rover based on pattern."""
        dt = self.PHYSICS_TIMESTEP
        
        if self._rover_pattern == 'stationary':
            return
            
        elif self._rover_pattern == 'linear':
            omega = self._rover_speed / self._rover_amplitude
            target_x = self._rover_amplitude * math.sin(omega * self._time)
            self._rover_pos[0] += 2.0 * (target_x - self._rover_pos[0]) * dt
            
        elif self._rover_pattern == 'circular':
            omega = self._rover_speed / self._rover_amplitude
            target_x = self._rover_amplitude * math.cos(omega * self._time)
            target_y = self._rover_amplitude * math.sin(omega * self._time)
            self._rover_pos[0] += 2.0 * (target_x - self._rover_pos[0]) * dt
            self._rover_pos[1] += 2.0 * (target_y - self._rover_pos[1]) * dt
            
        elif self._rover_pattern == 'square':
            segment = int(self._time / 3) % 4
            corners = [(1, 1), (-1, 1), (-1, -1), (1, -1)]
            target = corners[segment]
            target_x = target[0] * self._rover_amplitude
            target_y = target[1] * self._rover_amplitude
            dx = target_x - self._rover_pos[0]
            dy = target_y - self._rover_pos[1]
            dist = math.sqrt(dx**2 + dy**2)
            if dist > 0.01:
                self._rover_pos[0] += self._rover_speed * dx / dist * dt
                self._rover_pos[1] += self._rover_speed * dy / dist * dt
        
        # Update rover position in simulation
        p.resetBasePositionAndOrientation(
            self._rover_id,
            [self._rover_pos[0], self._rover_pos[1], self._rover_pos[2]],
            [0, 0, 0, 1]
        )

    def _get_telemetry(self) -> Dict[str, Any]:
        """Get current drone state."""
        pos, orn = p.getBasePositionAndOrientation(self._drone_id)
        vel, ang_vel = p.getBaseVelocity(self._drone_id)
        euler = p.getEulerFromQuaternion(orn)
        
        # Landing pad detection
        dx = self._rover_pos[0] - pos[0]
        dy = self._rover_pos[1] - pos[1]
        dz = self._rover_pos[2] - pos[2]
        horizontal_dist = math.sqrt(dx**2 + dy**2)
        vertical_dist = -dz
        
        landing_pad = None
        if vertical_dist > 0 and vertical_dist < 10.0:
            fov_rad = math.radians(self.CAMERA_FOV / 2)
            max_horizontal = vertical_dist * math.tan(fov_rad)
            if horizontal_dist < max_horizontal:
                landing_pad = {
                    "relative_x": dx,
                    "relative_y": dy,
                    "distance": vertical_dist,
                    "confidence": 1.0 - (horizontal_dist / max_horizontal)
                }
        
        return {
            "altimeter": {"altitude": pos[2], "vertical_velocity": vel[2]},
            "velocity": {"x": vel[0], "y": vel[1], "z": vel[2]},
            "imu": {
                "orientation": {"roll": euler[0], "pitch": euler[1], "yaw": euler[2]},
                "angular_velocity": {"x": ang_vel[0], "y": ang_vel[1], "z": ang_vel[2]}
            },
            "landing_pad": landing_pad,
            "rover_position": {"x": self._rover_pos[0], "y": self._rover_pos[1]}
        }

    def _run_controller(self) -> None:
        """Simple landing controller."""
        telemetry = self._get_telemetry()
        
        altimeter = telemetry['altimeter']
        altitude = altimeter['altitude']
        vertical_vel = altimeter['vertical_velocity']
        
        velocity = telemetry['velocity']
        vel_x = velocity['x']
        vel_y = velocity['y']
        
        landing_pad = telemetry['landing_pad']
        
        # Altitude control
        target_alt = 0.0
        Kp_z, Kd_z = 0.5, 0.3
        thrust = Kp_z * (target_alt - altitude) - Kd_z * vertical_vel
        
        # Horizontal control
        Kp_xy, Kd_xy = 0.3, 0.2
        
        if landing_pad is not None:
            target_x = landing_pad['relative_x']
            target_y = landing_pad['relative_y']
            pitch = Kp_xy * target_x - Kd_xy * vel_x
            roll = Kp_xy * target_y - Kd_xy * vel_y
        else:
            pitch = -Kd_xy * vel_x
            roll = -Kd_xy * vel_y
        
        # Clamp
        thrust = max(-1.0, min(1.0, thrust))
        pitch = max(-0.5, min(0.5, pitch))
        roll = max(-0.5, min(0.5, roll))
        
        self._command = {'thrust': thrust, 'pitch': pitch, 'roll': roll, 'yaw': 0.0}

    def _apply_controls(self) -> None:
        """Apply control commands to drone."""
        pos, orn = p.getBasePositionAndOrientation(self._drone_id)
        rot_matrix = p.getMatrixFromQuaternion(orn)
        local_up = [rot_matrix[2], rot_matrix[5], rot_matrix[8]]
        
        total_thrust = self.HOVER_THRUST + (self._command['thrust'] * self.THRUST_SCALE)
        total_thrust = max(0, total_thrust)
        
        thrust_force = [
            local_up[0] * total_thrust,
            local_up[1] * total_thrust,
            local_up[2] * total_thrust
        ]
        
        p.applyExternalForce(
            self._drone_id, -1,
            forceObj=thrust_force,
            posObj=pos,
            flags=p.WORLD_FRAME
        )
        
        p.applyExternalTorque(
            self._drone_id, -1,
            torqueObj=[
                self._command['pitch'] * self.PITCH_TORQUE_SCALE,
                self._command['roll'] * self.ROLL_TORQUE_SCALE,
                self._command['yaw'] * self.YAW_TORQUE_SCALE
            ],
            flags=p.LINK_FRAME
        )

    def _check_landing(self) -> bool:
        """Check if drone landed on rover."""
        contacts = p.getContactPoints(bodyA=self._drone_id, bodyB=self._rover_id)
        if len(contacts) > 0:
            vel, _ = p.getBaseVelocity(self._drone_id)
            speed = math.sqrt(vel[0]**2 + vel[1]**2 + vel[2]**2)
            return speed < 0.5
        return False


def main():
    parser = argparse.ArgumentParser(description='Standalone Drone Simulation')
    parser.add_argument('--pattern', '-p', type=str, default='stationary',
                       choices=['stationary', 'linear', 'circular', 'square'],
                       help='Rover movement pattern')
    parser.add_argument('--speed', '-s', type=float, default=0.5,
                       help='Rover speed in m/s')
    parser.add_argument('--amplitude', '-a', type=float, default=2.0,
                       help='Rover movement amplitude in meters')
    args = parser.parse_args()
    
    sim = StandaloneSimulation(
        rover_pattern=args.pattern,
        rover_speed=args.speed,
        rover_amplitude=args.amplitude
    )
    sim.run()


if __name__ == '__main__':
    main()