
import os
import warnings
import urllib3
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


from scipy import constants
from scipy.interpolate import RectSphereBivariateSpline,SmoothSphereBivariateSpline

import cdsapi
import numpy as np
import cartopy.crs as ccrs
import pandas as pd
import scipy
import scipy.integrate
import scipy.interpolate
import xarray as xr

from ddeq.smartcarb import DOMAIN
import ddeq


g = constants.g
R = 287.058


def get_wind_at_location(winds, lon0, lat0):
    """\
    Obtain wind at location given by lon0 and lat0 in wind field (xr.Dataset)
    using nearest neighbor interpolation.
    """
    lat = winds['lat'].data
    lon = winds['lon'].data

    if np.ndim(lat) == 2:
        dist = (lat - lat0)**2 + (lon - lon0)**2
        i, j = np.unravel_index(dist.argmin(), dist.shape)
    else:
        i = np.argmin(np.abs(lat-lat0))
        j = np.argmin(np.abs(lon-lon0))

    return np.array([float(winds['U'][i,j]), float(winds['V'][i,j])])


def read_smartcarb(time, lon, lat, radius=None, data_path='.', method='linear',
                   average=False):
    """\
    Read SMARTCARB winds at given location at given `time`. The location is
    interpolated from SMARTCARB model grid to given `lon` and `lat`. If `lon`
    and `lat` are given as scalar, it is possible to provide a radius around
    location for which are extracted.

    radius :: radius around location given in rotated degrees
    data_path :: path to SMARTCARB wind fields
    method :: interpolation method (used by xr.DataArray.interp method)
    average :: average extracted winds

    Return xr.Dataset with wind components `U` and `V` as well as `wind speed`
    and `direction`.
    """
    lon = np.atleast_1d(lon).flatten()
    lat = np.atleast_1d(lat).flatten()

    if lon.size > 1 and radius is not None:
        raise ValueError(
            'Please provide lon and lat as scalar if radius is given.'
        )

    rlon, rlat, _ = np.transpose(
        DOMAIN.proj.transform_points(ccrs.PlateCarree(), lon, lat)
    )

    filename = time.round('H').strftime('SMARTCARB_winds_%Y%m%d%H.nc')
    filename = os.path.join(data_path, filename)

    with xr.open_dataset(filename) as nc:
        if radius is None:
            u = [float(nc['U_GNFR_A'].interp(rlon=rlon[i], rlat=rlat[i],
                                             method=method))
                 for i in range(rlon.size)]

            v = [float(nc['V_GNFR_A'].interp(rlon=rlon[i], rlat=rlat[i],
                                             method=method))
                 for i in range(rlon.size)]

        else:
            distance = np.sqrt((nc.rlon - rlon)**2 + (nc.rlat - rlat)**2)

            if np.any(distance <= radius):
                region = nc.where(distance <= radius, drop=True)

                lon = region.lon.values.flatten()
                lat = region.lat.values.flatten()
                u = region['U_GNFR_A'].values.flatten()
                v = region['V_GNFR_A'].values.flatten()
            else:
                lon = lon
                lat = lat
                u = np.array([np.nan])
                v = np.array([np.nan])

    wind = xr.Dataset()
    wind['lon'] = xr.DataArray(lon, dims='index')
    wind['lat'] = xr.DataArray(lat, dims='index')
    wind['U'] = xr.DataArray(u, dims='index')
    wind['V'] = xr.DataArray(v, dims='index')

    if average:
        u_std = wind['U'].std()
        v_std = wind['V'].std()
        wind = wind.mean('index')
        wind['U_std'] = u_std
        wind['V_std'] = v_std

    wind['speed'] = np.sqrt(wind.U**2 + wind.V**2)
    wind['speed_precision'] = 1.0 # FIXME
    wind['direction'] = ddeq.wind.calculate_wind_direction(wind.U, wind.V)

    return wind.where(np.isfinite(wind['U']), drop=True)


def read_era5(time, lon, lat, data_path='.',timesteps=1):

    lon = np.atleast_1d(lon).flatten()
    lat = np.atleast_1d(lat).flatten()

    u = np.zeros_like(lon)
    v = np.zeros_like(lon)

    for i in range(lon.size):
        era5 = get_era5_wind_profile(time, lon[i], lat[i],
                                     data_path=data_path,
                                     timesteps=timesteps)

        # TODO: use weights
        if i == 0:
            wind = era5.mean('level').expand_dims("index")
        else:
            wind = xr.concat([wind, era5.mean('level')], dim="index")

    wind = wind.rename({"longitude": "lon",
                        "latitude": "lat",
                        "u": "U",
                        "v": "V"})
    wind = wind.drop_vars(["t", "z"])
    wind = wind.reset_coords()
    
    wind['speed'] = np.sqrt(wind.U**2 + wind.V**2)
    wind['speed_precision'] = (wind.U.dims, np.full_like(wind.U, 1.0)) # FIXME
    wind['direction'] = ddeq.wind.calculate_wind_direction(wind.U, wind.V)

    return wind.where(np.isfinite(wind['U']), drop=True)


