"""
Flowstone Management (F) - Sedimentary Memory parameter.
Handles measurement and analysis of flowstone and speleothem deposition rates.

Equation: F(t) = F₀ e^(-λt) + ∫₀ᵗ η(τ) e^(-λ(t-τ)) dτ

Author: Samir Baladi
"""

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

logger = logging.getLogger(__name__)

@dataclass
class FlowstoneMeasurement:
    """Represents a single flowstone measurement."""
    thickness: float  # Measured thickness (mm)
    timestamp: float  # Time of measurement (years)
    deposition_rate: float  # Current deposition rate (mm/yr)
    layer_count: int  # Number of visible layers
    precision: float  # Measurement precision (mm)
    station_id: str  # Survey station identifier
    notes: Optional[str] = None


class FlowstoneDeposition:
    """
    Handles flowstone deposition measurements and analysis.
    
    Attributes:
        F0: Initial deposition thickness (mm)
        lambda_decay: Erosion decay rate (yr⁻¹)
        eta: Deposition rate function
        measurements: List of flowstone measurements
    """
    
    def __init__(self, F0: float = 0.0, lambda_decay: float = 0.01):
        """
        Initialize flowstone deposition analyzer.
        
        Args:
            F0: Initial deposition thickness (mm)
            lambda_decay: Erosion decay rate (yr⁻¹, default 0.01 from paper)
        """
        self.F0 = F0
        self.lambda_decay = lambda_decay
        self.measurements = []
        self.deposition_history = []
        logger.info(f"Initialized FlowstoneDeposition with F₀={F0}mm, λ={lambda_decay}yr⁻¹")

    def add_measurement(self,
                       thickness: float,
                       deposition_rate: float,
                       layer_count: int,
                       timestamp: float,
                       station_id: str,
                       precision: float = 0.5) -> None:
        """
        Add a new flowstone measurement.
        
        Args:
            thickness: Measured thickness (mm)
            deposition_rate: Current deposition rate (mm/yr)
            layer_count: Number of visible layers
            timestamp: Time of measurement (years)
            station_id: Survey station identifier
            precision: Measurement precision (default 0.5mm from paper)
        """
        measurement = FlowstoneMeasurement(
            thickness=thickness,
            timestamp=timestamp,
            deposition_rate=deposition_rate,
            layer_count=layer_count,
            precision=precision,
            station_id=station_id
        )
        
        self.measurements.append(measurement)
        self.deposition_history.append((timestamp, deposition_rate))
        logger.info(f"Added flowstone measurement at {station_id}: {thickness:.2f}mm, rate={deposition_rate:.3f}mm/yr")

    def calculate_thickness(self, t: np.ndarray, eta_func: Optional[callable] = None) -> np.ndarray:
        """
        Calculate flowstone thickness over time using convolution integral.
        
        F(t) = F₀ e^(-λt) + ∫₀ᵗ η(τ) e^(-λ(t-τ)) dτ
        
        Args:
            t: Time array (years)
            eta_func: Deposition rate function η(τ) (mm/yr)
            
        Returns:
            Thickness time series (mm)
        """
        if eta_func is None:
            # Default constant deposition rate
            eta_func = lambda tau: 0.087  # Mean rate from paper
        
        F = np.zeros_like(t)
        
        for i, ti in enumerate(t):
            # Initial term
            F[i] = self.F0 * np.exp(-self.lambda_decay * ti)
            
            if ti > 0:
                # Convolution integral
                def integrand(tau):
                    return eta_func(tau) * np.exp(-self.lambda_decay * (ti - tau))
                
                integral, _ = integrate.quad(integrand, 0, ti)
                F[i] += integral
        
        return F

    def fit_deposition_model(self, 
                            t: np.ndarray,
                            F_measured: np.ndarray) -> Dict[str, float]:
        """
        Fit deposition model parameters to measurements.
        
        Args:
            t: Time array
            F_measured: Measured thickness values
            
        Returns:
            Dictionary with fitted parameters
        """
        # Simple exponential fit
        # F = A * (1 - exp(-B*t)) + C
        
        from scipy.optimize import curve_fit
        
        def model(t, A, B, C):
            return A * (1 - np.exp(-B * t)) + C
        
        try:
            popt, pcov = curve_fit(model, t, F_measured, p0=[0.1, 0.05, 0])
            
            return {
                'A': popt[0],
                'B': popt[1],
                'C': popt[2],
                'R2': self._calculate_r2(F_measured, model(t, *popt))
            }
        except:
            logger.warning("Curve fitting failed")
            return {}

    def _calculate_r2(self, y_actual: np.ndarray, y_pred: np.ndarray) -> float:
        """Calculate R² goodness of fit."""
        ss_res = np.sum((y_actual - y_pred) ** 2)
        ss_tot = np.sum((y_actual - np.mean(y_actual)) ** 2)
        return 1 - (ss_res / ss_tot) if ss_tot > 0 else 0

    def analyze_deposition_rate(self,
                               window_years: float = 10.0) -> Dict[str, float]:
        """
        Analyze deposition rate trends over time.
        
        Args:
            window_years: Time window for rate calculation (years)
            
        Returns:
            Dictionary with rate statistics
        """
        if len(self.deposition_history) < 2:
            return {'mean_rate': 0.087, 'trend': 'unknown'}
        
        times = np.array([t for t, _ in self.deposition_history])
        rates = np.array([r for _, r in self.deposition_history])
        
        # Sort by time
        idx = np.argsort(times)
        times = times[idx]
        rates = rates[idx]
        
        # Calculate statistics
        mean_rate = np.mean(rates)
        std_rate = np.std(rates)
        
        # Calculate trend
        if len(times) > 2:
            slope, _ = np.polyfit(times - times[0], rates, 1)
            trend = 'increasing' if slope > 0.001 else 'decreasing' if slope < -0.001 else 'stable'
        else:
            slope = 0
            trend = 'unknown'
        
        # Detect cycles
        if len(rates) > 10:
            # Detrend
            rates_detrend = rates - np.polyval([slope, rates[0]], times - times[0])
            
            # Find dominant frequency
            fft = np.fft.fft(rates_detrend)
            freqs = np.fft.fftfreq(len(rates_detrend), np.mean(np.diff(times)))
            
            positive_idx = freqs > 0
            if np.any(positive_idx):
                peak_idx = np.argmax(np.abs(fft[positive_idx]))
                peak_freq = freqs[positive_idx][peak_idx]
                cycle_period = 1 / peak_freq if peak_freq > 0 else np.inf
            else:
                cycle_period = np.inf
        else:
            cycle_period = np.inf
        
        return {
            'mean_rate': round(mean_rate, 3),
            'std_rate': round(std_rate, 3),
            'min_rate': round(np.min(rates), 3),
            'max_rate': round(np.max(rates), 3),
            'trend': trend,
            'trend_slope': round(slope, 4),
            'cycle_period': round(cycle_period, 1) if cycle_period != np.inf else 'none',
            'n_measurements': len(rates)
        }

    def calculate_sedimentary_memory(self,
                                    current_thickness: Optional[float] = None) -> Dict[str, Union[float, str]]:
        """
        Calculate sedimentary memory score based on deposition history.
        
        Args:
            current_thickness: Current measured thickness (uses latest if None)
            
        Returns:
            Dictionary with memory score and assessment
        """
        if current_thickness is None and self.measurements:
            current_thickness = self.measurements[-1].thickness
        elif current_thickness is None:
            return {'score': 50, 'assessment': 'Insufficient data'}
        
        # Expected thickness based on typical rates
        if self.measurements:
            avg_rate = np.mean([m.deposition_rate for m in self.measurements])
        else:
            avg_rate = 0.087  # From paper
        
        # Time span of record
        if len(self.measurements) >= 2:
            time_span = self.measurements[-1].timestamp - self.measurements[0].timestamp
            expected_thickness = avg_rate * time_span
        else:
            expected_thickness = current_thickness
        
        # Score based on preservation
        preservation_ratio = current_thickness / expected_thickness if expected_thickness > 0 else 1
        score = min(100, preservation_ratio * 100)
        
        if score >= 80:
            assessment = "Excellent sedimentary record"
            memory = "High preservation"
        elif score >= 60:
            assessment = "Good sedimentary record"
            memory = "Moderate preservation"
        elif score >= 40:
            assessment = "Partial sedimentary record"
            memory = "Some erosion evident"
        elif score >= 20:
            assessment = "Poor sedimentary record"
            memory = "Significant erosion"
        else:
            assessment = "Minimal sedimentary memory"
            memory = "Mostly eroded"
        
        return {
            'score': round(score, 1),
            'assessment': assessment,
            'memory': memory,
            'current_thickness': current_thickness,
            'expected_thickness': round(expected_thickness, 2),
            'preservation_ratio': round(preservation_ratio, 2)
        }

    def detect_climate_signal(self) -> Dict[str, Union[float, str]]:
        """
        Detect climate change signals in deposition rates.
        
        Returns:
            Dictionary with climate signal analysis
        """
        if len(self.deposition_history) < 5:
            return {'signal': 'Insufficient data for climate analysis'}
        
        times = np.array([t for t, _ in self.deposition_history])
        rates = np.array([r for _, r in self.deposition_history])
        
        # Sort by time
        idx = np.argsort(times)
        times = times[idx]
        rates = rates[idx]
        
        # Detrend for cycle analysis
        if len(times) > 2:
            coeffs = np.polyfit(times - times[0], rates, 1)
            trend = coeffs[0]
            rates_detrend = rates - np.polyval(coeffs, times - times[0])
        else:
            trend = 0
            rates_detrend = rates
        
        # Look for periodic signals
        if len(rates_detrend) > 10:
            # Spectral analysis
            freqs, power = signal.periodogram(rates_detrend, fs=1.0/np.mean(np.diff(times)))
            
            # Find significant peaks
            significant = power > 2 * np.mean(power)
            peak_freqs = freqs[significant]
            peak_powers = power[significant]
            
            if len(peak_freqs) > 0:
                main_cycle = 1 / peak_freqs[0] if peak_freqs[0] > 0 else np.inf
                climate_signal = f"{main_cycle:.1f}-year cycle detected"
            else:
                climate_signal = "No significant cycles"
        else:
            climate_signal = "Insufficient for cycle detection"
        
        # Temperature correlation proxy
        # Warmer temperatures generally decrease deposition
        if trend < -0.001:
            climate_interpretation = "Decreasing deposition - possibly warming trend"
        elif trend > 0.001:
            climate_interpretation = "Increasing deposition - possibly cooling/wetter trend"
        else:
            climate_interpretation = "Stable deposition - no clear climate signal"
        
        return {
            'trend': round(trend, 4),
            'trend_interpretation': climate_interpretation,
            'climate_signal': climate_signal,
            'mean_rate': round(np.mean(rates), 3),
            'rate_change': round(rates[-1] - rates[0], 3) if len(rates) > 1 else 0,
            'rate_change_percent': round((rates[-1] / rates[0] - 1) * 100, 1) if len(rates) > 1 and rates[0] > 0 else 0
        }

    def predict_future_deposition(self,
                                 years_ahead: float = 10,
                                 scenario: str = 'current') -> Dict[str, float]:
        """
        Predict future deposition rates.
        
        Args:
            years_ahead: Years to predict ahead
            scenario: 'current', 'warming', or 'cooling'
            
        Returns:
            Dictionary with predictions
        """
        if not self.measurements:
            base_rate = 0.087
            base_thickness = 0
        else:
            base_rate = np.mean([m.deposition_rate for m in self.measurements])
            base_thickness = self.measurements[-1].thickness
        
        # Scenario adjustments
        if scenario == 'warming':
            rate_change = -0.0034  # From paper (mm/yr²)
        elif scenario == 'cooling':
            rate_change = 0.002
        else:
            rate_change = 0
        
        # Predict future rates
        future_times = np.linspace(0, years_ahead, int(years_ahead) + 1)
        future_rates = base_rate + rate_change * future_times
        future_rates = np.maximum(future_rates, 0.01)  # Minimum positive rate
        
        # Integrate to get thickness
        future_thickness = base_thickness + np.cumsum(future_rates) * np.mean(np.diff(future_times))
        
        return {
            'years_ahead': years_ahead,
            'scenario': scenario,
            'final_rate': round(future_rates[-1], 3),
            'final_thickness': round(future_thickness[-1], 2),
            'total_growth': round(future_thickness[-1] - base_thickness, 2),
            'rate_reduction': round((1 - future_rates[-1] / base_rate) * 100, 1) if base_rate > 0 else 0
        }

    def summary(self) -> Dict[str, Union[float, int, str]]:
        """Get summary statistics of flowstone deposition."""
        if not self.measurements:
            return {'status': 'No data'}
        
        thicknesses = [m.thickness for m in self.measurements]
        rates = [m.deposition_rate for m in self.measurements]
        layers = [m.layer_count for m in self.measurements]
        
        return {
            'n_measurements': len(self.measurements),
            'mean_thickness': np.mean(thicknesses),
            'std_thickness': np.std(thicknesses),
            'min_thickness': np.min(thicknesses),
            'max_thickness': np.max(thicknesses),
            'latest_thickness': thicknesses[-1],
            'mean_rate': np.mean(rates),
            'mean_layers': np.mean(layers),
            'latest_station': self.measurements[-1].station_id,
            'lambda_decay': self.lambda_decay
        }


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Flowstone Deposition Analysis")
    print("="*60)
    
    fd = FlowstoneDeposition(F0=0.5)
    
    # Add measurements
    fd.add_measurement(2.5, 0.087, 15, 2015.5, "F-001")
    fd.add_measurement(3.2, 0.085, 18, 2020.5, "F-002")
    fd.add_measurement(3.8, 0.082, 21, 2024.5, "F-003")
    
    # Analyze deposition rate
    rates = fd.analyze_deposition_rate()
    print(f"\nDeposition Rate Analysis:")
    print(f"  - Mean Rate: {rates['mean_rate']} mm/yr")
    print(f"  - Trend: {rates['trend']}")
    print(f"  - Cycle Period: {rates['cycle_period']} years")
    
    # Calculate sedimentary memory
    memory = fd.calculate_sedimentary_memory()
    print(f"\nSedimentary Memory:")
    print(f"  - Score: {memory['score']}/100")
    print(f"  - Assessment: {memory['assessment']}")
    print(f"  - Memory: {memory['memory']}")
    
    # Detect climate signal
    climate = fd.detect_climate_signal()
    print(f"\nClimate Signal:")
    print(f"  - Trend: {climate['trend']:.4f} mm/yr²")
    print(f"  - Interpretation: {climate['trend_interpretation']}")
    print(f"  - Climate Signal: {climate['climate_signal']}")
    
    # Predict future
    prediction = fd.predict_future_deposition(years_ahead=25, scenario='warming')
    print(f"\nFuture Prediction (25 years, warming scenario):")
    print(f"  - Final Rate: {prediction['final_rate']} mm/yr")
    print(f"  - Total Growth: {prediction['total_growth']} mm")
    print(f"  - Rate Reduction: {prediction['rate_reduction']}%")
    
    print("\n" + "="*60)
