from filterpy.kalman import ExtendedKalmanFilter

class WindEstimator:
    def __init__(self):
        self.ekf = ExtendedKalmanFilter(dim_x=6, dim_z=3)
        
        # State: [x, y, z, vx, vy, vz] + wind [wx, wy]
        self.ekf.x = np.zeros(6)
        self.ekf.F = np.eye(6)
        self.ekf.H = np.array([
            [0, 0, 0, 1, 0, 0],
            [0, 0, 0, 0, 1, 0],
            [0, 0, 0, 0, 0, 1]
        ])
        
        self.ekf.P *= 10
        self.ekf.R = np.eye(3) * 0.1
        self.ekf.Q = np.eye(6) * 0.01
        
        self.wind_estimate = np.zeros(2)
        self.drag_coefficient = 0.5
        self.air_density = 1.225
        self.drone_mass = 1.0
    
    def predict(self, dt, control_input):
        """Predict step with drone dynamics."""
        # Update state transition
        self.ekf.F[0, 3] = dt
        self.ekf.F[1, 4] = dt
        self.ekf.F[2, 5] = dt
        
        self.ekf.predict()
    
    def update(self, measured_velocity, commanded_thrust):
        """Update step with velocity measurement."""
        # Compute expected wind force
        airspeed = measured_velocity[:2] - self.wind_estimate
        wind_force = 0.5 * self.air_density * self.drag_coefficient * airspeed
        
        # Measurement model
        z = measured_velocity
        self.ekf.update(z, self._measurement_function)
        
        # Extract wind estimate
        self.wind_estimate = self.ekf.x[3:5]
    
    def _measurement_function(self, state):
        return state[3:6]
    
    def get_wind_speed(self):
        return np.linalg.norm(self.wind_estimate)
    
    def get_wind_direction(self):
        return np.arctan2(self.wind_estimate[1], self.wind_estimate[0])


class AdaptiveWindCompensator:
    def __init__(self, wind_estimator: WindEstimator):
        self.estimator = wind_estimator
        self.base_pid_gains = {'kp': 1.0, 'ki': 0.1, 'kd': 0.5}
        self.adapted_gains = self.base_pid_gains.copy()
    
    def compute_feedforward(self, drone_mass, drag_coeff):
        """Compute feedforward compensation for estimated wind."""
        wind_speed = self.estimator.get_wind_speed()
        wind_dir = self.estimator.get_wind_direction()
        
        if wind_speed < 0.1:
            return np.zeros(2)
        
        # Force compensation
        force = 0.5 * 1.225 * drag_coeff * wind_speed**2
        direction = np.array([np.cos(wind_dir), np.sin(wind_dir)])
        
        return force * direction / drone_mass
    
    def adapt_pid_gains(self, wind_speed):
        """Adapt PID gains based on wind conditions."""
        if wind_speed < 3:
            # Light wind
            self.adapted_gains = self.base_pid_gains.copy()
        elif wind_speed < 7:
            # Moderate wind - increase gains
            factor = 1 + (wind_speed - 3) / 10
            self.adapted_gains = {k: v * factor for k, v in self.base_pid_gains.items()}
        else:
            # Strong wind - significantly increase gains
            factor = 1.5
            self.adapted_gains = {k: v * factor for k, v in self.base_pid_gains.items()}
        
        return self.adapted_gains

from typing import Dict, Optional

