#!/usr/bin/env python
#
# This script checks every single key for GRIB files within the destination
# folder you have typed in the command line
#
# AUTHOR: Barcelona Supercomputing Center-Centro Nacional de Supercomputacion
#         Earth Sciences Department
#         I. CERNICHARO: ivan.cernicharo@bsc.es
#

import os
import sys
import datetime
import json

from eccodes import GribFile

class LOG:

    def __init__(self):

        self.log = ""
        self.date = datetime.datetime.now().date().isoformat().replace("-","")

    def save_log(self):
        
        f_log = open('LOG_GRIB_KEYS-' + self.date + '.txt','w')
        f_log.write(self.log)
        f_log.close()


def list_grib_files_tree(path):
    """Function to list all GRIB files of a tree directory
    """
    grib_files = []
    
    for root, dirs, files in os.walk(path):
            for f in files:
                if (f.endswith("grib") or f.endswith("grb")) and "efas" not in f and "ecv" not in f:
                    grib_files.append(os.path.join(root,f))
    
    grib_files.sort()

    string = ""

    for i,f in enumerate(grib_files):
        string += "\n" + str(i)+ '> ' + f
    
    return string,grib_files

def list_grib_files_folder(path):
    """Function to list all GRIB files of a folder
    """
    grib_files = []
    
    for f in os.listdir(path):
        if (f.endswith("grib") or f.endswith("grb")) and "efas" not in f and "ecv" not in f:
            grib_files.append(os.path.join(path,f))
    
    grib_files.sort()

    string = ""

    for i,f in enumerate(grib_files):
        string += "\n" + str(i)+ '> ' + f
    
    return string,grib_files

class grib_mssg:
    
    def __init__(self,msg):
        self.msg = msg
        self.attrcolection = msg.keys()
        
    def return_keys(self):
        return self.attrcolection
    
def map_keys_gfiles(grib_files):
    """Map all available GRIB files within a folder or tree directory
    and extracts all structure keys for each message within each
    GRIB file listed
    """
    grib_keys = []

    for f in grib_files:
        g_f = GribFile(f)
        k_keys = []
        for i,m in enumerate(g_f):

            if i == 0:
                g_v = m["GRIBEditionNumber"]
                g_c = grib_mssg(m)
            if "seas_for" in f:
                type_file = "seasonal forecast"
            if "reanalysis" in f or f.endswith("grb"):
                if "uerra" in f:
                    type_file = "reanalysis uerra"
                else:
                    type_file = "reanalysis era5"
            g_keys = g_c.return_keys()

            for kk in g_keys:
                if kk not in k_keys: 
                    k_keys.append(kk)

        g_f.close()
        grib_keys.append([k_keys,type_file,g_v])
        
    return grib_keys

