#!/usr/bin/env python3

try:
    import argparse
    import sys
    import os
    import glob
    import time
    from scipy.interpolate import griddata
    import numpy as np
    import netCDF4 as nc4
    from netCDF4 import Dataset
    import matplotlib.pyplot as plt
    from matplotlib.colors import LogNorm
    import matplotlib.cm as cm
    import seaborn as sns
    from termcolor import colored

    print(colored("All modules loaded!\n", 'green'))
except ModuleNotFoundError:
    print(colored("Module not found: %s"%ModuleNotFoundError, 'red'))

folder='./incoming/L2__NO2___/Poland/'
output_dir='./incoming/croped/Poland/'

#obszar interpolacji
lon_pl_min = 13.0
lon_pl_max = 26.0
lat_pl_min = 48.0
lat_pl_max = 55.0
krok_siatki = 0.025

qa_threshold = 0.7

 
DEBUG = True
CONVERT_TO_Pmolec_cm2 = True

def set_parser():
    """ set custom parser """
    
    parser = argparse.ArgumentParser(description="")
    parser.add_argument("-df", "--date_from", type=str, required=True, 
                        help="Sensing date from YYYYMMDD ")                        
    parser.add_argument("-dt", "--date_to", type=str, required=True, 
                        help="Sensing date to YYYYMMDD ")   
    parser.add_argument("-s", "--save_each_raster", type=bool,  default=False,
                        help="Shall I produce output for each timestep? ")   
    return parser

def process_file(ncfile,PRODUCT_NAME = 'nitrogendioxide_tropospheric_column',method='nearest'):
    global grid_x, grid_y 
    fh = Dataset(ncfile)
    #loaging variables from netCDFa
    lons = (fh.groups['PRODUCT'].variables['longitude'][:][0,:,:])
    lats = (fh.groups['PRODUCT'].variables['latitude'][:][0,:,:])
    product = (fh.groups['PRODUCT'].variables[PRODUCT_NAME][0,:,:])
    qa = (fh.groups['PRODUCT'].variables['qa_value'][0,:,:])
    
    #2D arrays -> vectors
    lonsa = np.array(lons).flatten() 
    latsa = np.array(lats).flatten() 
    producta = np.array(product).flatten()
    qaa = np.array(qa).flatten()
    
    #generating target grid
    grid_x, grid_y = np.mgrid[lon_pl_min:lon_pl_max:krok_siatki, lat_pl_min:lat_pl_max:krok_siatki]
    points=np.squeeze(np.array(list(zip([lonsa,latsa])))).T
    
    #actual interpolation
    grid_product = griddata(points, producta, (grid_x, grid_y), method=method,fill_value=-1)
    grid_qa = griddata(points, qaa, (grid_x, grid_y), method=method,fill_value=-1)
    
    #rotating north up
    grid_product = (np.rot90(grid_product))
    grid_qa = (np.rot90(grid_qa))

    #masking low quality values
    too_low_quality = (grid_qa<qa_threshold)
    grid_product[too_low_quality]=-1

    return grid_product, grid_qa
    

def save_to_PNG(output_raster,nazwa_pliku_wynikowego):
    global grid_x, grid_y 
    if 'Basemap' not in sys.modules:
        from mpl_toolkits.basemap import Basemap
    
    lat_c=(lat_pl_min+lat_pl_max)/2
    lon_c=(lon_pl_min+lon_pl_max)/2
    fig = plt.figure(figsize=(10, 10), edgecolor='k')
    m = Basemap(width=1300000,height=1000000,
                resolution='h',projection='tmerc',\
                lat_ts=40,lat_0=lat_c,lon_0=lon_c,
                llcrnrlat=lat_pl_min-0.1, llcrnrlon=lon_pl_min-0.1,
                urcrnrlat=lat_pl_max+0.1, urcrnrlon=lon_pl_max+0.1)

    xi, yi = m(grid_x.flatten(), grid_y.flatten())


    # Masking bad values
    masked_array = np.ma.array (output_raster, mask=np.isnan(output_raster))
    cmap1 = cm.jet
    #cmap1 = cm.gist_ncar
    #

    if CONVERT_TO_Pmolec_cm2:
        cs = m.pcolor(xi.reshape( [grid_y.shape[1],grid_y.shape[0]],order='F'),
            yi.reshape( [grid_y.shape[1],grid_y.shape[0]],order='F'),
            np.flip(masked_array *6.02214*1e4,axis=0), cmap=cmap1,vmin=0.0,vmax=10.0)
        
        # color scale
        cbar = m.colorbar(cs, location='bottom', pad="10%")
        cbar.set_label('$NO_2$ tropospheric column $[10^{15}mol/cm^2]$')
        cmap1.set_bad('white',6.02214*1e4)
    else:
        # Drawing map in RAW units [mol/m2]
        cs = m.pcolor(xi.reshape( [grid_y.shape[1],grid_y.shape[0]],order='F'),
            yi.reshape( [grid_y.shape[1],grid_y.shape[0]],order='F'),
            np.flip(masked_array,axis=0), cmap=cmap1,vmin=0.0,vmax=0.00005)

        # color scale
        cbar = m.colorbar(cs, location='bottom', pad="10%")
        cbar.set_label('mol/m2')
        cmap1.set_bad('white',1.)

    # Siatka
    m.drawparallels(np.arange(-80., 81., 10.), labels=[1,0,0,0], fontsize=10)
    m.drawmeridians(np.arange(-180., 181., 10.), labels=[0,0,0,1], fontsize=10)

    # Wybrzeża, kraje etc.
    m.drawcoastlines()
    m.drawstates()
    m.drawcountries()

   
    # tytuł i zapis
    plt.title(nazwa_pliku_wynikowego)
    fig.tight_layout()
    plt.savefig(output_dir+'png_output/'+nazwa_pliku_wynikowego+'.png',dpi=300,box_layout='tight')

