#!/usr/bin/env python3
"""
1D Drone PID Simulation — Simulates altitude/attitude control in one dimension.

Usage:
    python3 pidsim.py
    python3 pidsim.py --mode angle  # test angle control
    python3 pidsim.py --mode altitude  # test altitude hold
"""

import argparse
import matplotlib.pyplot as plt
import numpy as np


class Drone1D:
    """Simple 1D drone model (Z-axis or pitch axis)."""
    
    def __init__(self, mode='angle'):
        self.mode = mode        # 'angle' or 'altitude'
        self.state = 0.0        # angle (rad) or altitude (m)
        self.velocity = 0.0     # angular rate (rad/s) or vertical vel (m/s)
        self.gravity = 9.81
        self.mass = 1.0         # kg
        self.I = 0.01           # moment of inertia (kg*m^2) for angle mode
        self.thrust_max = 20.0  # N or N*m
        
    def step(self, thrust, dt):
        """Apply thrust and simulate one timestep."""
        thrust = np.clip(thrust, -self.thrust_max, self.thrust_max)
        
        if self.mode == 'angle':
            acc = thrust / self.I
            self.velocity += acc * dt
            self.velocity *= 0.995  # air drag
            self.state += self.velocity * dt
        else:  # altitude
            acc = (thrust / self.mass) - self.gravity
            self.velocity += acc * dt
            self.velocity *= 0.98   # air drag
            self.state += self.velocity * dt
            
        return self.state, self.velocity


class PID:
    def __init__(self, Kp, Ki, Kd, setpoint=0.0):
        self.Kp = Kp
        self.Ki = Ki
        self.Kd = Kd
        self.setpoint = setpoint
        self.integral = 0.0
        self.prev_error = 0.0
        
    def update(self, measurement, dt):
        error = self.setpoint - measurement
        self.integral += error * dt
        derivative = (error - self.prev_error) / dt if dt > 0 else 0
        output = self.Kp * error + self.Ki * self.integral + self.Kd * derivative
        self.prev_error = error
        return output


def simulate_angle():
    """Simulate pitch/roll angle control."""
    drone = Drone1D(mode='angle')
    pid = PID(Kp=30.0, Ki=0.5, Kd=1.5, setpoint=0.0)
    
    dt = 0.005  # 200Hz
    t = 0.0
    ts, states, outputs = [], [], []
    
    # Add disturbance at t=1s
    disturbance_applied = False
    
    for _ in range(400):
        ts.append(t)
        
        # Measurement (with noise)
        meas = drone.state + np.random.randn() * 0.005
        
        # PID output
        thrust = pid.update(meas, dt)
        
        # Disturbance: external torque at t=1s
        dist = 0
        if 1.0 < t < 1.05 and not disturbance_applied:
            dist = -5.0
            disturbance_applied = True
        
        drone.step(thrust + dist, dt)
        
        states.append(drone.state * 57.3)  # convert to degrees
        outputs.append(thrust)
        t += dt
    
    return ts, states, outputs


def simulate_altitude():
    """Simulate altitude hold."""
    drone = Drone1D(mode='altitude')
    pid = PID(Kp=10.0, Ki=0.5, Kd=3.0, setpoint=5.0)  # hover at 5m
    
    dt = 0.02  # 50Hz
    t = 0.0
    ts, alts, thrusts = [], [], []
    
    for _ in range(500):
        ts.append(t)
        meas = drone.state + np.random.randn() * 0.02
        thrust = pid.update(meas, dt)
        
        # Gravity compensation
        thrust += drone.mass * drone.gravity
        
        drone.step(thrust, dt)
        alts.append(drone.state)
        thrusts.append(thrust)
        t += dt
    
    return ts, alts, thrusts


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--mode', default='angle', choices=['angle', 'altitude'])
    args = parser.parse_args()
    
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6))
    
    if args.mode == 'angle':
        ts, states, outputs = simulate_angle()
        ax1.set_ylabel('Angle (deg)')
        ax2.set_ylabel('Torque (N*m)')
        ax1.axhline(y=0, color='gray', linestyle='--')
    else:
        ts, alts, thrusts = simulate_altitude()
        ax1.set_ylabel('Altitude (m)')
        ax2.set_ylabel('Thrust (N)')
        ax1.axhline(y=5.0, color='gray', linestyle='--', label='Target')
        ax1.legend()
    
    ax1.plot(ts, states if args.mode == 'angle' else alts, 'b-')
    ax1.grid(True)
    ax1.set_title(f'{args.mode.capitalize()} PID Response')
    
    ax2.plot(ts, outputs if args.mode == 'angle' else thrusts, 'r-')
    ax2.set_xlabel('Time (s)')
    ax2.grid(True)
    
    plt.tight_layout()
    plt.show()


if __name__ == '__main__':
    main()