def check_distinct_keys(grib_keys):
    """This function takes as input a list of lists of grib keys and analyze which
    are the common keys for seasonal, reanalysis, ERA5, UERRA and returns as an
    output a dictornary with these keys. It also returns a LOG of what has been found
    """
    
    common_keys,non_common_keys,seasonal_keys,reanalysis_keys,all_keys = [],[],[],[],[]
    common_era5_uerra,non_common_era5_uerra,r_era5_keys,r_uerra_keys = [],[],[],[]
    common_seasonal,non_common_seasonal = [[],[]]
    g1,g2 = [],[]
    
    for ll in grib_keys:
        for a in ll[0]:

            counter,counter_r,counter_s = 0,0,0

            for kl in grib_keys:

                if a not in kl[0]:
                    counter += 1

                if "reanalysis" in kl[1] and "reanalysis" in ll[1] and a not in kl[0]:
                    counter_r += 1

                if "seasonal" in kl[1] and "seasonal" in ll[1] and a not in kl[0]:
                    counter_s += 1

            if a not in g1 and ll[2] == 1:
                g1.append(a)

            if a not in g2 and ll[2] == 2:
                g2.append(a)

            if counter == 0 and a not in common_keys:  
                common_keys.append(a)

            if counter != 0 and a not in non_common_keys:  
                non_common_keys.append(a)

            if counter_r == 0 and a not in common_era5_uerra and "reanalysis" in ll[1]:
                common_era5_uerra.append(a)

            if counter_r != 0 and a not in non_common_era5_uerra and "reanalysis" in ll[1]:
                non_common_era5_uerra.append(a)

            if counter_s == 0 and a not in common_seasonal and "seasonal" in ll[1]:
                common_seasonal.append(a)

            if counter_s != 0 and a not in non_common_seasonal and "seasonal" in ll[1]:
                non_common_seasonal.append(a)

            if 'seasonal forecast' in ll[1] and a not in seasonal_keys:
                seasonal_keys.append(a)

            if 'reanalysis' in ll[1] and a not in reanalysis_keys:
                reanalysis_keys.append(a)

            if ('era5' in ll[1] or 'ecv' in ll[1]) and a not in r_era5_keys:
                r_era5_keys.append(a)

            if 'uerra' in ll[1] and a not in r_uerra_keys:
                r_uerra_keys.append(a)

            if a not in all_keys:
                all_keys.append(a)
    
    # Saving LOG of number of distinct keys
    string = "\n" + 60*"#" + "\n"
    string += "\nLOG GRIB KEYS FOR SEASONAL AND REANALYSIS\n"
    string += "\n" + 60*"-" + "\n"
    string += "\n\t||GENERAL KEYS||\n"
    string += "\nAll Keys: %s \n" %len(all_keys)
    string += "\nCommon Keys: %s \n" %len(common_keys)
    string += "\nNon common keys for reanalysis and seasonal forecast: %s \n" %len(non_common_keys)
    string += "\n" + 60*"-" + "\n"
    string += "\n\t||SEASONAL||\n"
    string += "\nKeys for seasonal: %s \n" %len(seasonal_keys)
    string += "\nCommon Seasonal Keys: %s \n" %len(common_seasonal)
    string += "\nNon common Seasonal Keys: %s \n" %len(non_common_seasonal)
    string += "\n" + 60*"-" + "\n"
    string += "\n\t||REANALYSIS||\n"
    string += "\nKeys for reanalysis: %s \n" %len(reanalysis_keys)
    string += "\nReanalysis ERA5 Keys: %s \n" %len(r_era5_keys)
    string += "\nReanalysis UERRA Keys: %s \n" %len(r_uerra_keys)
    string += "\nCommon ERA5 and UERRA keys: %s \n" %len(common_era5_uerra)
    string += "\nNon common ERA5 and UERRA keys: %s \n" %len(non_common_era5_uerra)
    string += "\nCommon ERA5 and UERRA keys are the same as \nGENERAL common keys above? %s \n" %(common_era5_uerra == common_keys)
    string += "\n" + 60*"-" + "\n"
    string += "\n\t||BY GRIB EDITION||\n"
    string += "\nKeys for GRIB1: %s \n" %len(g1)
    string += "\nKeys for GRIB2: %s \n" %len(g2)
    string += "\n" + 60*"#"
        
    mapped_keys = {
        "general":{
            "all":all_keys,
            "common":common_keys
        },
        "seasonal":{
            "all":seasonal_keys,
            "common":common_seasonal
        },
        "reanalysis":{
            "all":reanalysis_keys,
            "era5":r_era5_keys,
            "uerra":r_uerra_keys
        },
        "by_grib_edition":{
            "1":g1,
            "2":g2
        }
    }
    
    return string,mapped_keys

def save_to_json(dictionary_name):
    
    string = ""
    for item in dictionary_name.keys():
        json_name = "GRIB_KEYS-" + item + ".json"
        
        string += "\n>Saving %s as %s\n" %(item,json_name)

        with open(json_name,"w") as fp:
            json.dump(dictionary_name[item],fp,indent=4,sort_keys=True)

        fp.close()

    string += "\n\n****FINISHED****\n"
    return string


def main(args):
    # Create a LOG class
    script_log = LOG()

    if len(args) > 1:
        if args[1] == "tree":
            # List all available GRIB file in the tree directory
            args = args[0]
            string,grib_files = list_grib_files_tree(args)
            script_log.log += string
        elif args[1] == "folder":
            # List all available GRIB file in the folder
            args = args[0]
            string,grib_files = list_grib_files_folder(args)
            script_log.log += string

    print(string)

    # Map all files
    grib_keys = map_keys_gfiles(grib_files)

    # Analyze the grib keys above
    string,mapped_keys = check_distinct_keys(grib_keys)
    script_log.log += string
    print(string)

    # Save keys above
    string = save_to_json(mapped_keys)
    script_log.log += string
    script_log.save_log()
    print(string)

if __name__ == "__main__":
    args = sys.argv[1:]

    print("\nMapping all GRIB files in %s \n" %args)

    sys.exit(main(args))
