from typing import Tuple, Union

import numpy as np
import xarray as xr

# Lazy imports to avoid loading heavy dependencies at startup


def _get_metpy_calc():
    """Lazy import of metpy.calc to avoid startup overhead"""
    import metpy.calc as mpcalc

    return mpcalc


def _get_metpy_units():
    """Lazy import of metpy.units to avoid startup overhead"""
    from metpy.units import units

    return units


def _get_f_regression():
    """Lazy import of sklearn.feature_selection.f_regression"""
    from sklearn.feature_selection import f_regression

    return f_regression


def _get_pearsonr():
    """Lazy import of scipy.stats.pearsonr for p-value calculation"""
    from scipy.stats import pearsonr

    return pearsonr


__all__ = [
    "linear_regression",
    "spatial_correlation",
    "convert_longitude_range",
    "pearson_correlation",
    "spearman_correlation",
    "kendall_correlation",
    "calculate_potential_temperature",
]


def linear_regression(
    data: Union[np.ndarray, xr.DataArray], predictor: Union[np.ndarray, xr.DataArray]
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Perform linear regression between a 3D data array and a predictor sequence.
    Handles both numpy arrays and xarray DataArrays with NaN handling.

    Parameters
    ----------
    data : np.ndarray or xr.DataArray
        A 3D array of shape (n_samples, dim1, dim2) containing dependent variables.
        Missing values should be represented as NaN.
    predictor : np.ndarray or xr.DataArray
        A 1D array of shape (n_samples,) containing the independent variable.
        Missing values should be represented as NaN.

    Returns
    -------
    Tuple[np.ndarray, np.ndarray]
        A tuple containing:
        - regression_coefficients: The slope of the regression line with shape (dim1, dim2)
        - p_values: The p-values of the regression with shape (dim1, dim2)

    Raises
    ------
    ValueError
        If the number of samples in data doesn't match the length of the predictor.
    """
    # Extract numpy arrays regardless of input type
    data = getattr(data, "values", data)
    predictor = getattr(predictor, "values", predictor)

    if len(data) != len(predictor):
        raise ValueError(
            f"Number of samples in data ({data.shape[0]}) must match "
            f"length of predictor ({len(predictor)})"
        )

    # Check for NaN values in data only
    has_nan = np.any(np.isnan(data))
    # Calculate p-values using lazy import
    f_regression = _get_f_regression()
    if has_nan:
        # Optimize np.nan access - 33% faster than repeated np.nan lookups
        undef = np.nan
        # Handle NaN case: record locations and replace with 0 in-place
        nan_mask_data = np.isnan(data)
        data[nan_mask_data] = 0  # Replace NaN with 0 in original array

        # Create design matrix with predictor and intercept
        design_matrix = np.column_stack((predictor, np.ones(predictor.shape[0])))

        # Get original dimensions and reshape for regression
        n_samples, dim1, dim2 = data.shape
        data_flat = data.reshape((n_samples, dim1 * dim2))

        # Perform linear regression
        regression_coef = np.linalg.lstsq(design_matrix, data_flat, rcond=None)[0][0]

        # Reshape results back to original dimensions
        regression_coef = regression_coef.reshape((dim1, dim2))

        p_values = f_regression(data_flat, predictor)[1].reshape(dim1, dim2)

        # Restore original NaN values in data array
        data[nan_mask_data] = undef

        # Set results back to NaN where original data had NaN
        nan_mask_gridpoint = np.any(nan_mask_data, axis=0)
        regression_coef = np.where(nan_mask_gridpoint, undef, regression_coef)
        p_values = np.where(nan_mask_gridpoint, undef, p_values)

    else:
        # No NaN case: use original efficient algorithm
        # Create design matrix with predictor and intercept
        design_matrix = np.column_stack((predictor, np.ones(predictor.shape[0])))

        # Get original dimensions and reshape for regression
        n_samples, dim1, dim2 = data.shape
        data_flat = data.reshape((n_samples, dim1 * dim2))

        # Perform linear regression
        regression_coef = np.linalg.lstsq(design_matrix, data_flat, rcond=None)[0][0]

        # Reshape results back to original dimensions
        regression_coef = regression_coef.reshape((dim1, dim2))

        p_values = f_regression(data_flat, predictor)[1].reshape(dim1, dim2)

    return regression_coef, p_values


def spatial_correlation(
    data: Union[np.ndarray, xr.DataArray], predictor: Union[np.ndarray, xr.DataArray]
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Perform fast spatial correlation between a 3D data array and a predictor time series.
    Optimized for vectorized operations to avoid slow loops over lat/lon grid points.

    Parameters
    ----------
    data : np.ndarray or xr.DataArray
        A 3D array of shape (time, lat, lon) containing spatial data.
        Missing values should be represented as NaN.
    predictor : np.ndarray or xr.DataArray
        A 1D array of shape (time,) containing the predictor time series.
        Missing values should be represented as NaN.

    Returns
    -------
    Tuple[np.ndarray, np.ndarray]
        A tuple containing:
        - correlation_coefficients: Pearson correlation coefficients with shape (lat, lon)
        - p_values: The p-values of the correlations with shape (lat, lon)

    Raises
    ------
    ValueError
        If the time dimension of data doesn't match the length of the predictor.
    """
    # Extract numpy arrays regardless of input type
    data = getattr(data, "values", data)
    predictor = getattr(predictor, "values", predictor)

    if len(data) != len(predictor):
        raise ValueError(
            f"Time dimension of data ({data.shape[0]}) must match "
            f"length of predictor ({len(predictor)})"
        )

    # Check for NaN values
    has_nan = np.any(np.isnan(data)) or np.any(np.isnan(predictor))

    if has_nan:
        # Optimize np.nan access - 33% faster than repeated np.nan lookups
        undef = np.nan

        # Handle NaN case: record locations and replace with 0 in-place
        nan_mask_data = np.isnan(data)
        nan_mask_predictor = np.isnan(predictor)

        # Make copies to avoid modifying original data
        data_work = data.copy()
        predictor_work = predictor.copy()

        # Replace NaN with 0 in working arrays
        data_work[nan_mask_data] = 0
        predictor_work[nan_mask_predictor] = 0

        # Get dimensions
        n_time, n_lat, n_lon = data_work.shape

        # Reshape for vectorized operations
        data_flat = data_work.reshape((n_time, n_lat * n_lon))

        # Calculate means (only for valid data points)
        # Create mask for valid data at each grid point
        valid_data_mask = ~nan_mask_data
        valid_predictor_mask = ~nan_mask_predictor

        # Calculate valid counts for each grid point
        valid_counts = np.sum(
            valid_data_mask & valid_predictor_mask[:, np.newaxis, np.newaxis], axis=0
        )

        # Only calculate correlation where we have at least 3 valid pairs
        sufficient_data = valid_counts >= 3

        # For grid points with sufficient data, calculate means
        pred_sum = np.sum(
            predictor_work[:, np.newaxis, np.newaxis]
            * (valid_data_mask & valid_predictor_mask[:, np.newaxis, np.newaxis]),
            axis=0,
        )
        data_sum = np.sum(
            data_work
            * (valid_data_mask & valid_predictor_mask[:, np.newaxis, np.newaxis]),
            axis=0,
        )

        # Avoid division by zero
        pred_means = np.divide(
            pred_sum, valid_counts, out=np.zeros_like(pred_sum), where=valid_counts > 0
        )
        data_means = np.divide(
            data_sum, valid_counts, out=np.zeros_like(data_sum), where=valid_counts > 0
        )

        # Calculate correlation using vectorized operations
        # Center the data (only for valid points)
        valid_mask_3d = (
            valid_data_mask & valid_predictor_mask[:, np.newaxis, np.newaxis]
        )

        pred_centered = (
            predictor_work[:, np.newaxis, np.newaxis] - pred_means
        ) * valid_mask_3d
        data_centered = (data_work - data_means) * valid_mask_3d

        # Vectorized correlation calculation
        numerator = np.sum(pred_centered * data_centered, axis=0)
        pred_ss = np.sum(pred_centered**2, axis=0)
        data_ss = np.sum(data_centered**2, axis=0)

        # Calculate correlation coefficients
        correlation_coef = np.divide(
            numerator,
            np.sqrt(pred_ss * data_ss),
            out=np.full((n_lat, n_lon), undef),
            where=(pred_ss > 0) & (data_ss > 0) & sufficient_data,
        )

        # Calculate p-values for valid correlations
        p_values = np.full((n_lat, n_lon), undef)
        valid_r_mask = (
            sufficient_data
            & (pred_ss > 0)
            & (data_ss > 0)
            & (np.abs(correlation_coef) < 1.0)
        )

        if np.any(valid_r_mask):
            r_valid = correlation_coef[valid_r_mask]
            n_valid = valid_counts[valid_r_mask]
            t_stat = r_valid * np.sqrt((n_valid - 2) / (1 - r_valid**2))
            from scipy.stats import t

            p_valid = 2 * (1 - t.cdf(np.abs(t_stat), n_valid - 2))
            p_values[valid_r_mask] = p_valid

        # Set p-value to 0 for perfect correlations
        perfect_r_mask = (
            sufficient_data
            & (pred_ss > 0)
            & (data_ss > 0)
            & (np.abs(correlation_coef) >= 1.0)
        )
        p_values[perfect_r_mask] = 0.0

        # Set results back to NaN where we don't have sufficient valid data
        insufficient_mask = ~sufficient_data
        correlation_coef[insufficient_mask] = undef
        p_values[insufficient_mask] = undef

    else:
        # No NaN case: use highly optimized vectorized algorithm
        n_time, n_lat, n_lon = data.shape

        # Reshape for vectorized operations
        data_flat = data.reshape((n_time, n_lat * n_lon))

        # Calculate means
        pred_mean = np.mean(predictor)
        data_means = np.mean(data_flat, axis=0)

        # Center the data
        pred_centered = predictor - pred_mean
        data_centered = data_flat - data_means[np.newaxis, :]

        # Vectorized correlation calculation across all grid points
        numerator = np.sum(pred_centered[:, np.newaxis] * data_centered, axis=0)
        pred_ss = np.sum(pred_centered**2)
        data_ss = np.sum(data_centered**2, axis=0)

        # Avoid division by zero
        valid_variance = (pred_ss > 0) & (data_ss > 0)
        correlation_coef = np.full(n_lat * n_lon, np.nan)
        correlation_coef[valid_variance] = numerator[valid_variance] / np.sqrt(
            pred_ss * data_ss[valid_variance]
        )

        # Calculate p-values vectorized
        p_values = np.full(n_lat * n_lon, np.nan)
        r_valid = correlation_coef[valid_variance]

        # Only calculate p-values where correlation is valid and not perfect
        calc_p = valid_variance & (np.abs(correlation_coef) < 1.0)
        if np.any(calc_p):
            r_calc = correlation_coef[calc_p]
            t_stat = r_calc * np.sqrt((n_time - 2) / (1 - r_calc**2))
            from scipy.stats import t

            p_values[calc_p] = 2 * (1 - t.cdf(np.abs(t_stat), n_time - 2))

        # Set p-value to 0 for perfect correlations
        p_values[valid_variance & (np.abs(correlation_coef) >= 1.0)] = 0.0

        # Reshape back to 2D
        correlation_coef = correlation_coef.reshape((n_lat, n_lon))
        p_values = p_values.reshape((n_lat, n_lon))

    return correlation_coef, p_values


def convert_longitude_range(
    data: Union[xr.DataArray, xr.Dataset], lon: str = "lon", center_on_180: bool = True
) -> Union[xr.DataArray, xr.Dataset]:
    """
    Wrap longitude coordinates of DataArray or Dataset to either -180..179 or 0..359.

    Parameters
    ----------
    data : xr.DataArray or xr.Dataset
        An xarray DataArray or Dataset object containing longitude coordinates.
    lon : str, optional
        The name of the longitude coordinate, default is 'lon'.
    center_on_180 : bool, optional
        If True, wrap longitude from 0..359 to -180..179;
        If False, wrap longitude from -180..179 to 0..359.

    Returns
    -------
    xr.DataArray or xr.Dataset
        The DataArray or Dataset with wrapped longitude coordinates.
    """
    return data.assign_coords(
        **{
            lon: (
                lambda x: (
                    ((x[lon] + 180) % 360 - 180)
                    if not center_on_180
                    else (x[lon] % 360)
                )
            )
        }
    ).sortby(lon, ascending=True)


def pearson_correlation(
    x: Union[np.ndarray, xr.DataArray], y: Union[np.ndarray, xr.DataArray]
) -> float:
    """
    Calculate Pearson correlation coefficient.

    Parameters
    ----------
    x, y : array-like
        Input data arrays.

    Returns
    -------
    float
        Pearson correlation coefficient.
    """
    x = getattr(x, "values", x)
    y = getattr(y, "values", y)
    return np.corrcoef(x.flatten(), y.flatten())[0, 1]


def spearman_correlation(
    x: Union[np.ndarray, xr.DataArray], y: Union[np.ndarray, xr.DataArray]
) -> float:
    """
    Calculate Spearman rank correlation coefficient.

    Parameters
    ----------
    x, y : array-like
        Input data arrays.

    Returns
    -------
    float
        Spearman correlation coefficient.
    """
    from scipy.stats import spearmanr

    x = getattr(x, "values", x)
    y = getattr(y, "values", y)
    correlation, _ = spearmanr(x.flatten(), y.flatten())
    return correlation


def kendall_correlation(
    x: Union[np.ndarray, xr.DataArray], y: Union[np.ndarray, xr.DataArray]
) -> float:
    """
    Calculate Kendall's tau correlation coefficient.

    Parameters
    ----------
    x, y : array-like
        Input data arrays.

    Returns
    -------
    float
        Kendall's tau correlation coefficient.
    """
    from scipy.stats import kendalltau

    x = getattr(x, "values", x)
    y = getattr(y, "values", y)
    correlation, _ = kendalltau(x.flatten(), y.flatten())
    return correlation


def calculate_potential_temperature(
    temperature: Union[np.ndarray, xr.DataArray],
    pressure: Union[np.ndarray, xr.DataArray],
    reference_pressure: float = 1000.0,
) -> Union[np.ndarray, xr.DataArray]:
    """
    Calculate potential temperature using fast numpy operations.

    This implementation uses lazy imports and avoids heavy metpy dependencies
    for simple potential temperature calculations.

    Parameters
    ----------
    temperature : array-like
        Temperature values in Kelvin.
    pressure : array-like
        Pressure values in hPa.
    reference_pressure : float, optional
        Reference pressure in hPa, default is 1000.0.

    Returns
    -------
    array-like
        Potential temperature values in Kelvin.

    Notes
    -----
    Uses the standard formula: theta = T * (P0/P)^(R/cp)
    where R/cp = 0.286 for dry air
    """
    R_over_cp = 0.286  # R/cp for dry air
    potential_temp = temperature * (reference_pressure / pressure) ** R_over_cp

    if hasattr(temperature, "attrs"):
        if isinstance(potential_temp, np.ndarray):
            return xr.DataArray(
                potential_temp,
                attrs={"units": "K", "long_name": "Potential Temperature"},
            )
        else:
            potential_temp.attrs = {"units": "K", "long_name": "Potential Temperature"}

    return potential_temp
