"""
Thruster Model Module - Viscosity-compensated propulsion
"""

import numpy as np
from typing import Optional, Dict


class ThrusterModel:
    """
    Thruster model with viscosity compensation for deep-sea operation
    
    Accounts for increased viscosity at depth (22% increase at 110 MPa)
    """
    
    def __init__(self,
                 viscosity: float = 1.8e-3,  # Pa·s at surface
                 density: float = 1025.0,  # kg/m³ at surface
                 propeller_diameter: float = 0.05,  # meters
                 num_thrusters: int = 4):
        """
        Initialize thruster model
        
        Args:
            viscosity: Dynamic viscosity (Pa·s)
            density: Water density (kg/m³)
            propeller_diameter: Propeller diameter (m)
            num_thrusters: Number of thrusters
        """
        self.viscosity = viscosity
        self.density = density
        self.propeller_diameter = propeller_diameter
        self.num_thrusters = num_thrusters
        
        # Thruster parameters
        self.Kt = 0.3  # Thrust coefficient
        self.Kq = 0.05  # Torque coefficient
        self.max_rpm = 3000
        self.efficiency = 0.7
        
        # Area for drag calculations
        self.vehicle_area = np.pi * (propeller_diameter / 2)**2 * num_thrusters
    
    def calculate_thrust(self,
                        rpm: float,
                        advance_ratio: float = 0) -> float:
        """
        Calculate thrust at given RPM
        
        Args:
            rpm: Propeller RPM
            advance_ratio: Advance ratio (J = V/(nD))
            
        Returns:
            Thrust in Newtons
        """
        # Convert RPM to RPS
        n = rpm / 60.0
        
        # Thrust = Kt * ρ * n² * D⁴
        thrust = (self.Kt * self.density * n**2 * 
                 self.propeller_diameter**4)
        
        # Correction for advance ratio
        if advance_ratio > 0:
            thrust *= (1 - advance_ratio / 2)
        
        return thrust * self.num_thrusters
    
    def viscosity_compensation(self,
                              required_velocity: float,
                              depth: float,
                              surface_viscosity: float = 1.8e-3) -> float:
        """
        Compensate for increased viscosity at depth
        
        Args:
            required_velocity: Desired velocity (m/s)
            depth: Operating depth (m)
            surface_viscosity: Viscosity at surface
            
        Returns:
            Compensated RPM
        """
        # Viscosity ratio
        viscosity_ratio = self.viscosity / surface_viscosity
        
        # Increased drag requires more thrust
        thrust_multiplier = viscosity_ratio**0.5  # Approximate
        
        # Calculate base RPM for required velocity
        # Simple linear approximation
        base_rpm = required_velocity * 1000  # Rough conversion
        
        compensated_rpm = base_rpm * thrust_multiplier
        
        return min(compensated_rpm, self.max_rpm)
    
    def power_required(self,
                      thrust: float,
                      efficiency: Optional[float] = None) -> float:
        """
        Calculate power required for given thrust
        
        Args:
            thrust: Required thrust (N)
            efficiency: Thruster efficiency (uses default if None)
            
        Returns:
            Power in Watts
        """
        if efficiency is None:
            efficiency = self.efficiency
        
        # Power = Thrust * Velocity / efficiency
        # Need to estimate velocity from thrust
        # Simplified: assume linear relationship
        velocity = thrust / (self.num_thrusters * 100)  # Rough estimate
        
        power = thrust * velocity / efficiency
        
        return max(0, power)
    
    def calculate_drag(self,
                      velocity: float,
                      depth: float,
                      area: Optional[float] = None) -> float:
        """
        Calculate hydrodynamic drag
        
        Args:
            velocity: Vehicle velocity (m/s)
            depth: Operating depth (m)
            area: Cross-sectional area (m²)
            
        Returns:
            Drag force in Newtons
        """
        if area is None:
            area = self.vehicle_area
        
        # Reynolds number dependent drag coefficient
        Re = (self.density * velocity * np.sqrt(area) / 
              self.viscosity)
        
        if Re < 1000:
            Cd = 1.0  # Laminar
        else:
            Cd = 0.4  # Turbulent
        
        drag = 0.5 * self.density * velocity**2 * area * Cd
        
        return drag
    
    def optimize_rpm(self,
                    required_thrust: float,
                    depth: float) -> Dict:
        """
        Find optimal RPM for required thrust
        
        Args:
            required_thrust: Desired thrust (N)
            depth: Operating depth
            
        Returns:
            Optimal operating parameters
        """
        # Search for RPM giving required thrust
        best_rpm = 0
        best_error = float('inf')
        
        for rpm in range(100, self.max_rpm, 10):
            thrust = self.calculate_thrust(rpm)
            error = abs(thrust - required_thrust)
            
            if error < best_error:
                best_error = error
                best_rpm = rpm
        
        # Calculate power at optimal RPM
        thrust_at_opt = self.calculate_thrust(best_rpm)
        power = self.power_required(thrust_at_opt)
        
        return {
            'rpm': best_rpm,
            'thrust_N': thrust_at_opt,
            'power_W': power,
            'error_percent': (best_error / required_thrust * 100) if required_thrust > 0 else 0
        }
    
    def estimate_velocity(self,
                         rpm: float,
                         depth: float) -> float:
        """
        Estimate vehicle velocity from RPM
        
        Args:
            rpm: Propeller RPM
            depth: Operating depth
            
        Returns:
            Estimated velocity in m/s
        """
        # Thrust at given RPM
        thrust = self.calculate_thrust(rpm)
        
        # Balance thrust and drag
        # Solve for velocity where thrust = drag
        # Simplified iterative approach
        
        v_guess = rpm / 1000  # Initial guess
        
        for _ in range(10):
            drag = self.calculate_drag(v_guess, depth)
            error = thrust - drag
            
            if abs(error) < 0.1:
                break
            
            # Adjust velocity
            v_guess += error * 0.001
        
        return v_guess
    
    def __repr__(self) -> str:
        return (f"ThrusterModel(μ={self.viscosity:.2e}Pa·s, "
                f"ρ={self.density:.0f}kg/m³, "
                f"max_rpm={self.max_rpm})")
