#!/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 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)
        print('ok')

        # For every dataset
        for prod in data:
            print('ok2')

            # 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()
                print(var)
                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
                        if prod['group'] == 'seasonal':
                            row['system'] = system
                            row['originating_centre'] = origin
                    # Drop other rows
                    else:
                        drops.append(ind)

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

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

                    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']))
                                if prod['group'] == 'seasonal':
                                    downloads[key]['leadtime_month'] = list(set(downloads[key]['leadtime_month'] + row['leadtime_month']))
                                else:
                                    downloads[key]['time'] = list(set(downloads[key]['time'] + row['time']))

                                # Pressure levels
                                if 'pressure_level' in downloads[key]:
                                    downloads[key]['pressure_level'] = sorted(list(set(downloads[key]['pressure_level'] + row['pressure_level'])), key = float)
                                    downloads[key]['pressure_level'] = sorted(downloads[key]['pressure_level'], key = float)

                            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'], system=info['system']
                        )
                    else:
                        outfile = "{dir}/{var}-{date}.{format}".format(
                            dir=outdir, var=info['variable'], date=k, format='grib', prod=prod['dataset']
                        )
                        
                    if prod['form']['format'] == 'grib': 
                        # Download perform and save logs
                        call_download(cred, outfile, info, prod['dataset'], logdir)

                    # Sample mode (for testing)
                    if 'sample' in prod:
                        break

                # Perform ONE download with ONE resolution (NetCDF sample of the converter)
                try:
                    get_netcdf_sample(info, cred, prod, outfile, logdir)
                except:
                    pass # CDS fails to convert sometimes

if __name__ == "__main__":
    main()
