#!/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)
#       Julian Rodrigo Berlin (julian.berlin@bsc.es)
#       Alejandro Alfonso Hernandez (alejandro.alfonso@bsc.es)

"""
Script: Generate a list of QAR ini configuration files
---
- Read from download lists
- Generate ini file to be used by the checker
- how to run:
  cd c3s512-wp1-datachecker
  python dqc_wrapper/QAR_dqc/generate_QARs_config.py
"""

import os
import json
import glob
import csv

from dqc_wrapper.utils.utils_forms import get_forms_dates


def generate_ini(rootd, pattern, variable_e, datatype, variable, units, res, gt, gmn,
                 ts, origin, system, dataset, is_s, dr, tg, com, lsm, lnd, vmin, vmax, cmsname):

    ini='''[general]
input = /shared/cds_downloads/{r}/{ds}/{v}
fpattern = {p}
log_dir = /shared/QAR_dqc/logs 
res_dir = /shared/QAR_dqc/results
forms_dir = /shared/c3s_repos/cds-forms-c3s
masks_dir = /shared/cds_downloads/masks
cms_name = {cms}

[dataset]
variable = {ve}
datatype = {dt}
cds_dataset = {ds}
cds_variable = {v}
cds_units = {un}
complete = {com}
origin = {o}
system = {s}

[file_format]

[standard_compliance_grib]

[standard_compliance_netcdf]

[spatial_completeness]
land_sea_mask = {lsm}
land_variable = {lnd}

[temporal_completeness]
is_seasonal = {is_s}

[spatial_consistency]
grid_interval = {res}
grid_type = {gt}
grid_mapping_name = {gmn}
dual_resolution = {dr}

[temporal_consistency]
time_step = {ts}
time_granularity = {tg}

[valid_ranges]
valid_min = {vmin}
valid_max = {vmax}
'''.format(
    p=pattern, ve=variable_e, dt=datatype, v=variable, is_s=is_s, dr=dr, gt=gt, com=com, gmn=gmn, un=units, cms=cmsname,
    o=origin, s=system, ds=dataset, ts=ts, r=rootd, res=res, tg=tg, lsm=lsm, lnd=lnd, vmin=vmin, vmax=vmax)

    return ini

"""
Get tolerance limits
"""
def get_tolerance_limits(tolerance_file):
    # Read and skip comments
    var_min_max = dict()
    with open(tolerance_file, 'r') as fp:
        rdr = csv.reader(filter(lambda row: row[0]!='#', fp))
        for row in rdr:
            if len(row):
                var_min_max[row[0]] = row[1:3]

    fp.close()
    return var_min_max

"""
Get temporal resolution (if needed) + land_sea_mask
"""
def get_temporal_res(var, temp_resolution):
    var = var.replace('_anomaly', '') # postprocessed
    for v in temp_resolution:
        variable, s, g, l = v.split(',')
        if variable.strip() == var.strip():
            return s, g, l

    return None, None, 'no'

"""
Get grid resolution of corresponding variable (for product_type=reanalysis),
return -1 if var is not found in the file
Note: if there are new variables in the CDS, the file should be updated accordingly
"""
def get_era5_sl_res(var, mask_variables):
    for v in mask_variables:
        variable, res, _, _ = v.split(',')
        if variable.rstrip() == var.rstrip():
            return res
    return -1

"""
Get lsm and lnd of corresponding variable, return None, None if var is not found in the file
Note: if there are new variables in the CDS, the file should be updated accordingly
"""
def get_era5_sl_lsm_lnd(var, mask_variables):
    for v in mask_variables:
        variable, _, lsm, lnd = v.split(',')
        if variable.rstrip() == var.rstrip():
            result = lsm, lnd
    return None, None

