#!/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
from datetime import datetime
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_daily.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']
            else:
                origin = prod['form']['origin'] if 'origin' in prod['form'] else ''

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

                # Perform download (2 months per year, those months change per variable(m))
                startM = len(var)%6 + 1
                endM = (startM + 6)

                # 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:
                        # UERRA models distinction
                        if origin != '':
                            condition = var in row['variable'] and origin in row['origin']
                        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 = {}
                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
                            #row['day'] = [ '{:02}'.format(d) for d in range(1,32) ] # small trick, accept all days
                            key = '{}{}01'.format(y,m)

                            # Changing parameters (unique values)                           
                            if key in downloads:
                                
                                # Some products do not have a product type
                                if 'product_type' in downloads[key]:
                                    downloads[key]['product_type'] = list(set(downloads[key]['product_type'] + row['product_type']))

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

                                # Time 
                                if 'time' in downloads[key]:
                                    downloads[key]['time'] = sorted(list(set(downloads[key]['time'] + row['time'])))

                                # All of them should have day
                                downloads[key]['day'] = sorted(list(set(downloads[key]['day'] + row['day'])))

                            else:
                                # First time
                                downloads[key] = row.to_dict()

                # Empty products
                if downloads == {}:
                    continue

                # Get time boundaries
                temporal_coverage = sorted(downloads.keys())
                dt0 = datetime.strptime(temporal_coverage[0], "%Y%m%d")
                dtn = datetime.strptime(temporal_coverage[-1], "%Y%m%d")

                # Keys to download (2 months per year rotation strategy)
                count = 0
                keys_to_download = []
                for yr in range(dt0.year, dtn.year+1):

                    # HISTORICAL: Round-Robin 2-months/year rotate with modulus
                    if yr < 2020:
                        mstart = 1 + (startM + count) % 12
                        mend = 1 + (endM + count) % 12
                        count += 1
                        k_start = '{}{:02}01'.format(yr, mstart)
                        k_end = '{}{:02}01'.format(yr, mend)

                        # Within the limits of the dataset
                        if k_start in downloads:
                            keys_to_download.append(k_start)
                        if k_end in downloads:
                            keys_to_download.append(k_end)

                    # PRODUCTION: from 2020 onwards everything is checked
                    else:
                        for mn in range(1,13):
                            k2020 = '{}{:02}01'.format(yr, mn)
                            if k2020 in downloads:
                                keys_to_download.append(k2020)

                # Download selection
                for k in keys_to_download:

                    # Round-Robin months rotate with modulus
                    info = downloads[k]

                    # 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:
                        # UERRA / SL
                        if 'origin' in info:
                            outfile = "{dir}/{var}-{date}.{format}".format(
                                dir=outdir, var=info['origin'][0], date=k, format='grib', prod=prod['dataset']
                            )
                        # REANALYSIS
                        else:
                            outfile = "{dir}/{var}-{date}.{format}".format(
                                dir=outdir, var=info['variable'], date=k, format='grib', prod=prod['dataset']
                            )
                            
                    if prod['form']['format'] == 'grib':        
                        # Perform download, gather statistics
                        call_download(cred, outfile, info, prod['dataset'], logdir)
                    else:
                        # We get only one file
                        get_netcdf_sample(info, cred, prod['dataset'], outfile, logdir)
                        break

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

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

if __name__ == "__main__":
    main()
