#!/usr/bin/env python

import sys
import time
import numpy as np
import numpy.linalg as la
from mpi4py import MPI

#
# pathrefiner refins an initial path using MFTP or MFEP
#
# input:  msys     -- user-defined object
#         initpath -- an array user-provided for initial path
#         param    -- a dictionary containing method parameters
# output: path     -- an array for "refined path" satisfying stopping criteria
#
def pathrefiner(msys, initpath, param):
    noiter = param['noiter']   # the maximum iterations users specify for a continuous run
    preiter = param['preiter'] # the number of iterations having been operated
    tau = param['tau']         # the pseudo time time square determining the moving speed
    TOL = param['TOL']         # the stopping criterion defined as d = max_{0<=j<=J}|x^current - x^previous|
    method = param['method']   # the MFEP or the MFTP. By defaul is the MFTP

    (norepl, nocv) = initpath.shape
    k_BT = (msys.temp/300.0)*0.59595

    # MPI parallel code starts here
    noproc = MPI.COMM_WORLD.Get_size()
    myrank = MPI.COMM_WORLD.Get_rank()
    noloop = (norepl-1)/noproc +1

    path = initpath
    maxiter = noiter + preiter

    d = -1.0

    # pathrefiner starts refining initial path
    for cycle in range(preiter+1, maxiter+1):
        
        start = time.time()      
        # initialize D, allD, F, and allF
        D = np.zeros((norepl,nocv,nocv),'float')
        F = np.zeros((norepl,nocv),'float')

        allD = np.zeros((noproc,noloop,nocv,nocv),'float')
        allF = np.zeros((noproc,noloop,nocv),'float')

        # the only interface with the user-defined module
        for j in range(myrank, norepl, noproc):

            [allD[myrank][j/noproc],allF[myrank][j/noproc]] = msys.mean_force(path[j],j)
        allD = MPI.COMM_WORLD.reduce(allD)
        allF = MPI.COMM_WORLD.reduce(allF)
        duration = time.time() - start
        
        if myrank == 0:
            
            oldpath = path

            for j in range(norepl):
                D[j] = allD[j%noproc][j/noproc]
                F[j] = allF[j%noproc][j/noproc]

            # choose method from MFEP, simplified string MFTP, implicit MFTP
            if param['method'] == 'MFTP_ISS':
                path = MFTP_ISS(D, F, path, tau, k_BT)
            elif param['method'] == 'MFEP_SS':
                path = MFEP_SS(D, F, path, tau)
            else:
                print 'Warning: no existing method matches your selection. Run implicit MFTP'
                path = MFTP_ISS(D, F, path, tau, k_BT)

            # normalization
            path = normalize(path)

            # pointwise distance between two paths
            diff = path - oldpath
            maxdis = np.array([np.sqrt(np.dot(diff[i],diff[i])) for i in range(norepl)])
            d = np.max(maxdis)/tau
          
            print 'It takes %f s to move %f in %d iteration.'\
                  %(duration, d, cycle)

            # save data into disk
            writepath(path,'path'+str(cycle)+'.dat')

        path = MPI.COMM_WORLD.bcast(path,root=0)
        d = MPI.COMM_WORLD.bcast(d,root=0)
        if d < TOL:
            break
    return [path,cycle]

