import numpy as np
import sys
import argparse
import subprocess
from tqdm import tqdm

R = 0.00831446261815324 # kJ / (mol * K)

def command(cmd, inp, out):
    """ Executes the command in the shell and raise exeption if command fails """
    process = subprocess.Popen(cmd, stdin = subprocess.PIPE, stdout = subprocess.PIPE, stderr = subprocess.PIPE, universal_newlines = True)
    if inp != None:
        process.stdin.write(inp)
    stdout, stderr = process.communicate()
    if process.returncode !=0:
        raise Exception(stderr)
    elif out:
        return stdout + stderr
    else:
        return None

def unpack_tpr(tprfiles, gmx):
    """ Read all tpr files, check for consistency and takes needed parameters """
    consistency = np.zeros((len(tprfiles), 8))
    for j in tqdm(range(0, len(tprfiles)), desc = "Loading .tpr files          "):
        tpr_dec = command([gmx, "dump", "-s", tprfiles[j]], None, True)
        tpr_dec = tpr_dec.split()
        for i in range(0, len(tpr_dec)):
            if tpr_dec[i] == "dt":
                dt = float(tpr_dec[i+2])*1000.0
                consistency[j,0] = dt
            elif tpr_dec[i] == "nsteps":
                nsteps = int(tpr_dec[i+2])
                consistency[j,1] = nsteps
            elif tpr_dec[i] == "pull-nstfout":
                pffreq = int(tpr_dec[i+2])
                consistency[j,2] = pffreq
            elif tpr_dec[i] == "init-step":
                start_steps = int(tpr_dec[i+2])
                consistency[j,3] = start_steps
            elif tpr_dec[i] == "pull-ncoords":
                nmol = int(tpr_dec[i+2])
                consistency[j,4] = nmol
            elif tpr_dec[i] == "ref-t:":
                T = float(tpr_dec[i+1])
                consistency[j,5] = T
            elif tpr_dec[i] == "nstxout-compressed":
                nstxout = int(tpr_dec[i+2])
                consistency[j,6] = nstxout
            elif tpr_dec[i] == "pull-nstxout":
                pxfreq = int(tpr_dec[i+2])
                consistency[j,7] = pxfreq
    print("Checking .tpr files...")
    problem = False
    what =("Time step of integration dt (fs)","Length of simulation (nsteps)","Output of forces (pull-nstfout)","Starting step for simulation (init-step)","Number of constrained molecules (pull-ncoords)","Temperature (K)","Output of positions (nstxout-compressed)","Output of positions (pull-nstxout)")
    for i in range(1, len(tprfiles)):
        for j in range(0, 8):
            if consistency[i-1,j] != consistency[i,j]:
                problem = inconsistency(what[j], tprfiles[i-1], tprfiles[i], consistency[i-1,j], consistency[i,j])
    if problem:
        raise Exception("Simulation data in tpr files is not consistent. Please check input files.")
    else:
        print("The following checked parameters are consistent: dt, nsteps, pull-nstfout, pull-nstxout, init-step, pullncoords, ref-t and nstxout-compressed.")
    return nmol, T, dt, pffreq, pxfreq, nsteps - start_steps, nstxout

def inconsistency(what, wherea, whereb, valuea, valueb):
    """ Notifies user of .tpr inconsistencies """
    print("{} inconsistency!".format(what))
    print("{} in file {}: {}".format(what, wherea, valuea))
    print("{} in file {}: {}\n".format(what, whereb, valueb))
    return True

def window_positions(pxfiles, nfiles, nmol):
    """ Read the pullx files and computes the average positions and variances """
    pos = np.zeros((nfiles,nmol))
    var = np.zeros((nfiles,nmol))
    for i in tqdm(range(0,len(pxfiles)), desc = "Loading pullx files         "):
        data = np.loadtxt(pxfiles[i], dtype = float, comments = ('#','@','&'))
        if len(data[0,:])-1 != nmol:
            raise Exception("Number of molecules in pullx files is different from number of constrained molecules in .tpr files.")
        for j in range(0, nmol):
            pos[i,j] = np.average(data[:,j + 1])
            var[i,j] = np.var(data[:,j + 1])
    return pos, var

def position_ACF(pos, pxfiles, pacfdir, savedir, nfiles, nmol, position_corr_time, dt, pxfreq, gmx):
    """ Computes position autocorrelation functions via gmx analyze, then loads them and prints them divided. """
    true_time = dt*pxfreq
    command(["mkdir", "./" + savedir + "/" + pacfdir], None, False)
    gmx_analyze_head = "Collection of gmx analyze commands and outputs"
    for i in tqdm(range(0, nfiles), desc = "Calculating position ACFs   "):
        gmx_analyze_head += "\n\n$> {} analyze -f {} -ac ./{}/{}/pacf_{}_tlength_{}.xvg -subav -nonormalize -acflen {}\n\n".format(gmx, pxfiles[i], savedir, pacfdir, i, position_corr_time, position_corr_time)
        gmx_analyze_head += command([gmx, "analyze", "-f", pxfiles[i], "-ac", "./{}/{}/pacf_{}_tlength_{}.xvg".format(savedir, pacfdir, i, position_corr_time), "-subav", "-nonormalize", "-acflen", str(position_corr_time)], None, True)
    text_file = open("./{}/{}/{}_pACF_analyze_output.txt".format(savedir, pacfdir, gmx), "wt")
    text_file.write(gmx_analyze_head)
    text_file.close()
    time = np.arange(0.0, position_corr_time*true_time/1000.0, true_time/1000.0)
    pACFs = np.zeros((position_corr_time, nfiles*nmol))
    for i in tqdm(range(0, nfiles), desc = 'Loading position ACFs       '):
        namexvg = "./{}/{}/pacf_{}_tlength_{}.xvg".format(savedir, pacfdir, i, position_corr_time)
        data = np.loadtxt(namexvg, dtype = float, comments = ('#','@','&'))
        for j in range(0, nmol):
            pACFs[:, j*nfiles + i] = data[position_corr_time*j:position_corr_time*(j + 1), 1]
    for j in range(0, nmol):
        head_a = "Window position and related pACF force given a molecule\n"
        head_b = "molecule : {}\n".format(j)
        head_b += "Corr time: {}ps\n".format(int(position_corr_time*dt*pxfreq/1000.0))
        head_b += " win (nm)  "
        head_c = "\n time (ps) "
        warner = False
        for i in range(0, nfiles):
            head_b += str("{:.4f}".format(pos[i,j])).center(13, ' ')
            head_c += ("window_{}".format(i)).center(13, ' ')
            max_start = max(np.abs(pACFs[:int(0.05*position_corr_time), j*nfiles+i]))
            max_end = max(np.abs(pACFs[int(0.95*position_corr_time):, j*nfiles+i]))
            if max_end/max_start >= 0.05:
                warner = True
                head_a += "Window: {} at position {:.2f}. WARNING: SLOW CONVERGENCE. Final values of pACF are ~{:.1f}% of initial ones. Please consider longer integration times.\n".format(i, pos[i,j], max_end/max_start*100.0)
            else:
                head_a += "Window: {} at position {:.2f}. Final values of pACF are ~{:.1f}% of initial ones.\n".format(i, pos[i,j], max_end/max_start*100.0)
        if warner:
            print(" -> Slow convergence in one or more windows of molecule {}. Check position ACFs in output and consider longer integration times.".format(j))
        np.savetxt("./{}/pACF_molecule_{}.txt".format(savedir, j), np.column_stack((time, pACFs[:,j*nfiles:(j + 1)*nfiles])), fmt = '%.6e', header = head_a + head_b + head_c)
    print("{} analyze output files saved in directory: ./{}/{}".format(gmx, savedir, pacfdir))
    print("Individual position ACFs profiles saved in directory: {}".format(savedir))
    return pACFs