def read_coco2_era5(time, lon, lat, suffix=None, method='linear'):
    """
    Read ERA-5 data from CoCO2 project on ICOS-CP.
    """
    warnings.warn('CoCO2-specific function will be removed in future.',
                  DeprecationWarning, stacklevel=2)

    if suffix is None:
        filename = time.strftime(
            '/project/coco2/jupyter/WP4/ERA5/ERA5-gnfra-%Y%m%dt%H00.nc'
        )
    else:
        filename = time.strftime(
            f'/project/coco2/jupyter/WP4/ERA5/ERA5-gnfra-%Y%m%dt%H00_{suffix}.nc'
        )

    lon = np.atleast_1d(lon)
    lat = np.atleast_1d(lat)

    with xr.open_dataset(filename) as nc:

        u = [float(nc['U_GNFR_A'].interp(longitude=lon[i], latitude=lat[i],
                                         method=method))
             for i in range(lon.size)]

        v = [float(nc['V_GNFR_A'].interp(longitude=lon[i], latitude=lat[i],
                                         method=method))
                 for i in range(lon.size)]

    wind = xr.Dataset()
    wind['lon'] = xr.DataArray(lon, dims='index')
    wind['lat'] = xr.DataArray(lat, dims='index')
    wind['U'] = xr.DataArray(u, dims='index')
    wind['V'] = xr.DataArray(v, dims='index')

    wind['speed'] = np.sqrt(wind.U**2 + wind.V**2)
    wind['speed_precision'] = 1.0 # FIXME
    wind['direction'] = ddeq.wind.calculate_wind_direction(wind.U, wind.V)

    return wind


def calculate_wind_direction(u, v):
    return (270.0 - np.rad2deg(np.arctan2(v,u))) % 360


def download_era5_winds(time, data_path='.', overwrite=False,
                        timesteps=1):
    """
    Download ERA-5 data of a given time step.
    
    timesteps :: number of hours to download on the given day.
        If 24, a full day will be returned.
    """

    if timesteps == 24:
        era5_filename = os.path.join(data_path,
                                     time.strftime('ERA5-wind-%Y%m%d.nc'))
    else:
        era5_filename = os.path.join(data_path,
                                     time.strftime('ERA5-wind-%Y%m%d_%H00.nc'))

    if os.path.exists(era5_filename) and not overwrite:
        return era5_filename

    cds = cdsapi.Client()

    query = {
        'product_type': 'reanalysis',
        'format': 'netcdf',
        'variable': [
            'geopotential', 'temperature', 'u_component_of_wind',
            'v_component_of_wind',
        ],
        'pressure_level': [
            '700', '750', '775',
            '800', '825', '850',
            '875', '900', '925',
            '950', '975', '1000',
        ],
        "area": [90, -180, -90, 180], # north, east, south, west
    }

    query["date"] = time.strftime('%Y-%m-%d')
    
    if timesteps == 1:
        query["time"] = time.round('H').strftime('%H:00')
    elif timesteps == 24:
        query['time'] = [f"{hour:02d}:00" for hour in range(0, timesteps, 1)]
    else:
        query['time'] = [f"{hour:02d}:00" for hour in range(
            max(0,time.round('H').hour+1-timesteps),
            time.round('H').hour+1,
            1)]
        
    r = cds.retrieve('reanalysis-era5-pressure-levels', query, era5_filename)
    
    return era5_filename


def get_era5_wind_profile(time, lon, lat, data_path='.',
                          timesteps=1):
    
    era5_filename = download_era5_winds(time, data_path,
                                        overwrite=False, timesteps=timesteps)

    with xr.open_dataset(era5_filename) as era5:
        era5 = era5.load().sel(longitude=lon, latitude=lat,
                               method='nearest')

    return era5



def calculate_pressure_levels(surface_pressure, nlevels=37, which='middle'):

    filename= os.path.join(os.path.dirname(__file__), 'data',
                           'ERA5 L137 model level definitions.csv')

    level_definitions = pd.read_csv(filename, index_col=0)

    a = np.array(level_definitions['a [Pa]'])
    b = np.array(level_definitions['b'])

    # pressure at boundary of model layers
    ph = a + b * float(surface_pressure)

    # TODO: how to interpolate?
    #ph = np.interp(np.linspace(0,138,nlevels+1), np.arange(138), ph)

    # pressure at middle of model layers
    pf = 0.5 * (ph[1:] + ph[:-1])

    if which == 'middle':
        return pf
    elif which == 'boundary':
        return ph
    else:
        raise ValueError


def calculate_level_heights(temperature, pressure, which='middle'):
    """
    temperature at middle of layers
    pressure at boundary of layers
    """
    g = 9.80665
    R = 287.06

    p1 = pressure[1:]
    p2 = pressure[:-1]

    dh = R * temperature / g * np.log(p1 / p2)

    hh = np.concatenate([[0.0], dh[::-1].cumsum()])
    hh = hh[::-1]

    hf = 0.5 * (hh[1:] + hh[:-1])

    if which == 'middle':
        return hf
    elif which == 'boundary':
        return hh
    else:
        raise ValueError


