"""
This python file takes output_parameters folder as input and provides 
asymmetric errors using the distribution of samples. 
"""
import glob
import pickle
import numpy as np
import pandas as pd
from statsmodels.stats import weightstats


def get_quantiles_on_best_val(samples, weights, best_val):
    """Generates lower and upper limit of errors for the best values.
    Sorts the samples and corresponding weights. 
    Gets value of samples such that they encompass 68.27% of the samples
    by weight on both sides of the best value.

    Args:
        samples (array): the sampled values for the parameter from dynesty
        weights (array): weights of the samples
        best_val (float): best value among the samples

    Returns:
        tuple: the lower and upper error on the best value.
    """

    weights = weights/np.sum(weights)
    sorter_sample = np.argsort(samples)

    sorted_weights = weights[sorter_sample]
    sorted_samples = samples[sorter_sample]

    best_val_idx = np.argmin(np.abs(sorted_samples-float(best_val)))

    # Faster method
    best_percentile = np.sum(sorted_weights[:best_val_idx])
    model = weightstats.DescrStatsW(sorted_samples, sorted_weights)
    lower_error, upper_error = model.quantile([(1-.6827)*best_percentile, best_percentile+(.6827*(
        1-best_percentile))], return_pandas=False)-sorted_samples[best_val_idx]

    """
    # Previous method (slower)
    # Shows more detail on how we are getting the errors
    
    total_sum = 0
    for i in range(best_val_idx, len(sorted_weights)):
        total_sum += sorted_weights[i]
        if total_sum >= .6827*np.sum(sorted_weights[best_val_idx:]):
            upper_error = sorted_samples[i]-best_val
            break

    total_sum = 0
    for i in range(best_val_idx, 0, -1):
        total_sum += sorted_weights[i]
        if total_sum >= .6827*np.sum(sorted_weights[:best_val_idx]):
            lower_error = sorted_samples[i]-best_val
            break"""

    return -np.abs(lower_error), np.abs(upper_error)

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#


def make_dict(val_dict, key, vals):
    """makes dictionary using the given values. if the key already 
    exists, then the values get appended. 

    Args:
        val_dict (dict): the dictionary where the values are to be added
        key (str): the key of the value
        vals (array): the values to be added
    """
    if key in val_dict:
        val_dict[key] = np.ndarray.flatten(np.append(val_dict[key], vals))
    else:
        val_dict[key] = np.ndarray.flatten(vals)


def q_to_u_err(q, q_err):
    """Converts error in q to errors in u parameters for quadratic limb-darkening

    Args:
        q (list): values of q0, q1
        q_err (list): errors on q0, q1

    Returns:
        list: errors on u0, u1
    """

    u_err = []
    u_err.append(np.sqrt(((q[1] * q_err[0])**2)/q[0] + 4 * q[0] * q_err[1]**2))
    u_err.append(np.abs(np.sqrt(
        (((1 - 2 * q[1]) * q_err[0])**2)/(0.25 * q[0]) + 4 * q[0] * q_err[1] ** 2)))

    return u_err


def q_to_u(q, q_err_l, q_err_u):
    """Converts q (Kipping parameters) to limb darkening parameters u.
    Handles only quadratic limb darkening currently.

    Args:
        q (list): values of q0, q1
        q_err_l (list): lower bound of errors on q0, q1
        q_err_u (list): upper bound of errors on q0, q1

    Returns:
        tuple: the values of u, lower bound on values, upper bound on values.
    """

    q_err_l = np.abs(np.array(q_err_l,dtype=float))
    q_err_u = np.abs(np.array(q_err_u,dtype=float))


    u = [2 * np.sqrt(q[0]) * q[1], np.sqrt(q[0]) * (1 - 2 * q[1])]

    u_err_l = q_to_u_err(q, q_err_l)
    u_err_u = q_to_u_err(q, q_err_u)

    return np.array(u,dtype=float), -np.abs(np.array(u_err_l,dtype=float)), np.abs(np.array(u_err_u,dtype=float))


