#!/data/miniconda3/envs/eqc/bin/python

from pathlib import Path
import os
import gc
import sys
import yaml
import time
import xarray as xr

from src.download import downloading
from src.transfer import transferring
from src.compute import computing
from src.plot import plotting
from src.tools import tools
t = tools()


tic = time.time()


# Parsing configuration file
with open('conf.yaml') as f:
    conf = yaml.load(f, Loader=yaml.FullLoader)

dataset_parameters = conf['dataset_parameters']

if len(sys.argv) > 1:
    varsys2assess_file = sys.argv[1]
    
    # Parsing list of variables and systems to be assessed
    with open(varsys2assess_file) as f:
        varsys2assess = yaml.load(f, Loader=yaml.FullLoader)
        
    var2ass = varsys2assess['variables2assess']
    centresystem2assess = varsys2assess['centresystem2assess']
        
else:
    var2ass = [dataset_parameters['variable']]
    centresystem2assess = {dataset_parameters['originating_centre']: dataset_parameters['system']}    


for var2ass_i in var2ass:
    for orig2ass_i in centresystem2assess.keys():
        sys2ass = centresystem2assess[orig2ass_i]
        
        if not isinstance(sys2ass, list):
            sys2ass = [sys2ass] 
            
        for sys2ass_i in sys2ass:
    
            tic_i = time.time()
        
            # Updating this dictionary will automatically update the loaded conf dictionary
            dataset_parameters['variable'] = var2ass_i
            dataset_parameters['originating_centre'] = orig2ass_i
            dataset_parameters['system'] = sys2ass_i

            print('\n\n________________________________________________________________\n')
            print('STARTING SCIENTIFIC ASSESSMENT\n\n')

    
            # DOWNLOADING..............
            ###########################

            if conf['download_parameters']['execute']:

                D = downloading(conf)

                # Seasonal forecast download requests
                # Monthly files
                if dataset_parameters['data_aggregation_sf'] == 'monthly':

                    for mm in range(dataset_parameters['start_months'][0], dataset_parameters['start_months'][1]+1):

                        D.start_months = [mm, mm]

                        t.thread(target = D.seasonal_forecast())

                # Single file
                else:
                    D.seasonal_forecast()
                    
                print('\nELAPSED TIME: {:.2f} hours'.format((time.time()-tic)/3600))


                # Reanalysis download requests
                # Monthly files
                if dataset_parameters['data_aggregation_ref'] == 'monthly':

                    for mm in range(dataset_parameters['start_months'][0], dataset_parameters['start_months'][1]+1):

                        D.start_months = [mm, mm]

                        t.thread(target = D.era5())

                # Single file
                else:
                    D.era5()
                    
                print('\nELAPSED TIME: {:.2f} hours'.format((time.time()-tic)/3600))



            # TRANSFERRING..............
            ###########################    

            if conf['transfer_parameters']['execute']:

                T = transferring(conf)

                # Extract Seasonal Forecast fields (individual members monthly means) from the datachecker archive.
                T.get_sf()

                # Extract ERA5 fields (monthly means) as reference from the datachecker archive.
                T.get_ref()
                
                # Clean garbage collector
                print('\nCleanning python garbage collector...\n')
                gc.collect()
                
                print('\nELAPSED TIME: {:.2f} hours'.format((time.time()-tic)/3600))



            # COMPUTING..............
            ###########################

            if conf['compute_parameters']['execute']:

                C = computing(conf)

                # Computations done independently for each level and start month
                levels = t.get_levels(dataset_parameters)
                sm_range = dataset_parameters['start_months']
                for level_i in levels:
                    for sm in range(sm_range[0], sm_range[1] + 1):

                        C.level = level_i
                        C.start_month = sm

                        # Automatic check if metrics are computed for given start month, allowing recomputation if required.
                        #if (not C.check_metrics()) or conf['compute_parameters']['recompute']: 

                        data_ok = C.get_data()                     # Prepare, combine and check data
                        d, load_ok = C.load_metrics()
                        if load_ok:
                            C.save_metrics(d)