def read_gnfra_wind(folder, date, UTC='11:00:00', latlims=None, lonlims=None, dataset='ERA5'):
    '''
    Parameters
    ----------
    folder : str
        Path to folder where wind data is located.
    date : str or datetime
        Date of the data.
    UTC : str, optional
        UTC time stamp as 'hh:mm:ss'. The default is '11:00:00'.
    latlims : tuple of float, optional
        Latitude limits to focus the data into a region. The default is None.
    lonlims : tuple of float, optional
        Latitude limits to focus the data into a region. The default is None.
    dataset : str, optional
        Dataset of wind data, choices are 'ERA5' or 'SMARTCARB'. The default is 'ERA5'.

    Returns
    -------
    ds : xarray dataset
        Wind data.
    '''
    if (latlims is None and lonlims is not None) or (latlims is not None and lonlims is None):
        raise ValueError('Provide either both or neither of longitude and latitude limits')

    # lims = True if both latlims and lonlims are provided
    lims = latlims is not None and lonlims is not None
    date_str = date.replace('-','') if isinstance(date, str) else date.strftime('%Y%m%d')
    utc_str = UTC.replace(':','')

    if dataset == 'ERA5':
        ds = xr.open_dataset(f'{folder}/ERA5-gnfra-{date_str}t{utc_str[:4]}.nc')
        if lims:
            wlon, wlat = ds['longitude'][:].data, ds['latitude'][:].data
            lat_slice = (min(latlims)-0.5 <= wlat) & (wlat <= max(latlims)+0.5)
            lon_slice = (min(lonlims)-0.5 <= wlon) & (wlon <= max(lonlims)+0.5)
            ds = ds.where(xr.DataArray(lat_slice[:,None] & lon_slice[None,:], dims=['latitude','longitude']), drop=True)

    elif dataset == 'SMARTCARB':
        # Read data
        ds = xr.open_dataset(f'{folder}/SMARTCARB_winds_{date_str}{utc_str[:2]}.nc',
                             drop_variables=['height_10m','U','V','U_10M','V_10M','HHL'])

        # Reindex latitude into decreasing order, the highest latitude at index 0 and the lowest at -1
        ds = ds.reindex(rlat=ds['rlat'][::-1].data)

        # If limits are given, slice data to the area of interest 
        if lims:
            wlon, wlat = ds['lon'][:,:].data.flatten(), ds['lat'][:,:].data.flatten()
            shape = (ds['rlat'].size, ds['rlon'].size)

            top_right = np.unravel_index(np.argmin((wlon-max(lonlims))**2+(wlat-max(latlims))**2),shape)
            bottom_right = np.unravel_index(np.argmin((wlon-max(lonlims))**2+(wlat-min(latlims))**2),shape)
            bottom_left = np.unravel_index(np.argmin((wlon-min(lonlims))**2+(wlat-min(latlims))**2),shape)
            top_left = np.unravel_index(np.argmin((wlon-min(lonlims))**2+(wlat-max(latlims))**2),shape)

            row_index = np.arange(0,shape[0],1,dtype=int)
            col_index = np.arange(0,shape[1],1,dtype=int)

            lat_slice = (min(top_left[0],top_right[0])-1 <= row_index) & (row_index <= max(bottom_left[0],bottom_right[0])+1)
            lon_slice = (min(top_left[1],bottom_left[1])-1 <= col_index) & (col_index <= max(top_right[1],bottom_right[1])+1)
            ds = ds.where(xr.DataArray(lat_slice[:,None] & lon_slice[None,:], dims=['rlat','rlon']), drop=True)

    return ds


def spatial_interpolation(data,wlon,wlat,lon,lat):
    '''
    Interpolates wind in spatial dimensions to given pixels lon-lat pixels.

    Parameters
    ----------
    data : 1d-array or 2d-array
        Wind data.
    wlon : 1d-array or 2d-array
        Wind longitude coordinates.
    wlat : 1d-array or 2d-array
        Wind latitude coordinates.
    lon : 1d-array
        Longitude coordinates to be interpolated.
    lat : 1d-array
        Latitude coordinates to be interpolated.

    Returns
    -------
    interp : 1d-array
        Interpolated data values.
    '''
    if wlon.ndim == 1 and wlat.ndim == 1:
        if data.ndim == 1:
            data = data.reshape((len(wlat),len(wlon)))
        # Interpolator in regular rectangular lat-lon grid
        interpolator = RectSphereBivariateSpline(np.deg2rad(90-wlat),np.deg2rad(wlon),data)
    elif wlon.ndim == 2 and wlat.ndim == 2:
        wlon, wlat = wlon.flatten(), wlat.flatten()
        if data.ndim == 2:
            data = data.flatten()
        # Same weights for each data value, the square sum of weights is exactly one
        weights = np.ones(len(data))/np.sqrt(len(data))
        # Interpolator on the irregular spherical lat-lon grid, requires smoothing factor s, for the chosen 
        # weights the function to to optimized is the root-mean-square distance between the data and the spline
        interpolator = SmoothSphereBivariateSpline(np.deg2rad(90-wlat.flatten()),np.deg2rad(wlon.flatten()),data,w=weights,s=1e-2*len(data))
    else:
        pass
    # Evaluate interpolator at requested coordinates
    interp = interpolator.ev(np.deg2rad(90-lat),np.deg2rad(lon))
    return interp


def read(time, sources, product='ERA5', data_path='.', radius=0.05,
         timesteps=1):
    """\
    Read wind dataset for sources.

    Product: ERA5, SMARTCARB, ...
    """
    if isinstance(time, xr.DataArray):
        time = pd.Timestamp(time.to_pandas())
    elif isinstance(time, np.datetime64):
        time = pd.Timestamp(time)

    winds = {}
    for name, source in sources.groupby('source', squeeze=False):
        if product in ['SC', 'SMARTCARB']:
            winds[name] = read_smartcarb(
                time, source.lon_o.values, source.lat_o.values,
                data_path=data_path,
                radius=radius, method='linear', average=True
            )
        elif product in ['ERA-5', 'ERA5']:
            wind = read_era5(
                time, source.lon_o.values, source.lat_o.values,
                data_path=data_path, timesteps=timesteps
            )
            winds[name] = wind.rename_dims({'index': 'source'})
        else:
            raise ValueError

    winds = xr.concat(winds.values(), dim='source')
    winds['source'] = xr.DataArray(sources.source, dims='source')

    return winds


