"""
Speleogenesis modeling - cave formation and evolution processes.
Analyzes passage development, dissolution rates, and speleogenetic history.

Author: Samir Baladi
"""

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

logger = logging.getLogger(__name__)

@dataclass
class PassagePhase:
    """Represents a phase in passage development."""
    phase_type: str  # 'phreatic', 'vadose', 'transition'
    start_time: float  # Start time (years BP)
    end_time: float  # End time (years BP)
    dissolution_rate: float  # Rate during phase (mm/yr)
    final_dimensions: Tuple[float, float, float]  # V, L, F at phase end


class Speleogenesis:
    """
    Models cave formation and evolution processes over geological time.
    """
    
    def __init__(self,
                 initial_porosity: float = 0.01,
                 rock_density: float = 2700,
                 co2_initial: float = 1000):
        """
        Initialize speleogenesis model.
        
        Args:
            initial_porosity: Initial rock porosity
            rock_density: Rock density (kg/m³)
            co2_initial: Initial CO₂ concentration (ppm)
        """
        self.initial_porosity = initial_porosity
        self.rock_density = rock_density
        self.co2_initial = co2_initial
        self.phases = []
        self.passage_history = []
        logger.info("Initialized Speleogenesis model")

    def add_phase(self,
                 phase_type: str,
                 start_time: float,
                 end_time: float,
                 dissolution_rate: float,
                 final_dimensions: Tuple[float, float, float]) -> None:
        """
        Add a speleogenetic phase.
        
        Args:
            phase_type: Type of phase ('phreatic', 'vadose', 'transition')
            start_time: Start time (years BP)
            end_time: End time (years BP)
            dissolution_rate: Dissolution rate during phase (mm/yr)
            final_dimensions: Final (V, L, F) at phase end
        """
        phase = PassagePhase(
            phase_type=phase_type,
            start_time=start_time,
            end_time=end_time,
            dissolution_rate=dissolution_rate,
            final_dimensions=final_dimensions
        )
        
        self.phases.append(phase)
        logger.info(f"Added {phase_type} phase: {start_time}-{end_time} years BP")

    def simulate_evolution(self,
                          total_time: float = 100000,  # 100 kyr
                          n_steps: int = 1000) -> Dict[str, np.ndarray]:
        """
        Simulate cave evolution over geological time.
        
        Args:
            total_time: Total simulation time (years)
            n_steps: Number of time steps
            
        Returns:
            Dictionary with evolution time series
        """
        times = np.linspace(0, total_time, n_steps)
        V = np.zeros(n_steps)
        L = np.zeros(n_steps)
        F = np.zeros(n_steps)
        
        # Initial conditions
        V[0] = 0.1  # Initial fracture width
        L[0] = 0.1
        F[0] = 0
        
        # Time-varying dissolution rate
        def dissolution_rate(t):
            # Base rate decreases over time as water becomes less aggressive
            base_rate = 0.1 * np.exp(-t / 50000)  # mm/yr
            return base_rate
        
        for i in range(1, n_steps):
            dt = times[i] - times[i-1]
            t = times[i]
            
            # Dissolution rate at this time
            rate = dissolution_rate(t)
            
            # Update dimensions (simplified)
            V[i] = V[i-1] + rate * dt * np.random.uniform(0.8, 1.2)
            L[i] = L[i-1] + rate * dt * np.random.uniform(0.8, 1.2)
            
            # Flowstone deposition (increases over time as passages open)
            if t > 20000:  # Deposition starts after significant opening
                F[i] = F[i-1] + 0.01 * dt / 1000  # mm/yr
            else:
                F[i] = F[i-1]
        
        self.passage_history = {
            'time': times,
            'V': V,
            'L': L,
            'F': F
        }
        
        return self.passage_history

    def identify_phreatic_vadose_transition(self,
                                           V_history: np.ndarray,
                                           L_history: np.ndarray,
                                           threshold_ratio: float = 0.5) -> List[Dict[str, float]]:
        """
        Identify transitions between phreatic and vadose conditions.
        
        Args:
            V_history: Vertical dimension history
            L_history: Lateral dimension history
            threshold_ratio: V/L ratio for transition identification
            
        Returns:
            List of transition events
        """
        if len(V_history) < 2:
            return []
        
        # Calculate aspect ratio
        aspect_ratio = V_history / (L_history + 1e-6)
        
        transitions = []
        current_phase = 'phreatic' if aspect_ratio[0] < threshold_ratio else 'vadose'
        
        for i in range(1, len(aspect_ratio)):
            if aspect_ratio[i] < threshold_ratio and current_phase != 'phreatic':
                transitions.append({
                    'time_index': i,
                    'transition_type': 'vadose → phreatic',
                    'aspect_ratio': float(aspect_ratio[i])
                })
                current_phase = 'phreatic'
            elif aspect_ratio[i] >= threshold_ratio and current_phase != 'vadose':
                transitions.append({
                    'time_index': i,
                    'transition_type': 'phreatic → vadose',
                    'aspect_ratio': float(aspect_ratio[i])
                })
                current_phase = 'vadose'
        
        return transitions

    def calculate_dissolution_history(self,
                                     time_points: np.ndarray,
                                     V: np.ndarray) -> Dict[str, np.ndarray]:
        """
        Calculate dissolution rate history.
        
        Args:
            time_points: Time points
            V: Vertical dimension history
            
        Returns:
            Dictionary with dissolution rates
        """
        if len(V) < 2:
            return {'rates': np.array([])}
        
        # Calculate rates
        rates = np.gradient(V, time_points)
        
        # Smooth rates
        from scipy.signal import savgol_filter
        if len(rates) > 10:
            window = min(11, len(rates) if len(rates) % 2 == 1 else len(rates) - 1)
            rates_smoothed = savgol_filter(rates, window, 2)
        else:
            rates_smoothed = rates
        
        return {
            'time': time_points,
            'rates': rates,
            'rates_smoothed': rates_smoothed,
            'mean_rate': float(np.mean(rates)),
            'max_rate': float(np.max(rates)),
            'min_rate': float(np.min(rates))
        }

    def estimate_cave_age(self,
                         current_dimensions: Tuple[float, float, float],
                         typical_rate: float = 0.05) -> Dict[str, float]:
        """
        Estimate cave age from current dimensions.
        
        Args:
            current_dimensions: Current (V, L, F) dimensions
            typical_rate: Typical dissolution rate (mm/yr)
            
        Returns:
            Dictionary with age estimates
        """
        V, L, F = current_dimensions
        
        # Convert to mm
        V_mm = V * 1000
        L_mm = L * 1000
        
        # Estimate age from each dimension
        age_V = V_mm / typical_rate
        age_L = L_mm / typical_rate
        
        # Flowstone gives minimum age
        if F > 0:
            age_F = F / 0.087  # Using mean deposition rate from paper
        else:
            age_F = 0
        
        return {
            'age_from_V_years': round(age_V),
            'age_from_L_years': round(age_L),
            'age_from_F_years': round(age_F),
            'mean_age': round((age_V + age_L + age_F) / 3),
            'min_age': round(min(age_V, age_L, age_F) if age_F > 0 else min(age_V, age_L)),
            'max_age': round(max(age_V, age_L, age_F))
        }

    def classify_speleogenesis(self,
                              V: float,
                              L: float,
                              aspect_ratio: Optional[float] = None) -> Dict[str, Union[str, float]]:
        """
        Classify speleogenetic type based on passage morphology.
        
        Args:
            V: Vertical dimension
            L: Lateral dimension
            aspect_ratio: V/L ratio (calculated if None)
            
        Returns:
            Dictionary with classification
        """
        if aspect_ratio is None:
            aspect_ratio = V / L if L > 0 else 0
        
        if aspect_ratio < 0.3:
            # Laterally elongated
            if V < 1.0:
                type_name = "Phreatic tube"
                description = "Formed under phreatic conditions, circular to elliptical"
            else:
                type_name = "Phreatic conduit"
                description = "Water-filled passage, dissolutionally enlarged"
        elif aspect_ratio > 2.0:
            # Vertically elongated
            if V > 10:
                type_name = "Vadose shaft"
                description = "Vertical shaft formed by free-falling water"
            else:
                type_name = "Vadose canyon"
                description = "Downcutting stream passage"
        else:
            # Mixed
            if V > 5:
                type_name = "Keyhole passage"
                description = "Phreatic tube with vadose trench"
            else:
                type_name = "Mixed morphology"
                description = "Multiple speleogenetic phases"
        
        return {
            'type': type_name,
            'description': description,
            'aspect_ratio': round(aspect_ratio, 2),
            'dominant_process': 'Phreatic' if aspect_ratio < 0.5 else 'Vadose' if aspect_ratio > 1.5 else 'Mixed'
        }

    def calculate_evolution_rate(self,
                                time_span: float,
                                dimension_change: float) -> float:
        """
        Calculate evolution rate over time span.
        
        Args:
            time_span: Time span (years)
            dimension_change: Change in dimension (m)
            
        Returns:
            Evolution rate (mm/yr)
        """
        if time_span <= 0:
            return 0
        
        return (dimension_change * 1000) / time_span

    def predict_future_evolution(self,
                                current_dimensions: Tuple[float, float, float],
                                time_ahead: float = 10000,
                                scenario: str = 'current') -> Dict[str, np.ndarray]:
        """
        Predict future passage evolution.
        
        Args:
            current_dimensions: Current (V, L, F)
            time_ahead: Years to predict
            scenario: 'current', 'accelerated', or 'reduced'
            
        Returns:
            Dictionary with predictions
        """
        V, L, F = current_dimensions
        
        # Future time points
        times = np.linspace(0, time_ahead, 100)
        
        # Rate adjustments by scenario
        if scenario == 'accelerated':
            rate_factor = 1.5
        elif scenario == 'reduced':
            rate_factor = 0.5
        else:
            rate_factor = 1.0
        
        # Dissolution rate decreases over time (exhaustion)
        base_rate_V = 0.05 * rate_factor * np.exp(-times / 50000)
        base_rate_L = 0.06 * rate_factor * np.exp(-times / 50000)
        base_rate_F = 0.087 * np.exp(-times / 10000)  # Deposition decreases
        
        # Integrate
        V_future = V + np.cumsum(base_rate_V) * np.mean(np.diff(times)) / 1000
        L_future = L + np.cumsum(base_rate_L) * np.mean(np.diff(times)) / 1000
        F_future = F + np.cumsum(base_rate_F) * np.mean(np.diff(times))
        
        return {
            'time': times,
            'V': V_future,
            'L': L_future,
            'F': F_future,
            'final_V': float(V_future[-1]),
            'final_L': float(L_future[-1]),
            'final_F': float(F_future[-1])
        }

    def summary(self) -> Dict[str, Union[float, str, List]]:
        """Get summary of speleogenetic analysis."""
        summary = {
            'initial_porosity': self.initial_porosity,
            'rock_density': self.rock_density,
            'co2_initial': self.co2_initial,
            'n_phases': len(self.phases)
        }
        
        if self.phases:
            # Phase statistics
            phase_types = [p.phase_type for p in self.phases]
            summary['phase_types'] = list(set(phase_types))
            summary['total_time_span'] = self.phases[-1].start_time - self.phases[0].end_time
        
        if self.passage_history:
            V = self.passage_history['V']
            summary['final_V'] = round(V[-1], 2)
            summary['final_L'] = round(self.passage_history['L'][-1], 2)
            summary['final_F'] = round(self.passage_history['F'][-1], 3)
        
        return summary


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Speleogenesis Modeling")
    print("="*60)
    
    sp = Speleogenesis()
    
    # Simulate evolution
    evolution = sp.simulate_evolution(total_time=100000)
    print(f"\nEvolution Simulation:")
    print(f"  - Final V: {evolution['V'][-1]:.2f} m")
    print(f"  - Final L: {evolution['L'][-1]:.2f} m")
    print(f"  - Final F: {evolution['F'][-1]:.3f} mm")
    
    # Identify transitions
    transitions = sp.identify_phreatic_vadose_transition(evolution['V'], evolution['L'])
    print(f"\nPhase Transitions: {len(transitions)} detected")
    
    # Classify speleogenesis
    classification = sp.classify_speleogenesis(V=2.8, L=3.2)
    print(f"\nSpeleogenetic Classification:")
    print(f"  - Type: {classification['type']}")
    print(f"  - Description: {classification['description']}")
    print(f"  - Dominant Process: {classification['dominant_process']}")
    
    # Estimate cave age
    age = sp.estimate_cave_age((2.8, 3.2, 0.087))
    print(f"\nEstimated Cave Age:")
    print(f"  - Mean Age: {age['mean_age']:,} years")
    print(f"  - Range: {age['min_age']:,} - {age['max_age']:,} years")
    
    print("\n" + "="*60)
