import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from collections import defaultdict
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from src.analysis.geospatial import haversine_distance, haversine_distance_vectorized
from src.config import config
from src.utils import get_analysis_radius_m
import logging

logger = logging.getLogger(__name__)

def filter_lightning_by_distance(lightning_df, centroid_lat, centroid_lng):
    """
    Filter lightning data to only include events within the farthest distance ring.
    """
    max_distance = get_analysis_radius_m()
    if len(lightning_df) == 0:
        return lightning_df.copy()
    dists = haversine_distance_vectorized(
        centroid_lat,
        centroid_lng,
        lightning_df['lat'].values,
        lightning_df['lng'].values,
    )
    mask = dists <= max_distance
    return lightning_df.loc[mask].copy()

def find_activity_periods(df, min_gap_minutes=30, min_events_per_period=10):
    """
    Find concentrated activity periods based on time gaps.
    """
    df_sorted = df.copy()
    if df_sorted['local_time'].dtype == 'object':
        df_sorted['local_time'] = pd.to_datetime(df_sorted['local_time'])
    df_sorted = df_sorted.sort_values('local_time').reset_index(drop=True)

    df_sorted['time_diff'] = df_sorted['local_time'].diff()
    gap = timedelta(minutes=min_gap_minutes)
    break_positions = df_sorted.index[df_sorted['time_diff'] > gap].tolist()

    periods = []
    start_pos = 0
    for break_pos in break_positions:
        period_slice = df_sorted.iloc[start_pos:break_pos]
        if len(period_slice) >= min_events_per_period:
            periods.append({
                'start_time': period_slice['local_time'].iloc[0],
                'end_time': period_slice['local_time'].iloc[-1],
                'data': period_slice,
                'event_count': len(period_slice)
            })
        start_pos = break_pos

    last_slice = df_sorted.iloc[start_pos:]
    if len(last_slice) >= min_events_per_period:
        periods.append({
            'start_time': last_slice['local_time'].iloc[0],
            'end_time': last_slice['local_time'].iloc[-1],
            'data': last_slice,
            'event_count': len(last_slice)
        })

    return periods

def create_minute_counts(period_data):
    """
    Create minute-by-minute counts of lightning events.
    """
    # Ensure local_time is datetime
    period_data = period_data.copy()
    if period_data['local_time'].dtype == 'object':
        period_data['local_time'] = pd.to_datetime(period_data['local_time'])
    
    # Round timestamps to nearest minute
    period_data['minute_rounded'] = period_data['local_time'].dt.floor('min')
    
    # Count events per minute for each p_type
    minute_counts = period_data.groupby(['minute_rounded', 'p_type']).size().unstack(fill_value=0)
    
    # Create complete minute range
    start_minute = period_data['minute_rounded'].min()
    end_minute = period_data['minute_rounded'].max()
    minute_range = pd.date_range(start=start_minute, end=end_minute, freq='min')
    
    # Reindex to include all minutes (fill missing with 0)
    minute_counts = minute_counts.reindex(minute_range, fill_value=0)
    
    return minute_counts

def find_peak_sub_periods(period_data, window_minutes=3):
    """
    Find peak sub-periods within an activity period.
    """
    minute_counts = create_minute_counts(period_data)
    
    # Calculate total events per minute
    if '0' in minute_counts.columns and '1' in minute_counts.columns:
        total_per_minute = minute_counts['0'] + minute_counts['1']
    elif '0' in minute_counts.columns:
        total_per_minute = minute_counts['0']
    elif '1' in minute_counts.columns:
        total_per_minute = minute_counts['1']
    else:
        return []
    
    # Use rolling window to find peak periods
    rolling_mean = total_per_minute.rolling(window=window_minutes, center=True).mean()
    mean_activity = total_per_minute.mean()
    std_activity = total_per_minute.std()
    
    # Find periods above mean + 1 standard deviation
    threshold = mean_activity + std_activity
    peak_mask = rolling_mean > threshold
    
    # Group consecutive peak minutes
    peak_periods = []
    in_peak = False
    start_time = None
    
    for i, (time, is_peak) in enumerate(zip(total_per_minute.index, peak_mask)):
        if is_peak and not in_peak:
            start_time = time
            in_peak = True
        elif not is_peak and in_peak:
            end_time = total_per_minute.index[i-1]
            peak_periods.append({
                'start': start_time,
                'end': end_time,
                'peak_rate': rolling_mean[start_time:end_time].max()
            })
            in_peak = False
    
    # Handle case where peak period extends to the end
    if in_peak:
        peak_periods.append({
            'start': start_time,
            'end': total_per_minute.index[-1],
            'peak_rate': rolling_mean[start_time:].max()
        })
    
    return peak_periods