class WeatherCompensator:
    def __init__(self):
        self.rain_intensity = 0.0  # 0-1 scale
        self.fog_density = 0.0     # 0-1 scale
        self.visibility = 1000.0   # meters
        
        self.sensor_degradation = {
            'camera': 1.0,
            'lidar': 1.0,
            'ultrasonic': 1.0,
            'gps': 1.0
        }
    
    def update_weather(self, rain_intensity: float, fog_density: float):
        self.rain_intensity = np.clip(rain_intensity, 0, 1)
        self.fog_density = np.clip(fog_density, 0, 1)
        self._update_sensor_degradation()
        self._update_visibility()
    
    def _update_sensor_degradation(self):
        """Compute sensor performance degradation."""
        # Camera degrades with fog and rain
        self.sensor_degradation['camera'] = max(0.1, 
            1.0 - 0.7 * self.fog_density - 0.3 * self.rain_intensity)
        
        # LiDAR degrades with rain (reflections)
        self.sensor_degradation['lidar'] = max(0.3,
            1.0 - 0.5 * self.rain_intensity)
        
        # Ultrasonic affected by rain noise
        self.sensor_degradation['ultrasonic'] = max(0.4,
            1.0 - 0.4 * self.rain_intensity)
        
        # GPS generally unaffected by weather
        self.sensor_degradation['gps'] = 0.95
    
    def _update_visibility(self):
        """Compute visibility based on fog density."""
        # Empirical relationship
        self.visibility = 1000 * np.exp(-3 * self.fog_density)
    
    def adapt_flight_parameters(self) -> Dict:
        """Return adapted flight parameters for current weather."""
        params = {
            'max_speed': 5.0 * (1 - 0.3 * self.rain_intensity),
            'max_altitude': 50.0 * (1 - 0.2 * self.fog_density),
            'position_accuracy': 0.5 / self.sensor_degradation['gps'],
            'obstacle_detection_range': 10.0 * self.sensor_degradation['lidar'],
            'camera_detection_range': 20.0 * self.sensor_degradation['camera'],
            'pid_gains_multiplier': 1.0 + 0.3 * self.rain_intensity
        }
        
        # Weather-based mission restrictions
        if self.rain_intensity > 0.7:
            params['recommendation'] = 'DELAY_MISSION'
            params['reason'] = 'Heavy rain - risk of water damage'
        elif self.fog_density > 0.8:
            params['recommendation'] = 'REDUCE_ALTITUDE'
            params['reason'] = 'Dense fog - limited visibility'
        elif self.visibility < 50:
            params['recommendation'] = 'LAND_IMMEDIATELY'
            params['reason'] = 'Near-zero visibility'
        else:
            params['recommendation'] = 'PROCEED_WITH_CAUTION'
        
        return params
    
    def fuse_sensor_data(self, measurements: Dict) -> Dict:
        """Fuse sensor data accounting for weather degradation."""
        fused = {}
        
        # Weight sensors by their degradation
        total_weight = sum(self.sensor_degradation.values())
        
        for sensor, measurement in measurements.items():
            weight = self.sensor_degradation.get(sensor, 0.5) / total_weight
            fused[sensor] = {
                'value': measurement,
                'weight': weight,
                'confidence': self.sensor_degradation.get(sensor, 0.5)
            }
        
        return fused

from scipy.signal import butter, filtfilt

class TurbulenceHandler:
    def __init__(self, sample_rate=100):
        self.sample_rate = sample_rate
        self.turbulence_intensity = 0.0
        self.buffer_size = 100
        self.accel_buffer = []
        
        # Turbulence detection thresholds
        self.detection_threshold = 2.0  # m/s^2
        self.severity_thresholds = {
            'light': 2.0,
            'moderate': 4.0,
            'severe': 8.0
        }
    
    def detect_turbulence(self, imu_data):
        """Detect turbulence from IMU accelerometer data."""
        accel = imu_data['acceleration']
        
        # Remove gravity
        accel_no_gravity = accel - np.array([0, 0, 9.81])
        
        # Compute acceleration magnitude
        accel_magnitude = np.linalg.norm(accel_no_gravity)
        
        # Add to buffer
        self.accel_buffer.append(accel_magnitude)
        if len(self.accel_buffer) > self.buffer_size:
            self.accel_buffer.pop(0)
        
        # Compute turbulence intensity (variance)
        if len(self.accel_buffer) > 10:
            self.turbulence_intensity = np.var(self.accel_buffer)
        
        # Classify severity
        severity = 'none'
        for level, threshold in self.severity_thresholds.items():
            if self.turbulence_intensity > threshold:
                severity = level
        
        return {
            'intensity': self.turbulence_intensity,
            'severity': severity,
            'acceleration': accel_magnitude
        }
    
    def filter_turbulence(self, control_signal):
        """Apply low-pass filter to reduce turbulence effects."""
        if self.turbulence_intensity > self.detection_threshold:
            # More aggressive filtering for stronger turbulence
            cutoff = max(1.0, 10.0 - self.turbulence_intensity)
            b, a = butter(2, cutoff / (self.sample_rate / 2), btype='low')
            
            if len(self.accel_buffer) > 30:
                filtered = filtfilt(b, a, control_signal)
                return filtered
        
        return control_signal
    
    def get_compensation_gains(self):
        """Get adapted controller gains for turbulence."""
        if self.turbulence_intensity < 1.0:
            return {'kp_mult': 1.0, 'kd_mult': 1.0, 'filter_strength': 0.1}
        elif self.turbulence_intensity < 4.0:
            return {'kp_mult': 1.2, 'kd_mult': 1.5, 'filter_strength': 0.3}
        else:
            return {'kp_mult': 1.5, 'kd_mult': 2.0, 'filter_strength': 0.5}

