"""
Lateral Conduit Control (L) - Hydrological Flexibility parameter.
Handles measurement and analysis of lateral passage dimensions and hydrological response.

Equation: L(t) = β₀ + k₁ΔQ + k₂Δh_w

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 LateralMeasurement:
    """Represents a single lateral conduit measurement."""
    width: float  # Measured width (m)
    timestamp: float  # Time of measurement (years)
    discharge: float  # Water discharge at time of measurement (m³/s)
    water_level: float  # Water level (m)
    precision: float  # Measurement precision (m)
    station_id: str  # Survey station identifier
    notes: Optional[str] = None


class LateralConduit:
    """
    Handles lateral conduit measurements and analysis.
    
    Attributes:
        beta0: Baseline conduit width (m)
        k1: Discharge response coefficient
        k2: Water level response coefficient
        measurements: List of lateral measurements
    """
    
    def __init__(self, 
                 beta0: float = 3.0,
                 k1: float = 0.15,
                 k2: float = 0.25):
        """
        Initialize lateral conduit analyzer.
        
        Args:
            beta0: Baseline width (default 3.0m from paper)
            k1: Discharge response coefficient
            k2: Water level response coefficient
        """
        self.beta0 = beta0
        self.k1 = k1
        self.k2 = k2
        self.measurements = []
        self.discharge_history = []
        logger.info(f"Initialized LateralConduit with β₀={beta0}m")

    def add_measurement(self,
                       width: float,
                       discharge: float,
                       water_level: float,
                       timestamp: float,
                       station_id: str,
                       precision: float = 0.1) -> None:
        """
        Add a new lateral conduit measurement.
        
        Args:
            width: Measured width (m)
            discharge: Water discharge at measurement time (m³/s)
            water_level: Water level (m)
            timestamp: Time of measurement (years)
            station_id: Survey station identifier
            precision: Measurement precision (default 0.1m from paper)
        """
        measurement = LateralMeasurement(
            width=width,
            timestamp=timestamp,
            discharge=discharge,
            water_level=water_level,
            precision=precision,
            station_id=station_id
        )
        
        self.measurements.append(measurement)
        self.discharge_history.append((timestamp, discharge))
        logger.info(f"Added lateral measurement at {station_id}: {width:.2f}m, Q={discharge:.2f}m³/s")

    def calculate_hydrological_response(self,
                                       discharge_change: float,
                                       water_level_change: float) -> float:
        """
        Calculate lateral width change in response to hydrological changes.
        
        ΔL = k₁ΔQ + k₂Δh_w
        
        Args:
            discharge_change: Change in discharge (m³/s)
            water_level_change: Change in water level (m)
            
        Returns:
            Predicted width change (m)
        """
        width_change = self.k1 * discharge_change + self.k2 * water_level_change
        return width_change

    def predict_width_from_hydrology(self,
                                    current_discharge: float,
                                    current_water_level: float,
                                    reference_time: Optional[float] = None) -> Dict[str, float]:
        """
        Predict current width based on hydrological conditions.
        
        Args:
            current_discharge: Current discharge (m³/s)
            current_water_level: Current water level (m)
            reference_time: Reference time for baseline (uses latest if None)
            
        Returns:
            Dictionary with predicted width and confidence
        """
        if not self.measurements:
            return {
                'predicted_width': self.beta0,
                'confidence': 0.5,
                'message': 'Using default baseline'
            }
        
        # Get reference state (latest measurement)
        ref = self.measurements[-1]
        
        # Calculate changes
        dQ = current_discharge - ref.discharge
        dh = current_water_level - ref.water_level
        
        # Predict new width
        predicted_width = ref.width + self.calculate_hydrological_response(dQ, dh)
        
        # Confidence decreases with magnitude of change
        confidence = max(0.3, 1.0 - 0.1 * (abs(dQ) + abs(dh)))
        
        return {
            'predicted_width': round(predicted_width, 2),
            'reference_width': ref.width,
            'discharge_change': dQ,
            'water_level_change': dh,
            'confidence': round(confidence, 2)
        }

    def calculate_adaptation_score(self,
                                  current_width: Optional[float] = None,
                                  current_discharge: Optional[float] = None) -> Dict[str, Union[float, str]]:
        """
        Calculate hydrological adaptation score.
        
        Args:
            current_width: Current measured width (uses latest if None)
            current_discharge: Current discharge (uses latest if None)
            
        Returns:
            Dictionary with adaptation score and assessment
        """
        if not self.measurements:
            return {'score': 50, 'assessment': 'Insufficient data'}
        
        if current_width is None:
            current_width = self.measurements[-1].width
        if current_discharge is None:
            current_discharge = self.measurements[-1].discharge
        
        # Normal flow condition from paper
        Q_normal = 0.15  # m³/s
        
        # Calculate how well width adapts to discharge
        expected_width = self.beta0 + self.k1 * (current_discharge - Q_normal)
        adaptation_error = abs(current_width - expected_width) / expected_width
        
        # Score (100 - scaled error)
        score = max(0, 100 * (1 - min(1, adaptation_error)))
        
        if score >= 80:
            assessment = "Excellent hydrological adaptation"
        elif score >= 60:
            assessment = "Good adaptation"
        elif score >= 40:
            assessment = "Moderate adaptation"
        elif score >= 20:
            assessment = "Poor adaptation"
        else:
            assessment = "Critical adaptation failure"
        
        return {
            'score': round(score, 1),
            'assessment': assessment,
            'current_width': current_width,
            'expected_width': round(expected_width, 2),
            'current_discharge': current_discharge,
            'adaptation_error': round(adaptation_error * 100, 1)
        }

    def analyze_flood_response(self,
                              flood_peak_discharge: float,
                              pre_flood_width: Optional[float] = None) -> Dict[str, float]:
        """
        Analyze passage response to flood pulse.
        
        Args:
            flood_peak_discharge: Peak discharge during flood (m³/s)
            pre_flood_width: Width before flood (uses baseline if None)
            
        Returns:
            Dictionary with predicted flood response
        """
        if pre_flood_width is None:
            pre_flood_width = self.beta0
        
        # Normal flow from paper
        Q_normal = 0.15
        
        # Calculate erosional widening during flood
        dQ_flood = flood_peak_discharge - Q_normal
        
        # Enhanced erosion during floods (k1 increases)
        k1_flood = self.k1 * 2.5  # 2.5x during floods
        
        width_increase = k1_flood * dQ_flood
        post_flood_width = pre_flood_width + width_increase
        
        # Post-flood recovery (partial)
        recovery = 0.7 * width_increase  # 70% recovery
        stabilized_width = pre_flood_width + recovery
        
        return {
            'pre_flood_width': pre_flood_width,
            'peak_width': round(post_flood_width, 2),
            'stabilized_width': round(stabilized_width, 2),
            'temporary_increase': round(width_increase, 2),
            'permanent_increase': round(recovery, 2),
            'widening_percent': round(width_increase / pre_flood_width * 100, 1)
        }

    def calculate_dissolution_rate(self,
                                  co2_concentration: float,
                                  water_velocity: float,
                                  temperature: float = 12.0) -> float:
        """
        Calculate limestone dissolution rate.
        
        Args:
            co2_concentration: CO₂ concentration (ppm)
            water_velocity: Water flow velocity (m/s)
            temperature: Water temperature (°C)
            
        Returns:
            Dissolution rate (mm/yr)
        """
        # Simplified dissolution kinetics
        # Based on PWP3D model from paper
        
        # Base rate
        k_diss = 0.08  # m³/mol from paper
        
        # CO₂ enhancement
        co2_factor = co2_concentration / 1000  # Normalized to 1000ppm
        
        # Velocity enhancement
        v_factor = np.sqrt(water_velocity / 0.1)  # Normalized to 0.1 m/s
        
        # Temperature effect (Arrhenius-type)
        T_ref = 12.0  # Reference temperature
        T_factor = np.exp(0.05 * (temperature - T_ref))
        
        dissolution_rate = k_diss * co2_factor * v_factor * T_factor
        
        return dissolution_rate * 1000  # Convert to mm/yr

    def get_route_flexibility(self, 
                             current_width: float,
                             team_size: int = 3) -> Dict[str, Union[str, float]]:
        """
        Assess route flexibility based on passage width.
        
        Args:
            current_width: Current passage width (m)
            team_size: Number of team members
            
        Returns:
            Dictionary with flexibility assessment
        """
        # Minimum width requirements
        min_single = 0.5  # m for single person
        min_side_by_side = 1.5  # m for two abreast
        min_equipment = 0.8  # m for equipment passage
        
        if current_width < min_single:
            flexibility = "Impassable"
            max_team = 1
            formation = "Crawling only"
        elif current_width < min_equipment:
            flexibility = "Very Limited"
            max_team = 1
            formation = "Single file, no equipment"
        elif current_width < min_side_by_side:
            flexibility = "Limited"
            max_team = 2
            formation = "Single file"
        elif current_width < 3.0:
            flexibility = "Moderate"
            max_team = 3
            formation = "Single file with passing"
        else:
            flexibility = "High"
            max_team = 5
            formation = "Multiple abreast"
        
        return {
            'flexibility': flexibility,
            'max_team_size': max_team,
            'recommended_formation': formation,
            'current_width': current_width,
            'team_comfort': 'Comfortable' if team_size <= max_team else 'Cramped'
        }

    def summary(self) -> Dict[str, Union[float, int, str]]:
        """Get summary statistics of lateral conduit."""
        if not self.measurements:
            return {'status': 'No data'}
        
        widths = [m.width for m in self.measurements]
        discharges = [m.discharge for m in self.measurements]
        
        return {
            'n_measurements': len(self.measurements),
            'mean_width': np.mean(widths),
            'std_width': np.std(widths),
            'min_width': np.min(widths),
            'max_width': np.max(widths),
            'latest_width': widths[-1],
            'latest_discharge': discharges[-1],
            'latest_station': self.measurements[-1].station_id,
            'baseline': self.beta0,
            'k1_discharge_coef': self.k1,
            'k2_water_level_coef': self.k2
        }


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Lateral Conduit Analysis")
    print("="*60)
    
    lc = LateralConduit()
    
    # Add measurements
    lc.add_measurement(3.2, 0.15, 0.45, 2021.5, "L-001")
    lc.add_measurement(3.4, 0.22, 0.62, 2022.5, "L-002")
    
    # Predict width from hydrology
    prediction = lc.predict_width_from_hydrology(0.35, 0.85)
    print(f"\nHydrological Prediction:")
    print(f"  - Predicted Width: {prediction['predicted_width']}m")
    print(f"  - Confidence: {prediction['confidence']:.0%}")
    
    # Calculate adaptation score
    score = lc.calculate_adaptation_score()
    print(f"\nAdaptation Score:")
    print(f"  - Score: {score['score']}/100")
    print(f"  - Assessment: {score['assessment']}")
    
    # Analyze flood response
    flood = lc.analyze_flood_response(flood_peak_discharge=5.2)
    print(f"\nFlood Response:")
    print(f"  - Peak Width: {flood['peak_width']}m")
    print(f"  - Stabilized Width: {flood['stabilized_width']}m")
    print(f"  - Widening: {flood['widening_percent']}%")
    
    # Route flexibility
    flexibility = lc.get_route_flexibility(3.2)
    print(f"\nRoute Flexibility:")
    print(f"  - Flexibility: {flexibility['flexibility']}")
    print(f"  - Max Team Size: {flexibility['max_team_size']}")
    print(f"  - Formation: {flexibility['recommended_formation']}")
    
    print("\n" + "="*60)
