#!/usr/bin/env python3
############!/data/miniconda3/envs/eqc/bin/python

'''
### #!/data/miniconda3/envs/eqc/bin/python
### #!/home/Earth/jpena/.conda/envs/xr/bin/python
'''

'''
#################################################################################
This software prototype has been developed under the Copernicus C3S_512 contract.

Institution: BSC
Author: Jesus Peña-Izquierdo
Email: jesus.pena@bsc.es
#################################################################################
'''


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

tic = time.time()


# Parsing configuration file
with open('conf.yaml') as f:
    conf = yaml.load(f, Loader=yaml.FullLoader)
if conf['automatic_pdf_key_strengths_and_limitations']['execute']:
    print('other modules incompatible with environment')
    from automatic_pdfs.automatic_text_specific import script
    s = script()
else:
    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()

dataset_parameters = conf['dataset_parameters']

#if dataset_parameters['variable'] == 'varsys2assess': 
if len(sys.argv) > 1:
    varsys2assess_file = sys.argv[1]
    #varsys2assess_file = 'varsys2assess.yaml'
    
    # 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')
            print('\nStart time: {}'.format(datetime.datetime.now()))
    
            # DOWNLOADING..............
            ###########################

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

                D = downloading(conf)

                # Seasonal forecast download requests
                # Monthly files
                if  conf['download_parameters']['download_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, datetime.datetime.now()))


                # Reanalysis download requests
                # Monthly files
                if conf['download_parameters']['download_ref']:
                    if  conf['download_parameters']['download_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, datetime.datetime.now()))


            # TRANSFERRING..............
            ###########################    
            if conf['transfer_parameters']['prepare_files_special_variables']:
                T = transferring(conf)
                T.special()

            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.
                if conf['transfer_parameters']['transfer_ref']:
                    T.get_ref()
                
                # Clean garbage collector
                print('\nCleanning python garbage collector...\n')
                gc.collect()
                
                print('\nELAPSED TIME: {:.2f} hours [{}]'.format((time.time()-tic)/3600, datetime.datetime.now()))


            # 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()[0]) or (conf['compute_parameters']['recompute'] and not C.check_metrics()[1]): 

                            data_ok = C.get_data()                     # Prepare, combine and check data

                            if data_ok:                                # If data exists and it is ok

                                if conf['compute_parameters']['metrics']['bias']:
                                    bias = C.bias()                        # Computing mean bias
                                else:
                                    bias = None
                                if conf['compute_parameters']['metrics']['correlation']:
                                    corr, pval_corr = C.correlation()      # Computing correlation and pvalues
                                else:
                                    corr = None
                                    pval_corr = None
                                if conf['compute_parameters']['metrics']['frpss']:
                                    frpss, rpss, rps_sf, rps_ref_clim = C.frpss()            # Computing frpss and rpss (fair and non-fair score)
                                else:
                                    frpss = None
                                    rpss = None
                                    rps_sf = None
                                    rps_ref_clim = None
                                if conf['compute_parameters']['metrics']['fcrpss']:
                                    fcrpss, crpss, crps_sf, crps_ref_clim = C.fcrpss()       # Computing fcrpss and crpss (fair and non-fair score)
                                else:
                                    fcrpss = None
                                    crpss = None
                                    crps_sf = None
                                    crps_ref_clim = None

                                # Saving computed metrics
                                metrics2save_sm = {'bias': bias,
                                                   'correlation': corr,
                                                   'pvalues_corr': pval_corr,
                                                   'frpss': frpss,
                                                   'rpss': rpss,
                                                   'fcrpss': fcrpss,
                                                   'crpss': crpss,
                                                   'crps': crps_sf,
                                                   'crps_clim': crps_ref_clim,
                                                   'rps': rps_sf,
                                                   'rps_clim': rps_ref_clim
                                                  }
                                
                                # Save metrics and required attributes (units and shortname and longname)
                                C.save_metrics(metrics2save_sm)
                                
                                # Clean garbage collector
                                print('\nCleanning python garbage collector...\n')
                                gc.collect()

                            print('\nELAPSED TIME: {:.2f} hours [{}]'.format((time.time()-tic)/3600, datetime.datetime.now()))


            # PLOTTING..............
            ###########################

            if conf['plot_parameters']['execute']:
                
                P = plotting(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):
                        
                        P.level = level_i
                        P.start_month = sm

                        # Load computed metrics and add required variable attributes to conf
                        metrics2load, load_ok = P.load_metrics()
                        
                        if load_ok:                                         # If metrics are already computed

                            bias = metrics2load.bias                        # mean bias
                            correlation = metrics2load.correlation          # correlation
                            pvalues_corr = metrics2load.pvalues_corr        # pvalues from the correlation
                            frpss = metrics2load.frpss                      # frpss
                            fcrpss = metrics2load.fcrpss                    # fcrpss
                            rpss = metrics2load.rpss
                            crpss = metrics2load.crpss


                            P.plot_map(bias)                                # Plotting bias

                            P.plot_map((correlation, pvalues_corr))         # Plotting correlation

                            P.plot_map(frpss)                               # Plotting frpss

                            P.plot_map(fcrpss)                              # Plotting fcrpss

                            P.plot_map(rpss)

                            P.plot_map(crpss)


                        print('\nELAPSED TIME: {:.2f} hours [{}]'.format((time.time()-tic)/3600, datetime.datetime.now()))
                      
                # Clean garbage collector
                print('\nCleanning python garbage collector...\n')
                gc.collect()
                    
            print('\nAll figures for {} system-{} {}  plotted.'.format(orig2ass_i, sys2ass_i,var2ass_i))
            
            print('\nTOTAL EXECUTION TIME: {:.2f} hours [{}]'.format((time.time()-tic)/3600, datetime.datetime.now()))
                    
            # SCORECARDS FUNCTION........
            ###########################     
            if conf['scorecards_function']['execute']:
                print('\n#########################')
                print('SCORECARDS FUNCTION......')
                print('#########################\n')

                scorecard_varname=conf['dataset_parameters']['variable']
                scorecard_inputpath=Path(conf['folders']['metrics_data_path']).joinpath(conf['dataset_parameters']['dataset_sf'],
                                                                            conf['dataset_parameters']['variable']
                                                                            )
                scorecard_centersystem=conf['dataset_parameters']['originating_centre'] + 's' + str(conf['dataset_parameters']['system'])
                scorecard_centersystemname = '{}-S{}'.format(conf['dataset_parameters']['originating_centre'].upper(),conf['dataset_parameters']['system'])
                
                # read attributes (paramID, units)
                if 'pressure' in conf['dataset_parameters']['dataset_sf']:
                    level_datafilename = '-p' + str(conf['dataset_parameters']['levels'][0])
                else:
                    level_datafilename = ''
            
                filename = 'EQCmetrics-{}s{}-{}-{}to{}-s{:02d}{}.nc'.format(conf['dataset_parameters']['originating_centre'],
                                                              conf['dataset_parameters']['system'],
                                                              conf['dataset_parameters']['variable'],
                                                              conf['dataset_parameters']['years'][0],
                                                              conf['dataset_parameters']['years'][1],
                                                              2,
                                                                level_datafilename
                                                             )
                data = xr.open_dataset(scorecard_inputpath.joinpath(filename))

                attrs = t.extract_attrs(data,
                                    attrs_in = ['units', 'varname_short', 'varname_long']
                                   )

                scorecard_varname_short = attrs['varname_short']
                scorecard_varname_long = attrs['varname_long']

                scorecard_units = attrs['units']

                if scorecard_units in conf['others']['dimensionless_options']:
                    scorecard_units = 'dimensionless'#'[-]'
                    
                #scorecard_units = scorecard_units.replace("**-1", "$<sup>-1</sup>")
                #scorecard_units = scorecard_units.replace("**-2", "$<sup>-2</sup>$")
                #scorecard_units = scorecard_units.replace("**-3", "$^{-3}$")
                #scorecard_units = scorecard_units.replace("**1", "$^{1}$")
                #scorecard_units = scorecard_units.replace("**2", "$^{2}$")
                #scorecard_units = scorecard_units.replace("**3", "$^{3}$")

                import json
                cdsparamId = json.load(open('cds_datasets.json'))
                scorecard_paramID = cdsparamId[conf['dataset_parameters']['dataset_sf']]['variables'][conf['dataset_parameters']['variable']]['paramId']
                
                scorecard_outputpath = Path(conf['folders']['figures_path']).joinpath(conf['dataset_parameters']['dataset_sf'].replace("-","_"), 
                                                                           conf['dataset_parameters']['originating_centre'] \
                                                                           + '_s' \
                                                                           + str(conf['dataset_parameters']['system']),
                                                                           scorecard_paramID
                                                                          ).joinpath("scorecards")
                
                if conf['folders']['save_figure_copy_in_anomaly_folder']:
                    scorecard_outputpath_anomaly = Path(str(scorecard_outputpath).replace('seasonal_monthly', 'seasonal_postprocessed'))
                    save_copy = [1,2]
                else:
                    save_copy = [1]

                if not os.path.exists(scorecard_outputpath):
                    os.makedirs(scorecard_outputpath)

                print("scorecard_centersystem: " + scorecard_centersystem)
                print("scorecard_centersystemname: " + scorecard_centersystemname)
                print("scorecard_units: " + scorecard_units)
                print("scorecard_varname: " + scorecard_varname)
                print("scorecard_varname_short: " + scorecard_varname_short)
                print("scorecard_varname_long: " + scorecard_varname_long)
                print("scorecard_paramID: " + scorecard_paramID)
                print("scorecard_inputpath: " + str(scorecard_inputpath))
                print("scorecard_outputpath: " + str(scorecard_outputpath))
                sc_type = conf['scorecards_function']['type']
                years = conf['dataset_parameters']['years']
                print('--------------------------------')

                if sc_type == 'teleconnection':
                    scorecard_type = 4
                    scorecard_plevel = ''
                    t.scorecardfunction(str(years[0]), str(years[1]), str(scorecard_centersystem), str(scorecard_centersystemname), str(scorecard_units), str(scorecard_varname), str(scorecard_varname_long), str(scorecard_plevel), str(scorecard_inputpath), str(scorecard_outputpath), str(scorecard_paramID), scorecard_type)
                elif sc_type == 'save_teleconnection':
                    scorecard_type = 0
                    scorecard_plevel = ''
                    t.scorecardfunction(str(years[0]), str(years[1]), str(scorecard_centersystem), str(scorecard_centersystemname), str(scorecard_units), str(scorecard_varname), str(scorecard_varname_long), str(scorecard_plevel), str(scorecard_inputpath), str(scorecard_outputpath), str(scorecard_paramID), scorecard_type)
                else:                    
                    if 'pressure' in conf['dataset_parameters']['dataset_sf']:
                        plevels = conf['dataset_parameters']['levels']
                    else:
                        plevels = ['']
                    if sc_type == "all":
                        sc_type = [1,2,3]
                    elif sc_type == "do_2_and_3":
                        sc_type = [2,3]
                    else:
                        sc_type = [sc_type]

                    for ss in save_copy:
                        if ss == 1:
                            scorecard_save_outputpath = scorecard_outputpath
                        else:
                            scorecard_save_outputpath = scorecard_outputpath_anomaly
                        for scorecard_plevel in plevels:
                            for scorecard_type in sc_type:
                                print(scorecard_save_outputpath)
                                t.scorecardfunction(str(years[0]), str(years[1]), str(scorecard_centersystem), str(scorecard_centersystemname), str(scorecard_units), str(scorecard_varname), str(scorecard_varname_long), str(scorecard_plevel), str(scorecard_inputpath), str(scorecard_save_outputpath), str(scorecard_paramID), scorecard_type)            
            
            
            # AUTOMATIC PDFs OF KEY STRENGTHS AND LIMITATIONS   ........
            ###########################     
            if conf['automatic_pdf_key_strengths_and_limitations']['execute']:
                print('\n#########################')
                print('KEY STRENGTHS AND LIMITATIONS.. environment needed: eqc-pdf-indepth ....')
                print('#########################\n')

                autopdf_varname=conf['dataset_parameters']['variable']
                autopdf_inputpath=Path(conf['folders']['metrics_data_path']).joinpath(conf['dataset_parameters']['dataset_sf'],
                                                                            conf['dataset_parameters']['variable']
                                                                            )
                autopdf_centersystem=conf['dataset_parameters']['originating_centre'] + 's' + str(conf['dataset_parameters']['system'])
                autopdf_center = conf['dataset_parameters']['originating_centre'].upper()
                autopdf_center = autopdf_center.replace('METEO_FRANCE', 'Météo-France')
                autopdf_center = autopdf_center.replace('UKMO', 'UK Met Office')
                autopdf_centersystemname = '{} System {}'.format(autopdf_center,conf['dataset_parameters']['system'])
                
                # read attributes (paramID, units)
                if 'pressure' in conf['dataset_parameters']['dataset_sf']:
                    level_datafilename = '-p' + str(conf['dataset_parameters']['levels'][0])
                else:
                    level_datafilename = ''
            
                filename = 'EQCmetrics-{}s{}-{}-{}to{}-s{:02d}{}.nc'.format(conf['dataset_parameters']['originating_centre'],
                                                              conf['dataset_parameters']['system'],
                                                              conf['dataset_parameters']['variable'],
                                                              conf['dataset_parameters']['years'][0],
                                                              conf['dataset_parameters']['years'][1],
                                                              2,
                                                                level_datafilename
                                                             )
                data = xr.open_dataset(autopdf_inputpath.joinpath(filename))

                attrs = s.extract_attrs(data,
                                    attrs_in = ['units', 'varname_short', 'varname_long']
                                   )

                autopdf_varname_short = attrs['varname_short']
                autopdf_varname_long = attrs['varname_long']

                autopdf_units = attrs['units']

                if autopdf_units in conf['others']['dimensionless_options']:
                    autopdf_units = 'dimensionless'#'[-]'
                    
                import json
                cdsparamId = json.load(open('cds_datasets.json'))
                autopdf_paramID = cdsparamId[conf['dataset_parameters']['dataset_sf']]['variables'][conf['dataset_parameters']['variable']]['paramId']
                
                autopdf_outputpath = Path('/shared/data_for_scorecards_SF/suite/output_pdfs').joinpath(autopdf_varname) #Path(conf['folders']['output_path_key_strength_and_limitations']).joinpath(autopdf_varname)
                workdir = conf['folders']['working_directory_for_automatic_pdfs']
                               
                if not os.path.exists(autopdf_outputpath):
                    #oldmask = os.umask(000)
                    os.makedirs(autopdf_outputpath) #, 755)
                    #oldmask = os.umask(000)
                
                years = conf['dataset_parameters']['years']
                
                if 'pressure' in conf['dataset_parameters']['dataset_sf']:
                    plevels = conf['dataset_parameters']['levels']
                else:
                    plevels = ['']
                #for autopdf_plevel in plevels:
                if 850 in plevels:
                    p850='T'
                else:
                    p850='F'
                if 500 in plevels:
                    p500='T'
                else:
                    p500='F'
                if 200 in plevels:
                    p200='T'
                else:
                    p200='F'
                if 50 in plevels:
                    p50='T'
                else:
                    p50='F'
                
                if 'pressure' in conf['dataset_parameters']['dataset_sf']:
                    autopdf_plevel = 'yes'
                else:
                    autopdf_plevel='no'
                saved_file_name = s.autopdffunction(str(years[0]), str(years[1]), str(autopdf_centersystem), str(autopdf_centersystemname), str(autopdf_units), str(autopdf_varname), str(autopdf_varname_long), str(autopdf_plevel), p850, p500, p200, p50, workdir, str(autopdf_inputpath), str(autopdf_outputpath), str(autopdf_paramID))
                
                from shutil import copyfile
                outpoutpath_final_pdf = Path(conf['folders']['output_path_key_strength_and_limitations']).joinpath(autopdf_varname)
                if not os.path.exists(outpoutpath_final_pdf):
                    #os.umask(0)
                    #os.makedirs(outpoutpath_final_pdf, mode = 0o777)
                    os.makedirs(outpoutpath_final_pdf)
                    #os.chmod(outpoutpath_final_pdf, Oo777)
                fullpathname = autopdf_outputpath.joinpath(saved_file_name[0] + '.pdf')
                fullpathnewname = outpoutpath_final_pdf.joinpath(saved_file_name[0] + '.pdf')
                copyfile(fullpathname, fullpathnewname)
                print("Key strengths and limitations pdf saved at:")
                print(fullpathnewname)