
from datetime import datetime, timedelta
import glob
import os
import time

from pyproj import Proj
import gdal
import netCDF4
import numpy as np
import scipy.interpolate as sint
import scipy.ndimage as ndi

import smartcarb as sc


"""\
Script for:
    - creating daily MODIS albedo at given band/wavelength for a MODIS "region"
    - project MODIS data to rotpole grid using data from different regions (if required)
"""

def normalised_convolution(values, std, truncate=4.0, keep_mask=False):

    if not hasattr(values, 'mask'):
        values = np.ma.array(values, mask=~np.isfinite(values))

    values = values.copy()
    mask = values.mask.copy()
    dtype = values.dtype

    certainty = 1.0 - mask
    values = values.data.copy()

    values[mask] = 0.0
    new_values = ndi.gaussian_filter(values, std, truncate=truncate)  / ndi.gaussian_filter(certainty, std, truncate=truncate)

    if keep_mask:
        return np.ma.array(new_values, mask=mask, dtype=dtype)
    else:
        return np.ma.array(new_values, mask=np.isnan(new_values), dtype=dtype)


def fill_gaps_and_smooth(a):

    # smooth in time (8 days) and fill gaps
    anew = normalised_convolution(a, [8.0, 0.0, 0.0])

    # fill remaining gaps by linear spatial interpolation
    n_days, ny, nx = anew.shape
    x,y = np.meshgrid(np.arange(nx), np.arange(ny))

    for day in range(n_days):
        mask = ~np.isfinite(anew[day])
        afill = normalised_convolution(anew[day], [1.0, 1.0], truncate=10.0)
        anew[day, mask] = afill[mask]

        if np.any(np.isnan(afill)):
            print(day, 'has still nans')

    return anew




def lonlat2rotpole(lon, lat):
    c = ccrs.RotatedPole(sc.cfg.pollon, sc.cfg.pollat)
    res = c.transform_points(ccrs.PlateCarree(), lon, lat)[:,:,:]
    rlon, rlat = res[:,:,0], res[:,:,1]
    return rlon ,rlat



def compute_lonlat(ds):
    """\
    Convert MODIS sinusoidal projection to WGS-84.
    """
    p_modis_grid = Proj('+proj=sinu +R=6371007.181 +nadgrids=@null +wktext')

    width = ds.RasterXSize
    height = ds.RasterYSize

    gt = ds.GetGeoTransform()

    xmin = gt[0]
    ymin = gt[3] + width * gt[4] + height * gt[5]
    xmax = gt[0] + width * gt[1] + height * gt[2]
    ymax = gt[3]

    x = np.linspace(xmin, xmax, width+1)
    x = 0.5 * (x[1:] + x[:-1])

    y = np.linspace(ymax, ymin, height+1)
    y = 0.5 * (y[1:] + y[:-1])

    x,y = np.meshgrid(x,y)

    lon, lat = p_modis_grid(x,y, inverse=True)

    return lon, lat




def read_scale_and_offset(ds):
    """\
    Read scale (default 1) and offset (default 0) from HDF4 dataset.
    """
    scale = ds.GetMetadataItem('scale_factor')
    if scale is None:
        scale = 1
    else:
        scale = float(scale)

    offset = ds.GetMetadataItem('add_offset')
    if offset is None:
        offset = 0
    else:
        offset = float(offset)

    return scale, offset



def read_field(filename, name):
    """\
    Read dataset with `name` from HDF4 file. Returns
    center longitudes and latitudes, and field.
    """
    # open file and find dataset
    f = gdal.Open(filename)
    d = f.GetSubDatasets()
    sub_datasets = f.GetSubDatasets()
    names = [s[0].split(':')[-1] for s in sub_datasets]
    key = sub_datasets[names.index(name)][0]

    # read dataset
    ds = gdal.Open(key)

    # corner lon/lat
    lon, lat = compute_lonlat(ds)

    # field metadata
    fill_value = int(ds.GetMetadataItem("_FillValue"))
    scale, offset = read_scale_and_offset(ds)

    # read, scale and mask field
    field = ds.ReadAsArray()
    field = np.ma.array((field - offset) * scale, mask=field==fill_value)

    return lon, lat, field



def read_and_append(filenames, name):

    fields = []
    masks = []

    for filename in filenames:
        # load field
        lon, lat, field = read_field(filename, name)
        rlon, rlat = sc.misc.wgs84_to_rotpole2(lon, lat, sc.cfg.pollon, sc.cfg.pollat)

        # corners of COSMO grid
        rlonc, rlatc = sc.cfg.rlonc, sc.cfg.rlatc

        # bin to COSMO grid
        valids = ~field.mask
        H, _, _ = np.histogram2d(rlat[valids], rlon[valids], bins=[rlatc, rlonc],
                                 weights=field[valids])
        N, _, _ = np.histogram2d(rlat[valids], rlon[valids], bins=[rlatc, rlonc])

        fields.append(H / N)
        masks.append(N==0)

    return np.ma.array(fields, mask=np.array(masks))



def bin_single_day(date, name, path='/input/MODIS/MCD43A3.006/'):

    # read all data fields 
    pattern = 'MCD43A3.A%s.h??v??.006.*.hdf' % date.strftime('%Y%j')
    filenames = glob.glob(os.path.join(path, pattern))

    fields = read_and_append(filenames, name)
    field = fields.mean(0)

    return field


def fill_gaps(filename, name, newname):

    with netCDF4.Dataset(filename, 'a') as nc:
        a = nc.variables[name][:]
        a = fill_gaps_and_smooth(a)

        var = nc.createVariable(newname, a.dtype, ('time', 'rlat', 'rlon'))
        var[:] = a


def main(filename, start, stop, name):

    # TODO: netCDF compression

    study = 'ESA project - SMARTCARB'
    user_name = 'Gerrit Kuhlmann'
    user_email = 'gerrit.kuhlmann@empa.ch'
    user_affiliation = 'Empa Duebendorf, Switzerland'

    attrs = {
            'DESCRIPTION':  'MODIS surface reflectance on SMARTCARB domain',
            'DATAORIGIN':   'MODIS product MCD43A3.006 ',
            'CREATOR':      user_name,
            'EMAIL':        user_email,
            'AFFILIATION':  user_affiliation,
            'VERSION':      '0.1',
            'DATE CREATED': time.ctime(time.time()),
            'STUDY':        study,
    }
    times = []

    for i, date in enumerate(sc.misc.iter_times(start, stop, timedelta(days=1))):

        seconds = (date - datetime(2015,1,1)).total_seconds()
        times.append(seconds)

        field = bin_single_day(date, name)
        field = field.astype('f4')

        j = int(date.strftime('%j')) - 1
        print(date.strftime('%Y-%m-%d'), 'index =', j)

        if i == 0 and not os.path.exists(filename):
            sc.io.init_netcdf(filename, sc.cfg.rlon, sc.cfg.rlat, attrs=attrs)


        with netCDF4.Dataset(filename, 'a') as nc:
            if name in nc.variables:
                var = nc.variables[name]
            else:
                # TODO attributes
                var = nc.createVariable(name, field.dtype, ('time', 'rlat', 'rlon'))
            var[j] = field

    # add times
    with netCDF4.Dataset(filename, 'a') as nc:
        var = nc.variables['time']
        var[:] = np.array(times)



if __name__ == '__main__':

    filename = 'COSMO_2D_MCD43A3.006_WSA.nc'
    names = ['Albedo_WSA_Band2', 'Albedo_WSA_Band6']

    for name in names:
        main(filename, datetime(2015,1,1), datetime(2015,12,31), name)
        fill_gaps(filename, name, name + '_filled')