"""
Init parameters for a reanalysis dataset
"""
def init_params_reanalysis(dl, d, v, vinfo, mask_variables, tolerance_limits, daily, uerra, land):
    root = 'reanalysis'

    # Get variable info
    ve = vinfo['shortName']
    u = vinfo['c3sunit']

    # Some units we know how to set limits
    if u == '%':
        vmin = 0.0
        vmax = 100.0
    elif u == '(0 - 1)':
        vmin = 0.0
        vmax = 1.0
    elif u.lower() == 'degrees':
        vmin = 0.0
        vmax = 360.0
    else:
        vmin = ''
        vmax = ''

    # Try to get tolerance limits (WMO or based on world records)
    try:
        vmin = float(tolerance_limits[v][0].strip())
    except:
        vmin = ''
    try:
        vmax = float(tolerance_limits[v][1].strip())
    except:
        vmax = ''

    # Model distinction
    try:
        o = dl['form']['origin']  # uerra has model distinction
    except:
        o = ''

    # Time resolution granularity (complete or not, split dates or not)
    if daily:
        ts = 3
        com = 'no'
        tg = 'hour'
    else:
        ts = 1
        com = 'yes'
        tg = 'month'
    if uerra:
        if 'mescan' in o:
            ts = 12
        else:
            ts = 6
    if land:
        ts = 1
        if daily:
            tg = 'hour'
            com = 'no'

    # Grid resolution (two resolutions allowed) - 0.25 and 0.5
    gmn = ''
    if uerra:
        res = 11000
        if 'mescan' in o:
            res = 5500
        dr = 'no'
        gt = 'projection'
        gmn = 'lambert_conformal_conic'

        patt = '*.grib'
        if o != '':
            patt = '{}-*.grib'.format(o)
    else:
        model = 'era5'
        res = 0.25
        dr = 'yes'
        gt = 'lonlat'
        if land:
            model = 'era5_land'
            res = 0.1
            dr = 'no'
        patt = '{}-*.grib'.format(v)

    # Product type
    dt = 'an, fc, em'

    # Get lsm and lnd flags and also update grid_res for era5-single-levels if necessary
    if land:
        # era5-land
        lsm = 'yes'
        lnd = 'yes'

    elif uerra:
        # uerra-height, uerra-pressure
        lsm = 'no'
        lnd = 'no'
        if 'soil' in d:
            # uerra-soil
            lsm = 'yes'
            lnd = 'yes'
        if 'single' in d:
            # Only var surface_roughness has lsm (for uerra-single-levels)
            if v == 'surface_roughness':
                lsm = 'yes'
                lnd = 'yes'
    else:
        # era5-pressure-levels
        lsm = 'no'
        lnd = 'no'

        if 'single' in d:
            # Get lsm, lnd and res for era5-single-levels
            has_lsm, is_lnd = get_era5_sl_lsm_lnd(v, mask_variables)
            v_res = get_era5_sl_res(v, mask_variables)
            if v_res != -1:
                lsm = has_lsm
                lnd = is_lnd
                res = v_res
            else:
                # case when variable is not found
                print('ERROR: variable {} not found in dqc_wrapper/QAR_dqc/mask_variables_reanalysis'.format(v))
                # TODO: any error must go to a log file (maybe QAR_dqc_generation.log )


    # Naming ini file and cmsname according to naming convention
    if len(o):
        cmsname = '{}-{}-{}'.format(d.replace('-', '_'), o, v.replace('-', '_'))
    else:
        if uerra:
            # for UERRA height_levels or pressure_levels the origin is always uerra_harmonie
            cmsname = '{}-uerra_harmonie-{}'.format(d.replace('-', '_'), v.replace('-', '_'))

        else: # era_5 / era5_land
            cmsname = '{}-{}-{}'.format(d.replace('-', '_'), model, v.replace('-', '_'))

    qar_name = cmsname + '.ini'



    return qar_name, generate_ini(root, patt, ve, dt, v, u, res, gt, gmn, ts, o, '', d, 'no', dr, tg, com, lsm, lnd, vmin, vmax, cmsname)

"""
Init parameters for a seasonal dataset
"""
def init_params_seasonal(dl, d, v, vinfo, mask_variables, temp_resolution, tolerance_limits, daily):

    # Get variable info
    ve = vinfo['shortName']
    u = vinfo['c3sunit']

    # Try to get tolerance limits
    try:
        vmin = float(tolerance_limits[v][0].strip())
    except:
        vmin = ''
    try:
        vmax = float(tolerance_limits[v][1].strip())
    except:
        vmax = ''

    # Time resolution + Data type
    com = 'yes'
    lnd = 'no'
    if daily:
        com = 'no'
        is_s = 'daily'
        dt = 'fc, hc, em'
        if 'pressure' in d:
            ts, tg, lsm = '12', 'hour', 'no'
        else:
            ts, tg, lsm = get_temporal_res(v, temp_resolution)
    else:
        ts, tg = '1', 'month'
        is_s = 'monthly'
        dt = 'fcmean, hcmean, em'
        lsm = 'no'
        if not('pressure' in d):
            try:
                _, _, lsm = get_temporal_res(v, temp_resolution)
            except:
                print(d,v)
                exit()

    # Grid resolution
    res = 1
    dr = 'no'
    gt = 'lonlat'

    # Origin/System
    o = dl['form']['originating_centre']
    s = dl['form']['system']

    # File pattern
    patt = o + '-' + s + '*.grib'

    # Naming ini file and cmsname according to naming convention
    cmsname = '{}-{}-{}-{}'.format(d.replace('-','_'), o, s, v.replace('-','_'))
    cmsname = cmsname.replace('ukmo', 'uk_met_office')

    qar_name = cmsname + '.ini'

    return qar_name, generate_ini('seasonal', patt, ve, dt, v, u, res, gt, '', ts, o, s, d, is_s, dr, tg, com, lsm, lnd, vmin, vmax, cmsname)