def read_wind_field(wind_data, source, lon, lat, radius):
    """
    Read wind at a certain point from TROPOMI or MicroHH files.
    
    Parameters:
    ----------
    wind_data: netCDF
        A file containing the wind information
    source: str
        The name of the source for which the wind is read in
    lon: float
        Longitude of the source
    lat: float
        Latitude of the source
    radius: int
        Radius of pixels around the source.
        
    Returns:
    wind_src: netCDF
        A file with the mean wind information of the kernel with radius 'radius'
    """
    
    if "UM" in wind_data.data_vars:
        wind_data = wind_data.rename({
            "UM": "U",
            "VM": "V"
        })
    elif "longitude" in wind_data.data_vars:
        wind_data = wind_data.rename({
            "longitude": "lon",
            "latitude": "lat"
        })

    wind_data['speed'] = np.sqrt(wind_data.U**2 + wind_data.V**2)
    wind_data['direction'] = ddeq.wind.calculate_wind_direction(wind_data.U, wind_data.V)
    
    lon_closest, lat_closest, dist = ddeq.misc.find_closest(wind_data, ("lon", "lat"), (lon, lat)) # get values for closest coordinate
    
    wind_src = xr.Dataset(
        coords = {
            "source": (["source"], np.array([source], dtype=object))
        },
        data_vars = {
            "lon": (
                ["source"],
                [wind_data.lon[lon_closest, lat_closest].values]
            ),
            "lat": (
                ["source"],
                [wind_data.lat[lon_closest, lat_closest].values]
            ),
            "U": (
                ["source"],
                [wind_data.U[lon_closest-radius:lon_closest+radius+1,
                                lat_closest-radius:lat_closest+radius+1].median().values],
                {"units": "m s-1"}
            ),
            "V": (
                ["source"],
                [wind_data.V[lon_closest-radius:lon_closest+radius+1,
                                lat_closest-radius:lat_closest+radius+1].median().values],
                {"units": "m s-1"}
            ),
            "speed": (
                ["source"],
                [wind_data.speed[lon_closest-radius:lon_closest+radius+1,
                                   lat_closest-radius:lat_closest+radius+1].median().values],
                {"units": "m s-1"}),
            "speed_precision": (
                ["source"],
                [wind_data.speed[lon_closest-radius:lon_closest+radius+1,
                                   lat_closest-radius:lat_closest+radius+1].std().values],
                {"units": "m s-1"}),
            "direction": (
                ["source"],
                [scipy.stats.circmean(wind_data.direction[lon_closest-radius:lon_closest+radius+1,
                                                           lat_closest-radius:lat_closest+radius+1].values.flatten(),
                                     high = 360, nan_policy="omit")],
                {"units": "°"}),
            "direction_precision": (
                ["source"],
                [scipy.stats.circstd(wind_data.direction[lon_closest-radius:lon_closest+radius+1,
                                                           lat_closest-radius:lat_closest+radius+1].values.flatten(),
                                     high = 360, nan_policy="omit")],
                {"units": "°"}),
        },
        attrs=dict(radius=f"{radius} pixels around determined pixel")
    )

    return wind_src


def integrate_wind(data, curves, winds, time, crs, pixel_size):
    """
    Integrate wind backwards in time
    Pixel size: dimensions of pixels in m
    """
    
    if isinstance(time, xr.DataArray):
        time = pd.Timestamp(time.to_pandas())
    elif isinstance(time, np.datetime64) or isinstance(time, str):
        time = pd.Timestamp(time)

    src = 0
    for source in data.source.values:
        
        if curves[source] is None:
            continue
            
        s = 0
        in_plume = True
        time_plume = time
        winds_int = winds.sel(time=time, method='nearest').sel(source=source)
        
        if curves[source] is None:
            in_plume = False
            angle = np.nan
        else:
            angle = ddeq.misc.compute_angle_between_curve_and_wind(curves[source],
                                                                   winds_int.direction,
                                                                   crs=crs)

        winds_int["distance"] = ((), 0, {"units": "m"})
        winds_int["in_plume"] = ((), True)
        winds_int["angle_between_curve_and_wind"] = ((), float(angle), {"units": "°"})
        w = float(winds_int.speed)

        while in_plume and pd.Timestamp(min(winds.time.values)):

            s += w * 60**2 # compute distance the air parcel has travelled in 1h
            nrow, nob, dist = ddeq.misc.find_closest(data.sel(source=source), ("xp","yp"), (s,0))

            # test if the parcel is still inside the plume. Use this as last value for interpolation
            if dist > np.sqrt(pixel_size**2 + pixel_size**2):
                in_plume = False

            time_plume = time_plume - pd.Timedelta(hours=1)
            
            # get wind data at the source of the plume
            winds_plume = winds.sel(time=time_plume, method='nearest').sel(source=source)
            winds_plume["lon"] = data.isel(nrows=nrow, nobs=nob).lon.values
            winds_plume["lat"] = data.isel(nrows=nrow, nobs=nob).lat.values
            winds_plume["distance"] = ((), s, {"units": "m"})
            winds_plume["in_plume"] = ((), in_plume)
            angle = ddeq.misc.compute_angle_between_curve_and_wind(curves[source],
                                                                   winds_plume.direction,
                                                                   crs=crs)
            winds_plume["angle_between_curve_and_wind"] = ((), float(angle), {"units": "°"})
            w = float(winds_plume.speed)

            if w == 0.0 or np.isnan(w):
                in_plume = False

            winds_int = xr.concat([winds_int, winds_plume], dim="time")
        
        if src == 0:
            winds_int_all = winds_int
            src += 1
        else:
            winds_int_all = xr.concat([winds_int_all, winds_int], dim="source")
    
    return winds_int_all


