import h5py
import numpy as np
from scipy.interpolate import RegularGridInterpolator, griddata
import cartopy.crs as ccrs

from util.util import get_grid_values_all

# resample methods:
linear = 'linear'
cubic = 'cubic'
nearest = 'nearest'


def fill_missing(fld_a, fld_b, mask):
    min_val = min(np.nanmin(fld_a), np.nanmin(fld_b))
    max_val = max(np.nanmax(fld_a), np.nanmax(fld_b))
    random_values_a = np.random.uniform(min_val, max_val, size=fld_a.shape)
    random_values_b = np.random.uniform(min_val, max_val, size=fld_a.shape)

    fld_a[mask] = random_values_a[mask]
    fld_b[mask] = random_values_b[mask]


def get_projection(cartopy_map_name, cen_lat, cen_lon):
    if cartopy_map_name == "LambertAzimuthalEqualArea":
        projection = ccrs.LambertAzimuthalEqualArea(
            central_longitude=cen_lon, central_latitude=cen_lat
        )
    elif cartopy_map_name == "AlbersEqualArea":
        projection = ccrs.AlbersEqualArea(
            central_longitude=cen_lon, central_latitude=cen_lat
        )
    elif cartopy_map_name == "Sinusoidal":
        projection = ccrs.Sinusoidal(central_longitude=cen_lon)
    else:
        raise ValueError("Projection: " + cartopy_map_name + " is not supported")

    return projection


def resample_reg_grid(scalar_field, y, x, y_s, x_s, method='linear'):

    intrp = RegularGridInterpolator((y, x), scalar_field, method=method, bounds_error=False)

    xg, yg = np.meshgrid(x_s, y_s, indexing='xy')
    yg, xg = yg.flatten(), xg.flatten()
    pts = np.array([yg, xg])
    t_pts = np.transpose(pts)

    return np.reshape(intrp(t_pts), (y_s.shape[0], x_s.shape[0]))


def resample(scalar_field, y_d, x_d, y_t, x_t, method='linear'):
    # 2D target locations shape
    t_shape = y_t.shape
    # reproject scalar fields
    fld_repro = griddata((y_d.flatten(), x_d.flatten()), scalar_field.flatten(),(y_t.flatten(), x_t.flatten()), method=method)
    fld_repro = fld_repro.reshape(t_shape)
    return fld_repro


def reproject(fld_2d, lat_2d, lon_2d, proj, target_grid=None, grid_spacing=15000.0, method=linear):
    """
    :param fld_2d: the 2D scalar field to reproject
    :param lat_2d: 2D latitude of the scalar field domain
    :param lon_2d: 2D longitude of the scalar field domain
    :param proj: the map projection (Cartopy). Default: LambertAzimuthalEqualArea
    :param region_grid: the larger region grid that we pull the target grid from
    :param target_grid: the resampling target (y_map, x_map) where y_map and x_map are 2D. If None, the grid is created
           automatically. The target grid is always returned.
    :param grid_spacing: distance between the target grid points (in meters)
    :param method: resampling method: 'linear', 'nearest', 'cubic'
    :return: reprojected 2D scalar field, the target grid
    """

    data_xy = proj.transform_points(ccrs.PlateCarree(), lon_2d, lat_2d)[..., :2]

    # Generate a regular 2d grid extending the min and max of the xy dimensions with grid_spacing
    if target_grid is None:
        x_min, y_min = np.amin(data_xy, axis=(0, 1))
        x_max, y_max = np.amax(data_xy, axis=(0, 1))
        x_map = np.arange(x_min, x_max, grid_spacing)
        y_map = np.arange(y_min, y_max, grid_spacing)
        x_map_2d, y_map_2d = np.meshgrid(x_map, y_map)
    else:
        y_map_2d, x_map_2d = target_grid

    fld_reproj = resample(fld_2d, data_xy[..., 1], data_xy[..., 0], y_map_2d, x_map_2d, method=method)

    return fld_reproj, (y_map_2d, x_map_2d)


def bisect_great_circle(lon_a, lat_a, lon_b, lat_b):
    lon_a = np.radians(lon_a)
    lat_a = np.radians(lat_a)
    lon_b = np.radians(lon_b)
    lat_b = np.radians(lat_b)

    dlon = lon_b - lon_a

    Bx = np.cos(lat_b) * np.cos(dlon)
    By = np.cos(lat_b) * np.sin(dlon)

    lat_c = np.arctan2(np.sin(lat_a) + np.sin(lat_b), np.sqrt((np.cos(lat_a) + Bx) ** 2 + By ** 2))
    lon_c = lon_a + np.arctan2(By, np.cos(lat_a) + Bx)

    lon_c = np.degrees(lon_c)
    lat_c = np.degrees(lat_c)

    return lon_c, lat_c


def acspo_sst(filename, stride=2, grid_spacing=2000):
    h5f = h5py.File(filename, 'r')

    # operational
    sst = get_grid_values_all(h5f, 'sea_surface_temperature')
    sst = sst[0, :, :]
    lons = get_grid_values_all(h5f, 'lon')
    lats = get_grid_values_all(h5f, 'lat')

    # classic
    # sst = get_grid_values_all(h5f, 'sst_regression')
    # sst = sst[:, :]
    # lons = get_grid_values_all(h5f, 'longitude')
    # lats = get_grid_values_all(h5f, 'latitude')

    print('data dims: ', lons.shape, lats.shape)

    sst = sst[::stride, ::stride]
    lons = lons[::stride, ::stride]
    lats = lats[::stride, ::stride]

    ylen, xlen = lons.shape
    print('final dims: ', ylen, xlen)

    cen_lon = lons[ylen // stride, xlen // stride]
    cen_lat = lats[ylen // stride, xlen // stride]
    print('center latitude/longitude: ', cen_lat, cen_lon)

    proj = get_projection('LambertAzimuthalEqualArea', cen_lat, cen_lon)

    fld_reproj, (y_map_2d, x_map_2d) = reproject(sst, lats, lons, proj, grid_spacing=grid_spacing)

    return fld_reproj, proj, y_map_2d, x_map_2d