"""
Water chemistry analysis for karst systems.
Handles pH, Ca²⁺ concentration, and dissolution kinetics.

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 WaterSample:
    """Represents a water chemistry sample."""
    pH: float
    calcium: float  # Ca²⁺ concentration (mg/L)
    temperature: float  # °C
    alkalinity: float  # mg/L as CaCO₃
    conductivity: float  # μS/cm
    timestamp: float
    location: str


class WaterChemistry:
    """
    Analyzes water chemistry for dissolution kinetics and cave formation.
    """
    
    def __init__(self,
                 k_diss: float = 0.08,
                 k_dep: float = 0.003,
                 temp_ref: float = 12.0):
        """
        Initialize water chemistry analyzer.
        
        Args:
            k_diss: Dissolution rate constant (m³/mol, default 0.08 from paper)
            k_dep: Deposition rate constant (m³/mol, default 0.003 from paper)
            temp_ref: Reference temperature (°C)
        """
        self.k_diss = k_diss
        self.k_dep = k_dep
        self.temp_ref = temp_ref
        self.samples = []
        logger.info(f"Initialized WaterChemistry with k_diss={k_diss}, k_dep={k_dep}")

    def add_sample(self,
                  pH: float,
                  calcium: float,
                  temperature: float,
                  alkalinity: float,
                  conductivity: float,
                  timestamp: float,
                  location: str) -> None:
        """
        Add a water chemistry sample.
        
        Args:
            pH: pH value
            calcium: Ca²⁺ concentration (mg/L)
            temperature: Water temperature (°C)
            alkalinity: Alkalinity (mg/L as CaCO₃)
            conductivity: Electrical conductivity (μS/cm)
            timestamp: Sample time
            location: Sample location
        """
        sample = WaterSample(
            pH=pH,
            calcium=calcium,
            temperature=temperature,
            alkalinity=alkalinity,
            conductivity=conductivity,
            timestamp=timestamp,
            location=location
        )
        
        self.samples.append(sample)
        logger.debug(f"Added water sample from {location}: pH={pH}, Ca={calcium}mg/L")

    def calculate_saturation_index(self,
                                  calcium: float,
                                  pH: float,
                                  temperature: float) -> float:
        """
        Calculate calcite saturation index.
        
        SI = log(IAP / Ksp)
        
        Args:
            calcium: Ca²⁺ concentration (mg/L)
            pH: pH value
            temperature: Water temperature (°C)
            
        Returns:
            Saturation index (positive = supersaturated, negative = undersaturated)
        """
        # Convert mg/L to mol/L
        Ca_mol = calcium / 40080  # Molecular weight of Ca²⁺: 40.08 g/mol
        
        # Estimate carbonate concentration from pH (simplified)
        # In reality, need full alkalinity speciation
        CO3_mol = 10**(pH - 10.3) * 1e-3  # Rough approximation
        
        # Ion activity product
        IAP = Ca_mol * CO3_mol
        
        # Temperature-dependent Ksp (simplified)
        Ksp = 10**(-171.9065 - 0.077993*temperature + 2839.319/temperature + 71.595*np.log10(temperature))
        
        # Saturation index
        if Ksp > 0:
            SI = np.log10(IAP / Ksp)
        else:
            SI = 0
        
        return round(SI, 2)

    def calculate_dissolution_rate(self,
                                  calcium: float,
                                  pH: float,
                                  temperature: float,
                                  discharge: float,
                                  area: float) -> Dict[str, float]:
        """
        Calculate limestone dissolution rate.
        
        Args:
            calcium: Ca²⁺ concentration (mg/L)
            pH: pH value
            temperature: Water temperature (°C)
            discharge: Water discharge (m³/s)
            area: Reaction surface area (m²)
            
        Returns:
            Dictionary with dissolution rates
        """
        # Temperature correction
        temp_factor = np.exp(0.05 * (temperature - self.temp_ref))
        
        # pH effect (more acidic = faster dissolution)
        pH_factor = 10 ** (7 - min(pH, 7))  # Simplified
        
        # Calculate dissolution rate (mol/m²/s)
        # Rate = k_diss * (1 - SI) * pH_factor * temp_factor
        SI = self.calculate_saturation_index(calcium, pH, temperature)
        
        if SI < 0:  # Undersaturated = dissolution
            rate_mol = self.k_diss * abs(SI) * pH_factor * temp_factor
            process = "dissolution"
        else:  # Supersaturated = deposition
            rate_mol = self.k_dep * SI * pH_factor * temp_factor
            process = "deposition"
        
        # Convert to mm/yr (assuming density 2.7 g/cm³)
        # mol/m²/s -> mm/yr: multiply by (mol weight / density) * seconds per year
        mol_weight = 100.09  # g/mol for CaCO₃
        density = 2.7e6  # g/m³
        conversion = mol_weight / density * 365.25 * 24 * 3600 * 1000  # to mm/yr
        
        rate_mm_yr = rate_mol * conversion
        
        # Total mass transfer
        if area > 0:
            total_mol_s = rate_mol * area
            total_kg_yr = total_mol_s * mol_weight / 1000 * 365.25 * 24 * 3600
        else:
            total_kg_yr = 0
        
        return {
            'rate_mm_yr': round(rate_mm_yr, 4),
            'rate_mol_m2_s': round(rate_mol, 8),
            'process': process,
            'saturation_index': SI,
            'total_kg_yr': round(total_kg_yr, 1),
            'temperature_factor': round(temp_factor, 2),
            'pH_factor': round(pH_factor, 2)
        }

    def estimate_co2_pressure(self,
                             pH: float,
                             alkalinity: float,
                             temperature: float) -> float:
        """
        Estimate CO₂ partial pressure from pH and alkalinity.
        
        Args:
            pH: pH value
            alkalinity: Alkalinity (mg/L as CaCO₃)
            temperature: Water temperature (°C)
            
        Returns:
            CO₂ partial pressure (atm)
        """
        # Convert alkalinity to eq/L
        Alk_eq = alkalinity / 50000  # 50 g/eq for CaCO₃
        
        # Henry's constant for CO₂ (temperature dependent)
        KH = 10**(-1.11 + 0.016*temperature - 0.00007*temperature**2)  # mol/L/atm
        
        # H₂CO₃* concentration from pH and alkalinity (simplified)
        # [H2CO3*] = Alk_eq * 10^(6.3 - pH) / (1 + 2*10^(6.3 - pH))
        H2CO3 = Alk_eq * 10**(6.3 - pH) / (1 + 2 * 10**(6.3 - pH))
        
        # pCO₂ = [H2CO3*] / KH
        pCO2 = H2CO3 / KH if KH > 0 else 0
        
        return round(pCO2 * 1e6, 1)  # Return in ppm

    def analyze_flood_chemistry(self,
                               pre_flood_sample: WaterSample,
                               flood_sample: WaterSample) -> Dict[str, Union[str, float]]:
        """
        Analyze chemical changes during flood event.
        
        Args:
            pre_flood_sample: Sample before flood
            flood_sample: Sample during flood
            
        Returns:
            Dictionary with flood chemistry analysis
        """
        # Calculate changes
        delta_pH = flood_sample.pH - pre_flood_sample.pH
        delta_Ca = flood_sample.calcium - pre_flood_sample.calcium
        delta_temp = flood_sample.temperature - pre_flood_sample.temperature
        
        # Dilution factor (using conductivity as proxy)
        if pre_flood_sample.conductivity > 0:
            dilution = flood_sample.conductivity / pre_flood_sample.conductivity
        else:
            dilution = 1.0
        
        # Interpret changes
        if delta_pH < -0.5:
            pH_change = "Significant acidification"
        elif delta_pH < -0.2:
            pH_change = "Moderate acidification"
        elif delta_pH > 0.5:
            pH_change = "Significant alkalization"
        elif delta_pH > 0.2:
            pH_change = "Moderate alkalization"
        else:
            pH_change = "Stable pH"
        
        if delta_Ca < -10:
            Ca_change = "Strong dilution"
        elif delta_Ca < -5:
            Ca_change = "Moderate dilution"
        elif delta_Ca > 10:
            Ca_change = "Enhanced dissolution"
        elif delta_Ca > 5:
            Ca_change = "Slight increase"
        else:
            Ca_change = "Stable"
        
        return {
            'delta_pH': round(delta_pH, 2),
            'delta_Ca': round(delta_Ca, 1),
            'delta_temperature': round(delta_temp, 1),
            'dilution_factor': round(dilution, 2),
            'pH_interpretation': pH_change,
            'Ca_interpretation': Ca_change,
            'flood_impact': 'Acidification' if delta_pH < -0.3 else 'Dilution' if dilution < 0.8 else 'Mixed'
        }

    def calculate_ca_co2_relationship(self,
                                     samples: Optional[List[WaterSample]] = None) -> Dict[str, float]:
        """
        Calculate relationship between Ca and CO₂.
        
        Args:
            samples: List of samples (uses all if None)
            
        Returns:
            Dictionary with correlation statistics
        """
        if samples is None:
            samples = self.samples
        
        if len(samples) < 3:
            return {'correlation': 0, 'significance': 'Insufficient data'}
        
        # Extract data
        Ca = np.array([s.calcium for s in samples])
        pH = np.array([s.pH for s in samples])
        
        # Calculate correlation
        corr = np.corrcoef(Ca, pH)[0, 1]
        
        # Significance (simplified)
        if abs(corr) > 0.7:
            significance = "Strong"
        elif abs(corr) > 0.4:
            significance = "Moderate"
        elif abs(corr) > 0.2:
            significance = "Weak"
        else:
            significance = "None"
        
        return {
            'correlation': round(corr, 3),
            'significance': significance,
            'n_samples': len(samples),
            'interpretation': 'Higher Ca with lower pH' if corr < -0.3 else 'No clear relationship'
        }

    def calculate_deposition_potential(self,
                                      calcium: float,
                                      pH: float,
                                      temperature: float,
                                      flow_velocity: float) -> float:
        """
        Calculate speleothem deposition potential.
        
        Args:
            calcium: Ca²⁺ concentration (mg/L)
            pH: pH value
            temperature: Water temperature (°C)
            flow_velocity: Water flow velocity (m/s)
            
        Returns:
            Deposition potential (0-1)
        """
        SI = self.calculate_saturation_index(calcium, pH, temperature)
        
        if SI <= 0:
            return 0  # No deposition potential
        
        # Higher SI means higher potential
        # But very high velocity inhibits deposition
        velocity_factor = max(0, 1 - flow_velocity / 0.5)  # Decreases above 0.5 m/s
        
        potential = min(1, SI / 2 * velocity_factor)
        
        return round(potential, 2)

    def summary(self) -> Dict[str, Union[float, str]]:
        """Get summary of water chemistry."""
        if not self.samples:
            return {'status': 'No samples'}
        
        pHs = [s.pH for s in self.samples]
        Cas = [s.calcium for s in self.samples]
        temps = [s.temperature for s in self.samples]
        
        return {
            'n_samples': len(self.samples),
            'mean_pH': round(np.mean(pHs), 2),
            'pH_range': f"{min(pHs):.1f}-{max(pHs):.1f}",
            'mean_Ca': round(np.mean(Cas), 1),
            'Ca_range': f"{min(Cas):.0f}-{max(Cas):.0f}",
            'mean_temperature': round(np.mean(temps), 1),
            'k_diss': self.k_diss,
            'k_dep': self.k_dep
        }


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Water Chemistry Analysis")
    print("="*60)
    
    wc = WaterChemistry()
    
    # Add samples
    wc.add_sample(pH=7.2, calcium=85, temperature=11.5, 
                  alkalinity=210, conductivity=450, 
                  timestamp=2024.5, location="Main Conduit")
    
    # Calculate saturation index
    SI = wc.calculate_saturation_index(85, 7.2, 11.5)
    print(f"\nSaturation Index: {SI}")
    
    # Calculate dissolution rate
    rate = wc.calculate_dissolution_rate(85, 7.2, 11.5, discharge=0.15, area=100)
    print(f"\nDissolution Rate:")
    print(f"  - Process: {rate['process']}")
    print(f"  - Rate: {rate['rate_mm_yr']} mm/yr")
    print(f"  - SI: {rate['saturation_index']}")
    
    # Estimate CO₂ pressure
    pCO2 = wc.estimate_co2_pressure(7.2, 210, 11.5)
    print(f"\nEstimated pCO₂: {pCO2} ppm")
    
    print("\n" + "="*60)
