import sys
import numpy as np
import subprocess
import os
from os import path
from scipy.integrate import simps
from matplotlib import pyplot as plt
from copy import deepcopy
from scipy.special import erf
from .utils import data2gauss, gauss_func

__all__ = ['read_dgdl_files', 'integrate_dgdl',
           'ks_norm_test', 'plot_work_dist']

def find_longest_dgdl_file(lst):
    '''Takes a list of dgdl.xvg files and returns the longest one.
    The length is measured by the timestamp of the last recorded datapoint.
    
    Parameters
    ----------
    lst : list
        list containing the paths to the dgdl.xvg files.

    Returns
    -------
    longest_idx : int
        index of the longest dgdl file in lst
    t_list : list
        List of last timestamps of the dhdl files
    '''
    #find last line of every file in lst
    last_t=0.0 #length of simultaion in ps
    longest_idx=-1 #index of first full-length (longest) simultaion
    t_list=[0.0]*len(lst)
    for idx in range(len(lst)):
        if(os.path.exists(lst[idx])):
            result = subprocess.run(['tail', '-n', '1', lst[idx]],\
                                    stdout=subprocess.PIPE)
            try:
                s=result.stdout.split()
                t=float(s[0])
                t_list[idx] = t
                if t > last_t:
                    last_t = t
                    longest_idx = idx
            except:
                continue #if can't convert to float or split line
    if(longest_idx<0):
        raise RuntimeError("No valid dgdl file found.")
        
    return(longest_idx, t_list)


def read_dgdl_files(lst, lambda0=0, invert_values=False, verbose=True,\
                    sigmoid=0.0):
    '''Takes a list of dgdl.xvg files and returns the integrated work values.

    Parameters
    ----------
    lst : list
        list containing the paths to the dgdl.xvg files.
    lambda0 : [0,1]
        whether the simulations started from lambda 0 or 1. Default is 0.
    invert_values : bool
        whether to invert the sign of the returned work value.

    Returns
    -------
    w : array
        array of work values.
    lst : list
        sorted list of input dgdl.avg files corresponding to the work values
        in w.
    '''

    # check lambda0 is either 0 or 1
    assert lambda0 in [0, 1]
    
    #Start with first full-length (longest) simultaion.
    #Everything before is too short.
    good=False
    idx, t_list=find_longest_dgdl_file(lst) 
    last_t=t_list[idx]
    first_w=0
    ndata=0
    while(not good and idx<len(lst)):
        if(t_list[idx] >= last_t): #only check full length simulations
            try:
                _check_dgdl(lst[idx], lambda0, verbose=verbose)
                first_w, ndata = integrate_dgdl(lst[idx], lambda0=lambda0,\
                                            invert_values=invert_values,\
                                            sigmoid=sigmoid)
                good=True
            except:
                print(' !! Error in checking %s' % (lst[idx]))
                good=False
                idx+=1
        else:
            idx+=1
    if(not good):
        raise RuntimeError("No good dgdl files provided.")

    w_list = [first_w]
    for idx, f in enumerate(lst[idx+1:]):
        if verbose is True:
            sys.stdout.write('\r    Reading %s' % f)
            sys.stdout.flush()

        w, _ = integrate_dgdl(f, ndata=ndata, lambda0=lambda0,
                              invert_values=invert_values, sigmoid=sigmoid)
        if w is not None:
            w_list.append(w)

    if verbose is True:
        print('\n')
    return w_list


def integrate_dgdl(fn, ndata=-1, lambda0=0, invert_values=False, sigmoid=0.0):
    '''Integrates the data in a dgdl.xvg file.

    Parameters
    ----------
    fn : str
        the inpur dgdl.xvg file from Gromacs.
    ndata : int, optional
        number of datapoints in file. If -1, then ??? default is -1.
    lambda0 : [0,1]
        whether the simulations started from lambda 0 or 1. Default is 0.
    invert_values : bool
        whether to invert the sign of the returned work value.

    Returns
    -------
    integr : float
        result of the integration performed using Simpson's rule.
    ndata : int
        number of data points in the input file.
    '''

    # check lambda0 is either 0 or 1
    assert lambda0 in [0, 1]

    try:
        lines = open(fn, encoding="ISO-8859-1").readlines()
    except:
        return None, None
    if not lines:
        return None, None

    # extract dgdl datapoints into r
    # TODO: we removed the check for file integrity. We could have an
    # optional files integrity check before calling this integration func

    lines = [l for l in lines if l[0] not in '#@&']
    r=[]
    try:
        r = list(map(lambda x: float(x.split()[1]), lines))
    except:
        print(' !! Error in reading %s' % (fn))
        return None, None

    if ndata != -1 and len(r) != ndata:
        try:
            print(' !! Skipping %s ( read %d data points, should be %d )'
                  % (fn, len(r), ndata))
        except:
            print(' !! Skipping %s ' % (fn))
        return None, None
    # convert time to lambda
    ndata = len(r)
    dlambda = 1./float(ndata)
    if lambda0 == 1:
        dlambda *= -1

    # arrays for the integration
    # --------------------------
    # array of lambda values
    x = [lambda0+i*dlambda for i, dgdl in enumerate(r)]
