#!/usr/bin/env python
#
# Data Quality Checker also known as DQC
#
# Read the README.md for more detailed information
#
#       Barcelona Supercomputing Center - Centro Nacional de Supercomputacion
#       Earth Sciences Department - Computational Earth Sciences
#     
#     Authors: 
#       Joan Sala Calero (joan.sala@bsc.es)

import requests
import yaml, json
from datetime import datetime

# Local imports
from dqc_wrapper.utils.utils_grib import *

"""
Given a dataset/variable a.k.a QAR
Return list of mandatory dates to find
"""
def get_forms_dates(dataset, variable, origin, system, forms_dir):
    constraints = {}
    
    # Load constraints file into dictionary
    try:            
        with open(os.path.join(forms_dir, dataset, 'constraints.json'), 'r') as cdf:
            constraints = json.load(cdf)                
    except:
        raise Exception('Problem loading constraints from cds-forms directory. Please check your configuration')

    # In particular dates for a given (dataset, variable)
    md = []
    try:
        dates = []
        for c in constraints:

            # Filter by variable
            if variable in c['variable']:

                # Seasonal needs filter by origin/system
                if 'seasonal' in dataset:
                    if not(origin in c['originating_centre']) or not(system in c['system']):
                        continue

                # Parse all dates
                for y in c['year']:
                    for m in c['month']:
                        if 'day' in c:
                            for d in c['day']:
                                dates.append(datetime(int(y), int(m), int(d)))
                        else:
                            dates.append(datetime(int(y), int(m), int(1)))

        # Get unique ordered values
        md = sorted(list(set(dates)))
    except Exception as e:
        print(e)
        raise Exception('Problem loading constraints from cds-forms directory. Please check your dataset name')

    return md

"""
Given a cds variable name (long-name)
Return a parameter identifier in the mapping.json
"""
def get_param_id(dataset, variable, forms_dir):
    mapping = {}
    code = None

    # Load constraints file into dictionary
    try:            
        with open(os.path.join(forms_dir, dataset, 'mapping.json'), 'r') as cdf:
            mapping = json.load(cdf)                
    except:
        raise Exception('Problem loading mapping from cds-forms directory. Please check your configuration')
 
    # Fetch given variable code
    try:            
        code = mapping['remap']['variable'][variable]              
    except:
        raise Exception('Could not find variable {} in dataset {}'.format(variable, dataset))


    return code

"""
Given a ECMWF url parameter database
Find shortname in the HTML (get them from the ECMWF parameter databse a.k.a web)
Function Not used
"""
def get_shortname_url(url):

    r = requests.get(url, stream=True)
    next_line = False

    for line in r.iter_lines():

        # Try to parse variable name
        l=str(line)
        if next_line:
            l = l[55:]
            l = l.replace('").appendTo(row);\'', '')
            return l
        
        if 'Short Name' in l:
            next_line = True
        else:
            next_line = False

    return ''

"""
Given a code return grib shortname for that variable
GRIB1 try first
"""
def get_shortname_grib(grib1_def, grib2_def, code):
    try:
        if code in grib1_def:
            return grib1_def[code]['cf']
        if code in grib2_def:
            return grib2_def[code]['cf']
    except:
        return None

"""
Given a cds dataset (long-name)
Return variables and shortnames (get them from the ECMWF parameter databse a.k.a web)
"""
def get_dataset_information(dataset, forms_dir, grib1_def, grib2_def):
    mapping = {}
    form = {}
    code = None

    # Load mapping/form file into dictionary
    try:            
        with open(os.path.join(forms_dir, dataset, 'mapping.json'), 'r') as cdf:
            mapping = json.load(cdf)
        with open(os.path.join(forms_dir, dataset, 'form.json'), 'r') as cdff:
            form = json.load(cdff)
    except:
        raise Exception('Problem loading mapping/form from cds-forms directory. Please check your configuration')

    # Load variables yaml file into dictionary (some in abstract.yaml, some in other yaml file)
    try:
        yaml_file = os.path.join(forms_dir, dataset, 'variables.yaml')
        if not(os.path.exists(yaml_file)):
            yaml_file = os.path.join(forms_dir, dataset, 'variablesanddescription.yaml')
        if not(os.path.exists(yaml_file)):
            yaml_file = os.path.join(forms_dir, dataset, 'abstract.yaml')
    except:
        raise Exception('Problem loading variables.yaml from cds-forms directory. Please check your configuration')

    # Try to load yaml
    with open(yaml_file, 'r') as cdyml:
        yvar = yaml.load(cdyml)

    # Get c3s variables info
    c3svinfo = get_variables_info_c3s(form)

    # Fetch given variables
    try:            
        var_dict = mapping['remap']['variable']             
    except:
        raise Exception('Could not find variables in dataset {}'.format(dataset))

    # Fetch product types if available
    prod_types = []
    if 'product_type' in mapping['remap']:
        prod_types = list(mapping['remap']['product_type'].values())
    else:
        if 'type' in mapping['force']:
            prod_types = mapping['force']['type']

    var_codes = dict()
    for k,c in var_dict.items(): 
        try:
            var_codes[k] = {
                'paramId' : c,
                'shortName': get_shortname_grib(grib1_def, grib2_def, c),
                'c3sname': c3svinfo[k],
                'c3sunit': yvar['variables'][c3svinfo[k]]['units']
            }
            if var_codes[k] is None:
                # Get Shortname from the website if not found
                url = 'https://apps.ecmwf.int/codes/grib/param-db/?id={}'.format(int(c))
                code = get_shortname_url(url)
                var_codes[k] = code
        except ValueError:
            var_codes[k] = c
        except KeyError:           
            print('Key not found, could not find shortname for variable={}, dataset={}'.format(k,dataset))

    return var_codes, prod_types

"""
Given a combination of dataset/variable
Return GRIB metadata available
"""
def get_metadata(forms_dir, dataset, variable):

    grib_info = load_grib_metadata()
    paramId = get_param_id(dataset, variable, forms_dir)

    try:
        return grib_info[str(paramId)]
    except:
        return { "varname": variable, "unit": "n/a", "grib_codes": {} }    

"""
Simulate Xpath in python dictionary
"""
def xpath_get(mydict, path):
    elem = mydict
    try:
        for x in path.strip("/").split("/"):
            elem = elem.get(x)
    except:
        pass
    return elem

"""
Return variables section from a forms.json cds-c3s-forms file
"""
def get_variables_info_c3s(form):
    kvars = {}
    for ff in form:
        if ff['label'] == 'Variable':
            # Depending if variable is grouped or not
            elems = xpath_get(ff, '/details/groups/values')
            if elems is None:
                kvars = xpath_get(ff, '/details/labels')
            else:
                for e in elems:
                    for ky,vl in e['labels'].items():
                        kvars[ky] = vl

    return kvars
