#!/usr/bin/env python
#
# Data Quality Checker also known as DQC
#
# Read the README.md for more detailed information
#
# Barcerlona Supercomputing Center - Centro Nacional de Supercomputacion
#        Earth Sciences Department - Computational Earth Sciences
#     
#     Authors: Jesús Peña-Izquierdo (jesus.pena@bsc.es)
#             -
#       
import xarray
import matplotlib.pyplot as plt
import numpy as np
import os
import sys
import pandas as pd
from scipy.stats import kurtosis, skew
#from statsmodels.stats.stattools import medcouple
import time
import datetime
import logging
import json


class valid_ranges_test:

    def __init__(self,
                 path,
                 fileformat = None,
                 nbins = 50,
                 OutIn = True,
                 xbin_initial_factor = 2
                ):
        
        self.fileformat = fileformat
        self.path = path
        self.varname = self.find_varname()
        self.datetime_tag = time.strftime("%Y%m%d_%H%M")
        self.log_path = '/data/DQC_LOGS'
        self.files = self.list_files_folder()
        self.nbins = nbins
        self.OutIn = OutIn
        self.xbin_initial_factor = xbin_initial_factor
        self.stats_varnames = ['Minimum', 'Maximum',
                    'Mean', 'Median',
		    'Kurtosis', 'Skewness', 
                    'Std', 'Std_menos', 'Std_mas', 
                    'Mad', 'Mad_menos', 'Mad_mas',
                    'MC', 'Q1', 'Q3', 'IQD']
        self.stats_k_varnames = ['k_sd_min', 'k_sd_max',
                    'k_sd_menos_min', 'k_sd_mas_max',
                    'k_mad_min', 'k_mad_max',
                    'k_mad_menos_min', 'k_mad_mas_max']
    
    def find_varname(self):
        islash = [i for i, ltr in enumerate(self.path) if ltr == '/'][-1]
        varname = self.path[islash+1:]
        
        return varname
    
    
    def list_files_folder(self):
        list_files = []
        
        #print(self.fileformat)
        
        if os.path.exists(self.path) and os.path.isdir(self.path):
            for root, dirs, files in os.walk(self.path):
                for f in files:
                    #print(f)
                    #print(self.fileformat)
                    #print(f.endswith('.nc') , self.fileformat == 'netcdf')
                    #print(f.endswith('.nc') and self.fileformat == 'netcdf')
                    if (self.fileformat == 'netcdf' and f.endswith('.nc')):                        
                        list_files.append(os.path.join(root,f))
                    elif (self.fileformat == 'grib' and f.endswith('.grib')):
                        list_files.append(os.path.join(root,f))
            
            return sorted(list_files)
        else:
            raise Exception('\n\tGiven path is not a valid path\n')
        
        return list_files
        
    def initiate_logging(self):
        # Initiate the logging
        
        path_logs = self.log_path

        if not os.path.exists(path_logs):
            os.makedirs(path_logs)
            
        formatter = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'

        log_filename = '_'.join(['LOG_TEST06_VALID_RANGE',
                                 (self.fileformat).upper(),
                                 self.datetime_tag,
                                 self.varname]) + '.log'

        log_filename_path = os.path.join(
            path_logs,
            log_filename)

        logging.shutdown()
        logging.basicConfig(filename=log_filename_path, level="DEBUG", format=formatter,
                        filemode="w")
        
        stringA = '\n' + '#'*80 + '\n' + '#'*80
        stringA += '\n RUNNING STANDARD COMPLIANCE TEST FOR FOLDER:'
        stringA += '\n ' + self.path
        stringA += '\n IN ' + (self.fileformat).upper() + ' FILE FORMAT '
        stringA += '\n EXECUTION TIME: ' + self.datetime_tag
        stringA += '\n' + '#'*80

        logging.info(stringA)

     
    #Loading data depending on fileformat
    def load_datafile(self,
                      filepath,
                      dt = 24
                     ):
        if filepath.endswith('.nc'):
            if dt is not None:
                data = xarray.open_dataset(filepath)
                Nt = len(data.time)
                for i in range(0, Nt, dt):
                    data_i = data.isel(time=range(i,i+dt)).to_array().values
                    data_i = data_i[~np.isnan(data_i)].flatten()
                    yield data_i, dt
                    
        if filepath.endswith('.grib'):
            print('WILL BE READY SOON')
    
    #Producing a list of statistical estimators
    def Stats(self,
              data
             ):
        #data to be a flatten numpy array

        Mean = np.mean(data)
        Median = np.median(data)
        Std = np.std(data)

        #MC = iter_medcouple(data, ni = 1000)
        MC = np.nan #skew(data)
        Q1 = np.quantile(data,0.25)
        Q3 = np.quantile(data,0.75)
        IQD = Q3 - Q1

        #Mad
        diff = np.abs(data - Median)
        Mad = np.median(diff)

        diff_menos = np.abs(data[data <= Median] - Median)
        #Mad_menos = np.median(diff_menos)
        Mad_menos = np.mean(diff_menos)

        diff_mas = np.abs(data[data >= Median] - Median)
        #Mad_mas = np.median(diff_mas)
        Mad_mas = np.mean(diff_mas)

        Std_menos =  np.sqrt( np.mean( np.square(data[data<=np.mean(data)] - np.mean(data)))) #np.mean( np.abs(data[data<=np.mean(data)] - np.mean(data)))
        Std_mas = np.sqrt( np.mean( np.square(data[data>=np.mean(data)] - np.mean(data)))) #np.mean( np.abs(data[data>=np.mean(data)] - np.mean(data)))

        Kurt = kurtosis(data)
        Skew = skew(data)

        stats_list = [Mean, Median, Kurt, Skew, Std, Std_menos, Std_mas, Mad, Mad_menos, Mad_mas, MC, Q1, Q3, IQD]

        return stats_list


    # Monitoring the absolute minimum and maximun values during online computation
    def minmax_monitor(self,
                       data,
                       minmax_record
                      ):

        if len(minmax_record)==2:
            min_record, max_record = minmax_record

            min_record = min([np.min(data), min_record])
            max_record = max([np.max(data), max_record])

        else:
            min_record = np.min(data)
            max_record = np.max(data)    

        minmax_record = np.array([min_record, max_record])
        return minmax_record

    # Weighted average for online computation 
    def weight_avg(self,
                   N__Ni,
                   Var_n__Var_i
                  ):

        N, Ni = N__Ni
        Var_n, Var_i = Var_n__Var_i

        Var_M = []

        for var_n_k, var_i_k in zip(Var_n, Var_i):
            var_m_k = (var_n_k*N + var_i_k*Ni) / (N + Ni)
            Var_M.append(var_m_k)

        return Var_M

    # Computing data distribution online
    def online_distribution(self,
                            data,
                            limits_record,
                            distribution_record,
                            nbins,
                            OutIn = True
                            ):

        dlimits = np.diff(limits_record) / self.xbin_initial_factor
        limits_record_ext = (int(limits_record[0] - dlimits) , int(limits_record[1] + dlimits))
        
        n_i, xbin = np.histogram(data, bins = nbins, range = limits_record_ext)

        #For accumulating outside xbin values within the edges
        if OutIn:
            n_outmin = len(np.where(data < limits_record_ext[0])[0])
            n_outmax = len(np.where(data > limits_record_ext[1])[0])
            n_i[0] = n_i[0] + n_outmax
            n_i[-1] = n_i[-1] + n_outmin

        distribution_record = distribution_record + n_i

        return distribution_record, xbin

    # Averaging statistical estimators online
    def avg_online_stats(self):

        iMean = 0
        iMed = 1
        iStd = 4
        iMad = 7

        files = self.files
        
        Nt = 0
        minmax_record = []

        for ifile in files[:12]:
            
            print('Analysing... {}'.format(ifile))
            logging.info('Loading file and computing stats... {}'.format(ifile))

            #Loading single files in chunks
            for data, Nt_i in self.load_datafile(ifile, dt=24):
                
                #MinMax record
                minmax_record =  self.minmax_monitor(data, minmax_record)

                #Stats
                stats_data_i = self.Stats(data)
                if Nt == 0:
                    stats_data_n = [0,] * len(stats_data_i)
                    limits_record = minmax_record
                    distribution_record = np.zeros((self.nbins,))

                stats_data_n = self.weight_avg( (Nt, Nt_i), (stats_data_n, stats_data_i))

                distribution_record, xbin = self.online_distribution(data, limits_record, distribution_record, 
                                                                nbins = self.nbins, OutIn = self.OutIn)
                
            logging.info('File loaded and stats computed.')

        #Proportion between Mean, sd and min/max
        k_sd = np.empty((3,2))

        #Proportion between Mads and (Min/Max - Median)
        k_mad = np.empty((3,2))
        for i in range(3):
            if stats_data_n[iMad + i] == 0:
                k_mad[i,:] = np.zeros((2,))
            else:
                k_mad[i,:] = np.abs(minmax_record - stats_data_n[iMed]) / stats_data_n[iMad + i]
            k_sd[i,:] = np.abs(minmax_record - stats_data_n[iMean]) / stats_data_n[iStd + i]

        k_sd[1,1] = k_sd[2,1] 
        k_sd = k_sd[:2,:].flatten()

        k_mad[1,1] = k_mad[2,1] #First row (Min/Max - Median), second row (Min/Max - Median_menos/Median_mas)
        k_mad = k_mad[:2,:].flatten()        


        Nt += Nt_i

        stats_var = np.concatenate((minmax_record, stats_data_n))
        stats_k =np.concatenate((k_sd, k_mad))

        logging.info('{} analysis completed.'.format(ifile))
            
           
        return stats_var, stats_k, distribution_record, xbin
    
    
    def check(self,
              k_sd_REF = 10
             ):
        iMin = 0
        iMax = 1
        iMean = 2
        iStd = 6
        
        self.initiate_logging()
        
        result_base_filename = 'TEST06_VALID_RANGE_RESULTS'
        
        stats_var, stats_k, _, _ = self.avg_online_stats()
                      
        result_filename = os.path.join(self.log_path,
                                        '_'.join([result_base_filename, self.varname,self.datetime_tag])) + '.txt'
        #Valid range estimation
        vr = []
        vr.append( stats_var[iMean] - k_sd_REF*stats_var[iStd] )
        vr.append( stats_var[iMean] + k_sd_REF*stats_var[iStd] )
        
        #Check
        k_sd = stats_k[0:2]
        if k_sd.any() > k_sd_REF:
            result1 = 'Test NOT PASSED: [ Min, Max ] = [{} , {}]'.format(str(stats_var[iMin]), str(stats_var[iMax]))
        else:
            result1 = 'Test PASSED: [ Min, Max ] = [{} , {}]'.format(str(stats_var[iMin]), str(stats_var[iMax]))
        
        result2 = 'Estimated valid range (k_sd_REF = {}) [Min_valid, Max_valid] = [{} , {}]'.format(str(k_sd_REF), str(vr[0]), str(vr[1]))
        
        logging.info(result1) 
        logging.info(result2) 
        logging.info('Writting extended results to {}'.format(result_filename))  
        
        stats_dict = {n: s for n, s in zip(self.stats_varnames, stats_var)}
        stats_k_dict = {n_k: s_k for n_k, s_k in zip(self.stats_k_varnames, stats_k)}
        
        #Creating result text file
        with open(result_filename, "w") as rf:
            json.dump('Results for the VALID RANGE TEST for {}'.format(self.path), rf,indent = 4)
            rf.write('\n')
            json.dump(result1, rf, indent = 4)
            rf.write('\n')
            json.dump(result2, rf, indent = 4)
            rf.write('\n')
            json.dump(stats_dict , rf, indent = 4)
            json.dump(stats_k_dict , rf, indent = 4)
            
        print('Results stored in {}'.format(result_filename))