def force_ACF(pos, pffiles, facfdir, savedir, nfiles, nmol, corr_time, dt, pffreq, gmx):
    """ Computes force autocorrelation functions via gmx analyze, then loads them and prints them divided. """
    true_time = dt*pffreq
    command(["mkdir", "./" + savedir + "/" + facfdir], None, False)
    gmx_analyze_head = "Collection of gmx analyze commands and outputs"
    for i in tqdm(range(0, nfiles), desc = "Calculating force ACFs      "):
        gmx_analyze_head += "\n\n$> {} analyze -f {} -ac ./{}/{}/facf_{}_tlength_{}.xvg -subav -nonormalize -acflen {}\n\n".format(gmx, pffiles[i], savedir, facfdir, i, corr_time, corr_time)
        gmx_analyze_head += command([gmx, "analyze", "-f", pffiles[i], "-ac", "./{}/{}/facf_{}_tlength_{}.xvg".format(savedir, facfdir, i, corr_time), "-subav", "-nonormalize", "-acflen", str(corr_time)], None, True)
    text_file = open("./{}/{}/{}_fACF_analyze_output.txt".format(savedir, facfdir, gmx), "wt")
    text_file.write(gmx_analyze_head)
    text_file.close()
    time = np.arange(0.0, corr_time*true_time/1000.0, true_time/1000.0)
    fACFs = np.zeros((corr_time, nfiles*nmol))
    for i in tqdm(range(0, nfiles), desc = 'Loading force ACFs          '):
        namexvg = "./{}/{}/facf_{}_tlength_{}.xvg".format(savedir, facfdir, i, corr_time)
        data = np.loadtxt(namexvg, dtype = float, comments = ('#','@','&'))
        for j in range(0, nmol):
            fACFs[:, j*nfiles + i] = data[corr_time*j:corr_time*(j + 1), 1]
    for j in range(0, nmol):
        head_a = "Window position and related fACF force given a molecule\n"
        head_b = "molecule : {}\n".format(j)
        head_b += "Corr time: {}ps\n".format(int(corr_time*dt*pffreq/1000.0))
        head_b += " win (nm)  "
        head_c = "\n time (ps) "
        warner = False
        for i in range(0, nfiles):
            head_b += str("{:.4f}".format(pos[i,j])).center(13, ' ')
            head_c += ("window_{}".format(i)).center(13, ' ')
            max_start = max(np.abs(fACFs[:int(0.05*corr_time), j*nfiles+i]))
            max_end = max(np.abs(fACFs[int(0.95*corr_time):, j*nfiles+i]))
            if max_end/max_start >= 0.05:
                warner = True
                head_a += "Window: {} at position {:.2f}. WARNING: SLOW CONVERGENCE. Final values of fACF are ~{:.1f}% of initial ones. Please consider longer integration times.\n".format(i, pos[i,j], max_end/max_start*100.0)
            else:
                head_a += "Window: {} at position {:.2f}. Final values of fACF are ~{:.1f}% of initial ones.\n".format(i, pos[i,j], max_end/max_start*100.0)
        if warner:
            print(" -> Slow convergence in one or more windows of molecule {}. Check force ACFs in output and consider longer integration times.".format(j))
        np.savetxt("./{}/fACF_molecule_{}.txt".format(savedir, j), np.column_stack((time, fACFs[:,j*nfiles:(j + 1)*nfiles])), fmt = '%.6e', header = head_a + head_b + head_c)
    print("{} analyze output files saved in directory: ./{}/{}".format(gmx, savedir, facfdir))
    print("Individual force ACFs profiles saved in directory: {}".format(savedir))
    return fACFs

def fACF_diffusion(pos, fACFs, savedir, nmol, nfiles, T, dt, pffreq, corr_time):
    """ Calculates diffusion with the integration of force autocorrelation functions via the Kubo relation """
    print("Processing individual perpendicular diffusion profiles via Kubo relation...")
    true_time = dt*pffreq
    diff = np.zeros((nfiles, nmol, 2))
    for i in range(0, nmol):
        for j in range(0, nfiles):
            diff[j,i,0] = 10.0*(R*T)**2/np.trapz(fACFs[:,i*nfiles + j], dx = true_time)
            diff[j,i,1] = 0.0
        head = "Window position and related perpendicular diffusion coefficient via Kubo given a molecule\n"
        head += "molecule : {}\n".format(i)
        head += "Corr time: {}ps\n".format(int(corr_time*dt*pffreq/1000.0))
        head += "  win (nm)   D_perp (cm**2/s)   std_dev (cm**2/s)"
        np.savetxt("./{}/Dperp_kubo_individual_molecule_{}.txt".format(savedir, i), np.stack((pos[:,i], diff[:,i,0], diff[:,i,1])).transpose(), fmt = '%.6e', header = head)
    print("Individual Kubo diffusion profiles saved in directory: {}".format(savedir))
    return diff

def pACF_diffusion(pos, pACFs, savedir, nmol, nfiles, dt, pxfreq, position_corr_time, var):
    """ Calculates diffusion with the integration of position autocorrelation functions via the pACF method """
    print("Processing individual perpendicular diffusion profiles via pACF...")
    true_time = dt*pxfreq
    diff = np.zeros((nfiles, nmol, 2))
    for i in range(0, nmol):
        for j in range(0, nfiles):
            diff[j,i,0] = 10.0*var[j,i]**2/np.trapz(pACFs[:,i*nfiles + j], dx = true_time)
            diff[j,i,1] = 0.0
        head = "Window position and related perpendicular diffusion coefficient via pACF given a molecule\n"
        head += "molecule : {}\n".format(i)
        head += "Corr time: {}ps\n".format(int(position_corr_time*dt*pxfreq/1000.0))
        head += "  win (nm)   D_perp (cm**2/s)   std_dev (cm**2/s)"
        np.savetxt("./{}/Dperp_pACF_individual_molecule_{}.txt".format(savedir, i), np.stack((pos[:,i], diff[:,i,0], diff[:,i,1])).transpose(), fmt = '%.6e', header = head)
    print("Individual pACF diffusion profiles saved in directory: {}".format(savedir))
    return diff

def global_profile(pos, delta_com, bin_width, nfiles, nmol, incoming, odd):
    """ Computes the global symmetrized and unsymmetrized profiles from the individual molecule profiles """
    count = 0
    local_error = 0
    values = np.zeros(nmol*nfiles)
    nwindows = int(np.ceil(max(np.abs(pos.min()), np.abs(pos.max()))/bin_width + 1))
    profile_unsym = np.zeros((nwindows*2+1, 3))
    for i in range(0, nwindows*2+1):
        a_min = -bin_width*nwindows + i*bin_width - bin_width/2.0
        a_max = -bin_width*nwindows + i*bin_width + bin_width/2.0
        values[:] = 0.0
        count = 0
        local_error = 0
        for j in range(0, nfiles):
            for k in range(0, nmol):
                if pos[j,k] > a_min:
                    if pos[j,k] < a_max:
                            count += 1
                            values[int(count-1)] = incoming[j,k,0]
                            local_error = incoming[j,k,1]
        if count > 0:
            profile_unsym[i,1] = np.average(values[:count])
            if count == 1:
                profile_unsym[i,2] = local_error
            else:
                profile_unsym[i,2] = np.std(values[:count])/np.sqrt(count)
        profile_unsym[i,0] = bin_width*(i-nwindows)
    shift_pos = pos-delta_com
    shift_nwindows = int(np.ceil(max(np.abs(shift_pos.min()), np.abs(shift_pos.max()))/bin_width + 1))
    profile_sym = np.zeros((shift_nwindows*2+1, 3))
    for i in range(0, shift_nwindows+1):
        a_min = i*bin_width - bin_width/2.0
        a_max = i*bin_width + bin_width/2.0
        values[:] = 0.0
        count = 0
        local_error = 0
        for j in range(0, nfiles):
            for k in range(0, nmol):
                if abs(shift_pos[j,k]) > a_min:
                    if abs(shift_pos[j,k]) < a_max:
                            count += 1
                            local_error = incoming[j,k,1]
                            if (odd and (shift_pos[j,k] < 0)):
                                values[int(count - 1)] = -incoming[j,k,0]
                            else:
                                values[int(count - 1)] = incoming[j,k,0]
        if count > 0:
            profile_sym[shift_nwindows + i,1] = np.average(values[:count])
            if count == 1:
                profile_sym[shift_nwindows + i,2] = local_error
            else:
                profile_sym[shift_nwindows + i,2] = np.std(values[:count])/np.sqrt(count)
        profile_sym[shift_nwindows + i,0] = i*bin_width + delta_com
        profile_sym[shift_nwindows - i,0] = -i*bin_width + delta_com
        if odd:
            profile_sym[shift_nwindows - i,1] = -profile_sym[shift_nwindows + i,1]
        else:
            profile_sym[shift_nwindows - i,1] = profile_sym[shift_nwindows + i,1]
        profile_sym[shift_nwindows - i,2] = profile_sym[shift_nwindows + i,2]
    return profile_unsym, profile_sym

