import pandas as pd
import numpy as np
from datetime import datetime
from src.analysis.geospatial import haversine_distance, haversine_distance_vectorized
from src.config import config
from src.utils import get_analysis_radius_m, parse_period_string_to_datetime
import logging

logger = logging.getLogger(__name__)

def calculate_lightning_statistics(lightning_df, centroid_lat, centroid_lng, start_date=None, end_date=None):
    """
    Calculate lightning statistics for the report with detailed breakdown by distance rings.
    
    Args:
        lightning_df: DataFrame containing lightning data
        centroid_lat: Center latitude
        centroid_lng: Center longitude
        start_date: Start date in 'DD-MM-YYYY' format (optional)
        end_date: End date in 'DD-MM-YYYY' format (optional)
    """
    max_distance = get_analysis_radius_m()
    
    # Filter lightning data by distance (vectorized)
    if len(lightning_df) == 0:
        return {
            'intercloud_by_day': {},
            'cloud_to_ground_by_day': {},
            'total_lightning_per_km2': 0,
            'daily_lightning_per_km2': 0,
            'period_days': 0.0,
            'total_events': 0,
            'area_km2': 0,
            'max_distance_km': max_distance / 1000,
            'lightning_by_distance_rings': {},
            'daily_lightning_by_rings': {}
        }

    dists_all = haversine_distance_vectorized(
        centroid_lat,
        centroid_lng,
        lightning_df['lat'].values,
        lightning_df['lng'].values,
    )
    mask = dists_all <= max_distance
    filtered_df = lightning_df.loc[mask].copy()
    
    if len(filtered_df) == 0:
        return {
            'intercloud_by_day': {},
            'cloud_to_ground_by_day': {},
            'total_lightning_per_km2': 0,
            'daily_lightning_per_km2': 0,
            'period_days': 0.0,
            'total_events': 0,
            'area_km2': 0,
            'max_distance_km': max_distance / 1000,
            'lightning_by_distance_rings': {},
            'daily_lightning_by_rings': {}
        }
    
    # Convert local_time to datetime if needed
    if filtered_df['local_time'].dtype == 'object':
        filtered_df['local_time'] = pd.to_datetime(filtered_df['local_time'])
    
    # Add distance column
    filtered_df['distance_km'] = (dists_all[mask] / 1000)
    
    # Add date column
    filtered_df['date'] = filtered_df['local_time'].dt.date
    
    # 1. Inter-cloud lightnings by day (outermost ring only)
    intercloud_df = filtered_df[filtered_df['p_type'].astype(str) == '1'].copy()
    intercloud_by_day = {}
    if len(intercloud_df) > 0:
        intercloud_by_day = intercloud_df['date'].value_counts().to_dict()
        # Convert date objects to strings for JSON serialization
        intercloud_by_day = {date.strftime("%d-%m-%Y"): count for date, count in intercloud_by_day.items()}
    
    # 2. Cloud-to-ground lightnings by day (outermost ring only)
    cloud_to_ground_df = filtered_df[filtered_df['p_type'].astype(str) == '0'].copy()
    cloud_to_ground_by_day = {}
    if len(cloud_to_ground_df) > 0:
        cloud_to_ground_by_day = cloud_to_ground_df['date'].value_counts().to_dict()
        # Convert date objects to strings for JSON serialization
        cloud_to_ground_by_day = {date.strftime("%d-%m-%Y"): count for date, count in cloud_to_ground_by_day.items()}
    
    # 3. Lightning counts by distance rings (ranges instead of cumulative)
    lightning_by_distance_rings = {}
    for i, ring_distance in enumerate(config.distance_rings):
        ring_km = ring_distance / 1000
        
        # Define distance range for this ring
        if i == 0:
            # First ring: 0 to ring_km
            min_distance = 0
            ring_max_distance = ring_km
            ring_name = f"0-{ring_km:.1f}km"
        else:
            # Other rings: previous ring to current ring
            prev_ring_km = config.distance_rings[i-1] / 1000
            min_distance = prev_ring_km
            ring_max_distance = ring_km
            ring_name = f"{prev_ring_km:.1f}-{ring_km:.1f}km"
        
        # Count lightning within this distance range
        ring_lightning = filtered_df[
            (filtered_df['distance_km'] > min_distance) & 
            (filtered_df['distance_km'] <= ring_max_distance)
        ]
        
        # Count by type
        intercloud_count = len(ring_lightning[ring_lightning['p_type'].astype(str) == '1'])
        cloud_to_ground_count = len(ring_lightning[ring_lightning['p_type'].astype(str) == '0'])
        total_count = len(ring_lightning)
        
        lightning_by_distance_rings[ring_name] = {
            'intercloud': intercloud_count,
            'cloud_to_ground': cloud_to_ground_count,
            'total': total_count
        }
    
    # 4. Daily lightning breakdown by distance rings
    daily_lightning_by_rings = {}
    
    # Get unique dates
    unique_dates = sorted(filtered_df['date'].unique())
    
    for date in unique_dates:
        date_str = date.strftime("%d-%m-%Y")
        daily_lightning_by_rings[date_str] = {}
        
        # Get lightning for this date
        daily_lightning = filtered_df[filtered_df['date'] == date]
        
        for i, ring_distance in enumerate(config.distance_rings):
            ring_km = ring_distance / 1000
            
            # Define distance range for this ring
            if i == 0:
                # First ring: 0 to ring_km
                min_distance = 0
                ring_max_distance = ring_km
                ring_name = f"0-{ring_km:.1f}km"
            else:
                # Other rings: previous ring to current ring
                prev_ring_km = config.distance_rings[i-1] / 1000
                min_distance = prev_ring_km
                ring_max_distance = ring_km
                ring_name = f"{prev_ring_km:.1f}-{ring_km:.1f}km"
            
            # Count lightning within this distance range for this date
            ring_lightning = daily_lightning[
                (daily_lightning['distance_km'] > min_distance) & 
                (daily_lightning['distance_km'] <= ring_max_distance)
            ]
            
            # Count by type
            intercloud_count = len(ring_lightning[ring_lightning['p_type'].astype(str) == '1'])
            cloud_to_ground_count = len(ring_lightning[ring_lightning['p_type'].astype(str) == '0'])
            total_count = len(ring_lightning)
            
            daily_lightning_by_rings[date_str][ring_name] = {
                'intercloud': intercloud_count,
                'cloud_to_ground': cloud_to_ground_count,
                'total': total_count
            }
    
    # 5. Calculate area and lightning density (outermost ring)
    area_km2 = np.pi * (max_distance / 1000) ** 2
    total_events = len(filtered_df)
    total_lightning_per_km2 = total_events / area_km2 if area_km2 > 0 else 0
    
    if start_date and end_date:
        start_dt = parse_period_string_to_datetime(start_date)
        end_dt = parse_period_string_to_datetime(end_date)
        if start_dt is not None and end_dt is not None:
            delta = end_dt - start_dt
            period_seconds = delta.total_seconds()
            period_days = period_seconds / 86400.0 if period_seconds > 0 else 1.0
            logger.info(f"Analysis period duration: {period_days:.4f} days ({start_date} to {end_date})")
        else:
            period_days = 1.0
            if len(filtered_df) > 0:
                first_date = filtered_df['local_time'].min()
                period_days = float(pd.Timestamp(first_date.year, first_date.month, 1).days_in_month)
            logger.warning(f"Failed to parse date range ({start_date}, {end_date}). Using period_days={period_days}")
    elif len(filtered_df) > 0:
        first_date = filtered_df['local_time'].min()
        period_days = float(pd.Timestamp(first_date.year, first_date.month, 1).days_in_month)
        logger.info(f"Using month as period: {period_days} days")
    else:
        period_days = 1.0

    daily_lightning_per_km2 = total_lightning_per_km2 / period_days if period_days > 0 else 0

    return {
        'intercloud_by_day': intercloud_by_day,
        'cloud_to_ground_by_day': cloud_to_ground_by_day,
        'total_lightning_per_km2': total_lightning_per_km2,
        'daily_lightning_per_km2': daily_lightning_per_km2,
        'period_days': period_days,
        'total_events': total_events,
        'area_km2': area_km2,
        'max_distance_km': max_distance / 1000,
        'lightning_by_distance_rings': lightning_by_distance_rings,
        'daily_lightning_by_rings': daily_lightning_by_rings
    } 