##### VG VG VG ####
    if sigmoid != 0.0:
        x = [1.0/(1.0+np.exp( -1.0*(lambda0+i*dlambda-0.5)*sigmoid ) ) for i, dgdl in enumerate(r)]
        x[0] = lambda0
        x[-1] = 1.0-lambda0
#    for i in x:
#        print(i)
##### VG VG VG ####
    # array of dgdl
    y = r

    if lambda0 == 1:
        x.reverse()
        y.reverse()

    if invert_values is True:
        integr = simps(y, x) * (-1)
        return integr, ndata
    else:
        integr = simps(y, x)
        return integr, ndata


def ks_norm_test(data, alpha=0.05, refks=None):
    '''Performs a Kolmogorov-Smirnov test of normality.

    Parameters
    ----------
    data : array_like
        a one-dimensional array of values. This is the distribution tested
        for normality.
    alpha : float
        significance level of the statistics. Default if 0.05.
    refks : ???
        ???

    Returns
    -------
    Q : float
    lam0 : float
    check : float
    bOk : bool
    '''

    def ksref():
        f = 1
        potent = 10000
        lamb = np.arange(0.25, 2.5, 0.001)
        q = np.zeros(len(lamb), float)
        res = []
        for k in range(-potent, potent):
            q = q + f*np.exp(-2.0*(k**2)*(lamb**2))
            f = -f
        for i in range(len(lamb)):
            res.append((lamb[i], q[i]))
        return res

    def ksfunc(lamb):
        f = 1
        potent = 10000
        q = 0
        for k in range(-potent, potent):
            q = q + f*np.exp(-2.0*(k**2)*(lamb**2))
            f *= -1
        return q

    def edf(dg_data):
        edf_ = []
        ndata = []
        data = deepcopy(dg_data)
        data.sort()
        N = float(len(data))
        cnt = 0
        for item in data:
            cnt += 1
            edf_.append(cnt/N)
            ndata.append(item)
        ndata = np.array(ndata)
        edf_ = np.array(edf_)
        return ndata, edf_

    def cdf(dg_data):
        data = deepcopy(dg_data)
        data.sort()
        mean = np.average(data)
        sig = np.std(data)
        cdf = 0.5*(1+erf((data-mean)/float(sig*np.sqrt(2))))
        return cdf

    N = len(data)
    nd, ed = edf(data)
    cd = cdf(data)
    siglev = 1-alpha
    dval = []
    for i, val in enumerate(ed):
        d = abs(val-cd[i])
        dval.append(d)
        if i:
            d = abs(ed[i-1]-cd[i])
            dval.append(d)
    dmax = max(dval)
    check = np.sqrt(N)*dmax
    if not refks:
        refks = ksref()
    lst = list(filter(lambda x: x[1] > siglev, refks))
    lam0 = lst[0][0]
    if check >= lam0:
        bOk = False
    else:
        bOk = True

    q = ksfunc(check)
    return (1-q), lam0, check, bOk