"""
Generate a ini file (datachecker configuration)
for each of dataset/variable combination
"""
def get_qar_names(forms_dir, jsonf, codes_var, mask_variables, temp_resolution, tolerance_limits):
    qar_names = []
    inis = []

    with open(jsonf, 'r') as f:
        dd=json.load(f)

        for dl in dd:
            for v in dl['variables']:
                try:
                    # General
                    d = dl['dataset']
                    ve = codes_var[d]['variables'][v]
                except:
                    print('ERROR: dataset={} variable={} not found in cds_datasets.json'.format(d, v))
                    continue

                # SEASONAL
                if 'originating_centre' in dl['form'] and 'system' in dl['form']:
                    mandatory_dates = get_forms_dates(d, v, dl['form']['originating_centre'], dl['form']['system'], forms_dir)
                    qar_name, ini = init_params_seasonal(dl, d, v, ve, mask_variables, temp_resolution, tolerance_limits, 'daily' in jsonf)
                # REANALYSIS
                else:
                    # UERRA-SL / ORIGIN
                    origin = ''
                    try:
                        origin = dl['form']['origin']
                    except:
                        pass

                    mandatory_dates = get_forms_dates(d, v, origin, '', forms_dir)
                    qar_name, ini = init_params_reanalysis(dl, d, v, ve, mask_variables, tolerance_limits,
                                                           'daily' in jsonf, 'uerra' in jsonf, 'land' in jsonf)

                # Next + skip empty datasets
                if len(mandatory_dates):
                    qar_names.append(qar_name)
                    inis.append(ini)
                else:
                    print(qar_name)

    return zip(qar_names, inis)

def write_file(outf, string):
    with open(outf, 'w') as f:
        f.write(string)

def main():
    #TODO: Add glob expression to filter which type of file we want to regenerate
    #TODO: add external config files for all paths, everything is hardcoded

    # Config paths
    conf_dir="../../dqc_downloader/conf"
    qar_dir="./conf"
    datasets = "dqc_wrapper/grib/cds_datasets.json"
    forms_dir = "/data/cds-forms-c3s"

    # Output directory
    out_dir="/tmp/conf_generate"
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    else:
        os.system('rm -rf /tmp/conf_generate/')

    # Get codes
    codes_var = json.load(open(datasets,'r'))

    # Variables that require a land_sea_mask
    mask_variables = []
    with open('dqc_wrapper/QAR_dqc/mask_variables_reanalysis') as f:
        mask_variables = f.read().splitlines()

    # Variables that require a land_sea_mask
    temp_resolution = []
    with open('dqc_wrapper/QAR_dqc/temp_resolution_seasonal') as f:
        temp_resolution = f.read().splitlines()

    # Variables that have tolerance limits
    tolerance_limits = get_tolerance_limits('dqc_wrapper/QAR_dqc/tolerance_limits_wmo')

    # For downloaded datasets (daily + monthly)
    for json_conf in glob.glob(conf_dir+'/*ly*.json'):

        # Get file names
        qar_names = get_qar_names(forms_dir, json_conf, codes_var, mask_variables, temp_resolution, tolerance_limits)
        for data in qar_names:
            q,i = data
            qar=os.path.join(qar_dir, q)
            qar_out=os.path.join(out_dir, q)

            # Skip if file exists
            if os.path.exists(qar_out):
                print('Skipping ' + qar_out)
            else:
                write_file(qar_out, i)

            # Show differences
            if os.path.exists(qar):
                os.system('diff -b {} {}'.format(qar_out, qar))
            else:
                print('NEW: {}'.format(qar_out))

if __name__ == "__main__":
    main()