#     def monitor_multivar(self,
#                          filespath,
#                          vars,
#                          outpath = '/home/Earth/jpena/Suso/Coding/Python/',
#                         stats_fileout = 'stats.csv',
#                         distrib_fileout = 'distributions.csv',
#                         bins = 10, OutIn = True, Ni_files = 12):

#         cols_name = ['Minimum', 'Maximum',
#                     'Mean', 'Median',
#                     'Skewness', 'Kurtosis',
#                     'Std', 'Std_menos', 'Std_mas', 
#                     'Mad', 'Mad_menos', 'Mad_mas',
#                     'MC', 'Q1', 'Q3', 'IQD',
#                     'k_sd_min', 'k_sd_max',
#                     'k_sd_menos_min', 'k_sd_mas_max',
#                     'k_mad_min', 'k_mad_max',
#                     'k_mad_menos_min', 'k_mad_mas_max']

#         stats_vars = pd.DataFrame(columns = cols_name)
#         distribution_vars = pd.DataFrame() # = pd.Series()

#         for var_i in vars:

#             i_ = vars[0].find('_')
#             if i_>0:
#                 var_ii = var_i[:i_]
#             else:
#                 var_ii = var_i
#             filespath_var_i = filespath + var_i + '/'

#             print('Reading {}{}'.format(filespath_var_i,var_ii))