def float_average(pos, pffiles, nfiles, nmol, pffreq, dt, nsteps, savedir, ave_length):
    """ Load pullf files, calculates individual averages and saves them"""
    true_time = pffreq * dt
    filelen = int(nsteps/pffreq) + 1
    blocklen = max(1, int(filelen/ave_length))
    tot = int(filelen/blocklen)
    forces_time = np.zeros((tot, nfiles*nmol))
    for i in tqdm(range(0, nfiles), desc = "Loading pullf files         "):
        data = np.loadtxt(pffiles[i], dtype = float, comments = ('#','@','&'))
        for j in tqdm(range(0, nmol), leave = False, desc = " ---> Computing averages    "):
            for k in range(0, tot):
                forces_time[k, j*nfiles + i] = np.average(data[:blocklen*(1+k),j+1])
    time = np.arange(0.0, nsteps*dt/10**6, blocklen*true_time/10**6)
    for j in range(0, nmol):
        head_a = " win (nm)  "
        head_b = "\n time (ps) "
        for i in range(0, nfiles):
            head_a += str("{:.4f}".format(pos[i, j])).center(14, ' ')
            head_b += ("window_{}".format(i)).center(14, ' ')
        np.savetxt("./{}/force_time_molecule_{}.txt".format(savedir, j), np.column_stack((time, forces_time[:,j*nfiles:(j + 1)*nfiles])), fmt = '%.6e', header = head_a + head_b)
    return forces_time

def force_average(incoming, nfiles, nmol, aveF_len, pos, savedir):
    """ Calculates average of forces <F> """
    print ("Processing individual <F> for all the molecules and input files...")
    forces = np.zeros((nfiles, nmol, 2))
    for i in range(0, nmol):
        for j in range(0, nfiles):
            forces[j,i,0] = np.average(incoming[int((1.0-aveF_len)*len(incoming)):,i*nfiles + j])
            forces[j,i,1] = np.std(incoming[int((1.0-aveF_len)*len(incoming)):,i*nfiles + j])
        head = "Window position and related average of force <F> given a molecule\n"
        head += "molecule : {}\n".format(i)
        head += "Time average length: {}ps\n".format(aveF_len)
        head += "  win (nm)   <F> (kJ/(mol*nm))  std_dev (kJ/(mol*nm))"
        np.savetxt("./{}/force_average_molecule_{}.txt".format(savedir, i), np.stack((pos[:,i], forces[:,i,0], forces[:,i,1])).transpose(), fmt = '%.6e', header = head)
    print("Individual <F> profiles saved in directory: {}".format(savedir))
    return forces

def pmf_integration(pos, delta_com, bin_width, delta_g, incoming_unsym, incoming_sym):
    """ Integrates the average of force <F> via trapezium rule to calculate the PMF """
    nwindows = int(np.ceil(max(np.abs(pos.min()), np.abs(pos.max()))/bin_width + 1))
    pmf_unsym = np.zeros((nwindows*2+1, 7))
    for i in range(0, nwindows*2+1):
        pmf_unsym[i,0] = incoming_unsym[i,0]
        pmf_unsym[i,1] = np.trapz(incoming_unsym[:i+1,1], incoming_unsym[:i+1,0]) + delta_g
        pmf_unsym[i,2] = bin_width*np.sqrt(np.sum(incoming_unsym[:i+1,2]**2))
        pmf_unsym[nwindows*2-i,3] = -np.trapz(incoming_unsym[nwindows*2-i:,1], incoming_unsym[nwindows*2-i:,0]) + delta_g
        pmf_unsym[nwindows*2-i,4] = bin_width*np.sqrt(np.sum(incoming_unsym[nwindows*2-i:,2]**2))
    for i in range(nwindows, nwindows*2+1):
        if (pmf_unsym[i,1] != 0.0 and pmf_unsym[i,3] != 0.0):
            pmf_unsym[:,3] -= pmf_unsym[i,3] - pmf_unsym[i,1]
            break
        elif (i == nwindows*2):
            print("Somethings weird in the PMF, check carefully the plots.")
    for i in range(0, nwindows*2+1):
        pmf_unsym[i,5] = (pmf_unsym[i,1] + pmf_unsym[i,3])/2.0
        pmf_unsym[i,6] = np.sqrt(pmf_unsym[i,2]**2 + pmf_unsym[i,4]**2)/2.0
    shift_pos = pos - delta_com
    shift_nwindows = int(np.ceil(max(np.abs(shift_pos.min()), np.abs(shift_pos.max()))/bin_width + 1))
    pmf_sym = np.zeros((shift_nwindows*2+1, 7))
    for i in range(0, shift_nwindows*2+1):
        pmf_sym[i,0] = incoming_sym[i,0]
        pmf_sym[i,1] = np.trapz(incoming_sym[:i+1,1], incoming_sym[:i+1,0]) + delta_g
        pmf_sym[i,2] = bin_width*np.sqrt(np.sum(incoming_sym[:i+1,2]**2))
        pmf_sym[shift_nwindows*2-i,3] = -np.trapz(incoming_sym[shift_nwindows*2-i:,1], incoming_sym[shift_nwindows*2-i:,0]) + delta_g
        pmf_sym[shift_nwindows*2-i,4] = bin_width*np.sqrt(np.sum(incoming_sym[shift_nwindows*2-i:,2]**2))
    for i in range(0, shift_nwindows*2+1):
        pmf_sym[i,5] = (pmf_sym[i,1] + pmf_sym[i,3])/2.0
        pmf_sym[i,6] = np.sqrt(pmf_sym[i,2]**2 + pmf_sym[i,4]**2)/2.0
    return pmf_unsym, pmf_sym

def resistivity(diff, pmf, T, bin_width):
    integrand = np.zeros(len(diff))
    err = 0.0
    for i in range(0, len(diff)):
        if diff[i,1] > 0:
            integrand[i] = np.exp(pmf[i,5]/(R*T))/diff[i,1]
            err += ((bin_width*integrand[i])**2)*((pmf[i,6]/(R*T))**2 + (diff[i,2]/diff[i,1])**2)
        else:
            continue
    Res = np.trapz(integrand, dx = bin_width)*10**(-7)
    dRes = np.sqrt(err)*10**(-7)
    return Res, dRes

def check_build_index(index, nmol, natoms, savedir, latdir):
    """ Checks if the number of atoms declared in the index is correct, creates a new index with split molecules """
    with open(index, 'r') as f:
        lines = f.read().splitlines()
        list_atoms = [lines[i].split() for i in range(1, len(lines))]
    check_atoms = 0
    for i in range(0, len(list_atoms)):
        check_atoms += len(list_atoms[i])
    if (nmol*natoms != check_atoms):
        wrong_index = "\nWrong index file and/or number of atoms per pull group."
        wrong_index += "\nIndex file                : {}".format(index)
        wrong_index += "\nAtoms in index files      : {}".format(check_atoms)
        wrong_index += "\nNumber of molecules (.tpr): {}".format(nmol)
        wrong_index += "\nNumber of atoms (-natoms) : {}".format(natoms)
        wrong_index += "\nNumber of atoms in index file must be the same of number of molecules times number of atoms."
        parser.error(wrong_index)
    new_index = ""
    for i in range(0, nmol):
        new_index += "[ PULL_GROUP_{} ]\n".format(i)
        for j in range(0, natoms):
            new_index += "{} ".format(i*natoms + j + 1)
        new_index += "\n"
    command(["mkdir", "./{}/{}".format(savedir, latdir)], None, False)
    command(["cp", index, "./{}/{}/original_index.ndx".format(savedir, latdir)], None, False)
    text_file = open("./{}/{}/local_index.ndx".format(savedir, latdir), "wt")
    text_file.write(new_index)
    text_file.close()