def interpolate_wind(data, winds, delta):
    """
    Interpolate wind to the same along plume distances as in 'data'.
    """
    
    src = 0
    for source in data.source.values:
        
        if "xp" not in data.sel(source=source).data_vars or source not in winds.source.values:
            continue

        dmax = np.nanmax(data.sel(source=source).xp.values[data.sel(source=source).detected_plume]) + delta/2
        new_along = np.insert(np.arange(delta/2, int(dmax), delta), 0, -7000)

        winds_interp = xr.Dataset(
            coords = {
                "source": (["source"], np.array([source], dtype=object)),
                "along": (["along"], np.array(new_along, dtype=int), {"units": "m"})
            })

        if not len(winds.time) == 0:
            for varname, da in winds.data_vars.items():
                if varname != "distance":
                    new_var = np.interp(
                        np.insert(np.arange(delta/2, int(dmax), delta), 0, 0),
                        winds.distance.values.flatten(),
                        da.values.flatten())
                    winds_interp[varname] = ("along", new_var, {"units": da.attrs.get('units', "")})
            
            int_time = np.interp(
                np.insert(np.arange(delta/2, int(dmax), delta), 0, 0),
                winds.distance.values.flatten(),
                abs((winds.time.values.flatten() - 
                     np.nanmax(winds.time.values)) / 1e9).astype(int))
            winds_interp["int_time"] = ("along", int_time, {"units": "s"})
                    
    if src == 0:
        winds_interp_all = winds_interp
        src += 1
    else:
        winds_interp_all = xr.concat([winds_interp_all, winds_interp], dim="source")
            
    return winds_interp_all


def read_field(filename, wind_alti_file=None, product='ERA5',
               altitude='GNFR-A', average_below=False):
    """\
    Return wind field from file for different products. 3D wind fields are
    taken at nearest altitude or averaged_below. If altitude is "GNFR_A", use
    vertically weighted wind field.
    """
    with xr.open_dataset(filename) as rfile:

        if product == 'SMARTCARB':

            lat = rfile.variables['lat'][:]
            lon = rfile.variables['lon'][:]

            if altitude == 'GNFR-A':
                u = np.squeeze(rfile['U_GNFR_A'][:])
                v = np.squeeze(rfile['V_GNFR_A'][:])
            else:
                u = np.squeeze(rfile['U'][:])
                v = np.squeeze(rfile['V'][:])
                h = np.squeeze(rfile['HHL'].values[:])  # altitude in m
                h = h[:] - h[-1]                        # above surface
                h = 0.5 * (h[1:] + h[:-1])               # at layer center

        elif product == 'ERA5':

            lat = rfile.variables['latitude'][:]
            lon = rfile.variables['longitude'][:]

            if altitude == 'GNFR-A':
                u = np.squeeze(rfile.variables['U_GNFR_A'][:])
                v = np.squeeze(rfile.variables['V_GNFR_A'][:])
            else:
                u = np.squeeze(rfile.variables['u'][:])
                v = np.squeeze(rfile.variables['v'][:])

                raise NotImplementedError
        else:
            raise ValueError(f'Unknown wind product "{product}".')

    # average wind below altitude
    if not isinstance(altitude, str):
        if average_below:
            u = np.nanmean(np.where(h < altitude, u, np.nan), axis=0)
            v = np.nanmean(np.where(h < altitude, v, np.nan), axis=0)
        else:
            # extract value at level
            k = np.argmin(np.abs(h - altitude), axis=0)

            unew = np.zeros(u.shape[1:])
            vnew = np.zeros(v.shape[1:])

            # quickest when itereating over small number of indices
            for index in set(k.flatten()):
                mask = (k == index)
                unew[mask] = u.values[index, mask]
                vnew[mask] = v.values[index, mask]

            u = xr.DataArray(unew, dims=u.dims[1:], attrs=u.attrs)
            v = xr.DataArray(vnew, dims=v.dims[1:], attrs=v.attrs)

    # Create wind dataset
    attrs = {}
    attrs['wind_product'] = product
    attrs['altitude'] = str(altitude)

    wind_field = xr.Dataset(attrs=attrs)

    wind_field['U'] = xr.DataArray(data=u)
    wind_field['V'] = xr.DataArray(data=v)
    wind_field['lat'] = xr.DataArray(data=lat)
    wind_field['lon'] = xr.DataArray(data=lon)

    return wind_field


#################################### Functions for ERA5 ####################################

# define some variables for the calculation of the weighted wind profile
lower_bounds = np.array([170,  310,  470,  710]) # height bound of GNFR-A weights
upper_bounds = np.array([310,  470,  710,  990]) # height bound of GNFR-A weights
values =       np.array([0.08, 0.46, 0.29, 0.17]) # GNFR-A weights
bounds = np.append(lower_bounds, upper_bounds[-1])
width = np.diff(bounds)