class ErrorLimits:
    """
    Initializes the error limit analysis.

    Args:
        output_parmeters (str): path of the output_parameters folder
                                generated by TransitFit
    """

    def __init__(self, output_parmeters):
        
        # The path to the output_parameters folder from the TransitFit output.
        self.OUTPUT_PARAMETERS_FOLDER = output_parmeters

        # This will be the name of the output file
        self.MODIFIED_SUMMARY_OUTPUT = 'modified_output.csv'

        # Summary_output file
        self.summary_output = glob.glob(
            self.OUTPUT_PARAMETERS_FOLDER+'/*summary_output.csv')
        if len(self.summary_output) == 0:
            #print('')
            raise RuntimeError(
                'Please check the pathname, it may not be correct.')

            # return

        # Searching for all the .pkl files
        self.files_pkl = glob.glob(
            self.OUTPUT_PARAMETERS_FOLDER+'/quicksaves/*results.pkl')
        self.files_pkl.sort()

        # Searching for all the .csv files
        self.output_csvs = glob.glob(
            self.OUTPUT_PARAMETERS_FOLDER+'/quicksaves/*output.csv')
        self.output_csvs.sort()

        # In folded mode, there is additional filter-dependent folder inside
        # output_parameters, checking if that is the case.
        self.folded_mode_csv = glob.glob(
            self.OUTPUT_PARAMETERS_FOLDER+'/*/quicksaves/*output.csv')

        if len(self.folded_mode_csv) != 0:
            print('Folded mode detected.')
            self.FOLDED_MODE = True
            self.folded_mode_pkl = glob.glob(
                self.OUTPUT_PARAMETERS_FOLDER+'/*/quicksaves/*results.pkl')
            self.folded_mode_pkl.sort()
            self.folded_mode_csv.sort()

            self.folded_mode_priors_pkl = glob.glob(
                self.OUTPUT_PARAMETERS_FOLDER+'/*/quicksaves/*priors.pkl')
            self.folded_mode_priors_pkl.sort()
            # if len(unfolded_priors_pkl)>0:
            with open(self.folded_mode_priors_pkl[0], 'rb') as handle:
                priors = pickle.load(handle)

            self.allow_ttv = priors.allow_ttv
            if self.allow_ttv:
                print('TTV mode detected.')

        else:
            self.allow_ttv = False
            self.FOLDED_MODE = False
            self.folded_params = []
            self.priors_pkl = glob.glob(
                self.OUTPUT_PARAMETERS_FOLDER+'/quicksaves/*priors.pkl')
            self.priors_pkl.sort()
            with open(self.priors_pkl[0], 'rb') as handle:
                priors = pickle.load(handle)

        # The list of required parameters to analyze
        self.required_params = ['P', 't0', 'a/r*', 'inc', 'rp/r*']
        self.limb_dark_coeffs = priors.limb_dark_coeffs#['q0','q1']#
        self.ld_method = priors.limb_dark #'quadratic' #
        self.values = {}

    def handle_ttv(self):
        """ Handles the case of TTV analysis.
        """
        all_epochs = np.empty(0)
        for i, fmpk in enumerate(self.folded_mode_priors_pkl):
            with open(fmpk, 'rb') as handle:
                priors = pickle.load(handle)

            fitting_params = priors.fitting_params  # [name, tidx,fidx, eidx]
            # filters = fitting_params[:,2]
            params = fitting_params[:, 0]

            df = pd.read_csv(self.folded_mode_csv[i])
            selected_df = df.loc[df['Parameter'] == 't0']
            epochs = selected_df['Epoch'].to_numpy(dtype=float)
            all_epochs = np.append(all_epochs, epochs)
            with open(self.folded_mode_pkl[i], 'rb') as handle:
                results = pickle.load(handle)

            for j, p in enumerate(params):
                if p == 't0':
                    make_dict(self.values, p+'_' +
                              str(int(epochs[j])), results.samples[:, j])
                    make_dict(self.values, p+'_' +
                              str(int(epochs[j]))+'_weights', results.weights)
                    self.values[p+'_'+str(int(epochs[j])) +
                                '_best'] = results.best[j]

        priors_pkl = glob.glob(
            self.OUTPUT_PARAMETERS_FOLDER+'/quicksaves/*priors.pkl')
        priors_pkl.sort()
        for i, pp in enumerate(priors_pkl):
            with open(pp, 'rb') as handle:
                priors = pickle.load(handle)
            fitting_params = priors.fitting_params  # [name, tidx,fidx, eidx]
            # filters=fitting_params[:,2]
            params = fitting_params[:, 0]

            with open(self.files_pkl[i], 'rb') as handle:
                results = pickle.load(handle)

            for j, p in enumerate(params):
                # elif params[p] in {'a','inc','rp','q0','q1'}:
                make_dict(self.values, p, results.samples[:, j])
                make_dict(self.values, p+'_weights', results.weights)
                self.values[p+'_best'] = results.best[j]

        all_epochs = np.unique(all_epochs)

        # Generating the output
        with open(self.OUTPUT_PARAMETERS_FOLDER+'/'+self.MODIFIED_SUMMARY_OUTPUT, 'w') as mso:
            mso.write(
                'Parameter,Filter,Epoch,Best,Lower_error,Upper_error\n')
            for e in all_epochs:
                param_ = 't0_'+str(int(e))
                errs = get_quantiles_on_best_val(
                    samples=self.values[param_], weights=self.values[param_+'_weights'], best_val=self.values[param_+'_best'])
                # (model.mean)
                mso.write(
                    f"t0,-,{e},{self.values[param_+'_best']},{errs[0]},{errs[1]}\n")

            q, q_low, q_up = [], [], []
            for p in params:
                param_ = p
                errs = get_quantiles_on_best_val(
                    samples=self.values[param_], weights=self.values[param_+'_weights'], best_val=self.values[param_+'_best'])
                # (model.mean)

                mso.write(
                    f"{p},-,-,{self.values[param_+'_best']},{errs[0]},{errs[1]}\n")

                if p in self.limb_dark_coeffs:
                    q.append(self.values[param_+'_best']),
                    q_low.append(errs[0])
                    q_up.append(errs[1])
            
            if self.ld_method=='quadratic':
                u, u_low, u_up = q_to_u(q, q_low, q_up)

                for i, p in enumerate(self.limb_dark_coeffs):
                    mso.write(f"u{str(i)},-,-,{u[i]},{u_low[i]},{u_up[i]}\n")
                    #mso.write(f"u0,-,-,{u[0]},{u_low[0]},{u_up[0]}\n")
                    #mso.write(f"u1,-,-,{u[1]},{u_low[1]},{u_up[1]}\n")

        print(
            f"Saved results in {self.OUTPUT_PARAMETERS_FOLDER+'/'+self.MODIFIED_SUMMARY_OUTPUT}")

        return

    def get_errors(self):
        """ Gets the upper and lower error bound on values.
        """

        # Initalising values dictionary and filters list.
        # selvalues = {}
        filters = []

        if self.allow_ttv:
            self.handle_ttv()
            return

        # Getting the values of 'P' and 't0' in case of folded mode.
        elif self.FOLDED_MODE:
            self.folded_params = self.required_params[:2]
            self.required_params = self.required_params[2:]
            for i, file in enumerate(self.folded_mode_pkl):
                with open(file, 'rb') as handle:
                    result = pickle.load(handle)

                df_csv = pd.read_csv(self.folded_mode_csv[i])

                for i, param in enumerate(self.folded_params):
                    selected_df = df_csv.loc[df_csv['Parameter'] == param]
                    index = selected_df.index

                    make_dict(self.values, param, result.samples[:, index])
                    make_dict(self.values, param+'_weights', result.weights)

        for i, file in enumerate(self.files_pkl):
            with open(file, 'rb') as handle:
                result = pickle.load(handle)

            df_csv = pd.read_csv(self.output_csvs[i])

            for i, param in enumerate(self.required_params):
                selected_df = df_csv.loc[df_csv['Parameter'] == param]
                index = selected_df.index
                if param in {'inc', 'rp/r*'}:
                    # To account for a/AU parameter present in the
                    # summary_output files.
                    index = index-1
                
                if '-' not in selected_df['Filter'].to_list():
                    batch_filters = selected_df['Filter'].to_list()
                    filters += batch_filters
                    for i, f in enumerate(batch_filters):
                        make_dict(self.values, param+'_'+f,
                                  result.samples[:, index[i]])
                        make_dict(self.values, param+'_'+f +
                                  '_weights', result.weights)
                        
                        for j,u in enumerate(self.limb_dark_coeffs):
                            i = len(self.limb_dark_coeffs)*i  
                            # limb darkenening coeffs are saved together in the .pkl files
                            
                            make_dict(self.values, 'q'+str(j)+'_'+f,
                                    result.samples[:, index[-1]+1+i])
                            make_dict(self.values, 'q'+str(j)+'_'+f +
                                    '_weights', result.weights)
                            
                            #if self.ld_method=='quadratic':
                            """make_dict(self.values, 'q0'+'_'+f,
                                    result.samples[:, index[-1]+1+i])
                            make_dict(self.values, 'q1'+'_'+f,
                                    result.samples[:, index[-1]+1+i+1])

                            make_dict(self.values, 'q0'+'_'+f +
                                    '_weights', result.weights)
                            make_dict(self.values, 'q1'+'_'+f +
                                    '_weights', result.weights)"""
                
                elif selected_df['Error'].iloc[0]=='-':
                    self.required_params.remove(param)
                
                else:

                    make_dict(self.values, param, result.samples[:, index])
                    make_dict(self.values, param+'_weights', result.weights)

                    if param == 'rp/r*':  # Folded mode
                        for j,u in enumerate(self.limb_dark_coeffs):

                            make_dict(self.values, 'q'+str(j),
                                    result.samples[:, index+1])
                            make_dict(self.values, 'q'+str(j)+'_weights', result.weights)

                            """make_dict(self.values, 'q1',
                                    result.samples[:, index+2])

                            make_dict(self.values, 'q'+str(j)+'_weights', result.weights)
                            make_dict(self.values, 'q1'+'_weights', result.weights)"""
        
        #if self.ld_method=='quadratic':
        self.required_params = self.folded_params + \
            self.required_params+self.limb_dark_coeffs#['q0', 'q1']
            

        filters = list(set(filters))
        filters = np.sort(np.array(filters, dtype=int))

        df_so = pd.read_csv(self.summary_output[0])

        # Getting the best_values from the summary_output.csv file
        for p in self.required_params:
            selected_df = df_so.loc[df_so['Parameter'] == p]
            batch_filters = selected_df['Filter']
            if len(selected_df) == 1 and '-' not in batch_filters:
                self.values[p+'_best'] = float(selected_df['Best'])

            else:
                #batch_filters = selected_df['Filter']
                bestvals = selected_df['Best'].to_numpy(dtype=float)
                for i, f in enumerate(batch_filters):
                    self.values[p+'_'+str(f)+'_best'] = float(bestvals[i])

        # Generating the output
        with open(self.OUTPUT_PARAMETERS_FOLDER+'/'+self.MODIFIED_SUMMARY_OUTPUT, 'w') as mso:
            mso.write('Parameter,Filter,Best,Lower_error,Upper_error\n')
            q, q_low, q_up = [], [], []
            q_dict = {}
            #print(self.values)

            for param in self.required_params:
                if param in ['rp/r*']+self.limb_dark_coeffs and len(filters) > 0:
                    for fil in filters:
                        param_ = param+'_'+str(int(fil))
                        try:
                            best_val=self.values[param_+'_best']
                            #errs = get_quantiles_on_best_val(
                            #    samples=self.values[param_], weights=self.values[param_+'_weights'], best_val=self.values[param_+'_best'])
                        except KeyError:
                            best_val=self.values[param+'_best']
                        
                        errs = get_quantiles_on_best_val(
                            samples=self.values[param_], weights=self.values[param_+'_weights'], best_val=best_val)

                        # (model.mean)
                        mso.write(
                            f"{param},{fil},{best_val},{errs[0]},{errs[1]}\n")

                else:
                    errs = get_quantiles_on_best_val(
                        samples=self.values[param], weights=self.values[param+'_weights'], best_val=self.values[param+'_best'])
                    mso.write(
                        f"{param},-,{self.values[param+'_best']},{errs[0]},{errs[1]}\n")

                    if param in self.limb_dark_coeffs:
                        q.append(self.values[param+'_best']),
                        q_low.append(errs[0])
                        q_up.append(errs[1])

            if self.ld_method=='quadratic':
                if len(q) > 0:
                    u, u_low, u_up = q_to_u(q, q_low, q_up)
                    
                    for j,p in enumerate(self.limb_dark_coeffs):
                        mso.write(f"u{str(j)},-,{u[j]},{u_low[j]},{u_up[j]}\n")
                        #mso.write(f"u1,-,{u[1]},{u_low[1]},{u_up[1]}\n")

                elif len(filters) > 0:
                    for fil in filters:
                        q, q_low, q_up = [], [], []
                        for param in self.limb_dark_coeffs:
                            param_ = param+'_'+str(int(fil))
                            try:
                                best_val=self.values[param_+'_best']
                            except:
                                best_val=self.values[param+'_best']
                            errs = get_quantiles_on_best_val(
                                samples=self.values[param_], weights=self.values[param_+'_weights'], best_val=best_val)

                            q.append(best_val),
                            q_low.append(errs[0])
                            q_up.append(errs[1])

                        u, u_low, u_up = q_to_u(q, q_low, q_up)

                        for j,p in enumerate(self.limb_dark_coeffs):
                            mso.write(
                                f"u{str(j)},{str(int(fil))},{u[j]},{u_low[j]},{u_up[j]}\n")
                            #mso.write(f"u1,{str(int(fil))},{u[1]},{u_low[1]},{u_up[1]}\n")

        print(
            f"Saved results in {self.OUTPUT_PARAMETERS_FOLDER+'/'+self.MODIFIED_SUMMARY_OUTPUT}")

        return