#
# maximum flux transition path
# (Semi-implicit scheme )
#
# input:  D    -- proto diffusion tensor matrix
#         F    -- mean force, i.e. gradient of free energy G
#         Z    -- string replicas before iteration
#         tau  -- timestep
#         k_BT -- the temperature in unit (kcal/mol)
# output: newZ -- string replicas after iteration
#
def MFTP_ISS(D, F, Z, tau, k_BT):

    (norepl,nocv) = F.shape

    # calculate the right hand sides

    # inverse proto diffusion tensors
    D_inv = np.zeros((norepl,nocv,nocv),'float')
    for j in range(norepl):
        D_inv[j] = la.inv(D[j])
        
    # precalculate forward and backward differences
    Zplus = np.zeros((norepl,nocv),'float')
    Zminus = np.zeros((norepl,nocv),'float')
    Dplus = np.zeros((norepl,nocv,nocv),'float')
    Dminus = np.zeros((norepl,nocv,nocv),'float')
    for j in range(1,norepl-1):
        Zplus[j] = Z[j+1]-Z[j]
        Zminus[j] = Z[j]-Z[j-1]
        Dplus[j] = D[j+1]-D[j]
        Dminus[j] = D[j]-D[j-1]

    # calculate the right hand sides
    S = np.zeros((norepl),'float')            # scalar 1/beta*Z^TD^{-1}Z
    T1 = np.zeros((norepl,nocv),'float')      # term D\nabla F
    T2 = np.zeros((norepl,nocv),'float')      # term D_sD^{-1}Z_s/beta*Z^TD^{-1}Z
    for j in range(1,norepl-1):
        S[j] = 2.0*k_BT/((np.dot(Zplus[j],np.dot(D_inv[j],Zplus[j])) + \
                np.dot(Zminus[j],np.dot(D_inv[j],Zminus[j]))))
        T1[j] = np.dot(D[j],F[j])
        T2[j] = S[j]*(np.dot(Dplus[j],np.dot(D_inv[j],Zplus[j])) + \
                np.dot(Dminus[j],np.dot(D_inv[j],Zminus[j])))/2.0
    # Discretize D_s = -D\nablaF+(1/beta*Z^TD^{-1}Z)*(Z_ss-D_sD^{-1}Z_s)
    RH = np.zeros((norepl,nocv),'float')
    for j in range(1,norepl-1):
        RH[j] = Z[j]- tau*(T1[j]+T2[j])

    # update two boundary points of the string, steepest descent
    RH[0] = Z[0] - tau*np.dot(D[0],F[0])
    RH[norepl-1] = Z[norepl-1] - tau*np.dot(D[norepl-1],F[norepl-1])

    # change the order of variables to transform block tridiagonal to block diagonal matrix
    RH = np.transpose(RH)

    # compose the mass matrix 
    a = -tau*S
    a[0] = 0
    a[norepl-1] = 0
    b = 1 + 2*tau*S
    b[0] = 1
    b[norepl-1] = 1
    c = -tau*S
    c[0] = 0
    c[norepl-1] = 0

    newZ = np.zeros((nocv,norepl),'float')

    # forward sweep of thomas method
    cp = thomas_c(a,b,c,norepl)
    # backward substitution of thomas method    
    for j in range(nocv):
        newZ[j] = thomas_d(a,b,cp,RH[j],norepl)

    # switch back to the original order of variables
    newZ = np.transpose(newZ)

    return newZ

#
# minimum free energy path
# (Luca Maragliano's paper)
#
# input:  D   -- diffusion matrix
#         F   -- mean force, i.e. gradient of free energy G
#         Z   -- string replicas before iteration
#         tau -- timestep
# output: Z   -- string replicas after iteration
#
def MFEP_SS(D, F, Z, tau):

    (norepl,nocv) = F.shape

    g = np.zeros((norepl,nocv),'float')
    for j in range(norepl):
        g[j] = - np.dot(D[j], F[j])

    newZ = Z + tau*g

    return newZ

#
# path is normalized using equal arc length, piecewise polynomial
#
# input:  path     -- path before normalization
# output: newpath  -- final path after normalization by equal arc length
#
def normalize(path):

    (norepl, nocv) = path.shape

    L = np.zeros((norepl),'float')
    for i in range(1,norepl): 
        diff = path[i] - path[i-1]
        L[i] = L[i-1] + np.sqrt(np.dot(diff,diff))

    newpath = np.zeros((norepl,nocv),'float')
    newpath[0] = path[0]
    newpath[norepl-1] = path[norepl-1]
    for i in range(1,norepl-1):
        s = (i/(float(norepl-1)))*L[norepl-1]
        j=0
        while s>=L[j]:
            j += 1
        diff = path[j] - path[j-1]
        newpath[i] = path[j-1]+(s-L[j-1])/np.sqrt(np.dot(diff,diff))*diff

    return newpath

#
# Thomas (tridiagonal matrix) algorithm
#
# LINEAR SYSTEM:
# a_i x_{i-1} + b_i x_i + c_i x_{i+1} = d_i
# where a_0 = 0 and c_n = 0
#
# input:  a: (n x 1) array
#         b: (n x 1) array
#         c: (n x 1) array
# output: cp: (n x 1) array
# 
def thomas_c(a,b,c,n):
    cp = np.zeros((n),'float')
    cp[0] = c[0]/b[0]
    for j in range(1,n-1):
        cp[j] = c[j]/(b[j] - cp[j-1]*a[j])

    return cp

#
# input:  a: (n x 1) array
#         b: (n x 1) array
#         cp: (n x 1) array
# output: x: (n x 1) array
# 
def thomas_d(a,b,cp,d,n):
    dp = np.zeros((n),'float')
    dp[0] = d[0]/b[0]
    for j in range(1,n):
        dp[j] = (d[j]-dp[j-1]*a[j])/(b[j]-cp[j-1]*a[j])

    #back substitution
    x = np.zeros((n),'float')
    x[n-1]=dp[n-1]
    for j in range(n-2,-1,-1):
        x[j]=dp[j]-cp[j]*x[j+1]
    return x

#
# The final path will be written on disk. The format of data file is: each replica
# has a separate line, and each component of replica is separated by blank.
#
# input: path, norepl*nocv array
#        filename, save into disc
#
def writepath(path, filename):
    path = np.array(path)
    (norepl, nocv) = path.shape
    file = open(filename, 'w')
    for i in range(norepl):
        for j in range(nocv):
            file.write("%f\t"%path[i][j])
        file.write("\n")
    file.close()