import numpy as np
import pandas as pd
from typing import Tuple
from .geospatial import haversine_distance
from ..config import config
import logging

logger = logging.getLogger(__name__)

def calculate_distance_matrix(turbine_coords: np.ndarray, lightning_coords: np.ndarray) -> np.ndarray:
    """
    Calculate distance matrix between turbines and lightning strikes using vectorized operations.
    
    Args:
        turbine_coords: Array of turbine coordinates (lat, lng)
        lightning_coords: Array of lightning coordinates (lat, lng)
        
    Returns:
        Distance matrix with shape (n_turbines, n_lightning)
    """
    # Extract coordinates
    turbine_lats = turbine_coords[:, 0]
    turbine_lons = turbine_coords[:, 1]
    lightning_lats = lightning_coords[:, 0]
    lightning_lons = lightning_coords[:, 1]
    
    # Use broadcasting to calculate all distances at once
    # This is much more efficient than nested loops
    lat_diff = lightning_lats[:, np.newaxis] - turbine_lats[np.newaxis, :]
    lon_diff = lightning_lons[:, np.newaxis] - turbine_lons[np.newaxis, :]
    
    # Haversine formula components
    lat_diff_rad = np.radians(lat_diff)
    lon_diff_rad = np.radians(lon_diff)
    turbine_lats_rad = np.radians(turbine_lats)[np.newaxis, :]
    lightning_lats_rad = np.radians(lightning_lats)[:, np.newaxis]
    
    a = (np.sin(lat_diff_rad/2)**2 + 
         np.cos(turbine_lats_rad) * np.cos(lightning_lats_rad) * np.sin(lon_diff_rad/2)**2)
    c = 2 * np.arcsin(np.sqrt(a))
    
    # Convert to kilometers
    distances_km = 6371 * c
    
    return distances_km.T  # Transpose to get (n_turbines, n_lightning)

def calculate_turbine_risks_vectorized(turbine_df: pd.DataFrame, lightning_df: pd.DataFrame) -> pd.DataFrame:
    """
    Calculate turbine risks using vectorized operations for better performance.
    
    Args:
        turbine_df: DataFrame containing turbine data
        lightning_df: DataFrame containing lightning data
        
    Returns:
        DataFrame with added risk columns
    """
    logger.info("Starting vectorized risk calculation...")
    
    # Filter for cloud-to-ground lightning only
    cg_lightning_df = lightning_df[lightning_df['p_type'].astype(str) == '0'].copy()
    
    if len(cg_lightning_df) == 0:
        logger.warning("No cloud-to-ground lightning found")
        turbine_df = turbine_df.copy()
        turbine_df['risk_score'] = 0.0
        turbine_df['risk_log'] = 0.0
        return turbine_df
    
    # Extract coordinates as numpy arrays
    turbine_coords = turbine_df[['lat', 'lng']].values
    lightning_coords = cg_lightning_df[['lat', 'lng']].values
    
    # Calculate distance matrix
    logger.info(f"Calculating distance matrix for {len(turbine_df)} turbines and {len(cg_lightning_df)} lightning strikes...")
    distance_matrix = calculate_distance_matrix(turbine_coords, lightning_coords)
    
    # Get risk parameters
    P_0 = config.risk_params['P_0']
    alpha = config.risk_params['alpha']
    current_weight = config.risk_params['current_weight']
    
    # Calculate current magnitudes
    current_magnitudes = np.abs(cg_lightning_df['current'].values)
    
    # Calculate risk matrix using vectorized operations
    # Shape: (n_turbines, n_lightning)
    current_factor = 1 + current_weight * current_magnitudes / 10000
    distance_factor = np.exp(-alpha * distance_matrix)
    risk_matrix = P_0 * current_factor[np.newaxis, :] * distance_factor
    
    # Sum risks for each turbine
    turbine_risks = np.sum(risk_matrix, axis=1)
    
    # Add risk columns to DataFrame
    turbine_df = turbine_df.copy()
    turbine_df['risk_score'] = turbine_risks
    turbine_df['risk_log'] = np.log10(turbine_risks + 1)
    
    logger.info(f"Risk calculation completed. Risk range: {turbine_risks.min():.2f} - {turbine_risks.max():.2f}")
    
    return turbine_df

def calculate_turbine_risks(turbine_df: pd.DataFrame, lightning_df: pd.DataFrame) -> pd.DataFrame:
    """
    Calculate turbine risks (legacy function for backward compatibility).
    Now uses vectorized implementation for better performance.
    """
    try:
        from sklearn.neighbors import BallTree
        logger.info("Starting BallTree-based risk calculation...")
        cg_lightning_df = lightning_df[lightning_df['p_type'].astype(str) == '0'].copy()
        if len(cg_lightning_df) == 0:
            result_df = turbine_df.copy()
            result_df['risk_score'] = 0.0
            result_df['risk_log'] = 0.0
            return result_df

        import numpy as np
        earth_km = 6371.0
        max_radius_km = max(config.distance_rings) / 1000.0
        radius_rad = max_radius_km / earth_km

        lightning_coords_rad = np.radians(cg_lightning_df[['lat', 'lng']].values)
        tree = BallTree(lightning_coords_rad, metric='haversine')

        turbine_coords_rad = np.radians(turbine_df[['lat', 'lng']].values)
        ind_list, dist_list = tree.query_radius(turbine_coords_rad, r=radius_rad, return_distance=True, sort_results=True)

        P_0 = config.risk_params['P_0']
        alpha = config.risk_params['alpha']
        current_weight = config.risk_params['current_weight']
        currents = np.abs(cg_lightning_df['current'].values)

        risk_scores = np.zeros(len(turbine_df), dtype=float)
        for i in range(len(turbine_df)):
            idxs = ind_list[i]
            if idxs.size == 0:
                continue
            dists_km = dist_list[i] * earth_km
            current_mag = currents[idxs]
            current_factor = 1.0 + current_weight * current_mag / 10000.0
            distance_factor = np.exp(-alpha * dists_km)
            risk_scores[i] = np.sum(P_0 * current_factor * distance_factor)

        result_df = turbine_df.copy()
        result_df['risk_score'] = risk_scores
        result_df['risk_log'] = np.log10(risk_scores + 1.0)
        logger.info("BallTree-based risk calculation completed")
        return result_df
    except Exception as e:
        logger.warning(f"Falling back to vectorized matrix risk calculation due to: {e}")
        return calculate_turbine_risks_vectorized(turbine_df, lightning_df)