from typing import Dict

class AdaptiveFlightController:
    def __init__(self):
        self.wind_estimator = WindEstimator()
        self.weather_compensator = WeatherCompensator()
        self.turbulence_handler = TurbulenceHandler()
        
        self.base_pid = {
            'position': {'kp': 1.0, 'ki': 0.1, 'kd': 0.5},
            'velocity': {'kp': 0.5, 'ki': 0.2, 'kd': 0.1},
            'attitude': {'kp': 2.0, 'ki': 0.5, 'kd': 1.0}
        }
        
        self.adapted_pid = self.base_pid.copy()
    
    def update(self, imu_data, gps_data, weather_data):
        """Update all adaptation modules."""
        # Wind estimation
        self.wind_estimator.predict(0.01, np.zeros(3))
        self.wind_estimator.update(
            gps_data['velocity'],
            imu_data.get('thrust', 0)
        )
        
        # Weather update
        self.weather_compensator.update_weather(
            weather_data.get('rain', 0),
            weather_data.get('fog', 0)
        )
        
        # Turbulence detection
        turbulence = self.turbulence_handler.detect_turbulence(imu_data)
        
        # Adapt PID gains
        wind_speed = self.wind_estimator.get_wind_speed()
        self._adapt_gains(wind_speed, turbulence)
        
        return {
            'wind': {
                'speed': wind_speed,
                'direction': self.wind_estimator.get_wind_direction(),
                'compensation': self.wind_estimator.wind_estimate
            },
            'weather': self.weather_compensator.adapt_flight_parameters(),
            'turbulence': turbulence,
            'adapted_gains': self.adapted_pid
        }
    
    def _adapt_gains(self, wind_speed, turbulence):
        """Adapt PID gains based on conditions."""
        # Base adaptation for wind
        wind_factor = 1.0 + 0.1 * min(wind_speed, 10)
        
        # Turbulence adaptation
        turb_gains = self.turbulence_handler.get_compensation_gains()
        
        # Weather adaptation
        weather_params = self.weather_compensator.adapt_flight_parameters()
        weather_factor = weather_params['pid_gains_multiplier']
        
        # Combine factors
        total_factor = wind_factor * turb_gains['kp_mult'] * weather_factor
        
        for controller in self.adapted_pid:
            self.adapted_pid[controller]['kp'] = self.base_pid[controller]['kp'] * total_factor
            self.adapted_pid[controller]['kd'] = self.base_pid[controller]['kd'] * turb_gains['kd_mult']
    
    def compute_control(self, target, current_state):
        """Compute control with weather adaptation."""
        # Wind feedforward
        wind_comp = self.wind_estimator.wind_estimate
        feedforward = self.wind_compensator.compute_feedforward(
            drone_mass=1.0, drag_coeff=0.5
        )
        
        # PID control with adapted gains
        error = target - current_state[:3]
        
        # Position control
        kp = self.adapted_pid['position']['kp']
        kd = self.adapted_pid['position']['kd']
        
        velocity_cmd = kp * error - kd * current_state[3:6] + feedforward
        
        return velocity_cmd

from dataclasses import dataclass
from typing import Optional

@dataclass
class WeatherCondition:
    wind_speed: float = 0.0
    wind_direction: float = 0.0
    rain_intensity: float = 0.0
    fog_density: float = 0.0
    temperature: float = 20.0
    humidity: float = 50.0