def get_coordinates(lat_diff, nfiles, savedir, latdir, gmx, index, tprfiles, xtcfiles, nmol):
    """ Reduces .xtc and .tpr files via gmx and extract coordinates """
    if lat_diff == "x":
        noaxis = "-nox"
    elif lat_diff == "y":
        noaxis = "-noy"
    elif lat_diff == "z":
        noaxis = "-noz"
    for i in tqdm(range(0, nfiles), desc = "Reducing xtc and tpr        "):
        tpr = "./{}/{}/red_{}.tpr".format(savedir, latdir, i)
        xtc = "./{}/{}/red_{}.xtc".format(savedir, latdir, i)
        command([gmx, "convert-tpr", "-n", index, "-s", tprfiles[i], "-o", tpr], None, False)
        command([gmx, "trjconv", "-f", xtcfiles[i], "-s", tpr, "-o", xtc], "0", False)
        for j in tqdm(range(0, nmol), leave = False, desc = " ---> Extracting coordinates"):
            command([gmx, "traj", "-f", xtc, "-s", tpr, "-ox", "./{}/{}/coord_file_{}_mol_{}.xvg".format(savedir, latdir, i, j), noaxis, "-com", "-nojump", "-n", "./{}/{}/local_index.ndx".format(savedir, latdir)],"{}".format(j), False)

def einstein(nsteps, nstxout, ave_mol, perc, nfiles, nmol, savedir, latdir, dt, pos, dlat_len):
    """ Calculates mean squared displacements and applies standard Einstein's MSD formula, returns calculated diffusion coefficients """
    coord_steps = int(nsteps/nstxout + 1)
    tot_steps = int(coord_steps/ave_mol)
    real_steps = int(perc*tot_steps)
    MSD = np.zeros((nfiles, nmol, real_steps))
    for i in tqdm(range(0, nfiles), desc = "Loading coordinate files    "):
        for j in tqdm(range(0, nmol), leave = False, desc = " ---> Calculating MSD...    "):
            data = np.loadtxt("./{}/{}/coord_file_{}_mol_{}.xvg".format(savedir, latdir, i, j), dtype = float, comments = ('#','@','&'))
            for k in range(1, real_steps):
                for l in range(0, ave_mol):
                    MSD[i,j,k] += ((data[k+l*tot_steps,1] - data[l*tot_steps,1])**2 + (data[k+l*tot_steps,2] - data[l*tot_steps,2])**2)/k
    MSD = (10.0*MSD)/(2.0*2.0*ave_mol*dt*nstxout)
    print("Processing individual lateral diffusion profiles...")
    out_MSD = np.zeros((real_steps, nfiles+1))
    ein_diff = np.zeros((nfiles, nmol, 2))
    for i in range(0, nmol):
        for j in range(0, real_steps):
            out_MSD[j,0] = j*dt*nstxout*10**(-6)
            out_MSD[j,1:] = MSD[:,i,j]
        head_a = "Time limit of Mean Squared Displacement as per Einsten formula given a molecule\n"
        head_a += "Molecule number   : {}\n".format(i)
        head_a += "Averaged molecules: {}\n".format(ave_mol)
        head_a += "Data per molecule : {}\n".format(perc)
        head_a += " win (nm)  "
        head_c = "\n time (ns)"
        head_b = "Window position and related lateral diffusion coefficient given a molecule\n"
        head_b += "Molecule number   : {}\n".format(i)
        head_b += "Averaged molecules: {}\n".format(ave_mol)
        head_b += "Data per molecule : {}\n".format(perc)
        head_b += "Length of average : {}\n".format(dlat_len)
        head_b += "win (nm)  D_lat (cm**2/s)  std_dev (cm**2/s)"
        for k in range(0, nfiles):
            head_a += str(" {:.4f} ".format(pos[k,i])).center(13, ' ')
            head_c += ("window_{}".format(k)).center(13, ' ')
            ein_diff[k,i,0] = np.average(out_MSD[int((1.0-dlat_len)*real_steps):,k+1])
            ein_diff[k,i,1] = np.std(out_MSD[int((1.0-dlat_len)*real_steps):,k+1])/np.sqrt(ave_mol)
        np.savetxt("./{}/MSD_molecule_{}.txt".format(savedir, i), out_MSD, fmt = '%.6e', header = head_a + head_c)
        np.savetxt("./{}/Dlat_individual_molecule_{}.txt".format(savedir, i),np.stack((pos[:,i], ein_diff[:,i,0], ein_diff[:,i,1])).transpose(), fmt = '%.6e', header = head_b)
    return ein_diff

