#!/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)

from pathlib import Path
import pandas as pd
import numpy as np
import sys

# Custom packages
from dqc_downloader.utils_download import *

def main():

    # Params
    if len(sys.argv) != 3:
        print('USAGE: python download_monthly.py <config_file> <credentials_file>')
        exit()
    else:
        conf_path = sys.argv[1]
        cred_path = sys.argv[2]

    # General configuration
    with open('./conf/download_options.json') as json_file:
        conf = json.load(json_file)
        outputdir = conf['outdir']
        logdir = conf['logdir']
        cds_forms_c3s_location = Path(conf['formsdir'])

    # URL and key to download
    with open(cred_path) as json_file:
        cred = json.load(json_file)

    # Read list of products to download
    with open(conf_path) as json_file:
        data = json.load(json_file)

        # For every dataset
        for prod in data:

            # Get constraints for this dataset
            constraints = get_constraints(prod['dataset'], cds_forms_c3s_location, "constraints")
            if prod['group'] == 'seasonal':
                origin = prod['form']['originating_centre']
                system = prod['form']['system']

            # For every variable
            for var in prod['variables']:

                # Load into pandas table
                df = pd.DataFrame()
                for c in constraints:
                    df = df.append(c, ignore_index=True)

                # Select variable/origin/system
                drops = []
                for ind, row in df.iterrows():
                    # Filtering differs by product family
                    if prod['group'] == 'seasonal':
                        condition = origin in row['originating_centre'] and system in row['system'] and var in row['variable']
                    else:
                        condition = var in row['variable']

                    # We do not need the rest
                    if condition:                        
                        row['variable'] = var
                    # Drop other rows
                    else:
                        drops.append(ind)

                # Drop at once
                df = df.drop(drops)

                # Group by year/month
                downloads = {}
                for ind, row in df.iterrows():
                    years, months = get_row_dates(row)

                    for y in years:
                        for m in months:
                            # Forced parameters
                            row['month'] = m
                            row['year'] = y
                            key = '{}{}01'.format(y,m)

                            # Changing parameters (unique values)
                            if key in downloads:
                                downloads[key]['product_type'] = list(set(downloads[key]['product_type'] + row['product_type']))
                            else:
                                downloads[key] = row.to_dict()


                # Perform download (GRIB)
                for k, info in downloads.items():
                    # Output directory                    
                    outdir = os.path.join(outputdir, '{}/{}/{}'.format(prod['group'], prod['dataset'], info['variable']))                    
                    if not(os.path.exists(outdir)):
                        os.makedirs(outdir)

                    # Output file
                    if prod['group'] == 'seasonal':
                        outfile = "{dir}/{origin}-{system}-{date}.{format}".format(
                            dir=outdir, date=k, format='grib', prod=prod['dataset'],
                            origin=info['originating_centre'][0], system=info['system'][0]
                        )
                    else:
                        outfile = "{dir}/{var}-{date}.{format}".format(
                            dir=outdir, var=info['variable'], date=k, format='grib', prod=prod['dataset']
                        )

                    # Download perform and save logs
                    call_download(cred, outfile, info, prod, logdir)

                # Perform ONE download with ONE resolution (NetCDF sample of the converter)
                info['format'] = 'netcdf'
                info['product_type'] = info['product_type'][0]
                call_download(cred, outfile.replace('.grib', '.nc'), info, prod, logdir)
             

if __name__ == "__main__":
    main()