"""
Lithology classification and rock property analysis for karst systems.
Handles rock type identification, purity analysis, and dissolution susceptibility.

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 RockSample:
    """Represents a rock sample from the cave system."""
    lithology: str  # Rock type (limestone, dolomite, gypsum, etc.)
    caco3_purity: float  # CaCO₃ percentage (0-100)
    fracture_density: float  # Fractures per meter
    porosity: float  # Rock porosity (0-1)
    hardness: float  # Mohs hardness
    location: str  # Sample location
    timestamp: float  # Collection time


class LithologyClassifier:
    """
    Classifies lithology and analyzes rock properties for karst systems.
    """
    
    LITHOLOGY_TYPES = {
        'limestone': {
            'caco3_range': (90, 100),
            'porosity_range': (0.05, 0.15),
            'hardness_range': (3, 4),
            'dissolution_rate': 1.0,
            'color': '🟤'
        },
        'dolomite': {
            'caco3_range': (50, 90),
            'porosity_range': (0.03, 0.10),
            'hardness_range': (3.5, 4.5),
            'dissolution_rate': 0.6,
            'color': '⚪'
        },
        'gypsum': {
            'caco3_range': (0, 10),
            'porosity_range': (0.10, 0.25),
            'hardness_range': (2, 2.5),
            'dissolution_rate': 2.5,
            'color': '⬜'
        },
        'marble': {
            'caco3_range': (95, 100),
            'porosity_range': (0.01, 0.05),
            'hardness_range': (3, 4),
            'dissolution_rate': 0.3,
            'color': '⚪'
        },
        'chalk': {
            'caco3_range': (95, 100),
            'porosity_range': (0.20, 0.40),
            'hardness_range': (1, 2),
            'dissolution_rate': 1.8,
            'color': '⬜'
        }
    }
    
    def __init__(self):
        """Initialize lithology classifier."""
        self.samples = []
        self.formations = {}
        logger.info("Initialized LithologyClassifier")

    def add_sample(self,
                  lithology: str,
                  caco3_purity: float,
                  fracture_density: float,
                  porosity: float,
                  hardness: float,
                  location: str,
                  timestamp: float) -> None:
        """
        Add a rock sample.
        
        Args:
            lithology: Rock type
            caco3_purity: CaCO₃ percentage
            fracture_density: Fractures per meter
            porosity: Rock porosity
            hardness: Mohs hardness
            location: Sample location
            timestamp: Collection time
        """
        sample = RockSample(
            lithology=lithology,
            caco3_purity=caco3_purity,
            fracture_density=fracture_density,
            porosity=porosity,
            hardness=hardness,
            location=location,
            timestamp=timestamp
        )
        
        self.samples.append(sample)
        
        # Update formation statistics
        if lithology not in self.formations:
            self.formations[lithology] = {
                'samples': [],
                'mean_purity': [],
                'mean_fracture': []
            }
        
        self.formations[lithology]['samples'].append(sample)
        logger.info(f"Added {lithology} sample from {location}")

    def classify_rock(self,
                     caco3_purity: float,
                     porosity: float,
                     hardness: float,
                     acid_reaction: str = 'strong') -> Dict[str, Union[str, float]]:
        """
        Classify rock type based on properties.
        
        Args:
            caco3_purity: CaCO₃ percentage
            porosity: Rock porosity
            hardness: Mohs hardness
            acid_reaction: Reaction to HCl ('strong', 'weak', 'none')
            
        Returns:
            Dictionary with classification
        """
        # Initial guess based on CaCO₃
        if caco3_purity > 90:
            if porosity > 0.2:
                candidate = 'chalk'
            elif hardness > 3.5:
                candidate = 'marble'
            else:
                candidate = 'limestone'
        elif caco3_purity > 50:
            candidate = 'dolomite'
        else:
            candidate = 'gypsum'
        
        # Adjust based on acid reaction
        if acid_reaction == 'strong':
            confidence = 0.9 if candidate in ['limestone', 'chalk', 'marble'] else 0.6
        elif acid_reaction == 'weak':
            confidence = 0.9 if candidate == 'dolomite' else 0.5
        else:
            confidence = 0.9 if candidate == 'gypsum' else 0.4
        
        # Get reference properties
        ref = self.LITHOLOGY_TYPES.get(candidate, self.LITHOLOGY_TYPES['limestone'])
        
        return {
            'lithology': candidate,
            'confidence': round(confidence, 2),
            'color': ref['color'],
            'dissolution_rate': ref['dissolution_rate'],
            'description': self._get_description(candidate),
            'typical_uses': self._get_typical_uses(candidate)
        }

    def _get_description(self, lithology: str) -> str:
        """Get description of lithology type."""
        descriptions = {
            'limestone': "Pure calcium carbonate, well-bedded, medium-hard",
            'dolomite': "Magnesium-rich, slower dissolving, harder",
            'gypsum': "Evaporite, very soluble, soft",
            'marble': "Metamorphosed limestone, crystalline, dense",
            'chalk': "Soft, porous, microcrystalline limestone"
        }
        return descriptions.get(lithology, "Unknown lithology")

    def _get_typical_uses(self, lithology: str) -> str:
        """Get typical uses in speleology."""
        uses = {
            'limestone': "Primary karst former, good for passage development",
            'dolomite': "Forms distinctive karst, less cavernous",
            'gypsum': "Forms rapid dissolution features, unstable",
            'marble': "Forms tight joints, limited karst",
            'chalk': "Forms small conduits, high porosity"
        }
        return uses.get(lithology, "Unknown application")

    def calculate_dissolution_potential(self,
                                       lithology: str,
                                       caco3_purity: float,
                                       fracture_density: float,
                                       water_aggressiveness: float = 1.0) -> float:
        """
        Calculate dissolution potential for given rock.
        
        Args:
            lithology: Rock type
            caco3_purity: CaCO₃ percentage
            fracture_density: Fractures per meter
            water_aggressiveness: Water aggressiveness factor (1.0 = normal)
            
        Returns:
            Dissolution potential (0-1)
        """
        # Base rate from lithology
        base_rate = self.LITHOLOGY_TYPES.get(lithology, {}).get('dissolution_rate', 1.0)
        
        # Purity factor (higher purity = faster dissolution)
        purity_factor = caco3_purity / 100
        
        # Fracture factor (more fractures = more surface area)
        fracture_factor = min(1.0, fracture_density / 10)
        
        # Combined potential
        potential = base_rate * purity_factor * (1 + fracture_factor) * water_aggressiveness
        
        return min(1.0, potential)

    def analyze_formation(self, lithology: str) -> Dict[str, Union[float, str]]:
        """
        Analyze statistics for a specific formation.
        
        Args:
            lithology: Rock type to analyze
            
        Returns:
            Dictionary with formation statistics
        """
        if lithology not in self.formations:
            return {'status': f'No samples for {lithology}'}
        
        samples = self.formations[lithology]['samples']
        
        purities = [s.caco3_purity for s in samples]
        fractures = [s.fracture_density for s in samples]
        porosities = [s.porosity for s in samples]
        
        return {
            'n_samples': len(samples),
            'mean_purity': round(np.mean(purities), 1),
            'purity_range': f"{min(purities):.1f}-{max(purities):.1f}",
            'mean_fracture_density': round(np.mean(fractures), 2),
            'fracture_range': f"{min(fractures):.1f}-{max(fractures):.1f}",
            'mean_porosity': round(np.mean(porosities), 3),
            'dissolution_rate': self.LITHOLOGY_TYPES[lithology]['dissolution_rate'],
            'color': self.LITHOLOGY_TYPES[lithology]['color']
        }

    def calculate_homogeneity(self) -> Dict[str, float]:
        """
        Calculate lithological homogeneity of the system.
        
        Returns:
            Dictionary with homogeneity metrics
        """
        if len(self.samples) < 5:
            return {'homogeneity': 0.5, 'confidence': 'low'}
        
        purities = [s.caco3_purity for s in self.samples]
        
        # Coefficient of variation
        cv = np.std(purities) / np.mean(purities) if np.mean(purities) > 0 else 1
        
        # Homogeneity score (0-1, higher = more homogeneous)
        homogeneity = 1 / (1 + cv)
        
        # Dominant lithology
        lithologies = [s.lithology for s in self.samples]
        dominant = max(set(lithologies), key=lithologies.count)
        dominant_pct = lithologies.count(dominant) / len(lithologies) * 100
        
        return {
            'homogeneity': round(homogeneity, 2),
            'dominant_lithology': dominant,
            'dominant_percentage': round(dominant_pct, 1),
            'variability': round(cv, 2),
            'n_samples': len(self.samples)
        }

    def get_correction_factor(self, lithology: str, purity: float) -> float:
        """
        Get correction factor for model predictions based on lithology.
        
        Args:
            lithology: Rock type
            purity: CaCO₃ purity
            
        Returns:
            Correction factor for VLF model
        """
        # Base correction from paper
        base_correction = {
            'limestone': 1.0,
            'dolomite': 0.7,
            'gypsum': 1.5,
            'marble': 0.5,
            'chalk': 1.3
        }.get(lithology, 1.0)
        
        # Purity adjustment
        purity_adjustment = purity / 90  # Normalized to 90% reference
        
        return base_correction * purity_adjustment

    def summary(self) -> Dict[str, Union[float, str, Dict]]:
        """Get summary of lithology analysis."""
        if not self.samples:
            return {'status': 'No samples'}
        
        formations = {}
        for lith in self.formations:
            formations[lith] = {
                'samples': len(self.formations[lith]['samples']),
                'color': self.LITHOLOGY_TYPES[lith]['color']
            }
        
        homogeneity = self.calculate_homogeneity()
        
        return {
            'total_samples': len(self.samples),
            'formations': formations,
            'homogeneity': homogeneity['homogeneity'],
            'dominant_lithology': homogeneity['dominant_lithology'],
            'variability': homogeneity['variability']
        }


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Lithology Classification")
    print("="*60)
    
    lc = LithologyClassifier()
    
    # Add samples
    lc.add_sample('limestone', 95, 2.5, 0.08, 3.5, "Main Passage", 2024.5)
    lc.add_sample('dolomite', 75, 1.8, 0.05, 4.0, "Side Passage", 2024.5)
    
    # Classify rock
    classification = lc.classify_rock(caco3_purity=92, porosity=0.07, hardness=3.2)
    print(f"\nRock Classification:")
    print(f"  - Type: {classification['lithology']} {classification['color']}")
    print(f"  - Confidence: {classification['confidence']:.0%}")
    print(f"  - Description: {classification['description']}")
    
    # Calculate dissolution potential
    potential = lc.calculate_dissolution_potential('limestone', 95, 2.5)
    print(f"\nDissolution Potential: {potential:.2f}")
    
    # Calculate homogeneity
    homogeneity = lc.calculate_homogeneity()
    print(f"\nSystem Homogeneity:")
    print(f"  - Score: {homogeneity['homogeneity']:.2f}")
    print(f"  - Dominant: {homogeneity['dominant_lithology']}")
    print(f"  - {homogeneity['dominant_percentage']:.0f}% of samples")
    
    print("\n" + "="*60)