def main(argv):
    whoami = "This python script calculates the potential of the mean force (PMF), the perpendicular (D_perp) and the lateral (D_lat) diffusion coefficients for a (set of) molecule(s) constrained/restrained in a lipid membrane. "
    whoami += "The script is intended for usage with output coming from GROMACS, in particular when the pull code is applied with a constrain or umbrella potential. "
    whoami += "It is implicitly supposed that the molecules are all the same and that the constraint type (and strength) are all the same. "
    whoami += "\n\n"
    whoami += " > \033[4mGROMACS binaries\033[0m"
    whoami += "\n"
    whoami += "Different tools from GROMACS library are utilized, therefore it is expected that GROMACS is installed and the corresponding binaries are sourced. The standard command used by this script is \'gmx\'. "
    whoami += "Please check if the binaries are correctly sourced by calling \'gmx --h\'. If not, source them before running this script or give the full path by specifying the option -gmx (e.g. -gmx /usr/bin/GMRX). "
    whoami += "The script has been developed with GROMACS2018.1, but should work for any version since 5.0 and after. If GROMACS was compiled with PMI, specify -gmx gmx_mpi and source the binary or specify the full path. "
    whoami += "\n\n"
    whoami += " > \033[4mPerpendicular diffusion coefficient profile via Kubo relation\033[0m"
    whoami += "\n"
    whoami += "To calculate the perpendicular diffusion profile via the Kubo relation (force ACF integration), specify the option -diff_perp_kubo. The required inputs are the following lists of files: tpr (-it), pullx (-ix) and pullf (-if). "
    whoami += "The diffusion coefficients are calculated via the autocorrelation function of the force acting on the centre of mass of the molecules "
    whoami += "[S.J. Marrink and H.J. Berendsen, The Journal of Physical Chemistry, vol. 98, no. 15, pp. 4155-4168, 1994]. "
    whoami += "The code takes the input parameters in the tpr files and checks their consistency. Then it loads the pullx files and computes the average position of the molecules, that is the window positions for the profile. "
    whoami += "Finally, it loads the pullf files and computes the autocorrelation functions (ACFs) of the force acting on the centre of mass (COM) of the molecules. Force ACFs are computed via GROMACS' \'gmx analyze\' tool and the results stored. "
    whoami += "The force ACFs are then loaded and integrated with the trapezium rule and the individual diffusion values per window are saved. "
    whoami += "The values are collected into a profile, that is saved both as is and symmetrized with respect to the 0 of the window positions. "
    whoami += "\n\n"
    whoami += "Errors are calculated by taking average and deviation if more than one value for the diffusion profile falls in one window. In case only one value is present in a window, no error estimate is given. "
    whoami += "\n\n"
    whoami += "Force ACFs are integrated along the first 10000 points, i.e. a time length of integration equal to 10000 times the time step of integration of the simulation and the output timestep. The number of points can be changed by specifying -fACFlen. "
    whoami += "The individual calculated values are binned to obtain a global profile. The bin width is 0.1 nm by default. It should roughly correspond to the original distance of the windows. It can be changed by specifying -bw and a value in nm. "
    whoami += "The symmetry point for the profiles is 0 by default, it can be changed by specifying -com and a value in nm. When saving the force ACFs, the code checks the ratio among the maximum values of the fACFs at the beginning and at the end of the integration time. "
    whoami += "The values of the ratios are store in the corresponding output files. If the ratio is more than 5%, the code warns the user about a slow convergence in the fACFs. In this case, please do consider longer integration times. "
    whoami += "\n\n"
    whoami += " > \033[4mPerpendicular diffusion coefficient profile with position autocorrelation function method\033[0m"
    whoami += "\n"
    whoami += "To calculate the perpendicular diffusion profile via pACF integration, specify the option -diff_perp_pACF. The required inputs are the following lists of files: tpr (-it) and pullx (-ix). "
    whoami += "The diffusion coefficients are calculated via the autocorrelation function of the position of the centre of mass of the molecules "
    whoami += "[G. Hummer, New Journal of Physics, vol. 7, no. 1, p. 34, 2005]. "
    whoami += "The code takes the input parameters in the tpr files and checks their consistency. Then it loads the pullx files and computes the average position of the molecules, that is the window positions for the profile. "
    whoami += "Then it computes the variance of the positions and the position autocorrelation functions of the coordinate of the centre of mass (COM) of the molecules. Position ACFs are computed via GROMACS' \'gmx analyze\' tool and the results stored. "
    whoami += "The position ACFs are then loaded and integrated with the trapezium rule and the individual diffusion values per window are saved. "
    whoami += "The values are collected into a profile, that is saved both as is and symmetrized with respect to the 0 of the window positions. "
    whoami += "\n\n"
    whoami += "Errors are calculated by taking average and deviation if more than one value for the diffusion profile falls in one window. In case only one value is present in a window, no error estimate is given. "
    whoami += "\n\n"
    whoami += "Position ACFs are integrated along the first 10000 points, i.e. a time length of integration equal to 10000 times the time step of integration of the simulation and the output timestep. The number of points can be changed by specifying -pACFlen. "
    whoami += "The individual calculated values are binned to obtain a global profile. The bin width is 0.1 nm by default. It should roughly correspond to the original distance of the windows. It can be changed by specifying -bw and a value in nm. "
    whoami += "The symmetry point for the profiles is 0 by default, it can be changed by specifying -com and a value in nm. When saving the position ACFs, the code checks the ratio among the maximum values of the pACFs at the beginning and at the end of the integration time. "
    whoami += "The values of the ratios are store in the corresponding output files. If the ratio is more than 5%, the code warns the uses about a slow convergence in the pACFs. In this case, please do consider longer integration times. "
    whoami += "\n\n"
    whoami += " > \033[4mPotential of the mean force profile via potential of the mean constraint force method\033[0m"
    whoami += "\n"
    whoami += "To calculate the potential of the mean force (PMF), specify the option -pmf. The required inputs are the following lists of files: tpr (-it), pullx (-ix) and pullf (-if). "
    whoami += "The PMF profile is calculated by using the average of the force acting on the centre of mass of the molecules, aka position of the mean constraint force method "
    whoami += "[H.J. Berendsen and S.J. Marrink, Pure and applied chemistry, vol. 65, no. 12, pp. 2513-2520, 1993]. "
    whoami += "The code takes the input parameters in the tpr files and checks their consistency. Then it loads the pullx files and computes the average position of the molecules, that is the window positions for the profile. "
    whoami += "Finally, it loads the pullf files and computes the average of the force (<F>) acting on the centre of mass (COM) of the molecules. The individual <F> vs time curves for the molecules are computed and saved. "
    whoami += "The final values for <F> are derived by averaging on the last part of the <F> vs time plots, where it is expected that the average of the force has reached a steady value. "
    whoami += "The values are collected into a force profile, that is saved both as is and symmetrized with respect to the 0 of the window positions. If more than one value is collected per window, the average and its deviation are saved as value and error. "
    whoami += "If only one value is present in a window, then the associated error is the deviation of the average done on the secondo half of <F>(t) plot. "
    whoami += "The PMF is calculated by integrating the force profile. PMF calculations are carried out and saved both for the symmetrized and unsymmetrized force inputs. "
    whoami += "Errors on the are calculated by variance propagation from the force averages. In particular, errors are calculated by propagating the error on the integral carried on from the left and the right. "
    whoami += "The final output files for the sym/unsym profiles contain the following: window position, PMF integrated from left and associated error, PMF calculated from right and associated error and average with propagated error. "
    whoami += "\n\n"
    whoami += "The <F> vs time plots are calculated with the input force values divided in 1000 blocks, i.e. the output are 1000 couples of (t,<F(z)>(t)) points. The length can be changed by specifying -avelen and a positive integer. "
    whoami += "Similarly, the final value of <F(z)> is calculated on the second half of the (t,<F(z)>(t)) curve. It can be changed by specifying -flen and the normalized amount of data that is wanted, e.g. 0.60 to use 60% of the curve for <F(z)> calculation. "
    whoami += "The individual calculated force values are binned to obtain a global profile. The bin width is 0.1 nm by default. It should roughly correspond to the original distance of the windows. It can be changed by specifying -bw and a value in nm. "
    whoami += "The symmetry point for the profiles is 0 by default, it can be changed by specifying -com and a value in nm. "
    whoami += "The PMF profile is computed as is and is not shifted vertically. This can be done by specifying -dG and an energy value in kJ/mol. Please note that the shift is eventually going to change the resistivity/permeability values, if calculated. "
    whoami += "\n\n"
    whoami += " > \033[4mResistivity and permeability\033[0m"
    whoami += "\n"
    whoami += "If both the perpendicular diffusion coefficient and potential of the mean force profiles calculations are specified with -diff_perp_kubo/-diff_perp_pACF and -pmf, the script calculates the resistivity and the permeability with the inhomogeneous solubility diffusion model "
    whoami += "[S.J. Marrink and H.J. Berendsen, The Journal of Physical Chemistry, vol. 98, no. 15, pp. 4155-4168, 1994]. "
    whoami += "In case the perpendicular diffusion profile with both Kubo and pACF methods are requested, the code evaluates resistivity and permeability with both diff profiles. "
    whoami += "Results are printed on screen and saved in the summary file. Errors are calculated by propagating the variance from diff and pmf profiles in the integration of the inhomogenous solubility diffusion model. "
    whoami += "\n\n"
    whoami += " > \033[4mLateral diffusion coefficient profile via MSD\033[0m"
    whoami += "\n"
    whoami += "To calculate the lateral diffusion profile via MSD method, specify the option -diff_lat. The required inputs are the following lists of files: tpr (-it), pullx (-ix) and xtc (-xtc). "
    whoami += "Moreover, the following parameters are needed: the number of atoms per group (-natoms) and a special GROMACS index file (-n). "
    whoami += "The number of atoms per group is the number of atoms of the pull groups of interest. For example, in an all-atoms simulation where one or more water molecules are constrained/restrained at a given distance from a lipid bilayer, you should specify -natoms 3 (two hydrogens and one oxygen). "
    whoami += "The index file given as an input should be built by the user via \'gmx make_ndx\' and must contain only one group with all the atoms of the pulled molecules. For example, in an all-atoms simulation where five water molecules are constrained the group must contain the 15 (5x3) indexes of all the water atoms. "
    whoami += "The lateral diffusion coefficients are calculated via the Einstein formula as the time limit of the mean squared displacement (MSD) of the centre of mass (COM) of the molecules. "
    whoami += "The code takes the input parameters in the tpr files and checks their consistency. Then it loads the pullx files and computes the average position of the molecules, that is the window positions for the profile. "
    whoami += "Then, the script reduces the tpr and xtc files via \'gmx convert-tpr\' and \'gmx trjconv\' commands and extracts the coordinates of the molecules of interest from the xtc files via \'gmx traj\'. "
    whoami += "The MSD of the molecules is then computed from the coordinate files and the corresponding time limits are calculated and saved per individual molecule. "
    whoami += "The lateral diffusion coefficients are derived by averaging on the final part of the MSD time limit plot and are then saved individually per molecule. "
    whoami += "The values are collected into a profile, that is saved both as is and symmetrized with respect to the 0 of the window positions. "
    whoami += "If more than one value per bin is found, the arithmetic average is returned. Error is calculated as the standard deviation of the mean. If only one value is present, the error is calculated as the standard deviation of the force mean squared displacement vs time curve. "
    whoami += "\n\n"
    whoami += "The individual calculated values are binned to obtain a global profile. The bin width is 0.1nm by default. It should roughly correspond to the original distance of the windows. It can be changed by specifying -bw and a value in nm. "
    whoami += "The symmetry point for the profiles is 0 by default, it can be changed by specifying -com and a value in nm. "
    whoami += "The script assumes that the plane of interest for the lateral diffusion calculation is perpendicular to the z axis. To choose another axis, specify -lateral {x,y,z} (default is -lateral z). "
    whoami += "To enhance statistics, it is possible to split a single run for a molecule into n sub runs. The default is to consider just one molecule, to change this specify -nave followed by the number of virtual molecules to be considered. "
    whoami += "If -nave is greater than 1, it is better to separate sequential MSD calculations to decorrelate the values. By default, the code considers all of the data. To change this, specify -e and a value in between 0 (discard all data) and 1 (consider all data). "
    whoami += "To calculate the individual lateral diffusion coefficients, an average of the MSD vs time plot is taken. Since the starting values are usually noisy, the first 50% (i.e. 0.5) of the plot is not considered. To change this, specify -dllen and a value in between 0 (average on whole data) and 1 (discard all). "

    parser=argparse.ArgumentParser(description = whoami, formatter_class = argparse.RawDescriptionHelpFormatter)
    parser.add_argument('-it', '--tpr', nargs = 1, type=str, required = True, help = 'List of .tpr files in order.')
    parser.add_argument('-ix', '--pullx', nargs = 1, type=str, required = True, default = None, help = 'List of pullx.xvg files in order.')
    parser.add_argument('-if', '--pullf', nargs = 1, type=str, required = False, default = None, help = 'List of pullf.xvg files in order.')
    parser.add_argument('-diff_perp_kubo', '--comp_diff_perp_kubo', default = False, required = False, action='store_true', help = 'Specify to calculate the perpendicular diffusion profile via Kubo relation with input data.')
    parser.add_argument('-pmf', '--comp_pmf', default = False, required = False, action='store_true', help = 'Specify to calculate the potential of the mean force profile with input data via potential of the mean constraint force.')
    parser.add_argument('-fACFlen', '--fACFlength', type = int, required = False, default = 10000, help = 'Step length of force autocorrelation function. Default is 10000.')
    parser.add_argument('-avelen', '--avelength', type = int, required = False, default = 1000, help = 'Step length of force average calculation. Default is 1000.')
    parser.add_argument('-flen', '--aveFlength', type = float, required = False, default = 0.50, help = 'Amount of data used to calculate average for <F> values in PMF calculation. Default is 0.50.')
    parser.add_argument('-bw', '--binwidth', type = float, required = False, default = 0.10, help = 'Bin width (nm) for binning profiles. Default is 0.1 nm.')
    parser.add_argument('-com', '--deltacom', type = float, required = False, default = 0.0, help = 'Shift of the COM (nm) for symmetrisation of profiles. Default is 0 nm.')
    parser.add_argument('-dG', '--deltag', type = float, required = False, default = 0.0, help = 'Shift of the PMF (kJ/mol). Default is 0 kJ/mol.')
    parser.add_argument('-diff_lat', '--comp_diff_lat', default = False, required = False, action='store_true', help = 'Specify to calculate the lateral diffusion profile with input data.')
    parser.add_argument('-xtc', '--coord', nargs = 1, type=str, required = False, default = None, help = 'List of .xtc files in order.')
    parser.add_argument('-n', '--index', nargs = 1, type = str, required = False, default = None, help = 'Index file in GROMACS format.')
    parser.add_argument('-natoms', '--num_atoms', type = int, required = False, default = 0, help = 'Number of atoms per pull group.')
    parser.add_argument('-lateral', '--lateral_diff', type = str, required = False, default = "z", choices = ["x", "y", "z"], help = 'Direction perpendicular to plane of diffusion. Default is z.')
    parser.add_argument('-nave', '--num_ave_mol', type = int, required = False, default = 1, help = 'Number of molecules to subdivide the mean squared displacement calculation.')
    parser.add_argument('-e', '--end_ave', type = float, required = False, default = 1.0, help = 'Amount of data used to calculate MSD. Default is the whole dataset.')
    parser.add_argument('-dllen', '--dlat_length', type = float, required = False, default = 0.50, help = 'Amount of data used to calculate average for lateral diffusion coefficients. Default is 0.50.')
    parser.add_argument('-gmx', '--gmx_binary', type = str, required = False, default = "gmx", help = 'GROMACS binary to run the commands.')
    parser.add_argument('-diff_perp_pACF', '--comp_diff_perp_pACF', default = False, required = False, action='store_true', help = 'Specify to calculate the perpendicular diffusion profile with input data via pACF.')
    parser.add_argument('-pACFlen', '--pACFlength', type = int, required = False, default = 10000, help = 'Step length of positions autocorrelation function. Default is 10000.')


    args = parser.parse_args()

    title = "\n"
    title += "*************************************************************************************************************************************\n"
    title += "****************   This python script calculates the potential of the mean force and the lateral and perpendicular   ****************\n"
    title += "****************             diffusion profiles for restrained/constrained molecules in a lipid bilayer.             ****************\n"
    title += "****************                                Call --h for help and documentation.                                 ****************\n"
    title += "*************************************************************************************************************************************\n"
    print(title)

    with open(args.tpr[0]) as f:
        tprfiles = [line.rstrip() for line in f]
    nfiles = len(tprfiles)

    with open(args.pullx[0]) as f:
        pxfiles = [line.rstrip() for line in f]

    calc_diff_perp_kubo = args.comp_diff_perp_kubo
    calc_pmf = args.comp_pmf
    corr_time = args.fACFlength
    ave_length = args.avelength
    aveF_len = args.aveFlength
    bin_width = args.binwidth
    delta_com = args.deltacom
    delta_g = args.deltag
    calc_diff_lat = args.comp_diff_lat
    index = args.index
    natoms = args.num_atoms
    lat_diff = args.lateral_diff
    ave_mol = args.num_ave_mol
    perc = args.end_ave
    dlat_len = args.dlat_length
    gmx = args.gmx_binary
    calc_diff_perp_pACF = args.comp_diff_perp_pACF
    position_corr_time = args.pACFlength

    if (calc_diff_perp_kubo or calc_pmf):
        if (args.pullf != None):
            with open(args.pullf[0]) as f:
                pffiles = [line.rstrip() for line in f]
        else:
            parser.error('Perpendicular diffusion with Kubo relation and/or PMF profiles requested with -diff_perp_kubo and/or -pmf, please specify -if.')
        if (len(pxfiles) != nfiles or len(pffiles) != nfiles or nfiles <= 0):
            parser.error("Number of input files is different or zero: check -ix, -if and -it file lists.")
        if ((aveF_len < 0.0) or (aveF_len >= 1.0)):
            parser.error('The amount of data to keep for <F> calculation specified with -flen should be in between 0 and 1. You inserted the value: {}.'.format(aveF_len))

    if (calc_diff_perp_pACF):
        if (len(pxfiles) != nfiles or nfiles <= 0):
            parser.error("Number of input files is different or zero: check -ix and -it file lists.")

    if calc_diff_lat:
        if ((args.coord != None) and (args.index != None) and (natoms > 0)):
            with open(args.coord[0]) as f:
                xtcfiles = [line.rstrip() for line in f]
            index = args.index[0]
        else:
            parser.error('Lateral diffusion profile requested with -diff_lat, please specify -n, -xtc and -natom.')
        if ((len(xtcfiles) != nfiles) or (len(pxfiles) != nfiles) or (nfiles <= 0)):
            parser.error("Number of input files is different or zero: check -xtc, -ix and -it file lists.")
        if ((perc <= 0.0) or (perc > 1.0)):
            parser.error('The amount of data to keep for MSD calculation specified with -e should be in between 0 and 1. You inserted the value: {}.'.format(perc))
        if ((dlat_len <= 0.0) or (dlat_len > 1.0)):
            parser.error('The amount of data to keep for lateral diffusion coefficients calculation specified with -dllen should be in between 0 and 1. You inserted the value: {}.'.format(dlat_len))
        if (ave_mol <= 0):
            parser.error('The number of molecules to use for average in lateral diffusion coefficients calculation specified with -nave should be at least 1. You inserted the value: {}.'.format(ave_mol))

    nmol, T, dt, pffreq, pxfreq, nsteps, nstxout = unpack_tpr(tprfiles,gmx)

    savedir = "script_output"
    if calc_diff_perp_kubo:
        facfdir = "fACF_{}_analyze_tlength_{}ps".format(gmx,int(corr_time*dt*pffreq/1000.0))
    pacfdir = ""
    if calc_diff_perp_pACF:
        pacfdir += "pACF_{}_analyze_tlength_{}ps".format(gmx,int(position_corr_time*dt*pxfreq/1000.0))
    if calc_diff_lat:
        latdir = "lateral_diffusion"
    if (not (calc_diff_perp_kubo or calc_diff_perp_pACF or calc_pmf or calc_diff_lat)):
        parser.error('No profile to compute requested. Please specify -diff_perp_kubo, -diff_perp_pACF, -pmf and/or -diff_lat.')
    command(["mkdir", savedir], None, False)

    head_all = "\n"
    head_all += "--- This is your list of parameters ---\n"
    head_all += "\n"
    head_all += "Results saved in        : {}\n".format(savedir)
    head_all += "Simulation time (ns)    : {}\n".format(nsteps*dt/10**6)
    head_all += "Number of molecules     : {}\n".format(nmol)
    head_all += "Number of pull files    : {}\n".format(nfiles)
    head_all += "Bin width (nm)          : {}\n".format(bin_width)
    head_all += "Temperature (K)         : {}\n".format(T)
    head_all += "Time step (fs)          : {}\n".format(dt)
    head_all += "Nsteps of force output  : {}\n".format(pffreq)
    head_all += "Nsteps of pos output    : {}\n".format(pxfreq)
    head_all += "PMF profile             : {}\n".format(calc_pmf)
    if calc_pmf:
        head_all += " -> <F> step length     : {}\n".format(ave_length)
        head_all += " -> <F> average length  : {}\n".format(aveF_len)
    head_all += "Perp diffusion profile  : {}\n".format(calc_diff_perp_kubo)
    if calc_diff_perp_kubo:
        head_all += " -> fACF step length    : {}\n".format(corr_time)
        head_all += " -> fACF time length(ps): {}\n".format(corr_time*dt*pffreq/1000.0)
    if calc_diff_perp_pACF:
        head_all += " -> pACF step length    : {}\n".format(position_corr_time)
        head_all += " -> pACF time length(ps): {}\n".format(position_corr_time*dt*pxfreq/1000.0)
    if delta_com > 0.0:
        head_all += "Symmetrization shift(nm): {}\n".format(delta_com)
    if delta_com > 0.0:
        head_all += "PMF shift       (kJ/mol): {}\n".format(delta_g)
    head_all += "Lat diffusion profile   : {}\n".format(calc_diff_lat)
    if calc_diff_lat:
        head_all += " -> Index file          : {}\n".format(index)
        head_all += " -> Atoms per pull group: {}\n".format(natoms)
        head_all += " -> Perpendicular to    : {} axis\n".format(lat_diff)
        head_all += " -> Nsteps coord output : {}\n".format(nstxout)
        head_all += " -> N mol subdivision   : {}\n".format(ave_mol)
        head_all += " -> Length of MSD data  : {}\n".format(perc)
        head_all += " -> Length of average   : {}\n".format(dlat_len)

    print(head_all)
    print("Starting computation...")

    pos, var = window_positions(pxfiles,nfiles,nmol)

    if calc_diff_perp_kubo:
        print("\n--> Calculating perpendicular diffusion coefficient profile via Kubo relation...\n")
        fACFs = force_ACF(pos,pffiles,facfdir,savedir,nfiles,nmol,corr_time,dt,pffreq,gmx)
        fACF_diff = fACF_diffusion(pos,fACFs,savedir,nmol,nfiles,T,dt,pffreq,corr_time)
        print("Processing and saving perpendicular diffusion profile...")
        fACF_diff_unsym, fACF_diff_sym = global_profile(pos,delta_com,bin_width,nfiles,nmol,fACF_diff,False)
        head_pdu = "Window position, perpendicular diffusion coefficient via Kubo and corresponding error (unsymmetrized)"
        head_pdu += "\nCorrelation time length: {}ps".format(int(corr_time*dt*pffreq/1000.0))
        head_pdu += "\n win (nm)  D_perp(cm**2/s)  sigma_D(cm**2/s)"
        np.savetxt("./{}/Dperp_kubo_unsym.txt".format(savedir), fACF_diff_unsym, fmt = '%.6e', header = head_pdu)
        head_pds = "Window position, perpendicular diffusion coefficient via Kubo and corresponding error (symmetrized)"
        head_pds += "\nCorrelation time length: {}ps".format(int(corr_time*dt*pffreq/1000.0))
        head_pds += "\n win (nm)  D_perp(cm**2/s)  sigma_D(cm**2/s)"
        np.savetxt("./{}/Dperp_kubo_sym.txt".format(savedir), fACF_diff_sym, fmt = '%.6e', header = head_pds)

    if calc_diff_perp_pACF:
        print("\n--> Calculating perpendicular diffusion coefficient profile via pACF...\n")
        pACFs = position_ACF(pos,pxfiles,pacfdir,savedir,nfiles,nmol,position_corr_time,dt,pxfreq,gmx)
        pACF_diff = pACF_diffusion(pos,pACFs,savedir,nmol,nfiles,dt,pxfreq,position_corr_time,var)
        print("Processing and saving pACF perpendicular diffusion profile...")
        pACF_diff_unsym, pACF_diff_sym = global_profile(pos,delta_com,bin_width,nfiles,nmol,pACF_diff,False)
        head_ppdu = "Window position, perpendicular diffusion coefficient via pACF and corresponding error (unsymmetrized)"
        head_ppdu += "\nCorrelation time length: {}ps".format(int(position_corr_time*dt*pxfreq/1000.0))
        head_ppdu += "\n win (nm)  D_perp(cm**2/s)  sigma_D(cm**2/s)"
        np.savetxt("./{}/Dperp_pACF_unsym.txt".format(savedir), pACF_diff_unsym, fmt = '%.6e', header = head_ppdu)
        head_ppds = "Window position, perpendicular diffusion coefficient via pACF and corresponding error (symmetrized)"
        head_ppds += "\nCorrelation time length: {}ps".format(int(position_corr_time*dt*pxfreq/1000.0))
        head_ppds += "\n win (nm)  D_perp(cm**2/s)  sigma_D(cm**2/s)"
        np.savetxt("./{}/Dperp_pACF_sym.txt".format(savedir), pACF_diff_sym, fmt = '%.6e', header = head_ppds)

    if calc_pmf:
        print("\n--> Calculating potential of the mean force profile...\n")
        forces_time = float_average(pos,pffiles,nfiles,nmol,pffreq,dt,nsteps,savedir,ave_length)
        forces = force_average(forces_time,nfiles,nmol,aveF_len, pos, savedir)
        print("Processing and saving <F> profile...")
        force_unsym, force_sym = global_profile(pos,delta_com,bin_width,nfiles,nmol,forces,True)
        head_fu = "Window position, average of force and corresponding error (unsymmetrized)"
        head_fu += "\n win (nm)   <F> (kJ/mol/nm)   sigma_<F> (kJ/mol/nm)"
        np.savetxt("./{}/force_unsym.txt".format(savedir), force_unsym, fmt = '%.6e', header = head_fu)
        head_fs = "Window position, average of force and corresponding error (symmetrized)"
        head_fs += "\n win (nm)   <F> (kJ/mol/nm)   sigma_<F> (kJ/mol/nm)"
        np.savetxt("./{}/force_sym.txt".format(savedir), force_sym, fmt = '%.6e', header = head_fs)
        print("Processing and saving PMF profile...")
        pmf_unsym, pmf_sym = pmf_integration(pos,delta_com,bin_width,delta_g,force_unsym,force_sym)
        head_pu = "Window position, PMF and corresponding error for curve integrated from left, right and their average respectively (unsymmetrized)"
        head_pu += "\n win (nm)    PMF_l (kJ/mol)  sigma_PMF_l (kJ/mol)  PMF_r (kJ/mol)  sigma_PMF_r (kJ/mol)  PMF_ave (kJ/mol)  sigma_PMF_ave (kJ/mol)"
        np.savetxt("./{}/PMF_unsym_control.txt".format(savedir), pmf_unsym, fmt = '%.6e', header = head_pu)
        head_ps = "Window position, PMF and corresponding error for curve integrated from left, right and their average respectively (symmetrized)"
        head_ps += "\n win (nm)    PMF_l (kJ/mol)  sigma_PMF_l (kJ/mol)  PMF_r (kJ/mol)  sigma_PMF_r (kJ/mol)  PMF_ave (kJ/mol)  sigma_PMF_ave (kJ/mol)"
        np.savetxt("./{}/PMF_sym_control.txt".format(savedir), pmf_sym, fmt = '%.6e', header = head_ps)
        head_fpu = "Window position, PMF and corresponding error (unsymmetrized)"
        head_fpu += "\n win (nm)    PMF (kJ/mol)  sigma_PMF (kJ/mol)"
        np.savetxt("./{}/PMF_unsym.txt".format(savedir), np.stack((pmf_unsym[:,0],pmf_unsym[:,5],pmf_unsym[:,6])).transpose(), fmt = '%.6e', header = head_fpu)
        head_fps = "Window position, PMF and corresponding error (symmetrized)"
        head_fps += "\n win (nm)    PMF (kJ/mol)  sigma_PMF (kJ/mol)"
        np.savetxt("./{}/PMF_sym.txt".format(savedir), np.stack((pmf_sym[:,0],pmf_sym[:,5],pmf_sym[:,6])).transpose(), fmt = '%.6e', header = head_fps)
        if calc_diff_perp_kubo:
            fACF_Res_u, fACF_dRes_u = resistivity(fACF_diff_unsym,pmf_unsym,T,bin_width)
            head_all += "\nFinal results of resistivity (R) and permeability (P) and associated error for Kubo diffusion coefficients (unsymmetrized profiles)"
            head_all += "\nR = {:.2e} \u00B1 {:.2e}".format(fACF_Res_u,fACF_dRes_u)
            head_all += "\nP = {:.2e} \u00B1 {:.2e}\n".format(1.0/fACF_Res_u,fACF_dRes_u/fACF_Res_u**2)
            fACF_Res_s, fACF_dRes_s = resistivity(fACF_diff_sym,pmf_sym,T,bin_width)
            head_all += "\nFinal results of resistivity (R) and permeability (P) and associated error for Kubo diffusion coefficients (symmetrized profiles)"
            head_all += "\nR = {:.2e} \u00B1 {:.2e}".format(fACF_Res_s,fACF_dRes_s)
            head_all += "\nP = {:.2e} \u00B1 {:.2e}\n".format(1.0/fACF_Res_s,fACF_dRes_s/fACF_Res_s**2)
            print("\nFinal results of resistivity (R) and permeability (P) and associated error (Kubo diffusion)"+"\nR = {:.2e} \u00B1 {:.2e}".format(fACF_Res_s,fACF_dRes_s)+"\nP = {:.2e} \u00B1 {:.2e}\n".format(1.0/fACF_Res_s,fACF_dRes_s/fACF_Res_s**2))
        if calc_diff_perp_pACF:
            pACF_Res_u, pACF_dRes_u = resistivity(pACF_diff_unsym,pmf_unsym,T,bin_width)
            head_all += "\nFinal results of resistivity (R) and permeability (P) and associated error for pACF diffusion coefficients (unsymmetrized profiles)"
            head_all += "\nR = {:.2e} \u00B1 {:.2e}".format(pACF_Res_u,pACF_dRes_u)
            head_all += "\nP = {:.2e} \u00B1 {:.2e}\n".format(1.0/pACF_Res_u,pACF_dRes_u/pACF_Res_u**2)
            pACF_Res_s, pACF_dRes_s = resistivity(pACF_diff_sym,pmf_sym,T,bin_width)
            head_all += "\nFinal results of resistivity (R) and permeability (P) and associated error for pACF diffusion coefficients (symmetrized profiles)"
            head_all += "\nR = {:.2e} \u00B1 {:.2e}".format(pACF_Res_s,pACF_dRes_s)
            head_all += "\nP = {:.2e} \u00B1 {:.2e}\n".format(1.0/pACF_Res_s,pACF_dRes_s/pACF_Res_s**2)
            print("\nFinal results of resistivity (R) and permeability (P) and associated error (pACF diffusion)"+"\nR = {:.2e} \u00B1 {:.2e}".format(pACF_Res_s,pACF_dRes_s)+"\nP = {:.2e} \u00B1 {:.2e}\n".format(1.0/pACF_Res_s,pACF_dRes_s/pACF_Res_s**2))

    if calc_diff_lat:
        print("\n--> Calculating lateral diffusion coefficient profile...\n")
        check_build_index(index,nmol,natoms,savedir,latdir)
        get_coordinates(lat_diff,nfiles,savedir,latdir,gmx,index,tprfiles,xtcfiles,nmol)
        ein_diff = einstein(nsteps,nstxout,ave_mol,perc,nfiles,nmol,savedir,latdir,dt,pos,dlat_len)
        ein_unsym, ein_sym = global_profile(pos,delta_com,bin_width,nfiles,nmol,ein_diff,False)
        head_ldu = "Window position, lateral diffusion coefficient and corresponding error (unsymmetrized)\n"
        head_ldu += "Averaged molecules: {}\n".format(ave_mol)
        head_ldu += "Data per molecule : {}\n".format(perc)
        head_ldu += "Length of average : {}\n".format(dlat_len)
        head_ldu += " win (nm)   D_lat(cm**2/s)  sigma_D (cm**2/s)"
        np.savetxt("./{}/Dlat_unsym.txt".format(savedir), ein_unsym, fmt = '%.6e', header = head_ldu)
        head_lds = "Window position, lateral diffusion coefficient and corresponding error (symmetrized)\n"
        head_lds += "Averaged molecules: {}\n".format(ave_mol)
        head_lds += "Data per molecule : {}\n".format(perc)
        head_lds += "Length of average : {}\n".format(dlat_len)
        head_lds += " win (nm)   D_lat(cm**2/s)  sigma_D (cm**2/s)"
        np.savetxt("./{}/Dlat_sym.txt".format(savedir), ein_sym, fmt = '%.6e', header = head_lds)

    head_all += "\nFiles used to calculate output listed in order\n"
    head_all += "->tpr".ljust(len(max(tprfiles,key=len)) + 2)
    if ((calc_diff_perp_kubo or calc_pmf) and (not calc_diff_lat)):
        head_all += "->px".ljust(len(max(pxfiles,key=len)) + 2) + "->pf".ljust(len(max(pffiles,key=len)) + 2)
        for i in range(0,nfiles):
            head_all += "\n"
            head_all += str(tprfiles[i]).ljust(len(max(tprfiles,key=len)) + 2) + str(pxfiles[i]).ljust(len(max(pxfiles,key=len)) + 2) + str(pffiles[i]).ljust(len(max(pffiles,key=len)) + 2)
    if ((calc_diff_perp_kubo or calc_pmf) and calc_diff_lat):
        head_all += "->px".ljust(len(max(pxfiles,key=len)) + 2) + "->pf".ljust(len(max(pffiles,key=len)) + 2) + "->xtc".ljust(len(max(xtcfiles,key=len)) + 2)
        for i in range(0,nfiles):
            head_all += "\n"
            head_all += str(tprfiles[i]).ljust(len(max(tprfiles,key=len)) + 2) + str(pxfiles[i]).ljust(len(max(pxfiles,key=len)) + 2) + str(pffiles[i]).ljust(len(max(pffiles,key=len)) + 2) + str(xtcfiles[i]).ljust(len(max(xtcfiles,key=len)) + 2)
    if (calc_diff_lat and (not calc_diff_perp_kubo) and (not calc_pmf)):
        head_all += "->px".ljust(len(max(pxfiles,key=len)) + 2) + "->xtc".ljust(len(max(xtcfiles,key=len)) + 2)
        for i in range(0,nfiles):
            head_all += "\n"
            head_all += str(tprfiles[i]).ljust(len(max(tprfiles,key=len)) + 2) + str(pxfiles[i]).ljust(len(max(pxfiles,key=len)) + 2) + str(xtcfiles[i]).ljust(len(max(xtcfiles,key=len)) + 2)
    if (calc_diff_perp_pACF and (not calc_pmf) and (not calc_diff_lat) and (not calc_diff_perp_kubo)):
        head_all += "->px".ljust(len(max(pxfiles,key=len)) + 2)
        for i in range(0,nfiles):
            head_all += "\n"
            head_all += str(tprfiles[i]).ljust(len(max(tprfiles,key=len)) + 2) + str(pxfiles[i]).ljust(len(max(pxfiles,key=len)) + 2)


    text_file = open("./{}/analysis_details.txt".format(savedir), "wt")
    text_file.write(head_all)
    text_file.close()

    print("\nCalculations terminated successfully. All results save in directory: {}\n".format(savedir))

if __name__ == '__main__':
    main(sys.argv)