#!/usr/bin/env python
#
# grib_checker.py checks Checks whatever you want to check in your life
#
# AUTHOR: BSC-CNS I. CERNICHARO
#

import os
import sys
import datetime
import time
import logging
import json
import warnings
import numpy as np
import eccodes as ecc

VERSION = "ALPHA 0.3"

###############################################################################
# DATA CHECKS FOR GRIB CONTENT CLASS

def check_data_values(msg):
    v_b = return_values_bounds(msg)
    # max_,min_,mean_,v_max,v_min,v_mean = v_b
    
    if v_b[0] == v_b[3] and v_b[1] == v_b[4] and v_b[2] == v_b[5]:
        return "OK"

    else:
        return "ERROR"

def check_coords(msg):
    c_b = return_coords_bounds(msg)
    # max_lat,min_lat,max_lat1,min_lat1,max_lon,min_lon,max_lon1,min_lon1 = c_b
    
    if c_b == None or c_b == "Error":
        return "ERROR"

    if c_b[4] == c_b[6] and c_b[5] == c_b[7]:
        long_cck = "OK"
    else:
        long_cck = "ERROR"

    if c_b[0] == c_b[2] and c_b[1] == c_b[3]:
        lat_cck = "OK"
    else:
        lat_cck = "ERROR"

    if lat_cck == "OK" and long_cck == "OK": 
        return "OK"
    else: 
        return "ERROR"

###############################################################################
# METADATA CHECKS FOR GRIB CONTENT CLASS
def check_metada_keys(msg):
    keys = return_metadata(msg)

    # Metadata
    fold_name = os.path.join(
        os.getcwd(),
        "utils/grib_checker",
        "structure_files/grib_keys.json"
    )
    with open(fold_name,"r") as j_doc:
        grib_meta = json.load(j_doc)
    j_doc.close()

    if "GRIBEditionNumber" in keys:
        ge = msg["GRIBEditionNumber"]
        grib_meta_man = grib_meta["gribEdition"][str(ge)]["mandatory"]
        #grib_meta_opt = grib_meta["gribEdition"][str(ge)]["optional"]

        counter = 0
        no_available = list()
        for item in grib_meta_man:
            if item not in keys:
                counter += 1
                no_available.append(item)
        return ["OK"] if counter == 0 else ["ERROR",no_available]

    else: return ["ERROR","GRIBEditionNumber"]

def check_grid_keys(msg):
    keys = return_metadata(msg)
    
    # Grid Information Metadata
    fold_name = os.path.join(
        os.getcwd(),
        "utils/grib_checker",
        "structure_files/grid_types_geography.json"
    )
    with open(fold_name,"r") as j_doc:
        grid_data = json.load(j_doc)
    j_doc.close()

    if "gridType" in keys:
        gt = msg["gridType"]

        if gt in grid_data.keys():
            msg_grid_keys = grid_data[gt]

            counter = 0
            no_available = list()
            for item in msg_grid_keys:
                if item not in keys:
                    counter += 1
                    no_available.append(item)
            return ["OK"] if counter == 0 else ["ERROR",no_available]

    else: return ["ERROR",["gridType"]]

###############################################################################ç
# GRIB MESSAGE CONTENT MANAGER

def return_coords_bounds(msg):
    keys = msg.keys()
    
    lat0 = 'latitudeOfFirstGridPointInDegrees'
    lat1 = 'latitudeOfLastGridPointInDegrees'
    lon0 = 'longitudeOfFirstGridPointInDegrees'
    lon1 = 'longitudeOfLastGridPointInDegrees'
    
    conditional = lon0 in keys and lat0 in keys and lon1 in keys and lat1 in keys

    if conditional:
        lat = msg['latitudes']
        bounds_lat = [
            msg[lat0],
            msg[lat1]
        ]

        lon = msg['longitudes']
        bounds_lon = [
            msg[lon0],
            msg[lon1]
        ]

        return [
            max(bounds_lat),
            min(bounds_lat),
            max(lat),
            min(lat),
            max(bounds_lon),
            min(bounds_lon),
            max(lon),
            min(lon),
        ]           

    else:
        gridtype = msg["gridType"]
        
        if gridtype == "lambert":
            return None
            
        else:
            return "Error"
    