def _build_histogram_figure_for_periods(periods_chunk, max_distance_km):
    """
    Build a Plotly figure for a subset of activity periods.
    
    Args:
        periods_chunk: List of period dictionaries to plot
        max_distance_km: Maximum distance in km for the title
        
    Returns:
        Plotly figure object
    """
    n_periods = len(periods_chunk)
    
    if n_periods == 0:
        return None
    
    # Determine layout: single column if <= 3 periods, otherwise 2 columns
    if n_periods <= 3:
        n_cols = 1
    else:
        n_cols = 2
    
    n_rows = (n_periods + n_cols - 1) // n_cols
    # Cap n_rows at 3 (max 2×3 grid = 6 periods per figure)
    n_rows = min(n_rows, 3)
    
    # Adjust periods if we capped rows
    if n_periods > n_rows * n_cols:
        periods_chunk = periods_chunk[:n_rows * n_cols]
        n_periods = len(periods_chunk)
    
    # Calculate spacing
    base_row_height = 350
    if n_rows > 1:
        vertical_spacing = 0.28  # larger gap between rows
    else:
        vertical_spacing = 0.15
    horizontal_spacing = 0.1 if n_cols == 2 else 0.08
    
    def format_period_title(i, p):
        start = p['start_time']
        end = p['end_time']
        event_count = p['event_count']
        
        # Format date
        date_str = start.strftime('%d-%m-%Y')
        
        # Determine period type
        if start.date() == end.date():
            period_str = f"{start.strftime('%H:%M')}-{end.strftime('%H:%M')}"
        else:
            period_str = f"{start.strftime('%d-%m-%Y %H:%M')}-{end.strftime('%d-%m-%Y %H:%M')}"
        
        # Use line breaks to fit within histogram width
        return f"Date: {date_str}<br>Period: {period_str}<br>Total lightnings: {event_count}"
    
    fig = make_subplots(
        rows=n_rows, 
        cols=n_cols,
        subplot_titles=[format_period_title(i, p) for i, p in enumerate(periods_chunk)],
        vertical_spacing=vertical_spacing,
        horizontal_spacing=horizontal_spacing,
        specs=[[{"secondary_y": False} for _ in range(n_cols)] for _ in range(n_rows)]
    )
    
    colors = {'0': '#FF6B6B', '1': '#4ECDC4'}  # Red for cloud-to-ground, Teal for intracloud
    
    for period_idx, period in enumerate(periods_chunk):
        row = (period_idx // n_cols) + 1
        col = (period_idx % n_cols) + 1
        
        minute_counts = create_minute_counts(period['data'])
        
        # Create time labels (minutes from start)
        start_time = minute_counts.index[0]
        minutes_from_start = [(t - start_time).total_seconds() / 60 for t in minute_counts.index]
        
        # Add bars for each p_type
        for p_type in ['0', '1']:
            if p_type in minute_counts.columns:
                counts = minute_counts[p_type].values
                
                p_type_name = "Cloud-to-Ground" if p_type == '0' else "Intercloud"
                
                fig.add_trace(
                    go.Bar(
                        x=minutes_from_start,
                        y=counts,
                        name=p_type_name if period_idx == 0 else None,
                        marker_color=colors[p_type],
                        opacity=0.7,
                        showlegend=(period_idx == 0),
                        legendgroup=f"p_type_{p_type}",
                        hovertemplate=f"{p_type_name}<br>Minute: %{{x:.0f}}<br>Count: %{{y}}<extra></extra>"
                    ),
                    row=row, col=col
                )
        
        # Find and highlight peak sub-periods
        peak_periods = find_peak_sub_periods(period['data'])
        for peak in peak_periods:
            peak_start_minutes = (peak['start'] - start_time).total_seconds() / 60
            peak_end_minutes = (peak['end'] - start_time).total_seconds() / 60
            
            fig.add_vrect(
                x0=peak_start_minutes, x1=peak_end_minutes,
                fillcolor="yellow", opacity=0.2,
                line_width=0,
                row=row, col=col
            )
        
        # Update axes for this subplot
        fig.update_xaxes(
            title_text="Minutes from start",
            title_standoff=6,
            row=row,
            col=col,
            tickfont=dict(size=18),
            title_font=dict(size=22),
        )
        # Y-axis labels only on leftmost column
        if col == 1:
            fig.update_yaxes(
                title_text="Lightning Count",
                row=row,
                col=col,
                tickfont=dict(size=18),
                title_font=dict(size=22),
            )
        else:
            fig.update_yaxes(
                title_text="",
                row=row,
                col=col,
                tickfont=dict(size=16),
                title_font=dict(size=22),
            )
    
    # Calculate figure height
    figure_height = base_row_height * n_rows
    
    # Update subplot title font size (subplot titles are rendered as annotations)
    fig.update_annotations(font=dict(size=22))
    
    fig.update_layout(
        height=figure_height,
        width=config.histogram_layout['plot_width'],
        barmode='stack',
        showlegend=True,
        font=dict(size=18),
        margin=dict(t=110, b=100, l=75, r=40),
        legend=dict(
            orientation="h",
            y=-0.1,
            x=0.5,
            xanchor="center",
            font=dict(size=18),
            title_font=dict(size=20),
        )
    )
    
    # Store metadata for render function
    fig.layout.meta = {'n_rows': n_rows, 'n_cols': n_cols}
    
    return fig

def create_lightning_histogram_pages(lightning_df, centroid_lat, centroid_lng):
    """
    Create multiple histogram figures for pagination when there are many activity periods.
    
    Args:
        lightning_df: DataFrame containing lightning data
        centroid_lat: Center latitude
        centroid_lng: Center longitude
        
    Returns:
        List of Plotly figure objects (empty list if no periods found)
    """
    # Filter lightning data by distance
    filtered_df = filter_lightning_by_distance(lightning_df, centroid_lat, centroid_lng)
    
    if len(filtered_df) == 0:
        return []
    
    # Find ALL activity periods (no truncation)
    periods = find_activity_periods(
        filtered_df, 
        min_gap_minutes=config.histogram_params['min_gap_minutes'], 
        min_events_per_period=config.histogram_params['min_events_per_period']
    )
    
    if len(periods) == 0:
        return []
    
    # Get max periods per figure from config
    max_periods_per_figure = config.histogram_params.get('max_periods_per_figure', 6)
    
    # Get max distance for title
    max_distance_km = max(config.distance_rings) / 1000
    
    # Split periods into chunks
    figures = []
    for i in range(0, len(periods), max_periods_per_figure):
        chunk = periods[i:i + max_periods_per_figure]
        fig = _build_histogram_figure_for_periods(chunk, max_distance_km)
        if fig:
            figures.append(fig)
    
    return figures

def create_lightning_histogram(lightning_df, centroid_lat, centroid_lng):
    """
    Create lightning activity histogram for the PDF report.
    
    This is a backward-compatible wrapper that returns the first page.
    For multi-page support, use create_lightning_histogram_pages instead.
    """
    pages = create_lightning_histogram_pages(lightning_df, centroid_lat, centroid_lng)
    return pages[0] if pages else None 