#!/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 os
import json

# Custom packages
from C3S512.cds import downloader

from dqc_wrapper.utils.utils_files import get_dates_from_files

"""
Downloader
"""
def call_download(cred, outfile, info, dataset, logdir):
    # Download or skip (do not overwrite)
    if not os.path.exists(outfile):
        print("DOWN = {}/{}".format(dataset, os.path.basename(outfile)))
        downloader(dataset, info, outfile, logpath=logdir, key=cred['key'], user=cred['url'])
    else:
        print("SKIP = {}/{}".format(dataset, os.path.basename(outfile)))

"""
Return full constraints/mapping object
"""
def get_constraints(dataset_name, cds_forms_c3s_location, json_name):

    try:
        json_path = cds_forms_c3s_location.joinpath("cds-forms-c3s", dataset_name, json_name + ".json")
        with open(json_path, "r") as cj:
            json_data = json.load(cj)
    except:
        print('Could not open: {}'.json_path)

    return json_data

"""
Return filtered constraints/mapping object
"""
def get_filtered_constraints(dataset_name, cds_forms_c3s_location, json_name, var, origin, system):

    res = []
    try:
        json_path = cds_forms_c3s_location.joinpath("cds-forms-c3s", dataset_name, json_name + ".json")
        with open(json_path, "r") as cj:
            json_data = json.load(cj)

        # Filter out variable
        for c in json_data:
            if var in c['variable'] and origin in c['originating_centre'] and system in c['system']:
                c['variable'] = var
                res.append(c)
    except:
        print('Could not open: {}'.json_path)

    return res

"""
Get netcdf sample from the C3s converter
"""
def get_netcdf_sample(info, cred, prod, outfile, logdir):

    # Hard-coded of course
    info['format'] = 'netcdf'
    outfile = outfile.replace('.grib', '.nc')

    # List of product types (we need to pick one)
    if 'product_type' in info:
        if isinstance(info['product_type'], type([])):
            info['product_type'] = info['product_type'][0]
        else:
            info['product_type'] = info['product_type']

    call_download(cred, outfile, info, prod, logdir)

"""
Get row (month, year)
"""
def get_row_dates(row):

    years = row['year']
    months = row['month']
    if isinstance(years, str):
        years = [years]
    if isinstance(months, str):
        months = [months]

    return years, months

"""
Generate a download json object compatible with the cdsapi
"""
def generate_download_obj(constraints, conf, file2dwn):

    # Extract date pattern from file
    dt = get_dates_from_files([file2dwn])
    yr, mn, d = dt[0].strftime('%Y'), dt[0].strftime('%m'), '01'

    # Filter out constraints - Get all product types associated with the constraint + month + variable/origin/system
    selection = {}
    for row in constraints:
        if 'seasonal' in conf.cds_dataset:
            condition = conf.origin in row['originating_centre'] and conf.system in row['system'] \
                        and conf.cds_variable in row['variable'] \
                        and mn in row['month'] and yr in row['year']
        else:
            condition = conf.cds_variable in row['variable'] and mn in row['month'] and yr in row['year']

        # If the condition is met, we need to store all "different" combinations, avoid repetitions
        if condition:
            if selection == {}:
                selection = row
            else:
                for k in row.keys():
                    if type(row[k]) is list:
                        if k in selection:
                            selection[k] = sorted(list(set(row[k]+selection[k])))
                    else:
                        selection[k] = row[k]

    # Fix variable/month/year + system/origin(seasonal)
    selection['variable'] = conf.cds_variable
    selection['month'] = mn
    selection['year'] = yr
    if 'seasonal' in conf.cds_dataset:
        selection['originating_centre'] = conf.origin
        selection['system'] = conf.system

    return selection