def return_values_bounds(msg):
    vals = msg['values']
    return [
        msg[ 'maximum'],
        msg['minimum'],
        msg['average'],
        max(vals),
        min(vals),
        np.mean(vals),
    ]

def return_metadata(msg):
    keys = msg.keys()
    return keys

###############################################################################
# FUNCTION TO RUN THE GRIB_CHECKER

def grib_cck_run(args):
    time0 = time.time()
    
    # Initiate the grib checks creating the grib_content object
    g_c = grib_content(args)

    # Start adding messages to the log
    string = "\n\tDATE AND TIME: " + datetime.datetime.now().isoformat() 
    string += "\n\tGRIB_CCK VERSION: " + VERSION 
    string += "\n\tFILENAME: " + g_c.pname
    logging.info(string)
    g_c.print_num_keys()
    
    # Metada checks iteration
    g_c.check_meta_grid()
    
    # Data checks iteration
    # Not useful right now. Have to be redefined
    # g_c.check_data_coords()

    dtime = time.time() - time0

    str_time = str(int(dtime//3600)) + ' hours, '
    str_time += str(int((dtime%3600)//60)) + ' minutes, '
    str_time += str(round((dtime%3600)%60,3)) + ' seconds '

    string = '\n\tTIME ELAPSED: "%s\n' %str_time

    logging.info(string)

###############################################################################
# GRIB CONTENT CLASS TO MANAGE THE CHECKS TO BE PERFORMED

class grib_content:
    """grib_content object class set a list of checks and properties where
    the content of a grib file will be evaluated. All it needs to initiate
    is the file path of the GRIB file we will evaluate
    """
    def __init__(self,filename):
        self.name = filename
        self.grib_msgs = self.return_all_msg()
        self.pname = filename.split("/")[-1]
        self.log = ''

    def return_all_msg(self):
        # This funtion returns all messages within a GRIB file
        all_mssg = list()
        content = ecc.GribFile(self.name)
        
        for i in range(len(content)):
            all_mssg.append(ecc.GribMessage(content))

        return all_mssg


    def print_num_keys(self):
        # It returns the file name without the path
        HEAD = "#NUM MSGs: "
        if len(HEAD) != 24: 
            HEAD += (24-len(HEAD))*" "
        mssg = HEAD +str(len(self.grib_msgs))+" messages"
        logging.info(mssg)

    def check_data_coords(self):
        # The following lists will be filled with the log report for each
        # message. Coordinates and data values will be tested
        coords_log = list()
        data_log = list()
        index = list()

        # Checking all messages. Here we call methods that are imported from
        # data_checks submodule
        for i,msg in enumerate(self.grib_msgs):
            coords_log.append(check_coords(msg))
            data_log.append(check_data_values(msg))
            index.append(i)
        
        err_coords = coords_log.count("ERROR")
        err_data = data_log.count("ERROR")

        if err_coords != 0:
            HEAD = "#VAR_COORDS >> WARNING: "
            if len(HEAD) != 24: 
                HEAD += (24-len(HEAD))*" "
            msg_coord = "\n" + "One or more messages have incorrect data"
            logging.warning(msg_coord)
        else:
            HEAD = "#VAR_COORDS >> INFO: "
            if len(HEAD) != 24: 
                HEAD += (24-len(HEAD))*" "
            msg_coord = "\n" + "Everything OK with coordinates"
            logging.info(msg_coord)

        if err_data != 0:
            HEAD = "#VAR_DATA >> WARNING: "
            if len(HEAD) != 24: 
                HEAD += (24-len(HEAD))*" "
            msg_data = "\n" + "One or more messages have incorrect data"
            logging.warning(msg_data)
        else:
            HEAD = "#VAR_DATA >> INFO: "
            if len(HEAD) != 24: 
                HEAD += (24-len(HEAD))*" "
            msg_data = "\n" + "Everything OK with data values"
            logging.info(msg_data)

    def check_meta_grid(self):
        # The following lists will be filled with the log report for each
        # message. Metadata of the header and of the grid description section
        # will be tested.
        meta_log = list()
        grid_log = list()
        index = list()
        # These lists, g_ed and grid_t will check the homogeneity of the 
        # messages for the GRIBEditionNumber and gridType keys.
        g_ed = list()
        grid_t = list()
        missing_meta = list()
        missing_grid = list()

        # Checking all messages. Here we call methods that are imported from
        # metadata_checks submodule

        for i,msg in enumerate(self.grib_msgs):

            index.append(i)

            g_ed.append(msg["GRIBEditionNumber"])
            grid_t.append(msg["gridType"])

            meta_cck = check_metada_keys(msg)
            grid_cck = check_grid_keys(msg)

            if len(meta_cck) > 1:
                for item in meta_cck[1]:
                    if item not in missing_meta:
                        missing_meta.append(item)

            if len(grid_cck) > 1:
                for item in grid_cck[1]:
                    if item not in missing_grid:
                        missing_grid.append(item)

            meta_log.append(meta_cck[0])
            grid_log.append(grid_cck[0])

        grib_editions = list(set(g_ed))
        grid_types = list(set(grid_t))
        
        err_meta = meta_log.count("ERROR")
        err_grid = grid_log.count("ERROR")

        if len(grib_editions) > 1:
            # There are more than 1 edition in gribEdition keys for GRIB messages
            if 1 not in grib_editions:
                # If edition 1 not in
                msg_editions = "\n" + "Grib Edition Check:"
                msg_editions += "\n" + "tFound different Grib editions in the messages"
                for item in grib_editions:
                    msg_editions += "\n" + "\t\tEdition: " + str(item)
                logging.warning(msg_editions)
            else:
                msg_editions = "\n" + "Grib Edition Check:"
                msg_editions += "\n" + "tFound different Grib editions in the messages"
                for item in grib_editions:
                    msg_editions += "\n" + "\t\tEdition: " + str(item)
                wrn = "Deprecated GRIB1 version found"
                msg_editions += "\n\t" + wrn
                warnings.warn(wrn)
                logging.warning(msg_editions)
        else:
            if 1 not in grib_editions:
                msg_editions = "\n" + "Grib Edition Check:"
                msg_editions += " Everything OK, one unique edition found"
                msg_editions += " GRIB" + str(grib_editions[0])
                logging.info(msg_editions)
            else:
                msg_editions = " Grib Edition Check: "
                msg_editions += "Deprecated GRIB1 version found"
                logging.warning(msg_editions)
                warn_msg = "#KEYS-GRIB_ED >> WARNING: " + msg_editions
                wrn = "Deprecated GRIB1 version found"
                warnings.warn(wrn)

        if len(grid_types) > 1:
            msg_grid_type = "\n" + "Grid Type Check:"
            msg_grid_type += "\n" + "\tFound different Grid Type in the messages"
            for item in grid_types:
                msg_grid_type += "\n" + "\t\tGrid type: " + str(item)
            logging.warning(msg_grid_type)
        else:
            msg_grid_type = "\n" + "Grid Type Check:"
            msg_grid_type += "\n" + "\tEverything OK, one unique Grid Type found"
            msg_grid_type += "\n" + "\t\tGrid type: " + str(grid_types[0])
            logging.info(msg_grid_type)

        if err_meta != 0:
            msg_meta = "\n" + "Metadata Check:"
            msg_meta += "\n" + "\tOne or more messages have missing KEYS"
            for item in missing_meta:
                msg_meta += "\n" + "\t\tKey: " + str(item)
            logging.warning(msg_meta)
        else:
            msg_meta = "\n" + "Metadata Check:"
            msg_meta += "\n" + "tEverything OK"
            logging.info(msg_meta)

        if err_grid != 0:
            msg_grid = "\n" + "Grid Metada Check:"
            msg_grid += "\n" + "\tOne or more messages have missing Keys"
            for item in missing_grid:
                msg_grid += "\n" + "\t\tKey: " + str(item)
            logging.warning(msg_grid)
        else:
            msg_grid = "\n" + "Grid Metada Check:"
            msg_grid += "\n" + "\tEverything OK with data values"
            logging.info(msg_grid)

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

    for item in args:
        grib_cck_run(item)