def era5_winds_ddeq(time, sources, query, timesteps=1, data_path=".", overwrite=False):
    """
    Read ERA5 wind dataset for given sources.
    
    Parameters:
    ----------
    time : datetime64 or xarray
        Time for which the wind data should be downloaded
    sources : xarray
        Dataset of sources with lat and lon
    query : str
        Indicate how the wind should be vertically averaged
            - complete: GNFR-A weighted wind on model level wind
            - pbl_mean: Mean wind in pbl 
            - pressure_levels: Mean wind of pressure levels 1000 to 700 hPa
    timesteps : int
        Number of timesteps before "time" for which the data should be downloaded
        If 24, a full day will be returned.
    data_path : str
        Path to which the files should be saved
    overwrite : bool
        Indicates if the file should be overwritten if it already exists
        
    Returns:
    --------
    winds : netCDF
        A file containing the u and v component of the wind, total wind speed and direction.
    """
    
    if isinstance(time, xr.DataArray):
        time = pd.Timestamp(time.to_pandas())
    elif isinstance(time, np.datetime64):
        time = pd.Timestamp(time)

    winds = {}
    for name, source in sources.groupby('source'):
        
        coords = (source.lat_o.values, source.lon_o.values)
        prefix = name
        
        wind = era5_winds(
            time, coords, query=query, prefix=prefix,
            timesteps=timesteps, data_path=data_path,
            overwrite=overwrite
        )
        
        wind = wind.reset_coords(names=["lat", "lon"])
        winds[name] = wind

    winds = xr.concat(winds.values(), dim='source')
    winds['source'] = xr.DataArray(sources.source, dims='source')
    

    return winds

def era5_winds(time, coords, query, prefix="", timesteps=1, data_path=".", overwrite=False):
    """
    Download ERA5 wind speed data for a given area. Calculate total wind speed and direction.
    
    Parameters:
    ----------
    time : datetime64 or xarray
        Time for which the wind data should be downloaded
    coords : str or tuple
        Indicate for which coordinates the wind data should be downloaded
        "global" for global coverage or a tuple of coordinates: (lat,lon)
    query : str
        Indicate how the wind should be vertically averaged
            - complete: GNFR-A weighted wind on model level wind
            - pbl_mean: Mean wind in pbl 
            - pressure_levels: Mean wind of pressure levels 1000 to 700 hPa
    prefix : str
        Prefix of the name of the downloaded data
    timesteps : int
        Number of timesteps before "time" for which the data should be downloaded
        If 24, a full day will be returned.
    data_path : str
        Path to which the files should be saved
    overwrite : bool
        Indicates if the file should be overwritten if it already exists
        
    Returns:
    --------
    winds : netCDF
        A file containing the u and v component of the wind, total wind speed and direction
    """
    
    if isinstance(time, xr.DataArray):
        time = pd.Timestamp(time.to_pandas())
    elif isinstance(time, np.datetime64):
        time = pd.Timestamp(time)
    
    if coords == "global":
        area = [90, -180, -90, 180]
        
    elif isinstance(coords, tuple):
        lat_o, lon_o = coords
        area = [lat_o - 0.251, lon_o - 0.251, lat_o + 0.251, lon_o + 0.251]
    
    if query == "complete":
        era5_filename_lvl, era5_filename_sfc = download_era5_model_lvl(time, area, prefix=prefix,
                                                                       data_path=data_path, overwrite=overwrite,
                                                                       timesteps=timesteps)
        
        weighted_wind = get_weighted_wind_fields(era5_filename_lvl, era5_filename_sfc, data_path=data_path)
        weighted_wind = weighted_wind[["u", "v"]]
        weighted_wind = weighted_wind.transpose(*("latitude", "longitude", "time"))

        wind = xr.Dataset(
            data_vars={
                'U': (['lat', 'lon', 'time'], weighted_wind.u.values, {'long_name': 'GNFR-A weighted u-wind', 'units': 'm/s'}),
                'V': (['lat', 'lon', 'time'], weighted_wind.v.values, {'long_name': 'GNFR-A weighted v-wind', 'units': 'm/s'}),
            },
            coords={
                'lat': weighted_wind['latitude'].values.flatten(),
                'lon': weighted_wind['longitude'].values.flatten(),
                'time': weighted_wind['time'],
            }
        )
        
        description = 'U- and V-wind component vertically weighted for GNFR-A emission profile (Brunner et al. 2019)'

        
    elif query == "pbl_mean":
        era5_filename_lvl, era5_filename_sfc = download_era5_model_lvl(time, area, prefix=prefix,
                                                                       data_path=data_path, overwrite=overwrite,
                                                                       timesteps=timesteps)
        
        area = [90, -180, -90, 180]
        era5_filename_pbl = ddeq.download_S5P.download_era5_pbl(time, area, prefix="", data_path=data_path, overwrite=overwrite, timesteps=24)
        
        era5_pbl = xr.open_dataset(era5_filename_pbl)
        wind_height = compute_height_levels(era5_filename_lvl, era5_filename_sfc, data_path=data_path)
        
        # interpolate pbl to era5 grid
        pbl_subset = era5_pbl.interp(latitude=wind_height.latitude, longitude=wind_height.longitude)
        
        wind = xr.merge([wind_height, pbl_subset[["blh"]]])
        
        # select data below pbl
        wind_pbl = wind.where(wind.h <= wind.blh)
        wind_pbl["blh"] = wind_pbl.blh.mean(dim="level", keep_attrs=True)
        
        wind_pbl = wind_pbl.mean('level')
        wind_pbl = wind_pbl.where(wind_pbl.blh > 400) # only use values where the blh is higher than 400m (based on GNFR-A profile)
        
        wind_pbl = wind_pbl[["u", "v", "blh"]]
        wind_pbl = wind_pbl.transpose(*("latitude", "longitude", "time"))

        wind = xr.Dataset(
            data_vars = {
                'U': (['lat', 'lon', 'time'], wind_pbl.u.values, {'long_name': 'mean u-wind in pbl', 'units': 'm/s'}),
                'V': (['lat', 'lon', 'time'], wind_pbl.v.values, {'long_name': 'mean u-wind in pbl', 'units': 'm/s'}),
                'blh': (['lat', 'lon', 'time'], wind_pbl.blh.values, {'long_name': 'boundary layer height', 'units': 'm'}),
            },
            coords = {
                'lat': wind_pbl['latitude'].values.flatten(),
                'lon': wind_pbl['longitude'].values.flatten(),
                'time': wind_pbl['time'],
            }
        )
        
        description = 'Vertical mean of U- and V-wind component inside the planetary boundary layer'
        
        
    elif query == "pressure_levels":
        era5_filename = download_era5_pressure_lvl(time, area, prefix=prefix,
                                                   data_path=data_path, overwrite=overwrite,
                                                   timesteps=timesteps)
            
        w = xr.open_dataset(era5_filename)

        w = w.mean('level')
        w =  w[["u", "v"]]
        w = w.transpose(*("latitude", "longitude", "time"))

        wind = xr.Dataset(
            data_vars = {
                'U': (['lat', 'lon', 'time'], w.u.values, {'long_name': 'unweighted u-wind', 'units': 'm/s'}),
                'V': (['lat', 'lon', 'time'], w.v.values, {'long_name': 'unweighted v-wind', 'units': 'm/s'}),
            },
            coords = {
                'lat': w['latitude'].values.flatten(),
                'lon': w['longitude'].values.flatten(),
                'time': w['time'],
            }
        )
        
        description = 'Mean of U- and V-wind component of the pressure levels [700, 750, 775, 800, 825, 850, 875, 900, 925, 950, 975, 1000]'
    
    attrs = {
        'CREATOR': 'Sandro Meier',
        'EMAIL': 'sandro.meier@empa.ch',
        'AFFILIATION': 'Empa Duebendorf, Switzerland',
        'DATE_CREATED': pd.Timestamp.now().strftime('%Y-%m-%d %H:%M'),
        'ORIGIN': 'ERA-5',
        'DESCRIPTION': description
    }
    
    wind.attrs.update(attrs)

    wind['speed'] = xr.DataArray(
        np.sqrt(wind.U**2 + wind.V**2),
        attrs = {'units': 'm s-1'}
    )        
    wind['speed_precision'] = xr.DataArray(
        np.full_like(wind.U, 1.0),
        dims = wind.speed.dims,
        attrs = {'units': 'm s-1'}
    )
    wind['direction'] = xr.DataArray(
        ddeq.wind.calculate_wind_direction(wind.U, wind.V),
        attrs = {'units': '°'}
    )
    
    if isinstance(coords, tuple):
        wind = wind.sel(lon=lon_o,
                        lat=lat_o,
                        method='nearest')

    return wind


