"""
Buoyancy Control Module - Depth-dependent buoyancy adjustment
"""

import numpy as np
from typing import Dict, Optional


class BuoyancyControl:
    """
    Buoyancy control system with pressure compensation
    
    Accounts for density increase at depth (4.7% at Challenger Deep)
    """
    
    def __init__(self,
                 vehicle_mass: float = 8.2,  # kg in water
                 displacement: float = 0.008,  # m³
                 foam_density: float = 600):  # kg/m³ (syntactic foam)
        """
        Initialize buoyancy control
        
        Args:
            vehicle_mass: Vehicle mass in water (kg)
            displacement: Vehicle displacement (m³)
            foam_density: Density of buoyancy foam
        """
        self.vehicle_mass = vehicle_mass
        self.displacement = displacement
        self.foam_density = foam_density
        
        # Water properties
        self.water_density_surface = 1025.0  # kg/m³
        self.water_density_deep = 1050.0  # kg/m³ at 11km
        
        # Buoyancy system
        self.ballast_volume = 0.001  # m³ (variable ballast)
        self.ballast_density = 1025.0  # kg/m³ (seawater)
        
        # Current state
        self.current_depth = 0.0
        self.current_buoyancy = 0.0  # N (positive = upward)
        self.ballast_fraction = 0.5  # 0-1
    
    def calculate_buoyancy(self,
                          depth: float,
                          water_density: Optional[float] = None) -> float:
        """
        Calculate buoyancy force at depth
        
        Args:
            depth: Current depth (m)
            water_density: Water density (calculated if None)
            
        Returns:
            Buoyancy force in Newtons (positive = upward)
        """
        if water_density is None:
            # Interpolate density based on depth
            depth_fraction = min(1.0, depth / 11000)
            water_density = (self.water_density_surface + 
                           depth_fraction * (self.water_density_deep - 
                                           self.water_density_surface))
        
        # Buoyancy force = ρ_water * V_displaced * g
        g = 9.81  # m/s²
        
        # Total displaced volume (hull + ballast)
        total_volume = self.displacement + self.ballast_volume * self.ballast_fraction
        
        buoyancy = water_density * total_volume * g
        
        # Weight
        weight = self.vehicle_mass * g
        
        net_buoyancy = buoyancy - weight
        
        self.current_depth = depth
        self.current_buoyancy = net_buoyancy
        
        return net_buoyancy
    
    def pressure_compensation(self,
                            pressure: float,
                            foam_density: Optional[float] = None) -> float:
        """
        Adjust for foam compression under pressure
        
        Args:
            pressure: Ambient pressure (MPa)
            foam_density: Foam density (uses default if None)
            
        Returns:
            Compensated buoyancy
        """
        if foam_density is None:
            foam_density = self.foam_density
        
        # Foam compression (simplified)
        # Syntactic foam compresses ~2% at 110 MPa
        compression_ratio = 1 - 0.0002 * pressure  # 0.02% per MPa
        
        # Effective foam density increases as volume decreases
        effective_foam_density = foam_density / compression_ratio
        
        # Adjust displacement
        compressed_displacement = self.displacement * compression_ratio
        
        # Recalculate buoyancy with compressed foam
        water_density = (self.water_density_surface + 
                        (pressure / 110) * (self.water_density_deep - 
                                          self.water_density_surface))
        
        g = 9.81
        
        # Foam contribution to buoyancy
        foam_buoyancy = (water_density - effective_foam_density) * compressed_displacement * g
        
        # Other components (assume incompressible)
        other_mass = self.vehicle_mass - foam_density * self.displacement
        other_buoyancy = water_density * self.displacement * (1 - compression_ratio) * g
        
        net_buoyancy = foam_buoyancy + other_buoyancy - other_mass * g
        
        return net_buoyancy
    
    def maintain_neutral(self,
                        depth: float,
                        control_rate: float = 0.1) -> float:
        """
        Adjust ballast to maintain neutral buoyancy
        
        Args:
            depth: Current depth
            control_rate: Control update rate (Hz)
            
        Returns:
            New ballast fraction
        """
        # Current buoyancy
        buoyancy = self.calculate_buoyancy(depth)
        
        # Target: neutral buoyancy (buoyancy = 0)
        # Adjust ballast fraction based on error
        if abs(buoyancy) < 0.1:  # Already neutral
            return self.ballast_fraction
        
        # Buoyancy change per unit ballast
        g = 9.81
        delta_per_unit = self.ballast_volume * self.ballast_density * g
        
        required_change = -buoyancy / delta_per_unit
        
        # Apply change (limited by control rate)
        max_change = control_rate * 0.1
        self.ballast_fraction += np.clip(required_change, -max_change, max_change)
        
        # Keep within bounds
        self.ballast_fraction = np.clip(self.ballast_fraction, 0, 1)
        
        return self.ballast_fraction
    
    def emergency_surface(self):
        """Execute emergency surfacing - maximum positive buoyancy"""
        self.ballast_fraction = 0.0  # Empty ballast
        return self.calculate_buoyancy(self.current_depth)
    
    def emergency_descent(self):
        """Execute emergency descent - maximum negative buoyancy"""
        self.ballast_fraction = 1.0  # Fill ballast
        return self.calculate_buoyancy(self.current_depth)
    
    def get_buoyancy_status(self) -> Dict:
        """Get buoyancy system status"""
        return {
            'depth_m': self.current_depth,
            'buoyancy_N': self.current_buoyancy,
            'ballast_fraction': self.ballast_fraction,
            'ballast_volume_L': self.ballast_volume * self.ballast_fraction * 1000,
            'is_neutral': abs(self.current_buoyancy) < 0.1,
            'is_positive': self.current_buoyancy > 0,
            'is_negative': self.current_buoyancy < 0
        }
    
    def __repr__(self) -> str:
        status = self.get_buoyancy_status()
        direction = "up" if status['is_positive'] else "down" if status['is_negative'] else "neutral"
        return (f"BuoyancyControl(depth={status['depth_m']:.0f}m, "
                f"buoyancy={status['buoyancy_N']:.1f}N, "
                f"direction={direction})")
