#!/usr/bin/env python3

import os,sys,gc
import glob, errno
import numpy as np
import time
import calendar
from netCDF4 import Dataset
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap, cm, maskoceans
from datetime import datetime, timedelta
from matplotlib import ticker
from flexpart.esonc.get_data_bwd import Get_data_bwd
from create_netCDF import make_nc3d

# ========== SETTINGS ==========
dirstat = 'COVID_BC/FLEXOUT/'
stat = '004/'
sp = 'spec_40'

# ============= READ FLEXPART OUT =================
outdir = '/xnilu_wrk/projects/ATMO-ACCESS/OUTDIR/RANDOM/MOSAIC/'
foot = np.zeros((0,360,720))
first_iteration = True
dat = []
for fname in sorted(glob.glob(outdir + '/*')):
    dat.append(datetime.strptime(fname.split('_')[-2], "%Y%m"))
    if first_iteration:
        lons = Get_data_bwd(fname,'001')[31]
        lats = Get_data_bwd(fname,'001')[32]
        first_iteration = False
    foo = Get_data_bwd(fname,'001')[0]
    foot = np.concatenate((foot, np.mean(foo[:],axis=0)[None,:,:]),axis=0)
    print('Accumulated shape of foot variable: ',foot.shape)
    del foo

### SAVE NETCDF ###
print('Store in netCDF file '+'FOOTPRINT_MONAVE.nc' + "foot variable of shape ",foot.shape)
make_nc3d('FOOTPRINT_MONAVE.nc',dat[0].year,lons[0],lons[-1],len(lons),lats[0],lats[-1],len(lats),foot,dat,'MANUAL','footprint',1,1,'seconds', misv = -999.)
del foot
gc.collect()