"""
Groundwater flow modeling in karst conduits.
Implements Darcy flow and conduit hydraulics.

Author: Samir Baladi
"""

import numpy as np
from typing import Dict, List, Tuple, Optional, Union
from dataclasses import dataclass
import logging

logger = logging.getLogger(__name__)

@dataclass
class FlowParameters:
    """Parameters for groundwater flow calculation."""
    hydraulic_conductivity: float  # K (m/s)
    cross_section_area: float  # A (m²)
    hydraulic_gradient: float  # dh/dl (dimensionless)
    porosity: float  # n (dimensionless)


class GroundwaterFlow:
    """
    Models groundwater flow in karst conduits using Darcy's law and
    more sophisticated conduit flow equations.
    """
    
    def __init__(self,
                 hydraulic_conductivity: float = 1e-4,
                 porosity: float = 0.08,
                 conduit_roughness: float = 0.01):
        """
        Initialize groundwater flow model.
        
        Args:
            hydraulic_conductivity: Base hydraulic conductivity (m/s)
            porosity: Rock porosity (default 0.08 from paper)
            conduit_roughness: Manning's n for conduits
        """
        self.K_base = hydraulic_conductivity
        self.porosity = porosity
        self.n = conduit_roughness
        self.measurements = []
        logger.info(f"Initialized GroundwaterFlow with K={hydraulic_conductivity}m/s")

    def darcy_velocity(self, 
                      K: float, 
                      gradient: float) -> float:
        """
        Calculate Darcy velocity.
        
        v = -K * dh/dl
        
        Args:
            K: Hydraulic conductivity (m/s)
            gradient: Hydraulic gradient (dimensionless)
            
        Returns:
            Darcy velocity (m/s)
        """
        return K * gradient

    def darcy_flux(self,
                  K: float,
                  gradient: float,
                  area: float) -> float:
        """
        Calculate Darcy flux (discharge).
        
        Q = -K * A * dh/dl
        
        Args:
            K: Hydraulic conductivity (m/s)
            gradient: Hydraulic gradient
            area: Cross-sectional area (m²)
            
        Returns:
            Discharge (m³/s)
        """
        return K * area * gradient

    def conduit_velocity(self,
                        discharge: float,
                        area: float) -> float:
        """
        Calculate mean velocity in conduit.
        
        v = Q / A
        
        Args:
            discharge: Flow rate (m³/s)
            area: Cross-sectional area (m²)
            
        Returns:
            Mean velocity (m/s)
        """
        if area <= 0:
            return 0
        return discharge / area

    def manning_flow(self,
                    area: float,
                    slope: float,
                    hydraulic_radius: float) -> float:
        """
        Calculate flow using Manning's equation for open channel.
        
        Q = (1/n) * A * R^(2/3) * S^(1/2)
        
        Args:
            area: Cross-sectional area (m²)
            slope: Energy slope (m/m)
            hydraulic_radius: Hydraulic radius (m)
            
        Returns:
            Discharge (m³/s)
        """
        if self.n <= 0 or hydraulic_radius <= 0:
            return 0
        
        return (1 / self.n) * area * (hydraulic_radius ** (2/3)) * np.sqrt(slope)

    def pipe_flow(self,
                 diameter: float,
                 slope: float,
                 full: bool = True) -> float:
        """
        Calculate flow in full pipe (Darcy-Weisbach).
        
        Args:
            diameter: Pipe diameter (m)
            slope: Hydraulic slope
            full: Whether pipe is flowing full
            
        Returns:
            Discharge (m³/s)
        """
        if not full:
            # Use Manning for partly full
            area = np.pi * diameter**2 / 4
            hydraulic_radius = diameter / 4
            return self.manning_flow(area, slope, hydraulic_radius)
        
        # Darcy-Weisbach for full pipe
        g = 9.81  # m/s²
        friction_factor = 0.02  # Typical for karst conduits
        area = np.pi * diameter**2 / 4
        
        velocity = np.sqrt(2 * g * diameter * slope / friction_factor)
        return area * velocity

    def calculate_travel_time(self,
                            distance: float,
                            velocity: float) -> float:
        """
        Calculate travel time for given distance.
        
        Args:
            distance: Flow distance (m)
            velocity: Flow velocity (m/s)
            
        Returns:
            Travel time (hours)
        """
        if velocity <= 0:
            return float('inf')
        return distance / velocity / 3600  # Convert to hours

    def estimate_parameters(self,
                          discharge: float,
                          water_level: float,
                          width: float) -> FlowParameters:
        """
        Estimate flow parameters from measurements.
        
        Args:
            discharge: Measured discharge (m³/s)
            water_level: Water level (m)
            width: Conduit width (m)
            
        Returns:
            FlowParameters object
        """
        # Estimate cross-sectional area (simplified rectangular)
        area = width * water_level
        
        # Estimate velocity
        velocity = self.conduit_velocity(discharge, area)
        
        # Estimate gradient from velocity (simplified)
        # Using Manning backwards
        if area > 0 and width > 0:
            hydraulic_radius = area / (width + 2 * water_level) if (width + 2 * water_level) > 0 else water_level/2
            if hydraulic_radius > 0:
                # S = (v * n / R^(2/3))^2
                gradient = (velocity * self.n / (hydraulic_radius ** (2/3))) ** 2
            else:
                gradient = 0.001  # Default
        else:
            gradient = 0.001
        
        # Estimate conductivity (K = v / gradient)
        K = velocity / gradient if gradient > 0 else self.K_base
        
        return FlowParameters(
            hydraulic_conductivity=K,
            cross_section_area=area,
            hydraulic_gradient=gradient,
            porosity=self.porosity
        )

    def predict_flood_wave_speed(self,
                                base_velocity: float,
                                flood_ratio: float = 2.0) -> float:
        """
        Predict flood wave celerity.
        
        Args:
            base_velocity: Base flow velocity (m/s)
            flood_ratio: Ratio of flood to base discharge
            
        Returns:
            Flood wave speed (m/s)
        """
        # Flood waves travel faster than water particles
        # Celerity ≈ 1.5 * velocity for kinematic waves
        celerity = 1.5 * base_velocity * np.sqrt(flood_ratio)
        return celerity

    def analyze_flow_regime(self,
                           velocity: float,
                           depth: float) -> Dict[str, str]:
        """
        Analyze flow regime (laminar/turbulent, subcritical/supercritical).
        
        Args:
            velocity: Flow velocity (m/s)
            depth: Flow depth (m)
            
        Returns:
            Dictionary with flow regime classification
        """
        # Reynolds number for turbulence
        kinematic_viscosity = 1e-6  # m²/s at 20°C
        Re = velocity * depth / kinematic_viscosity if depth > 0 else 0
        
        if Re < 500:
            turbulence = "Laminar"
        elif Re < 2000:
            turbulence = "Transitional"
        else:
            turbulence = "Turbulent"
        
        # Froude number for subcritical/supercritical
        g = 9.81
        Fr = velocity / np.sqrt(g * depth) if depth > 0 else 0
        
        if Fr < 1:
            critical = "Subcritical"
            if Fr < 0.5:
                description = "Slow, deep flow"
            else:
                description = "Near-critical"
        elif Fr > 1:
            critical = "Supercritical"
            description = "Fast, shallow flow"
        else:
            critical = "Critical"
            description = "Critical flow"
        
        return {
            'turbulence': turbulence,
            'critical_flow': critical,
            'description': description,
            'Reynolds': round(Re),
            'Froude': round(Fr, 2)
        }

    def calculate_recharge(self,
                          precipitation: float,
                          catchment_area: float,
                          infiltration_rate: float = 0.3) -> float:
        """
        Calculate groundwater recharge from precipitation.
        
        Args:
            precipitation: Precipitation amount (mm)
            catchment_area: Catchment area (km²)
            infiltration_rate: Fraction of precipitation that infiltrates
            
        Returns:
            Recharge volume (m³)
        """
        # Convert precipitation from mm to m
        precip_m = precipitation / 1000
        
        # Convert catchment area from km² to m²
        area_m2 = catchment_area * 1e6
        
        recharge_volume = precip_m * area_m2 * infiltration_rate
        return recharge_volume

    def summary(self) -> Dict[str, Union[float, str]]:
        """Get summary of groundwater flow model."""
        return {
            'base_conductivity': self.K_base,
            'porosity': self.porosity,
            'roughness': self.n,
            'model_type': 'Darcy + Manning',
            'applicability': 'Karst conduits, open channels'
        }


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Groundwater Flow Modeling")
    print("="*60)
    
    gw = GroundwaterFlow()
    
    # Calculate Darcy flux
    Q = gw.darcy_flux(K=1e-3, gradient=0.01, area=5.0)
    print(f"\nDarcy Flux:")
    print(f"  - Discharge: {Q:.3f} m³/s")
    
    # Calculate conduit velocity
    v = gw.conduit_velocity(Q, area=5.0)
    print(f"  - Velocity: {v:.3f} m/s")
    
    # Manning flow in conduit
    Q_manning = gw.manning_flow(area=5.0, slope=0.01, hydraulic_radius=0.5)
    print(f"\nManning Flow:")
    print(f"  - Discharge: {Q_manning:.3f} m³/s")
    
    # Analyze flow regime
    regime = gw.analyze_flow_regime(velocity=0.5, depth=1.0)
    print(f"\nFlow Regime:")
    print(f"  - Turbulence: {regime['turbulence']}")
    print(f"  - Critical Flow: {regime['critical_flow']}")
    print(f"  - Description: {regime['description']}")
    
    # Estimate parameters from measurements
    params = gw.estimate_parameters(discharge=0.5, water_level=1.2, width=3.0)
    print(f"\nEstimated Parameters:")
    print(f"  - K: {params.hydraulic_conductivity:.2e} m/s")
    print(f"  - Area: {params.cross_section_area:.1f} m²")
    print(f"  - Gradient: {params.hydraulic_gradient:.4f}")
    
    print("\n" + "="*60)