#             #stats_var_i, distribution_record, xbin = avg_online_stats_multi(filespath_var_i, var_ii, bins = bins, OutIn = OutIn, Ni_files = Ni_files)
#             stats_var_i, distribution_record, xbin = avg_online_stats(filespath_var_i, var_ii, bins = bins, OutIn = OutIn)


#             s_stats_var_i = pd.Series( stats_var_i, name = var_ii, index = cols_name)
#             stats_vars = stats_vars.append(s_stats_var_i)

#             s_distrib_var_i = pd.Series( distribution_record, name = var_ii) # = pd.Series( distribution_record, index = [var_i]*bins, xbin)
#             distribution_vars = distribution_vars.append(s_distrib_var_i) # = distribution_vars.append(s_distrib_var_i)
#             s_distrib_var_i_xbin = pd.Series( xbin, name = var_ii + '_xbin') # = pd.Series( distribution_record, index = [var_i]*bins, xbin)
#             distribution_vars = distribution_vars.append(s_distrib_var_i_xbin) # = distribution_vars.append(s_distrib_var_i)

#             print(' ')
#             print(var_i + ' added.')
#             print(' ')


#         stats_vars.to_csv(outpath + stats_fileout)
#         distribution_vars.to_csv(outpath + distrib_fileout)


#         return stats_vars, distribution_vars
            

    
            

    
    
if __name__ == '__main__':
    
    sys.exit()