class AllWeatherDroneSystem:
    def __init__(self):
        self.wind_estimator = WindEstimator()
        self.weather_comp = WeatherCompensator()
        self.turbulence_handler = TurbulenceHandler()
        self.adaptive_controller = AdaptiveFlightController()
        
        self.current_weather = WeatherCondition()
        self.flight_log = []
    
    def update_weather_sensors(self, weather_data: Dict):
        """Update weather from onboard sensors."""
        self.current_weather = WeatherCondition(
            wind_speed=weather_data.get('wind_speed', 0),
            wind_direction=weather_data.get('wind_direction', 0),
            rain_intensity=weather_data.get('rain', 0),
            fog_density=weather_data.get('fog', 0)
        )
        
        self.weather_comp.update_weather(
            self.current_weather.rain_intensity,
            self.current_weather.fog_density
        )
    
    def check_flight_safety(self) -> Dict:
        """Check if conditions are safe for flight."""
        safety = {
            'can_fly': True,
            'restrictions': [],
            'recommendations': []
        }
        
        # Wind check
        if self.current_weather.wind_speed > 15:
            safety['can_fly'] = False
            safety['restrictions'].append('Wind speed exceeds 15 m/s')
        elif self.current_weather.wind_speed > 10:
            safety['restrictions'].append('Reduced max speed due to wind')
            safety['recommendations'].append('Fly at lower altitude')
        
        # Rain check
        if self.current_weather.rain_intensity > 0.8:
            safety['can_fly'] = False
            safety['restrictions'].append('Heavy rain - water damage risk')
        elif self.current_weather.rain_intensity > 0.5:
            safety['restrictions'].append('Moderate rain - reduced sensor range')
        
        # Fog/visibility check
        visibility = 1000 * np.exp(-3 * self.current_weather.fog_density)
        if visibility < 50:
            safety['can_fly'] = False
            safety['restrictions'].append('Visibility below 50m')
        elif visibility < 200:
            safety['recommendations'].append('Use enhanced sensor fusion')
        
        return safety
    
    def get_adapted_trajectory(self, planned_trajectory):
        """Adapt trajectory for current weather."""
        adapted = planned_trajectory.copy()
        
        wind_speed = self.current_weather.wind_speed
        wind_dir = self.current_weather.wind_direction
        
        # Reduce speed in high wind
        speed_factor = max(0.5, 1.0 - 0.05 * wind_speed)
        adapted['max_speed'] = planned_trajectory.get('max_speed', 5.0) * speed_factor
        
        # Add wind compensation waypoints
        if wind_speed > 3:
            # Compute crosswind component
            for i, waypoint in enumerate(adapted.get('waypoints', [])):
                wind_compensation = self._compute_wind_compensation(
                    waypoint, wind_speed, wind_dir
                )
                adapted['waypoints'][i] = waypoint + wind_compensation
        
        return adapted
    
    def _compute_wind_compensation(self, waypoint, wind_speed, wind_dir):
        """Compute position offset to compensate for wind drift."""
        # Simple drift estimation
        drift_rate = 0.1 * wind_speed  # Simplified
        compensation = np.array([
            -drift_rate * np.cos(wind_dir),
            -drift_rate * np.sin(wind_dir),
            0
        ])
        return compensation
    
    def monitor_flight(self, state, dt=0.1):
        """Monitor flight and adjust in real-time."""
        # Update estimators
        imu_data = {'acceleration': state[6:9], 'gyroscope': state[9:12]}
        gps_data = {'position': state[:3], 'velocity': state[3:6]}
        weather_data = {
            'rain': self.current_weather.rain_intensity,
            'fog': self.current_weather.fog_density
        }
        
        adaptation_info = self.adaptive_controller.update(
            imu_data, gps_data, weather_data
        )
        
        # Log
        self.flight_log.append({
            'time': len(self.flight_log) * dt,
            'weather': self.current_weather.__dict__,
            'adaptation': adaptation_info
        })
        
        return adaptation_info
    
    def generate_weather_report(self):
        """Generate comprehensive weather report."""
        safety = self.check_flight_safety()
        
        report = {
            'current_conditions': self.current_weather.__dict__,
            'safety_assessment': safety,
            'sensor_status': self.weather_comp.sensor_degradation,
            'visibility': self.weather_comp.visibility,
            'recommended_action': safety['recommendations'][0] if safety['recommendations'] else 'Normal operations'
        }
        
        return report


# Usage Example
def demo_weather_adaptation():
    system = AllWeatherDroneSystem()
    
    # Simulate changing weather
    weather_sequence = [
        {'wind_speed': 2, 'rain': 0, 'fog': 0},
        {'wind_speed': 5, 'rain': 0.1, 'fog': 0},
        {'wind_speed': 8, 'rain': 0.3, 'fog': 0.2},
        {'wind_speed': 12, 'rain': 0.6, 'fog': 0.4},
    ]
    
    for weather in weather_sequence:
        system.update_weather_sensors(weather)
        report = system.generate_weather_report()
        
        print(f"\nWeather: Wind={weather['wind_speed']}m/s, "
              f"Rain={weather['rain']:.0%}, Fog={weather['fog']:.0%}")
        print(f"  Can fly: {report['safety_assessment']['can_fly']}")
        print(f"  Visibility: {report['visibility']:.0f}m")
        print(f"  Recommendation: {report['recommended_action']}")


if __name__ == "__main__":
    demo_weather_adaptation()