import numpy as np
import dask.array as da


R = 6371


def to_0_360(lon):
    """
    Convert longitudes to 0 360
    """
    return lon % 360


def to_180_180(lon):
    """
    Convert longitudes to -180 180
    """
    return ((lon+180) % 360) - 180


def convert_to_lat_lon(data):
    return data.rename({"latitude": "lat", "longitude": "lon"})


def extract_region(data, 
                   lon_ini, lon_fin, 
                   lat_ini, lat_fin, 
                   transform_lon=to_180_180):
    """
    Extract region and convert the longitudes to a definition in transform_lon
    It will also reorder the data if the input is regular grid
    """
    if hasattr(data, 'latitude'):
        data = convert_to_lat_lon(data)
    lon = transform_lon(data.lon)
    lon_ini = transform_lon(lon_ini)
    lon_fin = transform_lon(lon_fin)
    if lon_ini < lon_fin:
        lon_out = np.logical_and(lon >= lon_ini, lon <= lon_fin)
    elif lon_fin < lon_out:
        lon_out = np.logical_or(lon >= lon_fin, lon <= lon_ini)
    else:
        lon_out = np.ones_like(lon, dtype=bool)
        
    lat_out = np.logical_and(data.lat >= lat_ini, data.lat <= lat_fin)
    
    if len(data.lon.shape) == 2:
        # irregular grid
        selection = np.logical_and(lon_out, lat_out).data
        
        i_range, j_range = selection.nonzero()
        if i_range.size == 0:
            raise ValueError("No data points available in selected region")
        i_min, i_max = i_range.min(), i_range.max()
        j_min, j_max = j_range.min(), j_range.max()
        i_slice, j_slice = slice(i_min, i_max + 1), slice(j_min, j_max + 1)
        out_data = data[..., i_slice, j_slice]
        selection = selection[i_slice, j_slice]
        # Mask remaining coordinates outside region
        mask = da.broadcast_to(selection, out_data.shape)
        return out_data.where(mask)
    else:
        # regular grid
        lon_index = np.arange(len(lon_out))[lon_out]
        lat_index = np.arange(len(lat_out))[lat_out]
        out_data = data.isel(lon=lon_index, lat=lat_index)
        out_data['lon'] = lon[lon_out]
        return out_data.sortby(out_data.lon)


def meridional_gradient(data):
    """
    Compute meridional gradients
    """
    if hasattr(data, 'latitude'):
        data = convert_to_lat_lon(data)
    
    lshape = [1]*len(data.shape)
    if len(data.lat.shape) == 2:
        # irregular grid: lat (J, I) -> (1, ..., 1, J, I)
        lshape[-2:] = data.lat.shape
    else:
        # irregular grid: lat (J,) -> (1, ..., 1, J, 1)
        lshape[-2] = data.lat.shape[0]

    # get values lazy in the same shape    
    lat = data.lat.data.reshape(lshape)
    var = data.data
    
    # compute the gradient
    gradient = data.copy(data=np.empty_like(data))
    gradient[..., 0, :] = (var[..., 1, :] - var[..., 0, :])/(R*(np.radians(lat[..., 1, :] - lat[..., 0, :])))
    gradient[..., -1, :] = (var[..., -1, :] - var[..., -2, :])/(R*(np.radians(lat[..., -1, :] - lat[..., -2, :])))
    gradient[..., 1:-1, :] = (var[..., 2:, :] - var[..., :-2, :])/(R*(np.radians(lat[..., 2:, :] - lat[..., :-2, :])))
    
    return gradient


def zonal_gradient(data):
    """
    Compute zonal gradients
    """
    if hasattr(data, 'latitude'):
        data = convert_to_lat_lon(data)

    lshape = [1]*len(data.shape)
    if len(data.lat.shape) == 2:
        # irregular grid: lon (J, I) -> (1, ..., 1, J, I)
        # irregular grid: lat (J, I) -> (1, ..., 1, J, I)
        lshape[-2:] = data.lat.shape
        lat = data.lat.data.reshape(lshape)
        lon = data.lon.values[None, :, :]
        # sinus of latitude to weight
        slatr = np.sin(np.radians(lat[..., 1:-1]))
        slat0 = np.sin(np.radians(np.mean(lat[..., :2], axis=-1)))
        slat1 = np.sin(np.radians(np.mean(lat[..., -2:], axis=-1)))
    else:
        # irregular grid: lon (I,) -> (1, ..., 1, 1, I)
        # irregular grid: lat (J,) -> (1, ..., 1, J, 1)
        lshape[-2] = data.lat.shape[0]
        lat = data.lat.data.reshape(lshape)
        # sinus of latitude to weight
        slatr = np.sin(np.radians(lat))
        slat0 = slatr[:, :, 0]
        slat1 = slat0
        lshape[-2:] = (1, data.lon.shape[0])
    
    # get values lazy in the same shape
    lon = data.lon.data.reshape(lshape)
    var = data.data
    
    # compute the gradient
    gradient = data.copy(data=np.empty_like(data))
    gradient[..., 0] = (var[..., 1] - var[..., 0])/(R*(np.radians(lon[..., 1] - lon[..., 0]))*slat0)
    gradient[..., -1] = (var[..., -1] - var[..., -2])/(R*(np.radians(lon[..., -1] - lon[..., -2]))*slat1)
    gradient[..., 1:-1] = (var[..., 2:] - var[..., :-2])/(R*(np.radians(lon[..., 2:] - lon[..., :-2]))*slatr)
    
    return gradient