def download_era5_pressure_lvl(time, area, prefix="", data_path='.', overwrite=False,
                               timesteps=1):
    """
    Download ERA-5 data on pressure levels of a given time step.
    Documentation: https://confluence.ecmwf.int/display/CKB/ERA5%3A+data+documentation
    
    timesteps :: number of hours to download on the given day.
        If 24, a full day will be returned.
    """
    
    if prefix != "":
        prefix = prefix + "_"

    if timesteps == 24:
        era5_filename = os.path.join(data_path,
                                     time.strftime(f'{prefix}ERA5-wind-%Y%m%d.nc'))
    else:
        era5_filename = os.path.join(data_path,
                                     time.strftime(f'{prefix}ERA5-wind-%Y%m%d_%H00.nc'))

    if os.path.exists(era5_filename) and not overwrite:
        return era5_filename
    
    if timesteps == 1:
        t = time.round('H').strftime('%H:00')
    elif timesteps == 24:
        t = [f"{hour:02d}:00" for hour in range(0, timesteps, 1)]
    else:
        t = [f"{hour:02d}:00" for hour in range(
            max(0,time.round('H').hour+1-timesteps),
            time.round('H').hour+1,
            1)]

    cds = cdsapi.Client()

    query = {
        'product_type'     : 'reanalysis',
        'date'             : time.strftime('%Y-%m-%d'),
        'time'             : t,
        'format'           : 'netcdf',
        'variable'         : ['geopotential', 'temperature', 'u_component_of_wind', 'v_component_of_wind'],
        'pressure_level'   : ['700', '750', '775', '800', '825', '850', '875', '900', '925', '950', '975', '1000'],
        'area'             : area, # north, east, south, west
    }

    r = cds.retrieve('reanalysis-era5-pressure-levels', query, era5_filename)
    
    return era5_filename


