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

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

import xarray as xr
import numpy as np
from pathlib import Path
import logging
import os
import sys
import pandas as pd
import yaml

from dask.distributed import Client, LocalCluster
from dask.diagnostics import ProgressBar
import dask

from src.tools import tools
t = tools()


class computing:

    def __init__(self,
                 conf,
                 start_month = 1,
                 level = None,
                ):
        
        self.start_month = start_month
        self.level = level
        
        self.dataset_parameters = conf['dataset_parameters']
        
        self.variable = conf['dataset_parameters']['variable']
        self.variable_ref = conf['dataset_parameters']['variable_ref']
        self.originating_centre =conf['dataset_parameters']['originating_centre']
        self.system = conf['dataset_parameters']['system']
        self.start_months = conf['dataset_parameters']['start_months']
        self.years = conf['dataset_parameters']['years']
        self.leadtime_months = conf['dataset_parameters']['leadtime_months']
        self.leadtime_periods = conf['dataset_parameters']['leadtime_periods']
        self.leadtime_months_T = self.leadtime_months +[self.leadtime_months[-1]+i+1 for i in range(len(self.leadtime_periods))]
        
        if not self.variable_ref:
            self.variable_ref = self.variable
        
        self.dataset_sf = conf['dataset_parameters']['dataset_sf']
        self.data_aggregation_sf = conf['transfer_parameters']['data_aggregation_sf']
        self.dataset_ref = conf['dataset_parameters']['dataset_ref']
        self.data_aggregation_ref = conf['transfer_parameters']['data_aggregation_ref']
        
        self.file_format_used = conf['transfer_parameters']['file_format_used']
        self.time_dimname = conf['compute_parameters']['xarray_parameters']['time_dimname']
        self.starttime_dimname = conf['compute_parameters']['xarray_parameters']['starttime_dimname']
        self.leadtime_dimname = conf['compute_parameters']['xarray_parameters']['leadtime_dimname']
        self.member_dimname = conf['compute_parameters']['xarray_parameters']['member_dimname']
        self.lon_dimname = conf['compute_parameters']['xarray_parameters']['lon_dimname']
        self.lat_dimname = conf['compute_parameters']['xarray_parameters']['lat_dimname']
        self.stop_different_number_nans = conf['compute_parameters']['stop_different_number_nans']
        
        self.metrics = conf['compute_parameters']['metrics']
        
        self.chunks = conf['compute_parameters']['xarray_parameters']['chunks']
        self.dask_distributed = conf['compute_parameters']['xarray_parameters']['dask_distributed']
        self.n_workers = conf['compute_parameters']['xarray_parameters']['n_workers']
        self.threads_per_worker = conf['compute_parameters']['xarray_parameters']['threads_per_worker']
        self.n_cores = conf['compute_parameters']['xarray_parameters']['n_cores']
        self.split_factor = conf['compute_parameters']['xarray_parameters']['split_factor']
        self.n_leadtimes = conf['compute_parameters']['xarray_parameters']['n_leadtimes']
        self.memory_limit = conf['compute_parameters']['xarray_parameters']['memory_limit']
        self.processes = conf['compute_parameters']['xarray_parameters']['processes']
        self.silence_logs = conf['compute_parameters']['xarray_parameters']['silence_logs']
        self.dashboard_address = conf['compute_parameters']['xarray_parameters']['dashboard_address']
        
        self.data_path_sf = Path(conf['folders']['base_dataOUT_path']).joinpath('seasonal',
                                                                        self.dataset_sf
                                                                       )
        self.data_path_ref = Path(conf['folders']['base_dataOUT_path']).joinpath('reanalysis',
                                                                         self.dataset_ref
                                                                        )
        self.logs_path = Path(conf['folders']['logs_path'])
        self.metrics_data_path = Path(conf['folders']['metrics_data_path']).joinpath(self.dataset_sf,
                                                                                     self.variable
                                                                                    )
        
        if self.file_format_used == 'netcdf':
            self.file_format_ext = 'nc'
            self.engine = 'netcdf4'
        else:
            self.file_format_ext = 'grib'
            self.engine = 'cfgrib'
            
        
        formatter = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        log_error_file = 'log_error_compute_{}s{}_{}.log'.format(self.originating_centre,
                                                                  self.system,
                                                                  self.variable)
        logging.basicConfig(filename = self.logs_path.joinpath(log_error_file),
                            level = "WARNING",
                            format = formatter,
                            filemode = "w")
        
        
        
        print('\n#########################')
        print('COMPUTING................')
        print('#########################\n')
        
        if self.dask_distributed: 
            # Starting dask distributed scheduler scheduler client
            cluster = LocalCluster(n_workers = self.n_workers,
                                   threads_per_worker = self.threads_per_worker,
                                    memory_limit = self.memory_limit,
                                    processes = self.processes,
                                    #silence_logs = self.silence_logs,
                                    #dashboard_address = self.dashboard_address
                                   )
            client = Client(cluster)

        else:
            # Starting a single machine scheduler
            dask.config.set(scheduler = 'threads')

     
    def update_level(self):
        
        if 'pressure' in self.dataset_sf:
            self.level_ext = '-p' + str(self.level)
            self.level_ext2 = '-p' + str(self.level)
        else:
            self.level_ext = ''
            self.level_ext2 = ''
            self.level = None
            self.levels = [None]

    
    def read_data_sf(self):
        
        # This function preload the Seasonal Forecast data for the given start month
        
        read_ok = True  # Flag for correct data read 
        
        # FOR SEASONAL FORECAST DATA
        # ..........................
        
        # Update level
        self.update_level()
        
        #Monthly files
        if self.data_aggregation_sf == 'monthly':
            
            print('\nPreloading SF data from monthly files for level {} and start month {}...'.format(self.level,
                                                                                                      self.start_month
                                                                                                     ))
            
            self.data_path_m_sf = self.data_path_sf.joinpath('monthly_files', self.variable)
            
            #Checking all dates (files) are located
            ok, files, files_missed = t.check_files_dates(folder_path = self.data_path_m_sf,
                                                          months = [self.start_month, self.start_month],
                                                          years = self.years,
                                                          filename_start = '{}s{}-{}-'.format(self.originating_centre,
                                                                                       self.system,
                                                                                       self.variable),
                                                          filename_end = self.level_ext + '.' + self.file_format_ext
                                                         )
                                                       
            if ok:
                print('All SF files within select time period collected.')
                print('...Preloading SF files...')

                self.data_sf = xr.open_mfdataset(files,
                                                 concat_dim = self.time_dimname,
                                                 chunks = self.chunks,
                                                 combine='nested',
                                                 engine = self.engine,
                                                 drop_variables = self.leadtime_dimname
                                                ).chunk({self.time_dimname: -1})
                # Convert to datarray
                var = list(self.data_sf.data_vars)
                self.data_sf = self.data_sf[var[0]]
                
                print('...Monthly SF files loaded.')
                    
            else:
                e0 = '---------------------------------------------------\n'
                e1 = '-----------------------ERROR-----------------------\n'
                e2 = 'SF files missing: {}\n'.format(files_missed)
                e3 = '\nAssessment for {}s{}-{}:\n...level: {}\n...start_month: {}\n...SKIPPED!!!\n'.format(self.originating_centre,
                                                                                           self.system,
                                                                                           self.variable,
                                                                                           self.level,
                                                                                            self.start_month
                                                                                          )
                print(e0+e1+e2+e3+e0)
                logging.error(e0+e1+e0+e2)
                
                read_ok = False
          
        # Start_month file
        else:   
            print('\nPreloading SF data from one single file for level {} and start_month {}...'.format(self.level,
                                                                                                        self.start_month
                                                                                                       ))

            self.data_path_s_sf = self.data_path_sf.joinpath('startmonth_files', self.variable)


            year_limits = '{}to{}'.format(self.years[0], self.years[1])
            filename = '{}s{}-{}-{}-s{:02d}{}.{}'.format(self.originating_centre,
                                                    self.system,
                                                    self.variable,
                                                    year_limits,
                                                    self.start_month,
                                                         self.level_ext,
                                                       self.file_format_ext
                                                   )
            
            file = self.data_path_s_sf.joinpath(filename)
            
            
            if os.path.exists(file):
                print('...Preloading... {}'.format(file))
                
                self.data_sf = xr.open_dataarray(file, 
                                                 chunks = self.chunks,
                                                 engine = self.engine).chunk({self.time_dimname: -1})
                
                # Remove problematic valid_time variable  
                #self.data_sf = self.data_sf.drop('valid_time')
            
                print('...Single SF file loaded.')
                
            else:
                e0 = '---------------------------------------------------\n'
                e1 = '-----------------------ERROR-----------------------\n'
                e2 = 'SF file missing: {}\n'.format(file)
                e3 = '\nAssessment for {}s{}-{}:\n...level: {}\n...start_month: {}\n...SKIPPED!!!\n'.format(self.originating_centre,
                                                                                           self.system,
                                                                                           self.variable,
                                                                                           self.level,
                                                                                            self.start_month
                                                                                          )
                print(e0+e1+e2+e3+e0)
                logging.error(e0+e1+e0+e2)
                
                read_ok = False
                
        
        if read_ok:
            
            # Renaming time as start_time
            self.data_sf = self.data_sf.rename({self.time_dimname: self.starttime_dimname})
            # Assigning 1-based values to the leadtime dimension
            self.leadtime_coord0 = np.arange(len(self.data_sf.coords[self.leadtime_dimname]))+1
            self.data_sf = self.data_sf.assign_coords({self.leadtime_dimname: self.leadtime_coord0})
            self.data_sf = self.data_sf.sel({self.leadtime_dimname: self.leadtime_months})
            
            # Get required attributes (units and shortname)
            attrs_sf = t.extract_attrs(self.data_sf,
                                    attrs_in = ['GRIB_cfVarName', 'GRIB_units'],
                                    attrs_out = ['varname_short', 'units']
                                   )

             # Get variable long name from the CDS catalogue forms
            names_map = t.var2longname(self.dataset_sf)
            attrs_sf['varname_long'] = names_map[self.variable]

            self.units_sf = attrs_sf['units']
            self.varname_short_sf = attrs_sf['varname_short']
            self.varname_long_sf = attrs_sf['varname_long']
            
        return read_ok
            
            
    def read_data_ref(self):
            
        # This function preloads the reference data for the given start month
        
        read_ok = True  # Flag for correct data read 
        
        # Update level
        self.update_level()
        
        # Identifying the grid of the seasonal forecast
        lon0 = str(self.data_sf.coords[self.lon_dimname][0].values) + 'E'
        lat0 = str(t.lat2NS(self.data_sf.coords[self.lat_dimname][0].values))
        lonlat0 = lon0 + lat0
        
        #Monthly files
        if self.data_aggregation_ref == 'monthly':
            
            print('\nPreloading REF data from monthly files...')
            
            lon0 = str(self.data_sf.coords[self.lon_dimname][0].values) + 'E'
            lat0 = str(t.lat2NS(self.data_sf.coords[self.lat_dimname][0].values))
            lonlat0 = lon0 + lat0
            
            self.data_path_m_ref = self.data_path_ref.joinpath('monthly_files', self.variable)
            
            #Checking all dates (files) are located
            ok, files, files_missed = t.check_files_dates(folder_path = self.data_path_m_ref,
                                                          months = [1,12], #self.start_months,
                                                          years = [self.years[0], self.years[1] + 1],
                                                          filename_start = '{}-grid{}-'.format(self.variable,
                                                                                              lonlat0
                                                                                             ),
                                                          filename_end = self.level_ext + '.' + self.file_format_ext
                                                         )

            if ok:
                print('All REF files within select time period collected.')
                print('...Preloading REF files...')
                
              
                self.data_ref = xr.open_mfdataset(files,
                                                 concat_dim = self.time_dimname,
                                                 chunks = self.chunks,
                                                 combine='nested',
                                                  engine = self.engine
                                                ).chunk({self.time_dimname: -1})
                # Convert to datarray
                var = list(self.data_ref.data_vars)
                self.data_ref = self.data_ref[var[0]]
                
                print('...Monthly REF files loaded.')
                
            else:
                e0 = '---------------------------------------------------\n'
                e1 = '-----------------------ERROR-----------------------\n'
                e2 = 'REF files missing: {}\n'.format(files_missed)
                e3 = '\nAssessment for {}s{}-{}:\n...level: {}\n...start_month: {}\n...SKIPPED!!!\n'.format(self.originating_centre,
                                                                                           self.system,
                                                                                           self.variable,
                                                                                           self.level,
                                                                                           self.start_month
                                                                                          )
                print(e0+e1+e2+e3+e0)
                logging.error(e0+e1+e0+e2)
                
                read_ok = False
             
         
        # Single file
        else:
            print('\nPreloading REF data from one single file...')
            
            self.data_path_s_ref = self.data_path_ref.joinpath('single_file', self.variable)


            year_limits = '{}to{}'.format(self.years[0], self.years[1]+1)
            filename = '{}-grid{}-{}{}.{}'.format(self.variable,
                                                lonlat0,
                                                year_limits,
                                                  self.level_ext,
                                                self.file_format_ext
                                               )
            
            file = self.data_path_s_ref.joinpath(filename)

            print('...Preloading... {}'.format(file))
            if os.path.exists(file):
                self.data_ref = xr.open_dataarray(file,
                                                  chunks = self.chunks,
                                                  engine = self.engine
                                                 ).chunk({self.time_dimname: -1})

                print('...Single REF file loaded.')
            
            else:
                e0 = '---------------------------------------------------\n'
                e1 = '-----------------------ERROR-----------------------\n'
                e2 = 'REF file missing: {}\n'.format(file)
                e3 = '\nAssessment for {}s{}-{}:\n...level: {}\n...start_month: {}\n...SKIPPED!!!\n'.format(self.originating_centre,
                                                                                           self.system,
                                                                                           self.variable,
                                                                                           self.level,
                                                                                              self.start_month
                                                                                          )
                print(e0+e1+e2+e3+e0)
                logging.error(e0+e1+e0+e2)
                
                read_ok = False
        
        if read_ok:
            # Get required attributes (units and shortname)
            attrs_ref = t.extract_attrs(self.data_ref,
                                    attrs_in = ['GRIB_cfVarName', 'GRIB_units'],
                                    attrs_out = ['varname_short', 'units']
                                   )

            # Get variable long name from the CDS catalogue forms
            names_map = t.var2longname(self.dataset_ref)

            if self.variable == 'maximum_2m_temperature_in_the_last_24_hours':
                attrs_ref['varname_long'] = 'maximum 2m temperature since previous post processing' ######special_cases
            elif self.variable == 'minimum_2m_temperature_in_the_last_24_hours':
                attrs_ref['varname_long'] = 'minimum 2m temperature since previous post processing' ######special_cases
            else:
                attrs_ref['varname_long'] = names_map[self.variable_ref]
            
            self.units_ref = attrs_ref['units']
            self.varname_short_ref = attrs_ref['varname_short']
            self.varname_long_ref = attrs_ref['varname_long']
                

                
        return read_ok
            
            
            
    def reshape_ref(self):
        
        # This function reshapes the reference data splitting the time dimension in start_month and leadtime_month

        print('\nReshaping REF data splitting time dimension in start_months and leadtimes for level {} and start_month {}...'.format(self.level,
                                                                                                                                      self.start_month
                                                                                                                                     ))

        Nyears = int(np.diff(self.years)+1) # Number of years
        Nleadtime = len(self.leadtime_months) # Number of leadtimes
        sm = self.start_month # Start month
        
        # Index for all leadtime corresponding to the given start month
        itime_leadtimes = np.array([np.arange(Nleadtime) + k*12 + sm -1 for k in range(Nyears)]).flatten()
        
        x_start_months = self.data_sf.coords[self.starttime_dimname].values # Start months from the SF
        x_leadtime_months = np.array(self.leadtime_months) # Leadtime months

        d = self.data_ref.isel({self.time_dimname: itime_leadtimes}) # Selected times for the leadtimes
        if 'step' in d.coords:
            d = d.drop('step') ######special_cases

        # Multiindex for splitting the time dimension in to startmonth and leadtime
        ind_multi = pd.MultiIndex.from_product((x_start_months, x_leadtime_months),
                                               names = (self.starttime_dimname, self.leadtime_dimname))

        
        self.data_ref = d.assign_coords({self.time_dimname: ind_multi}).unstack(self.time_dimname)

        print('...REF reshaped.')
        
        
    def add_periods(self):
        
        # This function adds as extra leadtimes the average of the given monthly periods
        
        print('\nAdding average leadtime periods...'.format(self.leadtime_periods))
        
        print('...SF data...')
        self.data_sf = t.add_periods(self.data_sf,
                                      periods = self.leadtime_periods,
                                      leadtime_dimname = self.leadtime_dimname
                                     )
        print('...REF data...')
        self.data_ref = t.add_periods(self.data_ref,
                                      periods = self.leadtime_periods,
                                      leadtime_dimname = self.leadtime_dimname
                                     )
        
        print('...Average leadtime periods added.')
        
    
    def check_units(self):
        
        # This function checks and convert the units of seasonal forecast and reference datasets
        
        # Check units
        print('\nChecking units for Seasonal Forecast and Reference...')
        
        if self.units_sf != self.units_ref:

            print('...Seasonal Forecast and Reference for {} do not have same units.'.format(self.variable))
            print('...{} vs {}'.format(self.units_sf, self.units_ref))

            conversion_coef, units_ok = t.find_conversion(self.variable)
            
            if units_ok:
                self.data_ref = conversion_coef * self.data_ref
                self.units_ref = self.units_sf
                
                print('...Following unit conversion has been used: SF_data = {} x REF_data'.format(conversion_coef))

            else:
                e0 = '---------------------------------------------------\n'
                e1 = '-----------------------ERROR-----------------------\n'
                e2 = '\nCoefficient units conversion missing in ./src/units_conversion.yaml\n'

                print(e0+e1+e2+e0)
                logging.error(e0+e1+e2+e0)

        else:
            units_ok = True
            print('...Units of Seasonal Forecast match units of Reference: {}'.format(self.units_sf))
            
        print('...Units checked.')
            
        return units_ok
            
            
    def check_nans(self):
        
        # This function creates a mask for lon,lat, leadtime points with any nan's along the 
        # member and/or start_time dimensions. A dummy value is used to fill the nan's to avoid
        # problems in the diagnostic computations. Later, the mask is applied to these points.
        
        print("\nChecking presence of nan's...")
        
        nans_ok = True
        
        with ProgressBar():

            Lmembers = len(self.data_sf.coords[self.member_dimname])
            Lyears = len(self.data_sf.coords[self.starttime_dimname])
            Lleadtimes = len(self.data_sf.coords[self.leadtime_dimname])
            
            # Number of nan's per pixel
            Nnans_sf = self.data_sf.isnull().sum([self.member_dimname, 
                                                   self.starttime_dimname, 
                                                   self.leadtime_dimname]).compute()
            Nnans_ref = self.data_ref.isnull().sum([self.starttime_dimname,
                                                     self.leadtime_dimname]).compute()
                
                
            # Boolean total mask
            mask_sf = Nnans_sf > 0
            mask_ref = Nnans_ref > 0
            mask = mask_sf + mask_ref
            
            if mask.sum() > 0:
                
                self.nans_found = True
                
                m0 = "\n---------------------------\n"
                e1 = "-----------ERROR-----------\n"
                w1 = "-----------WARNING-----------\n"
                m11 = "...NAN's found!!!\n"
                m22 = "...NAN's found!!!\nSame number of nan's in each file."
                ew1 = "...Variable {} have different number of nan's for start month {}\n".format(self.variable,
                                                                                                 self.start_month
                                                                                                )
                ew2 = "...Check if variable {} should have nan's\n".format(self.variable)
                ew3 = '...or re-run the transfer in case it is related with creation of the .idx\n'
                w2 = "...Dummy values will be used to fill nan's\n"
                w3 = "...Final maps will be masked.\n"
                
                # Problematic case when there are nans only for some members, leadtimes and/or years 
                nan_problem_files_sf = ((Nnans_sf > 0) * (Nnans_sf < Lmembers*Lyears*Lleadtimes)).sum() > 0
                nan_problem_files_ref = ((Nnans_ref > 0) * (Nnans_ref < Lyears*Lleadtimes)).sum() > 0

                if nan_problem_files_sf or nan_problem_files_ref:

                    logging.error(m0+e1+m11+ew1+ew2+ew3)

                    if self.stop_different_number_nans:
                        nans_ok = False

                        print(m0+e1+m11+ew1+ew2+ew3)
                        #raise NameError(m0+e1+m1+e2+m3+ew)
                        
                    else:
                        print(m0+w1+m11+ew1+ew2+ew3+w2+w3)
                else:
                    print(m0+w1+m22+w2+w3)
                
                # Nan mask
                self.mask = xr.where(mask, np.nan, 1).compute()
                
                # Avoid nan's for the diagnostic computations including a dummy value.
                fill_value = self.data_ref.mean().compute()
                self.data_sf = self.data_sf.fillna(fill_value)
                self.data_ref = self.data_ref.fillna(fill_value)
                
                
            else:
                print("...No NAN's found.")
                self.nans_found = False
                
        return nans_ok
            


    
    def get_data(self):
        
        # This function combines several preparatory steps in order to combine the
        # seasonal forecast and reference data
        
        t.print_header(self.dataset_parameters)
        
        read_ref_ok = False
        units_ok = True
        nans_ok = True
        
        read_sf_ok = self.read_data_sf()
        if read_sf_ok:
            read_ref_ok = self.read_data_ref()
            
            if read_ref_ok:
                
                # Check units
                units_ok = self.check_units()
                
                # Prepare data
                self.reshape_ref()
                self.add_periods()
                
                # Check nans
                nans_ok = self.check_nans()

        if read_sf_ok and read_ref_ok and units_ok and nans_ok:
            ok = True
            
            print('\nData loaded and ready.')
        else:
            ok = False
        
        return ok

        
    
    def bias(self):
        
        # This function just computes the bias between the seasonal forecast and the reference
        
        if self.metrics['bias']:
            print('\nComputing Mean Bias...')

            with ProgressBar():
                data_sf_ensm = self.data_sf.mean(dim = self.member_dimname)    # Ensemble mean
                bias_sf = data_sf_ensm - self.data_ref
                bias_sf = bias_sf.mean(dim = self.starttime_dimname).compute()
                bias_sf = bias_sf.rename('bias')

            print('...Mean Bias computed.')
            return bias_sf
    
    
    
    
    def correlation(self):
        
        # This function computes the correlation between the seasonal forecast and the reference time series
        
        corr, pval_corr = None, None
        
        if self.metrics['correlation']:
            print('\nComputing Correlation...')

            with ProgressBar():
                # Ensemble mean
                data_sf_ensm = self.data_sf.mean(dim = self.member_dimname)   
                # Climatologies for the SF
                data_sf_clim = data_sf_ensm.mean(self.starttime_dimname).compute()  
                # Climatologies for the REF
                data_ref_clim = self.data_ref.mean(self.starttime_dimname).compute()

                #Anomalies
                anom_sf = data_sf_ensm - data_sf_clim
                anom_ref = self.data_ref - data_ref_clim

                #Correlation
                corr, pval_corr = t.corr(anom_sf, anom_ref, dim = self.starttime_dimname)
                corr = corr.rename('correlation')
                pval_corr = pval_corr.rename('pvalues_corr')

            print('...Correlation computed.')
        
        return corr, pval_corr
   




    def frpss(self,
              quantile_threshold = [0.333,0.667]
             ):
        
        # This function computes the frpss for all leadtimes. 
        # n_leadtime defines how many leadtimes are loaded in to memory in each loop
        
        if self.metrics['frpss']:
            
            frpss_list = []
            rpss_list = []
            rps_list = []
            rps_clim_list = []
            for i in range(0,len(self.leadtime_months_T), self.n_leadtimes):

                ll = self.leadtime_months_T[i: i + self.n_leadtimes] # Leadtimes to be loaded in this loop


                print('\nComputing fRPSS and RPSS for level {}, start_month {} and leadtime {}...'.format(self.level,
                                                                                                 self.start_month,
                                                                                                 ll))

                #Extracting coordinates
                coords_frpss = [self.data_sf.coords[self.lat_dimname],
                                self.data_sf.coords[self.lon_dimname],
                                self.data_sf.coords[self.leadtime_dimname][ll[0]-1:ll[-1]]
                               ]
                coords_rpss = coords_frpss
                coords_rps = coords_frpss
                coords_rps_clim = coords_frpss

                with ProgressBar():
                    # Converting to numpy ndarrays for using R2py
                    print('...Loading leadtimes...')
                    data_ref_ll = self.data_ref.sel({self.leadtime_dimname: ll}).compute()
                    data_sf_ll = self.data_sf.sel({self.leadtime_dimname: ll}).compute()

                    print('...Computing terciles...')

                    #Tercile computation for sf
                    q_sf = data_sf_ll.quantile(quantile_threshold, dim = [self.starttime_dimname, self.member_dimname])
                    cat_sf = xr.ones_like(data_sf_ll) * 2
                    cat_sf = cat_sf.where(data_sf_ll > q_sf.isel(quantile = 0), other = 1 )
                    cat_sf = cat_sf.where(data_sf_ll <= q_sf.isel(quantile = 1), other = 3 )

                    #Tercile computation for ref
                    q_ref = data_ref_ll.quantile(quantile_threshold, dim=[self.starttime_dimname])
                    cat_ref = xr.ones_like(data_ref_ll) * 2
                    cat_ref = cat_ref.where(data_ref_ll > q_ref.isel(quantile = 0), other = 1 )
                    cat_ref = cat_ref.where(data_ref_ll <= q_ref.isel(quantile = 1), other = 3 )

                    #Checking points with no variability (std=0). Here the RPS should be undefined (i.e. nan)
                    #cond = (data_sf_ll.std([self.starttime_dimname, self.member_dimname]) == 0) | \
                    #       (data_ref_ll.std([self.starttime_dimname]) == 0)
                    cond = (cat_sf.std([self.starttime_dimname, self.member_dimname]) == 0) | \
                           (cat_ref.std([self.starttime_dimname]) == 0)

                    if cond.sum().values > 0:
                        # Simply imposing tercile=2 to all those locations because it gives no erros in SpecsVerification 
                        # (all terciles = 1 gives error, instead)
                        # Later these points become nan's in the skill score since rps_sf / rps_ref = 0/0 = nan
                        cat_sf = xr.where(cond, 2, cat_sf)
                        cat_ref = xr.where(cond, 2, cat_ref)

                        print('......{} data points found with no variability: Filter implemented.'.format(cond.sum().values))
                        print('......NaNs will appear in the fRPSS.')
                        
                    # Assuring the shape of the numpy arrays are correct
                    cat_sf = cat_sf.transpose(self.starttime_dimname,
                                          self.member_dimname,
                                          self.lat_dimname,
                                          self.lon_dimname,
                                          self.leadtime_dimname
                                         ).values
                    cat_ref = cat_ref.transpose(self.starttime_dimname,
                                          self.lat_dimname,
                                          self.lon_dimname,
                                          self.leadtime_dimname
                                         ).values

                    Nyears_sf, Nmembers_sf, Nlat_sf, Nlon_sf, Nlead_sf = cat_sf.shape
                    Nyears_ref, Nlat_ref, Nlon_ref, Nlead_ref = cat_ref.shape

                    # Check if shapes are consistent
                    if Nyears_sf != Nyears_ref or Nlat_sf != Nlat_ref or Nlon_sf != Nlon_ref or Nlead_sf != Nlead_ref:

                        e1 = '\nERROR. \nINCORRECT DIMENSIONs IN TERCILE MATRIXES:'
                        e2 = '(Nyears_sf, Nlat_sf, Nlon_sf, Nlead_sf) = {} vs (Nyears_ref, Nlat_ref, Nlon_ref, Nlead_ref) = {}' \
                             .format((Nyears_ref, Nlat_ref, Nlon_ref, Nlead_ref),
                                     (Nyears_sf, Nlat_sf, Nlon_sf, Nlead_sf)
                                    )
                        e = e1 + e2
                        print(e)
                        sys.exit(e)
                        logging.error(e)

                print('...Computing fRPS and RPS for the forecast...')
                frps_sf = t.frps(cat_sf, cat_ref, self.n_cores).mean(axis = 0)
                rps_sf = t.rps(cat_sf, cat_ref, self.n_cores).mean(axis = 0)
                
                print('...Computing climatology for the observations...')
                cat_ref_ens_clim = t.obs2ens_clim(cat_ref, self.n_cores)

                print('...Computing fRPS and RPS for the climatology...')
                frps_ref_clim =t.frps(cat_ref_ens_clim, cat_ref, self.n_cores).mean(axis = 0)
                rps_ref_clim =t.rps(cat_ref_ens_clim, cat_ref, self.n_cores).mean(axis = 0)

                print('...Computing fRPSS and RPSS...')
                frpss_ll, frpss_std_ll = t.skillscore(frps_sf, frps_ref_clim, R=True)
                rpss_ll, rpss_std_ll = t.skillscore(rps_sf, rps_ref_clim, R=True)

                frpss_ll = xr.DataArray(frpss_ll, coords = coords_frpss, name = 'frpss')
                rpss_ll = xr.DataArray(rpss_ll, coords = coords_rpss, name = 'rpss')
                rps_sf = xr.DataArray(rps_sf, coords = coords_rps, name = 'rps')
                rps_ref_clim = xr.DataArray(rps_ref_clim, coords = coords_rps_clim, name = 'rps_ref_clim')

                print('...fRPSS and RPSS computed for leadtime/s {}.'.format(ll))

                # Adding to a list to concatenate later
                frpss_list.append(frpss_ll)
                rpss_list.append(rpss_ll)
                rps_list.append(rps_sf)
                rps_clim_list.append(rps_ref_clim)


            print('\nTotal fRPSS and RPSS computed.')
        
            # Concatenate all leadtimes
            frpss = xr.concat(frpss_list, dim = self.leadtime_dimname)
            rpss = xr.concat(rpss_list, dim = self.leadtime_dimname)
            rps_sf = xr.concat(rps_list, dim = self.leadtime_dimname)
            rps_ref_clim = xr.concat(rps_clim_list, dim = self.leadtime_dimname)

        
            return frpss, rpss, rps_sf, rps_ref_clim
        
        
        
        
    def fcrpss(self):
        
        # This function computes the fcrpss for all leadtimes. 
        # self.n_leadtime defines how many leadtimes are loaded in to memory in each loop
        
        if self.metrics['fcrpss']:
            fcrpss_list = []
            crpss_list = []
            crps_list = []
            crps_clim_list = []
            for i in range(0,len(self.leadtime_months_T), self.n_leadtimes):

                ll = self.leadtime_months_T[i: i + self.n_leadtimes] # Leadtimes to be loaded in this loop

                print('\nComputing fcRPSS and cRPSS for level {}, start_month {} and leadtime/s {}...'.format(self.level,
                                                                                                    self.start_month,
                                                                                                    ll))

                #Extracting coordinates
                coords_fcrpss = [self.data_sf.coords[self.lat_dimname],
                                 self.data_sf.coords[self.lon_dimname],
                                 self.data_sf.coords[self.leadtime_dimname][ll[0]-1:ll[-1]]
                                ]
                coords_crpss = coords_fcrpss
                coords_crps = coords_fcrpss
                coords_crps_clim = coords_fcrpss

                with ProgressBar():

                    # Computing anomalies and converting to numpy ndarrays
                    print('...Computing anomalies...')
                    data_ref_ll = self.data_ref.sel({self.leadtime_dimname: ll}).compute()
                    data_ref_ll = data_ref_ll - data_ref_ll.mean(self.starttime_dimname)
                    data_ref_ll = data_ref_ll.values

                    data_sf_ll = self.data_sf.sel({self.leadtime_dimname: ll}).compute()
                    data_sf_ll = data_sf_ll - data_sf_ll.mean(self.starttime_dimname)
                    data_sf_ll = data_sf_ll.values

                    # Reshaping ndarray for correct matching
                    if data_ref_ll.ndim == 3:
                        data_ref_ll = np.transpose(data_ref_ll, (2,0,1))
                    else:
                        data_ref_ll = np.transpose(data_ref_ll, (2,0,1,3))
                        data_sf_ll = np.transpose(data_sf_ll, (0,1,3,4,2))

                    data_ref_shape = data_ref_ll.shape

                    if self.originating_centre == 'ncep' or (self.originating_centre == 'ukmo' and self.start_month < 5):
                        data_sf_ll = np.swapaxes(data_sf_ll, 0,1)

                    data_sf_shape = data_sf_ll.shape
                    
                    Nyears_sf, Nmembers_sf, Nlat_sf, Nlon_sf = data_sf_shape[0], data_sf_shape[1], data_sf_shape[2], data_sf_shape[3]
                    Nyears_ref, Nlat_ref, Nlon_ref = data_ref_shape[0], data_ref_shape[1], data_ref_shape[2]

                    # Check if shapes are consistent
                    if Nyears_sf != Nyears_ref or Nlat_sf != Nlat_ref or Nlon_sf != Nlon_ref:

                        e1 = '\nERROR. \nINCORRECT DIMENSIONs IN TERCILE MATRIXES:'
                        e2 = '(Nyears_ref, Nlat_ref, Nlon_ref) = {} vs (Nyears_sf, Nlat_sf, Nlon_sf) = {}'.format((Nyears_ref, Nlat_ref, Nlon_ref),
                                                                                                                     (Nyears_sf, Nlat_sf, Nlon_sf)
                                                                                                                    )
                        e = e1 + e2
                        print(e)
                        sys.exit(e)
                        logging.error(e)

                print('...Computing fCRPS and CRPS for the forecast...')
                fcrps_sf = t.fcrps(data_sf_ll, data_ref_ll, self.n_cores).mean(axis = 0)
                crps_sf = t.crps(data_sf_ll, data_ref_ll, self.n_cores).mean(axis = 0)

                print('...Computing climatology for the observations...')
                data_ref_ens_clim = t.obs2ens_clim(data_ref_ll, self.n_cores)

                print('...Computing fCRPS and CRPS for the climatology...')
                fcrps_ref_clim =t.fcrps(data_ref_ens_clim, data_ref_ll, self.n_cores).mean(axis = 0)
                crps_ref_clim =t.crps(data_ref_ens_clim, data_ref_ll, self.n_cores).mean(axis = 0)

                print('...Computing fCRPSS and CRPSS...')
                fcrpss_ll, fcrpss_std_ll = t.skillscore(fcrps_sf, fcrps_ref_clim, R=True)
                crpss_ll, crpss_std_ll = t.skillscore(crps_sf, crps_ref_clim, R=True)

                fcrpss_ll = xr.DataArray(fcrpss_ll, coords = coords_fcrpss, name = 'fcrpss')
                crpss_ll = xr.DataArray(crpss_ll, coords = coords_crpss, name = 'crpss')
                crps_sf = xr.DataArray(crps_sf, coords = coords_crps, name = 'crps')
                crps_ref_clim = xr.DataArray(crps_ref_clim, coords = coords_crps_clim, name = 'crps_ref_clim')

                print('...fCRPSS and CRPSS computed for leadtime/s {}.'.format(ll))

                # Adding to a list to concatenate later
                fcrpss_list.append(fcrpss_ll)
                crpss_list.append(crpss_ll)
                crps_list.append(crps_sf)
                crps_clim_list.append(crps_ref_clim)


            print('\nTotal fCRPSS and CRPSS computed.')

            # Concatenate all leadtimes
            fcrpss = xr.concat(fcrpss_list, dim = self.leadtime_dimname)
            crpss = xr.concat(crpss_list, dim = self.leadtime_dimname)
            crps_sf = xr.concat(crps_list, dim = self.leadtime_dimname)
            crps_ref_clim = xr.concat(crps_clim_list, dim = self.leadtime_dimname)
        
            return fcrpss, crpss, crps_sf, crps_ref_clim
    
    

    def save_metrics(self,
                     data_dict
                    ):
        
        # This function stores the dataarrays containd in data_dict. To be used to store all the computed metrics.
        # It also extract some attributes from the dataarrays and cdsforms and includes them in the saved netcdf
        # and the class
        
        self.update_level()
        
        # Get required attributes (units and shortname)
        attrs = dict()
        attrs['units'] = self.units_sf
        attrs['varname_short'] = self.varname_short_sf
        attrs['varname_long'] = self.varname_long_sf
        
        # File to be saved.
        filename = 'EQCmetrics-{}s{}-{}-{}to{}-s{:02d}{}.nc'.format(self.originating_centre,
                                                              self.system,
                                                              self.variable,
                                                              self.years[0],
                                                              self.years[1],
                                                              self.start_month,
                                                                  self.level_ext2
                                                             )
        filenameflag = 'EQCmetrics-{}s{}-{}-{}to{}-s{:02d}{}.txt'.format(self.originating_centre,
                                                              self.system,
                                                              self.variable,
                                                              self.years[0],
                                                              self.years[1],
                                                              self.start_month,
                                                                  self.level_ext2
                                                             )
        # Check if folder exists, if not it creates it
        self.metrics_data_path = self.metrics_data_path
        if not os.path.exists(self.metrics_data_path):
            os.makedirs(self.metrics_data_path)
         
        # File path
        filename_path = self.metrics_data_path.joinpath(filename)
        
        # Remove if exist to avoid writting permission problems
        if os.path.exists(filename_path):
            os.remove(filename_path)
        
        print('\nSaving computed metrics in...{}'.format(filename_path))
        
        # Masking metrics in case of nan's
        if self.nans_found:
            print("...masking metrics in case of nan's...")
            for di in data_dict:
                data_dict[di] = data_dict[di]*self.mask
        
        d = xr.Dataset(data_dict,
                       attrs = attrs
                      )
        
        d.to_netcdf(filename_path)

        #filename_path = self.metrics_data_path.joinpath(filenameflag)
        #with open(filename_path, 'w') as fp: 
        #    pass
        
        print('\nFile saved.')
        
    
    
    def load_metrics(self):
        
        # This function loads previously computed metrics. To be used for debugging plotting scripts.
        # It also add required variable attributes
        
        # Update level
        self.update_level()
           
        # File to be saved.
        filename = 'EQCmetrics-{}s{}-{}-{}to{}-s{:02d}{}.nc'.format(self.originating_centre,
                                                              self.system,
                                                              self.variable,
                                                              self.years[0],
                                                              self.years[1],
                                                              self.start_month,
                                                                    self.level_ext2
                                                             )
        
        filename_path = self.metrics_data_path.joinpath(filename)
        
        if os.path.exists(filename_path):
            ok_load = True

            print('\nLoading previously computed metrics from...\n{}'.format(filename_path))

            data = xr.open_dataset(filename_path)

            # Extracting required attributes (longname, shortname and units)
            attrs = t.extract_attrs(data,
                                    attrs_in = ['units', 'varname_short', 'varname_long']
                                   )

            self.units = attrs['units']
            self.varname_short = attrs['varname_short']
            self.varname_long = attrs['varname_long']

            print('\nFile {} loaded.'.format(filename_path))
            
        else:
            ok_load = False
            data = None
            
            e0 = '\nMetrics file missing: {}\n'.format(filename_path)
            e1 = '\nPlotting for {}-s{}-{}:\n...level: {}\n...start_month: {}\n...SKIPPED!!!\n'.format(self.originating_centre,
                                                                                       self.system,
                                                                                       self.variable,
                                                                                        self.level,
                                                                                       self.start_month
                                                                                      )
            print(e0+e1)

        return data, ok_load
    
    
    def check_metrics(self):
        
        # This function checks if metrics for a given start_months exist / are already computed
        
        # Update level
        self.update_level()
        
        # File to be saved.
        filename = 'EQCmetrics-{}s{}-{}-{}to{}-s{:02d}{}.nc'.format(self.originating_centre,
                                                              self.system,
                                                              self.variable,
                                                              self.years[0],
                                                              self.years[1],
                                                              self.start_month,
                                                                  self.level_ext2
                                                             )
        filenameflag = 'EQCmetrics-{}s{}-{}-{}to{}-s{:02d}{}.txt'.format(self.originating_centre,
                                                              self.system,
                                                              self.variable,
                                                              self.years[0],
                                                              self.years[1],
                                                              self.start_month,
                                                                  self.level_ext2
                                                             )
        
        filename_path = self.metrics_data_path.joinpath(filename)
        
        if os.path.exists(filename_path):
            ok_check = True
            
            print('\nMetrics already computed for level {} and start_month {}.'.format(self.level,
                                                                                       self.start_month
                                                                                      ))
            print('Saved at {}.'.format(filename))
        else:
            ok_check = False

        filename_path = self.metrics_data_path.joinpath(filenameflag)
        
        if os.path.exists(filename_path):
            ok_flag = True
        else:
            ok_flag = False
            
        return ok_check, ok_flag
        