
import glob
import os

import numpy as np
import pandas as pd
import ucat
import xarray as xr


def update_vcds(data, path, gas='NO2', max_iter=10):
    """
    Update vertical column densities (VCD) of `gas` by recomputing air mass
    factors (AMF) for Sentinel-5P/TROPOMI data. The method computes
    the mean concentration in the boundary layer from ERA-5 boundary layer
    height and TROPOMI `gas` enhancement above the background
    (i.e. `gas` - `gas`_estimated_background).

    The mean concentration is used to update the TM5 a priori profiles in
    the boundary layer and to compute the air mass factor using the
    averaging kernels. Since updating the AMFs will increase the enhancement,
    iteration is used till convergence.

    The method will only update AMFs for pixels that have been detected as
    enhanced above the background by the plume detection algorithm.

    The function will recompute {gas}_minus_estimated_background and the mass
    columns from the updated `gas` VCDS.
    """

    # if not valid data at detected pixels do not update
    if np.all(np.isnan(data[gas].values[data.is_hit])):
        return data

    time = pd.Timestamp(data.attrs['time'])
    pattern = time.strftime(
        'S5P_OPER_AUX_CTMANA_%Y%m%dT000000_????????T000000_????????T??????.nc'
    )
    # TODO: check if taken last makes sense
    aux_filename = sorted(glob.glob(os.path.join(path, pattern)))[-1]

    # open aux filename and interpolate to remote sensing image
    aux = xr.open_dataset(aux_filename)
    aux = aux.interp(lon=data.longitude, lat=data.latitude,
                     time=pd.Timestamp(data.attrs['time']),
                     method='linear')
    aux = aux.rename(lev='layer')

    # constants
    from scipy.constants import g
    R = 287.052874  # specific gas constant for dry air
    GAMMA = -0.0065 # lapse rate

    # pressure levels in middle and at interface
    pressure_mid = aux.hyam + aux.hybm * aux.ps
    pressure_int = aux.hyai + aux.hybi * aux.ps

    # correct surface pressure for topography (Zhou et al. 2009)
    ps_eff = aux.ps * (aux.t[0] / (aux.t[0] + GAMMA * (aux.surface_altitude - data.surface_altitude))) ** (-g / R / GAMMA)

    # effective pressure levels in middle and at interface
    pressure_mid_eff = aux.hyam + aux.hybm * ps_eff
    pressure_int_eff = aux.hyai + aux.hybi * ps_eff

    # layer thickness and mid height (TODO: use virtual temperature)
    dh = R * aux.t / g * np.log(pressure_int_eff[:-1].values / pressure_int_eff[1:].values)
    dh = xr.where(np.isfinite(dh), dh, 0.0)
    h = dh.cumsum(dim='layer') - 0.5 * dh

    # profile correction for topography
    topo_correction = (pressure_int_eff[:-1].values - pressure_int_eff[1:].values) / (pressure_int[:-1].values - pressure_int[1:].values)
    values = aux[gas.lower()] * topo_correction

    # Averaging kernels above tropopause to zero for tropospheric AMF
    if gas == 'NO2':
        amf_varname = 'air_mass_factor_troposphere'
        AK = xr.where(data.averaging_kernel.layer >= aux.tropopause_layer_index,
                      data.averaging_kernel, data.averaging_kernel)
    else:
        amf_varname = 'air_mass_factor'
        AK = data.averaging_kernel.copy()
    AMF = data[amf_varname].copy()

    # Convert TM5 to molar concentration (TODO: account for humidity)
    TM5_profiles = ucat.convert_points(values, 'mol mol-1', 'mol m-3',
                                       p=pressure_mid_eff, T=aux.t)

    data = data.rename(**{gas: f'{gas}_standard',
                          amf_varname: f'{amf_varname}_standard'})

    # Iterative update AMFs and VCDs
    VCD_current = data[f'{gas}_standard'].copy()

    for i in range(max_iter):

        # Compute concentration in boundary layer (mol / m³)
        inside_blh = (VCD_current - data[f'{gas}_estimated_background']) / data.blh

        profile = xr.where((h <= data.blh) & data.is_hit,
                           inside_blh, TM5_profiles)

        # Update AMFs and VCDs
        AMF_update = AMF * (AK * profile * dh).sum('layer') / (profile * dh).sum('layer')
        VCD_update = AMF * data[f'{gas}_standard'] / AMF_update

        change = np.nanmean(
            np.abs(VCD_update.values[data.is_hit] - VCD_current.values[data.is_hit])
            / VCD_current.values[data.is_hit]
        )
        if change <= 0.01:
            break

        VCD_current, VCD_update = VCD_update, None

    else:
        raise ValueError(f'Maximum number of iterations ({max_iter}) without convergence.')

    # Add updated VCDs and AMFs to remote sensing dataset
    data[gas] = xr.DataArray(
        VCD_update,
        dims=data[f'{gas}_standard'].dims,
        attrs=data[f'{gas}_standard'].attrs
    )
    data[gas].attrs['comment'] = 'AMF correction applied assuming well-mixed profile for pixels with enhanced columns.'

    data[amf_varname] = xr.DataArray(
        AMF_update,
        dims=data[f'{amf_varname}_standard'].dims,
        attrs=data[f'{amf_varname}_standard'].attrs
    )
    data[amf_varname].attrs['comment'] = 'AMF correction applied assuming well-mixed profile for pixels with enhanced columns.'


    # Update fields from plume detection and pre-processing
    # (i.e. {gas}_minus_estimated_background and *_mass)
    data[f'{gas}_minus_estimated_background'] = data[gas] - data[f'{gas}_estimated_background']

    data[f'{gas}_mass'] = xr.DataArray(
            ucat.convert_columns(data[gas], 'mol m-2', 'kg m-2', molar_mass=gas),
            dims=data[f'{gas}_mass'].dims, attrs=data[f'{gas}_mass'].attrs
    )
    data[f'{gas}_minus_estimated_background_mass'] = xr.DataArray(
            ucat.convert_columns(data[f'{gas}_minus_estimated_background'],
                                 'mol m-2', 'kg m-2', molar_mass=gas),
            dims=data[f'{gas}_minus_estimated_background_mass'].dims,
            attrs=data[f'{gas}_minus_estimated_background_mass'].attrs
    )

    return data