import xarray as xr
import glob,time
import numpy as np
import pandas as pd
from ILAMB.constants import mid_months,bnd_months
import os
from netCDF4 import Dataset

""" The observational data cannot be downloaded automatically. You
will need to register with the FluxCom Data Portal:

https://www.bgc-jena.mpg.de/geodb/projects/Home.php

and then with the username and password you will receive, download the
Artificial Neural Network (ANN) data found in the 'remote_source'
variable below.

In addition to copying the data into a single netCDF file for use in
ILAMB, we have found a number of values mostly near coastlines which
are hard zeros for all times defined in the dataset. We believe these
to be artifacts of the neural network and not true zeros. Thus we
apply an additional mask to areas which are a hard zero (less than
1e-15) for all times defined.
"""
def CheckFileList(V):
    v0 = V[0].split(".")
    y0 = int(v0[-2])
    vf = V[-1].split(".")
    yf = int(vf[-2])
    v0[-2] = "****"
    assert (yf-y0+1)==len(V)
    return ".".join(v0)
    
for cf_name,var_name in zip(["gpp","reco","hfss","hfls","rns"],
                            ["GPP","TER" ,"H"   ,"LE"  ,"Rn" ]):
    V = sorted(glob.glob("%s*.nc" % var_name))
    if len(V) == 0: continue
    remote_source = CheckFileList(V)
    dset = xr.concat([xr.open_dataset(f) for f in V],dim="time")
    data = np.ma.masked_invalid(dset[var_name])
    mask = np.ones(data.shape[0],dtype=bool)[:,np.newaxis,np.newaxis]*(np.abs(data)<1e-15).all(axis=0)
    data = np.ma.masked_array(data,mask=data.mask+mask)
    download_stamp = time.strftime('%Y-%m-%d', time.localtime(os.path.getmtime(V[0])))
    generate_stamp = time.strftime('%Y-%m-%d')
    dt = pd.to_datetime(dset.time.data)
    t  = np.array([(t.year-1850)*365+mid_months[t.month-1] for t in dt])
    tb = np.array([[(t.year-1850)*365+bnd_months[t.month-1],
                    (t.year-1850)*365+bnd_months[t.month]] for t in dt])
    lat = dset.lat
    lon = dset.lon
    with Dataset("%s.nc" % cf_name, mode="w") as oset:

        # dimensions
        oset.createDimension("time", size = t.size)
        oset.createDimension("lat", size = lat.size)
        oset.createDimension("lon", size = lon.size)
        oset.createDimension("nb", size = 2)
        
        # time
        T = oset.createVariable("time", t.dtype, ("time"))
        T[...] = t
        T.units = "days since 1850-01-01 00:00:00"
        T.calendar = "noleap"
        T.bounds = "time_bounds"
        
        # time bounds
        TB = oset.createVariable("time_bounds", t.dtype, ("time", "nb"))
        TB[...] = tb
    
        # latitude
        X = oset.createVariable("lat", lat.dtype, ("lat"))
        X[...] = lat
        X.standard_name = "latitude"
        X.long_name = "site latitude"
        X.units = "degrees_north"
    
        # longitude
        Y = oset.createVariable("lon", lon.dtype, ("lon"))
        Y[...] = lon
        Y.standard_name = "longitude"
        Y.long_name = "site longitude"
        Y.units = "degrees_east"
    
        # data
        missing = data[data.mask].data[0]
        D = oset.createVariable(cf_name, data.dtype, ("time", "lat", "lon"), zlib=True, fill_value=missing)
        D[...] = data
        D.units = dset[var_name].units.replace("gC","g")
        with np.errstate(invalid='ignore'):
            D.actual_range = np.asarray([data.min(),data.max()])

        oset.title = "FLUXCOM (RS+METEO) Global Land Carbon Fluxes using CRUNCEP climate data"
        oset.version = "1"
        oset.institutions = "Department Biogeochemical Integration, Max Planck Institute for Biogeochemistry, Germany"
        oset.source = """Data generated by Artificial Neural Networks and forced with CRUNCEPv6 meteorological data and MODIS (RS+METEO)
ftp://ftp.bgc-jena.mpg.de/pub/outgoing/FluxCom/CarbonFluxes_v1_2017/RS+METEO/CRUNCEPv6/raw/monthly/
ftp://ftp.bgc-jena.mpg.de/pub/outgoing/FluxCom/EnergyFluxes/RS_METEO/member/CRUNCEP_v8/monthly/"""
        oset.history = """
%s: downloaded %s;
%s: converted to netCDF, additionally we apply a mask where |var|<1e-15 for all time.""" % (download_stamp, remote_source, generate_stamp)
        oset.references  = """
@ARTICLE{Jung2019,
  author = {Jung, M., S. Koirala, U. Weber, K. Ichii, F. Gans, Gustau-Camps-Valls, D. Papale, C. Schwalm, G. Tramontana, and M. Reichstein},
  title = {The FLUXCOM ensemble of global land-atmosphere energy fluxes},
  journal = {Scientific Data},
  year = {2019},
  volume = {6},
  issue = {1},
  page = {74},
  doi = {https://doi.org/10.1038/s41597-019-0076-8}
}
@ARTICLE{Tramontana2016,
  author = {Tramontana, G., M. Jung, C.R. Schwalm, K. Ichii, G. Camps-Valls, B. Raduly, M. Reichstein, M.A. Arain, A. Cescatti, G. Kiely, L. Merbold, P. Serrano-Ortiz, S. Sickert, S. Wolf, and D. Papale},
  title = {Predicting carbon dioxide and energy fluxes across global FLUXNET sites with regression algorithms},
  journal = {Biogeosciences},
  year = {2016},
  number = {13},
  page = {4291-4313},
  doi = {https://doi.org/10.5194/bg-13-4291-2016}
}"""
        oset.comments = """
time_period: %d-%02d through %d-%02d; temporal_resolution: monthly; spatial_resolution: 0.5 degree; units: %s""" % (dt[0].year,dt[0].month,dt[-1].year,dt[-1].month,D.units)
        oset.convention = "CF-1.8"