def download_era5_model_lvl(time, area, prefix = "", data_path='.', overwrite=False,
                            timesteps=1): 
    """
    Download ERA-5 data on model levels of a given time step.
    
    timesteps :: number of hours to download on the given day.
        If 24, a full day will be returned.
    """
    
    if prefix != "":
        prefix = prefix + "_"

    if timesteps == 24:
        era5_filename_lvl = os.path.join(data_path,
                                     time.strftime(f'{prefix}ERA5-lvl-%Y%m%d.nc'))
    else:
        era5_filename_lvl = os.path.join(data_path,
                                     time.strftime(f'{prefix}ERA5-lvl-%Y%m%d_%H00.nc'))

    if os.path.exists(era5_filename_lvl) and not overwrite:
        return era5_filename_lvl, era5_filename_lvl.replace("lvl", "sfc")

    cds = cdsapi.Client()
    
    # data download specifications:
    cls     = "ea"
    expver  = "1"
    levtype = "ml"
    stream  = "oper"
    tp      = "an"
    date    = time.strftime("%Y-%m-%d")
    grid    = [0.25, 0.25]
    area    = area # north, east, south, west
    
    if timesteps == 1:
        time = time.round('H').strftime('%H:00')
    elif timesteps == 24:
        time = [f"{hour:02d}:00" for hour in range(0, timesteps, 1)]
    else:
        time = [f"{hour:02d}:00" for hour in range(
            max(0,time.round('H').hour+1-timesteps),
            time.round('H').hour+1,
            1)]

    # model levels
    query_lvl = {
        'class'   : cls,
        'date'    : date,
        'expver'  : expver,
        'levelist': '100/to/137/by/1', # 1 is top level, 137 the lowest model level in ERA5
        'levtype' : levtype,
        'param'   : '130/131/132/133', # temperature (t), u- and v-wind and specific humidity (q)
        'stream'  : stream,
        'time'    : time,
        'type'    : tp,
        'grid'    : grid, # Latitude/longitude grid: east-west (longitude) and north-south resolution (latitude). Default: 0.25 x 0.25°
        'area'    : area, #example: [60, -10, 50, 2], # North, West, South, East. Default: global
        'format'  : 'netcdf'
    }
        
    cds.retrieve('reanalysis-era5-complete', query_lvl, era5_filename_lvl)
    
    query_sfc = {
            'class'   : cls,
            'date'    : date,
            'expver'  : expver,
            'levelist': '1',       # Geopotential (z) and Logarithm of surface pressure (lnsp) are 2D fields, archived as model level 1
            'levtype' : levtype,
            'param'   : '129/152', # Geopotential (z) and Logarithm of surface pressure (lnsp)
            'stream'  : stream,
            'time'    : time,
            'type'    : tp,
            'grid'    : grid, # Latitude/longitude grid: east-west (longitude) and north-south resolution (latitude). Default: 0.25 x 0.25°
            'area'    : area, 
            'format'  : 'netcdf'
        }
    
    cds.retrieve('reanalysis-era5-complete', query_sfc, era5_filename_lvl.replace("lvl", "sfc"))
    
    
    return era5_filename_lvl, era5_filename_lvl.replace("lvl", "sfc")

def compute_height_levels(filename_lvl, filename_sfc, data_path="."):
    """
    Compute the geometric height of the ERA5 model levels using the hypsometric equation
    
    filename_lvl : str
        File name of the file containing the ERA5 data on model levels
    filename_sfc : str
        File name of the file containing the ERA5 data on the surface
    """
        
    lvl = xr.open_dataset(filename_lvl)
    sfc = xr.open_dataset(filename_sfc)
    
    era5 = xr.merge([lvl, sfc])
    era5["ps"] = np.exp(sfc['lnsp'])

    a,b = np.loadtxt(os.path.join(data_path,'L137_model_level_definitions.dat'), usecols=(1,2), unpack=True, skiprows=1)
    lower_value = np.nanmin(era5.level.values.astype(int))-1
    era5["a"] = (("level_bound"), a[np.insert(era5.level.values.astype(int), 0, lower_value)], {"units": "Pa"})
    era5["b"] = (("level_bound"), b[np.insert(era5.level.values.astype(int), 0, lower_value)], {"units": 1})

    # calculate pressure
    era5["p_bound"] = era5.a + era5.b * era5.ps
    era5["p_bound"] = era5["p_bound"].transpose(*('time', 'level_bound', 'latitude', 'longitude'))
    era5["p_mid"] = (era5.t.dims, 0.5 * (era5.p_bound.values[:,1:,:,:] + era5.p_bound.values[:,:-1,:,:]), {"units": "Pa"})

    # calculate virtual temperature
    era5["t_v"] = era5.t * (1. + 0.609133 * era5.q)
    era5["t_v"].attrs = {"units": "K",
                         "long_name": "virtual_temperature"}

    # calculate thickness of layer with hypsometric equation
    era5["dh"] = R * era5.t_v / g * np.log(era5.p_bound.values[:,1:,:,:] / era5.p_bound.values[:,:-1,:,:])
    era5["dh"].attrs = {"units": "m", "long_name": "thickness_of_model_level"}
    era5["h"] = era5["dh"][:,::-1,:,:].cumsum(dim='level') # height at upper boundary of level
    era5["h"] = era5.h-era5.dh/2 # height at the middle of the level
    era5["h"].attrs = {"units": "m", "long_name": "height_of_model_level"}
    
    return era5

def func(s):
    # get values on model level boundaries
    if bounds[0] <= s < bounds[-1]:
        i = np.searchsorted(bounds, s, side='right')
        return values[i-1]
    else:
        return 0.0


def get_weighted_wind_fields(filename_lvl, filename_sfc, data_path='.'):
    """
    Compute a vertically averaged wind speed based on a GNFR-A weighted wind
    """
    # compute heights
    wind = ddeq.wind.compute_height_levels(filename_lvl, filename_sfc, data_path=data_path)


    # compute weights (by quick interpolation)
    z2 = np.linspace(0,1000, 501)
    func_vec = scipy.vectorize(func)
    w2 = func_vec(z2)
    ifunc = scipy.interpolate.interp1d(z2, w2, bounds_error=False, fill_value=0.0)

    w2 = ifunc(wind.h.values)
    wind["weights"] = ((wind.h.dims), w2)
    wind["weights"] = wind.weights / wind.weights.sum(dim="level")

    # computed vertically-weighted winds
    wind["u"] = (wind.u * wind.weights).sum(dim="level") / wind.weights.sum(dim="level")
    wind["v"] = (wind.v * wind.weights).sum(dim="level") / wind.weights.sum(dim="level")
    wind["u"].attrs = {"units":"m/s", "long_name": "GNFR-A weighted u-wind"}
    wind["v"].attrs = {"units":"m/s", "long_name": "GNFR-A weighted v-wind"}
    
    return wind