def save_to_GeoTiff(output_array, output_filename):
    try:
        if 'gdal' not in sys.modules:
            from osgeo import gdal
        if 'gdal_array' not in sys.modules:
            from osgeo import gdal_array
        if 'osr' not in sys.modules:        
            from osgeo import osr
    except ModuleNotFoundError:
        print(colored("Gdal module not found: %s"%ModuleNotFoundError, 'red'))

    nrows,ncols = np.shape(output_array)
    xres = (lon_pl_max-lon_pl_min)/float(ncols)
    yres = (lat_pl_max-lat_pl_min)/float(nrows)
    geotransform=(lon_pl_min,xres,0,lat_pl_max,0, -yres) 

    #creating output file for every step if requested
    output_raster = gdal.GetDriverByName('GTiff').Create(output_dir+'tif_output/'+output_filename+'.tif',ncols, nrows, 1 ,gdal.GDT_Float64)  
    output_raster.SetGeoTransform(geotransform) 
        
    #creating CRS info
    srs = osr.SpatialReference()                
    srs.ImportFromEPSG(4326)                    
    output_raster.SetProjection( srs.ExportToWkt() )   

    #actual save
    if CONVERT_TO_Pmolec_cm2:
        output_raster.GetRasterBand(1).WriteArray(output_array*6.02214*1e4)   
        output_raster.FlushCache()

    else:
        output_raster.GetRasterBand(1).WriteArray(output_array)   
        output_raster.FlushCache()


def main():
    mean_raster = None
    counter = 0
    #tego chwilowo nie używamy
    #liczba_plikow_do_usredniania = 2
    os.system('mkdir -p '+output_dir+'tif_output/')
    os.system('mkdir -p '+output_dir+'png_output/')
    
    parser = set_parser()
    options = parser.parse_args()
    
    #TODO: DateRane support
    for ncfile in sorted(glob.glob(folder+"*.nc")):
        
        try:
            date_str = ncfile.split('_')[9]
        except IndexError:
            date_str = ncfile.split('T')[0][-8:]
            time_str = ncfile.split('T')[1][:4]
        
        if ((int(date_str)>=int(options.date_from)) and (int(date_str)<=int(options.date_to))):
            print("Processing data from date: "+str(date_str)+' time: '+str(time_str))
            
            start_time = time.time()
            result_array, result_qa = process_file(ncfile,method='linear')

            #Maska z True tam gdzie wartości mają sens
            bad_values = np.logical_or(result_array<0, result_array>1)
            result_array[bad_values] = np.nan
            nans = np.sum(np.isnan(result_array))
            total = result_array.shape[0]*result_array.shape[1]
            nans_percentage = nans/total*100
            
            if DEBUG:
                
                print("\t Processing took  "+"{:.2f}".format(time.time()-start_time)+" seconds")
                print("\t There are "+"{:.2f}".format(nans_percentage)+" % nans in result array ")

            if options.save_each_raster:
                #Saving current step
                save_to_PNG(result_array,'no2_QA_'+str(qa_threshold)+'_'+date_str+"T"+time_str)
                save_to_GeoTiff(result_array,'no2_QA_'+str(qa_threshold)+'_'+date_str+"T"+time_str)
                
            if nans_percentage <90:
                if mean_raster is None: 
                    result_array[bad_values] = np.nan
                    mean_raster = result_array
                else:
                    #Tam gdzie mamy braki pomiarów, uzupełniamy dotychczasową średnią
                    result_array[bad_values] = mean_raster[bad_values]
                    #sposób na liczenie średniej w miarę kolejnych dni
                    mean_raster = np.add(counter*mean_raster , result_array)/(counter+1)

                counter = counter +1
                
    
    save_to_GeoTiff(mean_raster,'no2_QA_'+str(qa_threshold)+"_average_"+options.date_from+'_'+options.date_to)
    save_to_PNG(mean_raster,"no2_QA_"+str(qa_threshold)+"_average_"+options.date_from+'_'+options.date_to)
    
    
if __name__ == "__main__":
    main()        