import numpy as np

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 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
    """
    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)
        
        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
        out_data[..., selection] = np.nan
        return out_data
    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
    """
    R = 6371
    if data.lat.shape == 2:
        # irregular grid: lat (J, I) -> (1, J, I)
        lat = data.lat.values[None, :, :]
    else:
        # regular grid: lat (J,) -> (1, J, 1)
        lat = data.lat.values[None, :, None]
    var = data.values
    
    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
