"""
CDZ Mapping - Spatial distribution and mapping of Creative Dissolution Zones.
Creates maps of CDZ locations and analyzes spatial patterns.

Author: Samir Baladi
"""

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

logger = logging.getLogger(__name__)

@dataclass
class CDZLocation:
    """Represents a mapped CDZ location."""
    x: float  # X coordinate (m)
    y: float  # Y coordinate (m)
    z: float  # Z coordinate (m)
    zone_type: str  # Type of CDZ
    size: float  # Approximate size (m)
    intensity: float  # CDZ intensity (0-1)
    passage_id: str  # Associated passage


class CDZMapper:
    """
    Maps spatial distribution of Creative Dissolution Zones.
    """
    
    def __init__(self):
        """Initialize CDZ mapper."""
        self.locations = []
        self.heatmap = None
        self.grid_size = 10  # meters
        logger.info("Initialized CDZMapper")

    def add_cdz_location(self,
                        x: float,
                        y: float,
                        z: float,
                        zone_type: str,
                        size: float,
                        intensity: float,
                        passage_id: str) -> None:
        """
        Add a CDZ location.
        
        Args:
            x, y, z: Coordinates
            zone_type: Type of CDZ
            size: Approximate size
            intensity: CDZ intensity
            passage_id: Associated passage
        """
        location = CDZLocation(
            x=x, y=y, z=z,
            zone_type=zone_type,
            size=size,
            intensity=intensity,
            passage_id=passage_id
        )
        
        self.locations.append(location)
        logger.debug(f"Added CDZ at ({x}, {y}, {z})")

    def create_heatmap(self,
                      x_range: Tuple[float, float],
                      y_range: Tuple[float, float],
                      resolution: int = 50) -> Dict[str, np.ndarray]:
        """
        Create a 2D heatmap of CDZ intensity.
        
        Args:
            x_range: (min_x, max_x)
            y_range: (min_y, max_y)
            resolution: Grid resolution
            
        Returns:
            Dictionary with heatmap data
        """
        if not self.locations:
            return {'status': 'No CDZ locations'}
        
        # Create grid
        x = np.linspace(x_range[0], x_range[1], resolution)
        y = np.linspace(y_range[0], y_range[1], resolution)
        X, Y = np.meshgrid(x, y)
        
        # Initialize heatmap
        heatmap = np.zeros_like(X)
        
        # Gaussian kernel for each CDZ
        for loc in self.locations:
            # Distance from grid points to CDZ
            dist = np.sqrt((X - loc.x)**2 + (Y - loc.y)**2)
            
            # Gaussian influence
            sigma = loc.size * 2
            influence = loc.intensity * np.exp(-(dist**2) / (2 * sigma**2))
            
            heatmap += influence
        
        # Normalize
        if np.max(heatmap) > 0:
            heatmap = heatmap / np.max(heatmap)
        
        self.heatmap = {
            'x': x,
            'y': y,
            'X': X,
            'Y': Y,
            'heatmap': heatmap
        }
        
        return self.heatmap

    def find_nearest_cdz(self,
                        x: float,
                        y: float,
                        z: float,
                        n_nearest: int = 3) -> List[Dict[str, float]]:
        """
        Find nearest CDZ locations to a point.
        
        Args:
            x, y, z: Query point coordinates
            n_nearest: Number of nearest CDZs to return
            
        Returns:
            List of nearest CDZ information
        """
        if not self.locations:
            return []
        
        # Build KD-tree
        points = np.array([[loc.x, loc.y, loc.z] for loc in self.locations])
        tree = KDTree(points)
        
        # Query nearest
        distances, indices = tree.query([x, y, z], k=min(n_nearest, len(points)))
        
        # Handle single vs multiple returns
        if n_nearest == 1:
            distances = [distances]
            indices = [indices]
        
        results = []
        for dist, idx in zip(distances, indices):
            loc = self.locations[idx]
            results.append({
                'distance': round(dist, 1),
                'x': loc.x,
                'y': loc.y,
                'z': loc.z,
                'zone_type': loc.zone_type,
                'intensity': loc.intensity,
                'passage_id': loc.passage_id
            })
        
        return results

    def analyze_spatial_pattern(self) -> Dict[str, Union[float, str]]:
        """
        Analyze spatial pattern of CDZs.
        
        Returns:
            Dictionary with spatial analysis
        """
        if len(self.locations) < 3:
            return {'status': 'Insufficient CDZs for spatial analysis'}
        
        # Extract coordinates
        xs = np.array([loc.x for loc in self.locations])
        ys = np.array([loc.y for loc in self.locations])
        zs = np.array([loc.z for loc in self.locations])
        
        # Calculate center of mass
        center_x = np.mean(xs)
        center_y = np.mean(ys)
        center_z = np.mean(zs)
        
        # Calculate spatial spread
        spread_x = np.std(xs)
        spread_y = np.std(ys)
        spread_z = np.std(zs)
        
        # Test for clustering (nearest neighbor analysis)
        if len(xs) > 1:
            points = np.column_stack([xs, ys, zs])
            tree = KDTree(points)
            
            # Mean nearest neighbor distance
            nn_distances = []
            for i in range(len(points)):
                dists, _ = tree.query(points[i], k=2)
                nn_distances.append(dists[1])  # Distance to nearest neighbor (excluding self)
            
            mean_nn = np.mean(nn_distances)
            
            # Expected mean NN distance for random distribution
            area = (np.max(xs) - np.min(xs)) * (np.max(ys) - np.min(ys))
            expected_nn = 0.5 / np.sqrt(len(xs) / area) if area > 0 else 0
            
            if expected_nn > 0:
                nn_ratio = mean_nn / expected_nn
                if nn_ratio < 0.8:
                    pattern = "Clustered"
                elif nn_ratio < 1.2:
                    pattern = "Random"
                else:
                    pattern = "Dispersed"
            else:
                pattern = "Unknown"
                nn_ratio = 0
        else:
            pattern = "Insufficient data"
            nn_ratio = 0
        
        # Group by zone type
        type_counts = {}
        for loc in self.locations:
            type_counts[loc.zone_type] = type_counts.get(loc.zone_type, 0) + 1
        
        return {
            'n_cdz': len(self.locations),
            'center': (round(center_x, 1), round(center_y, 1), round(center_z, 1)),
            'spread': (round(spread_x, 1), round(spread_y, 1), round(spread_z, 1)),
            'spatial_pattern': pattern,
            'nn_ratio': round(nn_ratio, 2),
            'type_distribution': type_counts,
            'depth_range': (round(np.min(zs), 1), round(np.max(zs), 1))
        }

    def identify_cdz_clusters(self,
                             max_cluster_distance: float = 50) -> List[Dict[str, Union[str, float, List]]]:
        """
        Identify clusters of CDZs.
        
        Args:
            max_cluster_distance: Maximum distance for cluster membership (m)
            
        Returns:
            List of clusters
        """
        if len(self.locations) < 2:
            return []
        
        # Simple clustering using distance threshold
        points = np.array([[loc.x, loc.y, loc.z] for loc in self.locations])
        
        # Assign clusters
        clusters = []
        assigned = set()
        
        for i in range(len(points)):
            if i in assigned:
                continue
            
            # Start new cluster
            cluster = [i]
            assigned.add(i)
            
            # Find all points within threshold
            for j in range(len(points)):
                if j in assigned:
                    continue
                
                # Check distance to any point in cluster
                for k in cluster:
                    dist = np.linalg.norm(points[j] - points[k])
                    if dist <= max_cluster_distance:
                        cluster.append(j)
                        assigned.add(j)
                        break
            
            if len(cluster) >= 2:  # Only keep clusters with multiple CDZs
                cluster_locs = [self.locations[idx] for idx in cluster]
                
                # Calculate cluster properties
                xs = [loc.x for loc in cluster_locs]
                ys = [loc.y for loc in cluster_locs]
                zs = [loc.z for loc in cluster_locs]
                
                clusters.append({
                    'size': len(cluster),
                    'center': (round(np.mean(xs), 1), round(np.mean(ys), 1), round(np.mean(zs), 1)),
                    'radius': round(max(np.std(xs), np.std(ys), np.std(zs)), 1),
                    'types': list(set(loc.zone_type for loc in cluster_locs)),
                    'mean_intensity': round(np.mean([loc.intensity for loc in cluster_locs]), 2),
                    'passages': list(set(loc.passage_id for loc in cluster_locs))
                })
        
        return clusters

    def calculate_cdz_density(self,
                             area: float,
                             volume: Optional[float] = None) -> Dict[str, float]:
        """
        Calculate CDZ density.
        
        Args:
            area: Survey area (m²)
            volume: Optional survey volume (m³)
            
        Returns:
            Dictionary with density metrics
        """
        if area <= 0:
            return {'error': 'Invalid area'}
        
        area_density = len(self.locations) / area * 10000  # per hectare
        
        result = {
            'n_cdz': len(self.locations),
            'area_density_per_ha': round(area_density, 2),
            'area': area
        }
        
        if volume and volume > 0:
            volume_density = len(self.locations) / volume * 1e6  # per million m³
            result['volume_density_per_Mm3'] = round(volume_density, 2)
            result['volume'] = volume
        
        return result

    def export_to_geojson(self, filename: str) -> None:
        """
        Export CDZ locations to GeoJSON format.
        
        Args:
            filename: Output filename
        """
        import json
        
        features = []
        for loc in self.locations:
            feature = {
                "type": "Feature",
                "geometry": {
                    "type": "Point",
                    "coordinates": [loc.x, loc.y, loc.z]
                },
                "properties": {
                    "zone_type": loc.zone_type,
                    "size": loc.size,
                    "intensity": loc.intensity,
                    "passage_id": loc.passage_id
                }
            }
            features.append(feature)
        
        geojson = {
            "type": "FeatureCollection",
            "features": features
        }
        
        with open(filename, 'w') as f:
            json.dump(geojson, f, indent=2)
        
        logger.info(f"Exported {len(features)} CDZs to {filename}")

    def summary(self) -> Dict[str, Union[float, str, Dict]]:
        """Get summary of CDZ mapping."""
        if not self.locations:
            return {'status': 'No CDZ locations'}
        
        spatial = self.analyze_spatial_pattern()
        clusters = self.identify_cdz_clusters()
        
        return {
            'n_cdz': len(self.locations),
            'spatial_pattern': spatial,
            'n_clusters': len(clusters),
            'clusters': clusters[:3] if clusters else [],  # Top 3 clusters
            'type_distribution': spatial.get('type_distribution', {}),
            'depth_range': spatial.get('depth_range', (0, 0))
        }


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - CDZ Mapping")
    print("="*60)
    
    mapper = CDZMapper()
    
    # Add some CDZ locations
    mapper.add_cdz_location(100, 200, -50, "Phreatic-Vadose", 25, 0.8, "P-7")
    mapper.add_cdz_location(150, 180, -45, "Phreatic-Vadose", 30, 0.9, "P-7")
    mapper.add_cdz_location(300, 400, -120, "Breakdown", 50, 0.6, "B-3")
    
    # Find nearest CDZs
    nearest = mapper.find_nearest_cdz(120, 190, -48)
    print(f"\nNearest CDZs:")
    for n in nearest:
        print(f"  - {n['zone_type']} at distance {n['distance']}m")
    
    # Analyze spatial pattern
    pattern = mapper.analyze_spatial_pattern()
    print(f"\nSpatial Pattern:")
    print(f"  - Pattern: {pattern.get('spatial_pattern', 'Unknown')}")
    print(f"  - Center: {pattern.get('center', (0,0,0))}")
    
    # Identify clusters
    clusters = mapper.identify_cdz_clusters()
    print(f"\nClusters: {len(clusters)} found")
    for i, cluster in enumerate(clusters):
        print(f"  Cluster {i+1}: {cluster['size']} CDZs at {cluster['center']}")
    
    print("\n" + "="*60)
