"""
Module: Crooks.py
Author: Miłosz Wieczór <milosz.wieczor@irbbarcelona.org>
License: GPL 3.0

Description:
    This module allows to run non-equilibrium alchemical
    free energy calculations, starting from alchemical topologies
    and equilibrated trajectories

Contents:
    Classes:
        CrooksPool:
            Creates an ensemble (pool) of workers that each run and analyze
            their own non-equilibrium MD from lambda 0->1 and 1->0
        Crooks:
            Represents a single worker running a fast-growth simulation

Usage:
    This module is intended to be employed as described in the tutorial; see:
        https://gitlab.com/KomBioMol/gromologist/-/wikis/Running-nonequilibrium-alchemical-free-energy-calculations

Notes:
    For advanced use, such as inserting mutations on-the-fly with pmx, contact the author.
"""

from concurrent.futures import ProcessPoolExecutor
import multiprocessing
import gromologist as gml
import os
from math import ceil
from shutil import copy

import numpy as np
try:
    from scipy.integrate import simps
    from scipy.optimize import fmin
    from sklearn.neighbors import KernelDensity
    from sklearn.model_selection import GridSearchCV
    from sklearn.model_selection import LeaveOneOut
except ImportError:
    pass


class CrooksPool:
    def __init__(self, struct, top, xtc0, xtc1, nruns=100, length=500, gmx=None, maxwarn=2, convergence=False,
                 alias='free', debug=False, offset=0, nmax=None, init_eq=50, temperature=300, stride=20,
                 plumed='', top2='', tmpi=True, mutate='', pmxff='', wt_only=False, mini=False, struct2='',
                 plumed2='', hbond=False, dt=2, weights=None, random=False, ncpu=None, **kwargs):
        """
        A pool of workers that will initialize all simulations
        and collect the results afterwards
        :param struct: str, a valid .gro or .pdb structure file for the system
        :param top: str, a .top topology file for the system
        :param xtc0: str, a trajectory capturing the equilibrium ensemble for state A (can also be .dcd or any other
        format accepted by MDTraj)
        :param xtc1: str, a trajectory capturing the equilibrium ensemble for state B
        :param nruns: int, number of A->B runs (overall, twice this many runs will be performed)
        :param length: int, length of a single trajectory (in ps)
        :param gmx: str, Gromacs executable (gmx by default)
        :param maxwarn: int, max number of warnings accepted by grompp
        :param convergence: bool, whether to do convergence analysis
        :param alias: str, a base name for all output files
        :param debug: bool, whether to print detailed convergence info
        :param offset: int, start numbering workers from this number (useful when more jobs are added post-factum)
        :param kwargs: dict of str:str pairs, will be converted into additional .mdp options
        """
        self.struct = struct
        self.struct2 = struct2 if struct2 else self.struct
        self.top = top
        self.top2 = top2 if top2 else self.top
        self.name = alias
        self.bandwidth = 0
        self.debug = debug
        self.convergence = convergence
        self.xtc = {0: xtc0, 1: xtc1}
        self.random = random
        self.nruns = nruns
        self.stride = stride
        self.mini = mini
        self.plumed = plumed
        self.plumed2 = plumed2 if plumed2 else self.plumed
        self.init_length = init_eq
        self.sim_length = length
        self.extra_args = kwargs
        self.lincs = 'all-bonds' if not hbond else 'h-bonds'
        self.maxwarn = maxwarn
        self.offset = offset
        self.temperature = temperature
        self.mutate = mutate
        self.pmxff = pmxff
        self.wt_only = wt_only
        self.dt = dt / 1000
        self.nmax = nmax if nmax is not None else self.nruns
        self.init_nst = self.init_length * int(1 / self.dt)
        self.nst = self.sim_length * int(1 / self.dt)
        if gmx is None:
            found_gmx = gml.find_gmx_dir()
            if found_gmx[1]:
                self.gmx = found_gmx
            else:
                raise RuntimeError("No valid gmx executable found, please specify it as gmx=/path/to/gmx")
        else:
            self.gmx = gmx
        self.gmx = 'gmx' if gmx is None else gmx
        self.tmpi = tmpi
        self.ncpu = self.ncpus() if ncpu is None else ncpu
        self.workers = [Crooks(self, n, lam) for n in range(offset, self.nruns + offset) for lam in [0, 1]]
        self.analysis_only = [Crooks(self, n, lam) for n in range(offset) for lam in [0, 1]]
        if weights is None:
            self.weights = [np.ones(self.nruns) / self.nruns, np.ones(self.nruns) / self.nruns]
        else:
            self.weights = [np.loadtxt(weights.replace('X', '0')), np.loadtxt(weights.replace('X', '1'))]
            for w in self.weights:
                w /= np.sum(w)

    def run(self):
        """
        The main routine: extracts initial frames,
        runs grompp and mdrun in parallel
        :return: None
        """
        exec = ProcessPoolExecutor
        ncpu = self.ncpus()
        self.drop_frames()
        if self.mutate:
            print('mutating structures...'.format())
            with exec(ncpu) as executor:
                executor.map(self.apply_pmx, [(self.mutate, self.wt_only, w.id, w.initlambda, self.pmxff)
                                              for w in self.workers])
        with exec(ncpu) as executor:
            print("preparing minimization...")
            executor.map(self.run_worker, [(w, 'mini', self.stride, self.temperature, self.init_nst, self.nst,
                                                self.lincs, self.extra_args, self.mini, self.dt) for w in self.workers])
        self.mdrun_multi('mini')
        with exec(ncpu) as executor:
            print("preparing equilibration...")
            executor.map(self.run_worker, [(w, 'eq', self.stride, self.temperature, self.init_nst, self.nst,
                                            self.lincs, self.extra_args, self.mini, self.dt) for w in self.workers])
        self.mdrun_multi('eq')
        with exec(ncpu) as executor:
            executor.map(self.run_worker, [(w, 'prod', self.stride, self.temperature, self.init_nst, self.nst,
                                            self.lincs, self.extra_args, self.mini, self.dt) for w in self.workers])
        self.mdrun_multi('prod')

    @staticmethod
    def run_worker(data):
        worker, runtype, stride, temperature, init_nst, nst, lincs, extra_args, mini, dt = data
        if runtype == 'mini':
            if mini:
                worker.prep_runs(runtype, stride, temperature, init_nst, nst, lincs, dt, extra_args)
                # We can minimize if required
                worker.log('running energy minimization...'.format())
                worker.grompp_me(runtype)
            else:
                worker.cp(f'run{worker.id}_l{worker.initlambda}/frame{worker.id}_l{worker.initlambda}.gro',
                          'mini.gro'.format())
        elif runtype == 'eq':
            # First there's equilibration:
            worker.prep_runs(runtype, stride, temperature, init_nst, nst, lincs, dt, extra_args)
            worker.log('running equilibration...')
            worker.grompp_me(runtype)
        elif runtype == 'prod':
            # The actual alchemistry:
            worker.prep_runs('prod', stride, temperature, init_nst, nst, lincs, dt, extra_args)
            worker.log('running grompp for all systems...')
            worker.grompp_me('prod')
            worker.log('running mdrun for all systems...')

    def analyze(self):
        """
        The analysis routine: performed after the run,
        reads all data files, processes them and
        prints out the results and/or convergence data
        :return:
        """
        print('analyzing results...')
        all_workers = self.analysis_only + self.workers
        for worker in all_workers:
            worker.analyze_me()
        works_0 = np.array([crk.work for crk in all_workers if crk.initlambda == 0]).reshape(-1, 1)
        # TODO sign
        works_1 = np.array([-crk.work for crk in all_workers if crk.initlambda == 1]).reshape(-1, 1)
        np.savetxt('work_0.dat', works_0, fmt='%10.5f')
        np.savetxt('work_1.dat', works_1, fmt='%10.5f')
        min_val = np.min(np.vstack([works_0, works_1]))
        max_val = np.max(np.vstack([works_0, works_1]))
        diff = max_val - min_val
        min_range = min_val - 0.1 * diff
        max_range = max_val + 0.1 * diff
        grid = np.linspace([min_range], [max_range], 200).reshape(-1, 1)
        dens0 = self.get_opt_kde(works_0, grid)
        dens1 = self.get_opt_kde(works_1, grid)
        np.savetxt('prob_0.dat', np.hstack([grid, dens0]), fmt='%10.5f')
        np.savetxt('prob_1.dat', np.hstack([grid, dens1]), fmt='%10.5f')
        result_bar = self.solve_bar(works_0, works_1)
        result_cgi = self.solve_cgi(works_0, works_1)
        result_kde = self.solve_kde(dens0, dens1, grid)
        np.savetxt('normalized_overlap_{}.dat'.format(self.name), self.bhattacharyya(dens0, dens1, grid),
                   fmt='%10.5f')
        if self.convergence:
            self.analyze_convergence(works_0, works_1, grid)
        with open(f'result_{self.name}.dat', 'w') as out:
            out.write(f'Result from CGI: {result_cgi:8.3f}\n')
            out.write(f'Result from KDE: {result_kde:8.3f}\n')
            out.write(f'Result from BAR: {result_bar:8.3f} (should be the most accurate)\n')

    @staticmethod
    def bhattacharyya(d0, d1, grid):
        """
        Calculates the Bhattacharyya (similarity) coefficient between distributions d0 and d1
        :param d0: np.array, rho_0, density of A->B work values
        :param d1: np.array, rho_1, density of B->A -work values
        :param grid: np.array, the corresponding X values
        :return: np.array of shape [1,1]
        """
        return simps(np.sqrt(d0.reshape(-1) * d1.reshape(-1)), grid.reshape(-1)).reshape(-1, 1)

    def analyze_convergence(self, w0, w1, grid):
        """
        Performs a series of bootstraps for a number of
        subsample sizes, then plots the resulting means/
        standard deviations to a file
        :param w0: np.array, work values for A->B
        :param w1: np.array, -work values for A->B
        :param grid: np.array, the X values, as in the original plot
        :return: None
        """
        samples = np.linspace(4, self.nruns, 10).astype(int)
        conv = {ns: [] for ns in samples}
        if self.debug:
            try:
                os.mkdir('bootstrap')
            except FileExistsError:
                pass
        for nsamp in samples:
            print('bootstrap for {} samples...'.format(nsamp))
            with ProcessPoolExecutor(max_workers=10) as executor:
                results = np.array(list(executor.map(self.boot, [(w0, w1, nsamp, grid, i)
                                                                 for i in range(10)]))).reshape(-1)
            conv[nsamp] = [np.mean(results), np.std(results)]

    def boot(self, params):
        """
        Performs a single bootstrap resampling,
        parallelized through ProcessPoolExecutor
        :param params: a 5-element tuple of input parameters:
        A->B work, B->A -work, number of subsamples to draw, original X values, random seed
        :return: float, the KDE solution for the intersection
        """
        w0, w1, nsamp, grid, n = params
        np.random.seed(n)
        w0_sampled = np.random.choice(w0.reshape(-1), nsamp, replace=True).reshape(-1, 1)
        w1_sampled = np.random.choice(w1.reshape(-1), nsamp, replace=True).reshape(-1, 1)
        dens0_sampled = self.get_opt_kde(w0_sampled, grid, self.bandwidth)
        dens1_sampled = self.get_opt_kde(w1_sampled, grid, self.bandwidth)
        result = self.solve_bar(w0_sampled, w1_sampled)
        if self.debug:
            np.savetxt(f'bootstrap/conv_{nsamp}_{n}_{result}.dat', np.hstack([grid, dens0_sampled, dens1_sampled]),
                       fmt='%10.5f')
        return result

    @staticmethod
    def solve_kde(d0, d1, grid):
        """
        Finds the intersection point based
        on KDE densities
        :param d0: np.array,
        :param d1: np.array,
        :param grid: np.array,
        :return: float, abscissa of intersection
        """
        diff = d0 - d1
        avg = 0.5 * (np.argmax(d0) + np.argmax(d1))
        solution = [x for x in np.where(diff[1:] * diff[:-1] < 0)[0]]
        arg = solution[np.argmin([np.abs(x - avg) for x in solution])]
        result = round(0.5 * (grid.reshape(-1)[arg] + grid.reshape(-1)[arg]), 3)
        return result

    @staticmethod
    def solve_cgi(w0, w1):
        """
        Finds the intersection point based on the CGI
        estimator (intersection of two Gaussians)
        :param w0: np.array, work values for A->B
        :param w1: np.array, -work values for B->A
        :return: float, abscissa of intersection
        """
        w0, w1 = w0.reshape(-1), w1.reshape(-1)
        m0, s0 = np.mean(w0), np.std(w0)
        m1, s1 = np.mean(w1), np.std(w1)
        x1 = m0 / s0 ** 2 - m1 / s1 ** 2
        den = 1 / s0 ** 2 - 1 / s1 ** 2
        sq = (m0 - m1) ** 2 / (s0 ** 2 * s1 ** 2) + 2 * den * np.log(s1 / s0)
        mid = (m0 + m1) / 2
        r1, r2 = (x1 - np.sqrt(sq)) / den, (x1 + np.sqrt(sq)) / den
        return round(r1, 3) if np.abs(r1 - mid) < np.abs(r2 - mid) else round(r2, 3)

    def solve_bar(self, w0, w1):
        return bar(w0, w1, self.temperature)

    def get_opt_kde(self, y, x, bandwidth=None):
        """
        Performs KDE (kernel density estimation) with a Gaussian
        kernel, and optimizes the bandwidth using the LeaveOneOut
        cross-validation if requested
        :param y: np.array, work values
        :param x: np.array, X values for the density
        :param bandwidth: if None, optimize; if a float, use this value
        :return: np.array, the estimated density
        """
        if bandwidth is None:
            bandwidths = 10 ** np.linspace([-1], [1], 20).reshape(-1)
            grid = GridSearchCV(estimator=KernelDensity(kernel='gaussian'), param_grid={'bandwidth': bandwidths},
                                cv=LeaveOneOut())
            grid.fit(y.reshape(-1, 1))
            kde = KernelDensity(kernel='gaussian', **grid.best_params_).fit(y.reshape(-1, 1))
            self.bandwidth = grid.best_params_['bandwidth']
        else:
            kde = KernelDensity(kernel='gaussian', bandwidth=bandwidth).fit(y.reshape(-1, 1))
        return np.exp(kde.score_samples(x)).reshape(-1, 1)

    def plot_results(self, w0='work_0.dat', w1='work_1.dat', p0='prob_0.dat', p1='prob_1.dat'):
        """
        Plots the two densities, their intersection
        and all individual work values at the bottom;
        also writes the precise solution to a file
        (this is admittedly a bit confusing)
        :param w0: np.array, work values for A->B
        :param w1: np.array, -work values for B->A
        :param p0: np.array, X:Y two columns for the first density
        :param p1: np.array, X:Y two columns for the second density
        :return: None
        """
        try:
            import matplotlib.pyplot as plt
        except ImportError:
            return
        fig, axes = plt.subplots(2, 1)
        axes[0].plot(*p0.T, c='C0')
        axes[0].plot(*p1.T, c='C1')
        result_bar = self.solve_bar(w0, w1)
        axes[0].axvline(result_bar, ls='--', c='k', lw=3)
        axes[0].fill_between(*p0.T, color='C0', alpha=0.25)
        axes[0].fill_between(*p1.T, color='C1', alpha=0.25)
        axes[0].plot(*p1.T, c='C1')
        axes[1].scatter(w0, np.zeros(len(w0)), c='C0')
        axes[1].scatter(w1, np.zeros(len(w1)), c='C1')
        axes[0].set_xlim(np.min(p0[:, 0]), np.max(p0[:, 0]))
        axes[1].set_xlim(np.min(p0[:, 0]), np.max(p0[:, 0]))
        datarange = np.max(p0[:, 0]) - np.min(p0[:, 0])
        axes[1].set_ylim(-datarange / 2, datarange / 2)
        axes[1].get_yaxis().set_visible(False)
        axes[0].get_xaxis().set_visible(False)
        axes[1].set_aspect(0.05)
        axes[1].spines['top'].set_visible(False)
        plt.subplots_adjust(hspace=-0.47)
        plt.savefig('{}.svg'.format(self.name))
        plt.close()

    def drop_frames(self):
        """
        Selects frames from the provided equilibrium trajectories, either in equal intervals or randomly
        :return: None
        """
        for i in [0, 1]:
            for worker in [w for w in self.workers if w.initlambda == i]:
                worker.mkdir()
                plu = self.plumed if i == 0 else self.plumed2
                if plu:
                    if plu.startswith('/'):
                        plumed = plu
                    else:
                        plumed = os.getcwd() + os.sep + plu
                    worker.cp(plumed, 'plumed.dat')
                    worker.plumed = True
            if not all([os.path.isfile('run{}_l{}/frame{}_l{}.gro'.format(num, i, num, i))
                        for num in range(self.offset, self.offset + self.nruns)]):
                struct = self.struct2 if i == 1 else self.struct
                nfr = gml.frames_count(self.xtc[i], self.gmx)
                if not self.random:
                    frames = np.linspace(0, nfr - 0.00001, self.nmax).astype(int)[
                             self.offset:self.offset + self.nruns]
                else:
                    frames = np.random.choice(nfr, size=self.nmax, replace=False)
                with open('frames.ndx', 'w') as ndxf:
                    ndxf.write(' [ frames ] \n')
                    for m in range(ceil(nfr/15.0)):
                        ndxf.write(' '.join([str(x + 1) for x in frames[15*m:15*(m+1)]]) + '\n')
                gml.gmx_command(self.gmx, 'trjconv', fr='frames.ndx', f=self.xtc[i], s=struct, sep=True,
                                o='frame.gro', pass_values=[0, 0])
                for fr in range(len(frames)):
                    os.rename(f'frame{fr}.gro', f'run{fr}_l{i}/frame{fr}_l{i}.gro')
            else:
                print("frames for lambda {} already present, skipping generating new ones".format(i))

    @staticmethod
    def apply_pmx(data):
        mutate, wt_only, num, lam, pmxff = data
        resp = {'A': 'ALA', 'C': 'CYS', 'D': 'ASP', 'E': 'GLU', 'F': 'PHE', 'G': 'GLY', 'H': 'HIS',
                'I': 'ILE', 'K': 'LYS', 'L': 'LEU', 'M': 'MET', 'N': 'ASN', 'P': 'PRO', 'Q': 'GLN',
                'R': 'ARG', 'S': 'SER', 'T': 'THR', 'V': 'VAL', 'W': 'TRP', 'Y': 'TYR'}
        resnums, targets = [], []

        try:
            import pmx
        except ImportError:
            raise RuntimeError("pmx is required for introducing alchemical mutation on-the-fly")

        for mut in mutate.split('/'):
            if mut[1] in '1234567890':
                if lam == 0:
                    target = resp[mut[-1]]
                elif lam == 1 and not wt_only:
                    target = resp[mut[0]]
                else:
                    target = None
                resnum = int(mut[1:-1])
            else:
                if lam == 0:
                    target = mut[-2:]
                elif lam == 1 and not wt_only:
                    target = mut[:2]
                else:
                    target = None
                resnum = int(mut[2:-2])
            resnums.append(resnum)
            targets.append(target)
        model = pmx.Model('run{}_l{}/frame{}_l{}.gro'.format(num, lam, num, lam))
        for resnum, target in zip(resnums, targets):  # TODO OP1, OP2 to O1P, O2P
            if target is not None:
                model = pmx.mutate(model, resnum, target, pmxff)
        model.writeGRO('run{}_l{}/frame{}_l{}.gro'.format(num, lam, num, lam))

    @staticmethod
    def grompp(data):
        crook, runtype = data
        crook.grompp_me(runtype)

    def mdrun_multi(self, runtype):
        """
        Runs gmx mdrun with the -multidir option for maximum efficiency
        :param runtype: str, 'mini', 'eq' or 'dyn'
        :return: None
        """
        if runtype == 'mini' and not self.mini:
            return
        ncpus = self.ncpus()
        if len(self.workers) % ncpus != 0:
            raise RuntimeError("Number of simulations should be divisible by the number of available CPUs, but {} and {} were specified")
        tpr = runtype if runtype in ['mini', 'eq'] else 'dyn'
        if self.plumed and runtype == 'eq':
            plu = ' -plumed plumed.dat '
        else:
            plu = ''
        multi = ['run{}_l{}'.format(w.id, w.initlambda) for w in self.workers]
        gmx_mpi = gml.find_gmx_dir(mpi=True)[1]
        if not all(['{}.gro'.format(tpr) in os.listdir('run{}_l{}'.format(w.id, w.initlambda)) for w in self.workers]):
            for multi_batch in [multi[ncpus*i : ncpus*(i+1)] for i in range(len(multi) // ncpus)]:
                print(gml.gmx_command(f'mpiexec {gmx_mpi}', 'mdrun', deffnm=tpr, v=plu, cpi=True,
                                      multidir=' '.join(multi_batch), answer=True))

    @staticmethod
    def ncpus():
        """
        Returns the total number of available CPUs across all nodes.
        """
        if 'SLURM_NTASKS' in os.environ:
            total_cpus = int(os.environ['SLURM_NTASKS'])
            return total_cpus
        elif 'PBS_NUM_NODES' in os.environ and 'PBS_NUM_PPN' in os.environ:
            nodes = int(os.environ['PBS_NUM_NODES'])
            cpus_per_node = int(os.environ['PBS_NUM_PPN'])
            total_cpus = nodes * cpus_per_node
            return total_cpus
        elif 'OMPI_UNIVERSE_SIZE' in os.environ:
            total_cpus = int(os.environ['OMPI_UNIVERSE_SIZE'])
            return total_cpus
        else:
            return multiprocessing.cpu_count()

    @staticmethod
    def read_results(crook):
        crook.analyze_me()


class Crooks:
    def __init__(self, master, num, initlambda):
        """
        A single worker instance
        :param master: a CrooksPool instance
        :param num: ID of the worker
        :param initlambda: initial alchemical lambda (0 or 1)
        """
        self.id = num
        self.master = master
        self.initlambda = initlambda
        self.plumed = False
        self.work = 0

    def mkdir(self):
        """
        Attempts to create a directory if not there yet
        :return: None
        """
        try:
            os.mkdir('run{}_l{}'.format(self.id, self.initlambda))
        except FileExistsError:
            pass

    def cp(self, filename, target_filename):
        copy(filename, 'run{}_l{}/{}'.format(self.id, self.initlambda, target_filename))

    def log(self, command):
        with open('run{}_l{}/crooks.log'.format(self.id, self.initlambda), 'a') as out:
            out.write(command.strip() + '\n')

    def prep_runs(self, runtype, stride, temperature, init_nst, nst, lincs, dt, extra_args):
        """
        Prepares the run by writing customized .mdp files
        :return: None
        """
        print('preparing runs in mode "{}"...'.format(runtype))
        mdp_defaults = {"nstcomm": 100, "nstenergy": 5000, "nstlog": 5000, "nstcalcenergy": 100,
                        "nstxout-compressed": int(1 / dt) * stride, "compressed-x-grps": "System",
                        "compressed-x-precision": 2500, "free-energy": "yes", "sc-alpha": 0.3, "sc-coul": "yes",
                        "sc-sigma": 0.25, "dt": dt, "constraints": lincs, "coulombtype": "PME",
                        "tcoupl": "v-rescale", "ref-t": temperature, "tau-t": 0, "ref-p": 1.0,
                        "rlist": 1.2, "rcoulomb": 1.2, "vdw-type": "Cut-off", "rvdw_switch": 0.8, "rvdw": 1.2,
                        "compressibility": "4.5e-5", "tau-p": 1.0, "tc-grps": "System"}
        mini_defaults = {"integrator": "steep", "free-energy": "yes", "sc-alpha": 0.3, "sc-coul": "yes",
                         "sc-sigma": 0.25, "nsteps": 1000, "emtol": 200, "emstep": 0.001, "nstlist": 10,
                         "pbc": "xyz", "coulombtype": "PME", "vdw-type": "Cut-off"}
        if runtype == 'eq':
            extra_defaults = {"gen-vel": "yes", "gen-temp": temperature, "pcoupl": "Berendsen",
                              "separate-dhdl-file": "no", "nsteps": init_nst, "nstxout": 10000, "nstvout": 10000}
        else:
            extra_defaults = {"gen-vel": "no", "pcoupl": "Parrinello-Rahman", "separate-dhdl-file": "yes",
                              "nsteps": nst, "nstxout": 100000, "nstvout": 100000}
        mdp_defaults.update(extra_defaults)
        mdp_defaults.update(extra_args)
        default = mini_defaults if runtype == 'mini' else mdp_defaults
        pref = 1 if self.initlambda == 0 else -1
        dl = 0 if runtype in ['eq', 'mini'] else pref / nst
        default.update({"init-lambda": self.initlambda, "delta-lambda": dl})
        mdp = '\n'.join(["{} = {}".format(param, value) for param, value in default.items()])
        name = "mini" if runtype == 'mini' else 'md'
        with open('run{}_l{}/{}.mdp'.format(self.id, self.initlambda, name), 'w') as outfile:
            outfile.write(mdp.format(nst=nst, init=self.initlambda, growth=pref / nst))

    def grompp_me(self, runtype):
        """
        Checks for the .tpr file, runs grompp if not present
        :return: None
        """
        tpr = 'eq' if runtype == 'eq' else 'mini' if runtype == 'mini' else 'dyn'
        mdp = 'mini' if runtype == 'mini' else 'md'
        trr = '' if runtype in ['mini', 'eq'] else ' -t eq.trr -time {} '.format(self.master.init_length)
        os.chdir('run{}_l{}'.format(self.id, self.initlambda))
        master_top = self.master.top if int(self.initlambda) == 0 else self.master.top2
        frame = 'mini.gro'.format(self.id, self.initlambda) if runtype == 'eq' \
            else 'frame{}_l{}.gro'.format(self.id, self.initlambda)
        if '{}.tpr'.format(tpr) not in os.listdir('.'):
            top = master_top if master_top.startswith('/') else '../' + master_top
            gml.gmx_command(self.master.gmx, f=f'{mdp}.mdp', p=top, c=frame, answer=True,
                            o=f'{tpr}.tpr', v=trr, maxwarn=self.master.maxwarn)
            # call('{gmx} grompp -f {mdp}.mdp -p {top} -c {frame} -o {tpr}.tpr {trr}'
            #      '-maxwarn {mw} >> gmp.log 2>&1'.format(gmx=self.master.gmx, l=self.initlambda, trr=trr,
            #                                             top=top, mw=self.master.maxwarn, tpr=tpr, mdp=mdp,
            #                                             frame=frame), shell=True)
            os.remove('mdout.mdp')

        os.chdir('..')

    def analyze_me(self):
        """
        Reads the .xvg file, removes potential duplicates (if the
        job was restarted), and integrates dH/dl to calculate work
        :return: None
        """
        dhdl = np.loadtxt('run{n}_l{l}/dyn.xvg'.format(n=self.id, l=self.initlambda), comments=['#', '@'])
        if not int(dhdl[-1, 0]) == int(self.master.sim_length):
            raise RuntimeError("In file run{}_l{}/dyn.xvg last line reads {}, less than the requested total "
                               "fime of {} ps".format(self.id, self.initlambda, dhdl[-1], self.master.sim_length))
        dhdl_dict = {i: j for i, j in dhdl}
        dhdl = np.array([dhdl_dict[i] for i in sorted(list(dhdl_dict.keys()))])  # avoid duplicates in case of restarts
        endlambda = 1 if self.initlambda == 0 else 0
        self.work = simps(dhdl, np.linspace([self.initlambda], [endlambda], len(dhdl)).reshape(-1))


def bar(forward_works, reverse_works, temperature):
    kb = 0.00831447215  # in kJ units
    forward_works = np.array(forward_works)
    reverse_works = np.array(reverse_works)
    temperature = float(temperature)
    nf = float(len(forward_works))
    nr = float(len(reverse_works))
    beta = 1. / (kb * temperature)
    fact = kb * temperature * np.log(nf / nr)

    def func(x, wf, wr):
        sf = 0
        for v in wf:
            sf += 1 / (1 + np.exp(beta * (fact + v - x)))
        sr = 0
        for v in wr:
            sr += 1 / (1 + np.exp(-beta * (fact + v - x)))
        r = sf - sr
        return r ** 2

    av_a = np.average(forward_works)
    av_b = np.average(reverse_works)
    x0 = (av_a + av_b) / 2.
    dg = fmin(func, x0=x0, args=(forward_works, reverse_works), disp=0)
    return float(dg)