def plot_work_dist(wf=[], wr=[], fname='Wdist.png', nbins=20, dG=None, dGerr=None,
                   units='kJ/mol', dpi=300, statesProvided='AB'):
    '''Plots forward and reverse work distributions. Optionally, it adds the
    estimate of the free energy change and its uncertainty on the plot.

    Parameters
    ----------
    wf : list
        list of forward work values.
    wr : list
        list of reverse work values.
    fname : str, optional
        filename of the saved image. Default is 'Wdist.png'.
    nbins : int, optional
        number of bins to use for the histogram. Default is 20.
    dG : float, optional
        free energy estimate.
    dGerr : float, optional
        uncertainty of the free energy estimate.
    units : str, optional
        the units of dG and dGerr. Default is 'kJ/mol'.
    dpi : int
        resolution of the saved image file.
    statesProvided: str
        work values for two states or only one

    Returns
    -------
    None

    '''

    def smooth(x, window_len=11, window='hanning'):

        if x.ndim != 1:
            raise ValueError("smooth only accepts 1 dimension arrays.")
        if x.size < window_len:
            raise ValueError("Input vector needs to be bigger than "
                             "window size.")
        if window_len < 3:
            return x
        if window not in ['flat', 'hanning', 'hamming',
                          'bartlett', 'blackman']:
            raise ValueError("Window is on of 'flat', 'hanning', 'hamming', "
                             "'bartlett', 'blackman'")
        s = np.r_[2*x[0]-x[window_len:1:-1], x, 2*x[-1]-x[-1:-window_len:-1]]
        # moving average
        if window == 'flat':
            w = np.ones(window_len, 'd')
        else:
            w = eval('np.' + window + '(window_len)')
        y = np.convolve(w/w.sum(), s, mode='same')
        return y[window_len-1:-window_len+1]

    plt.figure(figsize=(8, 6))
    x1 = list(range(len(wf)))
    x2 = list(range(len(wr)))
    if 'A' in statesProvided:
        x1 = list(range(len(wf)))
    if 'B' in statesProvided:
        x2 = list(range(len(wr)))
    if len(x1) > len(x2):
        x = x1
    else:
        x = x2
    if 'A' in statesProvided:
        mf, devf, Af = data2gauss(wf)
    if 'B' in statesProvided:
        mb, devb, Ab = data2gauss(wr)

    if 'AB' in statesProvided: 
        maxi = max(wf+wr)
        mini = min(wf+wr)
        plt.subplot(1, 2, 1)
        plt.plot(x1, wf, 'g-', linewidth=2, label="Forward (0->1)", alpha=.3)
        plt.plot(x2, wr, 'b-', linewidth=2, label="Backward (1->0)", alpha=.3)
        ### smoothing ###
        try:
            sm1 = smooth(np.array(wf))
            plt.plot(x1, sm1, 'g-', linewidth=3)
        except:
            print("Plotting: no smoothing for Wf")
        try:
            sm2 = smooth(np.array(wr))
            plt.plot(x2, sm2, 'b-', linewidth=3)
        except:
            print("Plotting: no smoothing for Wr")
    elif 'A' in statesProvided: 
        maxi = max(wf)
        mini = min(wf)
        plt.subplot(1, 2, 1)
        plt.plot(x1, wf, 'g-', linewidth=2, label="Forward (0->1)", alpha=.3)
        ### smoothing ###
        try:
            sm1 = smooth(np.array(wf))
            plt.plot(x1, sm1, 'g-', linewidth=3)
        except:
            print("Plotting: no smoothing for Wf")
    elif 'B' in statesProvided: 
        maxi = max(wr)
        mini = min(wr)
        plt.subplot(1, 2, 1)
        plt.plot(x2, wr, 'b-', linewidth=2, label="Backward (1->0)", alpha=.3)
        ### smoothing ###
        try:
            sm2 = smooth(np.array(wr))
            plt.plot(x2, sm2, 'b-', linewidth=3)
        except:
            print("Plotting: no smoothing for Wr")

    plt.legend(shadow=True, fancybox=True, loc='upper center',
               prop={'size': 12})
    plt.ylabel(r'W [kJ/mol]', fontsize=20)
    plt.xlabel(r'# Snapshot', fontsize=20)
    plt.grid(lw=2)
    plt.xlim(0, x[-1]+1)
    xl = plt.gca()
    for val in xl.spines.values():
        val.set_lw(2)
    plt.subplot(1, 2, 2)
    if 'A' in statesProvided:
        plt.hist(wf, bins=nbins, orientation='horizontal', facecolor='green',
             alpha=.75, density=True)
    if 'B' in statesProvided:
        plt.hist(wr, bins=nbins, orientation='horizontal', facecolor='blue',
             alpha=.75, density=True)

    x = np.arange(mini, maxi, .5)

    if 'AB' in statesProvided:
        y1 = gauss_func(Af, mf, devf, x)
        y2 = gauss_func(Ab, mb, devb, x)
        plt.plot(y1, x, 'g--', linewidth=2)
        plt.plot(y2, x, 'b--', linewidth=2)
        size = max([max(y1), max(y2)])
    elif 'A' in statesProvided:
        y1 = gauss_func(Af, mf, devf, x)
        plt.plot(y1, x, 'g--', linewidth=2)
        size = max(y1)
    elif 'B' in statesProvided:
        y2 = gauss_func(Ab, mb, devb, x)
        plt.plot(y2, x, 'b--', linewidth=2)
        size = max(y2)

    res_x = [dG, dG]
    res_y = [0, size*1.2]
    if dG is not None and dGerr is not None:
        plt.plot(res_y, res_x, 'k--', linewidth=2,
                 label=r'$\Delta$G = %.2f $\pm$ %.2f %s' % (dG, dGerr, units))
        plt.legend(shadow=True, fancybox=True, loc='upper center',
                   prop={'size': 12})
    elif dG is not None and dGerr is None:
        plt.plot(res_y, res_x, 'k--', linewidth=2,
                 label=r'$\Delta$G = %.2f %s' % (dG, units))
        plt.legend(shadow=True, fancybox=True, loc='upper center',
                   prop={'size': 12})
    else:
        plt.plot(res_y, res_x, 'k--', linewidth=2)

    plt.xticks([])
    plt.yticks([])
    xl = plt.gca()
    for val in xl.spines.values():
        val.set_lw(2)
    plt.subplots_adjust(wspace=0.0, hspace=0.1)
    plt.savefig(fname, dpi=dpi)


def _check_dgdl(fn, lambda0, verbose=True):
    '''Prints some info about a dgdl.xvg file.'''
    lines = open(fn, encoding="ISO-8859-1").readlines()
    if not lines:
        return None
    r = []
    for line in lines:
        if line[0] not in '#@&':
            r.append([float(x) for x in line.split()])
    ndata = len(r)
    dlambda = 1./float(ndata)
    if lambda0 == 1:
        dlambda *= -1

    if verbose is True:
        print('    First accepted file: %s' % fn)
        print('    # data points: %d' % ndata)
        print('    Length of trajectory: %8.3f ps' % r[-1][0])
        print('    Delta lambda: %8.5f' % dlambda)
