#!/usr/bin/env python
"""
mmarvin 20200221
(mmarvin@ed.ac.uk)

This Python script processes GFED4.1s files for input to GEOS-Chem.
It can be used to process standard emissions, as well as 
preliminary 'beta' emissions.

To start: enter year, whether this year requires beta emissions, 
whether this is a leap year, and the path to the original hdf5 
GFED4.1s source files.

"""

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

#Select options
yyyy = '2017'
isbeta = True
isleap = False
path2hdf5 = "./source_gfed_hdf5"

months = ['01','02','03','04','05','06','07','08','09','10','11','12']

#Open yearly HDF5 file
if isbeta == True: 
    fn = path2hdf5+"/GFED4.1s_"+yyyy+"_beta.hdf5"
else:
    fn = path2hdf5+"/GFED4.1s_"+yyyy+".hdf5" 

hf = h5py.File(fn, 'r')

lat = hf['lat']
lat = np.flip(lat[:,0])
nlat = np.size(lat)

lon = hf['lon']
lon = lon[0,:]
nlon = np.size(lon)

#Create monthly NetCDF files
for mm in months:
    
    if mm == '04' or mm == '06' or mm == '09' or mm == '11':
        dd = 30
    elif mm == '02'and isleap == True:
        dd = 29
    elif mm == '02' and isleap == False:
        dd = 28
    else:
        dd = 31
    
    time_mm = pd.date_range(yyyy+'-'+mm+'-01',yyyy+'-'+mm+'-01')
    time_dd = pd.date_range(yyyy+'-'+mm+'-01',yyyy+'-'+mm+'-'+str(dd))
    time_hh = pd.date_range(yyyy+'-'+mm+'-01T00:00:00',yyyy+'-'+mm+'-01T21:00:00',periods=8)
    
    #Read total DM emissions: kg/m2/month => kg/m2/s
    data = hf['emissions/'+mm+'/DM']
    data_dm = np.flip(data[:,:],axis=0).reshape(1,nlat,nlon)/(60*60*24*dd)
    
    #Calculate partitioned emissions from total DM
    data = hf['emissions/'+mm+'/partitioning/DM_AGRI']
    data_frac = np.flip(data[:,:],axis=0).reshape(1,nlat,nlon)
    data_agri = data_dm*data_frac
    
    data = hf['emissions/'+mm+'/partitioning/DM_BORF']
    data_frac = np.flip(data[:,:],axis=0).reshape(1,nlat,nlon)
    data_borf = data_dm*data_frac
    
    data = hf['emissions/'+mm+'/partitioning/DM_DEFO']
    data_frac = np.flip(data[:,:],axis=0).reshape(1,nlat,nlon)
    data_defo = data_dm*data_frac
    
    data = hf['emissions/'+mm+'/partitioning/DM_SAVA']
    data_frac = np.flip(data[:,:],axis=0).reshape(1,nlat,nlon)
    data_sava = data_dm*data_frac
    
    data = hf['emissions/'+mm+'/partitioning/DM_PEAT']
    data_frac = np.flip(data[:,:],axis=0).reshape(1,nlat,nlon)
    data_peat = data_dm*data_frac
    
    data = hf['emissions/'+mm+'/partitioning/DM_TEMF']
    data_frac = np.flip(data[:,:],axis=0).reshape(1,nlat,nlon)
    data_temf = data_dm*data_frac
    
    #Read daily scaling factors
    daily_fraction = np.zeros([dd,nlat,nlon])
    for ndays in np.arange(dd):
        day_string = str(ndays+1)
        frac = hf['emissions/'+mm+'/daily_fraction/day_'+day_string]
        dfrac = np.flip(frac[:,:],axis=0).reshape(1,nlat,nlon)
        daily_fraction[ndays,:,:] = dfrac*dd
        
    daily_fraction = np.where(np.sum(daily_fraction,axis=0) > 0,daily_fraction,1.0)
        
    #Read diurnal scaling factors
    hours = ['UTC_0-3h','UTC_3-6h','UTC_6-9h','UTC_9-12h','UTC_12-15h','UTC_15-18h','UTC_18-21h','UTC_21-24h']
    diurnal_fraction = np.zeros([8,nlat,nlon])
    for hrs in np.arange(len(hours)):
        frac = hf['emissions/'+mm+'/diurnal_cycle/'+hours[hrs]]
        dfrac = np.flip(frac[:,:],axis=0).reshape(1,nlat,nlon)
        diurnal_fraction[hrs,:,:] = dfrac*8
        
    diurnal_fraction = np.where(np.sum(diurnal_fraction,axis=0) > 0,diurnal_fraction,1.0)
    
    #Creat NetCDF datasets
    data_file = 'GFED4_gen.025x025.'+yyyy+mm+'.nc'
    dailyfrac_file = 'GFED4_dailyfrac_gen.025x025.'+yyyy+mm+'.nc'
    diurnalfrac_file = 'GFED4_3hrfrac_gen.025x025.'+yyyy+mm+'.nc'
    
    ds = xr.Dataset(coords={'time':time_mm,'lat':lat,'lon':lon})

    # "time", "lat", "lon"
    ds['time'].attrs['standard_name'] = 'time'
    ds['time'].attrs['axis'] = 'T'
    
    ds['lat'].attrs['standard_name'] = 'latitude'
    ds['lat'].attrs['long_name'] = 'latitude'
    ds['lat'].attrs['units'] = 'degrees_north'
    ds['lat'].attrs['axis'] = 'Y'

    ds['lon'].attrs['standard_name'] = 'longitude'
    ds['lon'].attrs['long_name'] = 'longitude'
    ds['lon'].attrs['units'] = 'degrees_east'
    ds['lon'].attrs['axis'] = 'X'
    
    #Assign data variables to dataset
    ds_data = ds.copy()
    ds_data['DM_TOTL'] = (['time','lat','lon'], data_dm)
    ds_data['DM_AGRI'] = (['time','lat','lon'], data_agri)
    ds_data['DM_BORF'] = (['time','lat','lon'], data_borf)
    ds_data['DM_DEFO'] = (['time','lat','lon'], data_defo)
    ds_data['DM_SAVA'] = (['time','lat','lon'], data_sava)
    ds_data['DM_PEAT'] = (['time','lat','lon'], data_peat)
    ds_data['DM_TEMP'] = (['time','lat','lon'], data_temf)
    
    ds_data = ds_data.astype('float32')
    for v in ds_data.data_vars.keys():
        ds_data[v].attrs['units'] = 'kg/m2/s'
        
    #Assign daily scale factors to dataset    
    ds_daily = ds.copy()
    ds_daily.coords['time'] = time_dd
    ds_daily['GFED_FRACDAY'] = (['time','lat','lon'], daily_fraction)
    ds_daily['GFED_FRACDAY'].attrs['units'] = '1'
    ds_daily = ds_daily.astype('float32')
    
    #Assign diurnal scale factors to dataset
    ds_diurnal = ds.copy()
    ds_diurnal.coords['time'] = time_hh
    ds_diurnal['GFED_FRAC3HR'] = (['time','lat','lon'], diurnal_fraction)
    ds_diurnal['GFED_FRAC3HR'].attrs['units'] = '1'
    ds_diurnal = ds_diurnal.astype('float32')

    #Create NetCDF files
    ds_data.to_netcdf(data_file)
    ds_daily.to_netcdf(dailyfrac_file)
    ds_diurnal.to_netcdf(diurnalfrac_file)