from __future__ import division, print_function, absolute_import
import numpy as np
import pandas as pd
import os
from copy import copy

databasepath = os.path.join(os.path.dirname(__file__), 'database')
databasepath += '/saftgamma_database.xlsx'

file = pd.ExcelFile(databasepath, engine='openpyxl')
df_groups = pd.read_excel(file, 'groups', index_col='groups')
df_mie_kl = pd.read_excel(file, 'unlikemie_kl')
df_asso_kl = pd.read_excel(file, 'unlikeasso_kl')
df_secondorder = pd.read_excel(file, 'secondmie')
df_secondasso = pd.read_excel(file, 'secondasso')

class GCdatabase(object):
    '''
    SAFT-Gamma-Mie group contribution database Object

    This object have implemeted methods to modify the database published
    groups and interactions parameters for SAFT-Gamma-Mie EoS.

    The included parameters are obtained from J. Chem. Eng. Data 2020,
    65, 12, 5862–5890. https://doi.org/10.1021/acs.jced.0c00746

    Parameters
    ----------
    df_groups: DataFrame
        DataFrame that includes group information
    df_mie_kl: DataFrame
        DataFrame that includes unlike Mie parameters
    df_asso_kl: DataFrame
        DataFrame that includes unlike association Parameters
    df_secondorder: DataFrame
        Dataframe that includes second order modification to Mie interactions
    df_scondasso: DataFrame
        Dataframe that includes second order modification to association
        interactions

    Attributes
    ----------
    group_list: List of groups on the database
    df_groups: DataFrame that includes group information
    df_mie_kl: DataFrame that includes unlike Mie parameters
    df_asso_kl: DataFrame that includes unlike association Parameters
    df_secondorder: Dataframe that includes second order modification
    to Mie interactions
    df_scondasso: Dataframe that includes second order modification to
    association interactions

    Methods
    -------
    add_group: method to add a new group
    new_interaction_mie: method to set unlike Mie interactions
    new_interaction_asso: method to add association parameters
    get_interactions: method to get the interactions between two groups
    restore_database: method to restore the database to its initial value
    '''
    def __init__(self, df_groups=df_groups, df_mie_kl=df_mie_kl,
                 df_asso_kl=df_asso_kl, df_secondorder=df_secondorder,
                 df_secondasso=df_secondasso):

        self.df_groups = df_groups
        self.df_mie_kl = df_mie_kl
        self.df_asso_kl = df_asso_kl
        self.df_secondorder = df_secondorder
        self.df_secondasso = df_secondasso

        self.df_groups_backup = copy(df_groups)
        self.df_mie_kl_backup = copy(df_mie_kl)
        self.df_asso_kl_backup = copy(df_asso_kl)
        self.df_secondorder_backup = copy(df_secondorder)
        self.df_secondasso_backup = copy(df_secondasso)

        self.group_list = list(self.df_groups.index)

    def add_group(self, name, vk=1., Sk=1., sigma=0., eps=0., lr=12., la=6.,
                  nH=0, ne1=0, ne2=0, charge=0., sigma_born=0., mw=0.,
                  overwrite=False):
        """
        add_group method

        Method that adds a new group to the database

        Parameters
        ----------
        name: string
            name of the group
        vk: float
            number os sites used by the group
        Sk: float
            shape factor of the group
        sigma: float
            lenght scale of the group, used in Mie potential [Amstrong]
        eps: float
            energy scale of the group, used in Mie potential [K]
        lr: float
            repulsive exponent of the group, used in Mie potential
        la: float
            attractive exponent of the group, used in Mie potential
        nH: int
            number of Hidrogen associative sites
        ne1: int
            number of e1 associative sites
        ne2: int
            number of e2 associative sites
        charge: int
            charge of the group (electron charge) [Adim.]
        sigma_born: float
            diameter used in Born contribution [Amstrong]
        mw: float
            molar weight of the group [g/mol]
        overwrite: bool, optional
            whether to overwrite or not current parameters of the database
        """
        Nst = np.count_nonzero([nH, ne1, ne2])
        group_included = name in self.group_list

        if group_included and not overwrite:
            raise Exception('group {} already included in database, set'.format(name) +
                            ' overwrite=True to overwrite the current parameters')
        elif group_included and overwrite:
            new_parameters = np.array([vk, Sk, sigma, eps, lr, la, Nst, nH,
                                       ne1, ne2, charge, sigma_born, mw])
            self.df_groups.loc[name] = new_parameters
        else:
            new_group = {'groups': name, 'vk*': vk, 'Sk': Sk,
                         'sigma_kk': sigma, 'eps_kk': eps,
                         'lr_kk': lr, 'la_kk': la, 'Nst_kk': Nst, 'nH_kk': nH,
                         'ne1_kk': ne1, 'ne2_kk': ne2, 'charge_kk': charge,
                         'sigma_born_kk': sigma_born, 'mw_kk': mw}
            groups_aux = self.df_groups.reset_index()
            groups_new = groups_aux.append(new_group, ignore_index=True)
            groups_new.set_index('groups', inplace=True)
            self.df_groups = groups_new
            self.group_list = list(self.df_groups.index)

    def new_interaction_mie(self, group_k, group_l, eps_kl=0., lr_kl='CR',
                            overwrite=False):
        """
        new_interaction_mie method

        Method that adds the unlike mie interactions between group k and
        group l

        Parameters
        ----------
        group_k: string
            name of the group k
        group_l: string
            name of the group l
        eps_kl: float
            unlike energy scale between the groups, used in Mie potential [K]
        lr_kl: float
            unlike repulsive exponent between the groups, used in Mie potential
        overwrite: bool, optional
            whether to overwrite or not current parameters of the database
        """
        bool_kk = self.df_mie_kl.group_k == group_k
        bool_ll = self.df_mie_kl.group_l == group_l
        bool_kl = self.df_mie_kl.group_k == group_l
        bool_lk = self.df_mie_kl.group_l == group_k

        index1 = self.df_mie_kl.index[bool_kk & bool_ll]
        len1 = index1.shape[0]

        index2 = self.df_mie_kl.index[bool_kl & bool_lk]
        len2 = index2.shape[0]

        if len1 == 1:
            index = index1
            n = len1
        elif len2 == 1:
            index = index2
            n = len2
        else:
            n = 0
        already_in = n > 0
        if already_in and not overwrite:
            raise Exception('Interaction parameters between group {} '.format(group_k) +
                            ' and group {}'.format(group_l) + ' are already in the database' +
                            ' set overwrite=True to overwrite the current parameters')

        elif already_in and overwrite:
            self.df_mie_kl.iloc[index, [2, 3]] = [eps_kl, lr_kl]
        else:
            new_interaction = {'group_k': group_k, 'group_l': group_l,
                               'eps_kl': eps_kl, 'lr_kl': lr_kl}
            df_mie_new = self.df_mie_kl.append(new_interaction, ignore_index=True)
            self.df_mie_kl = df_mie_new

    def new_interaction_asso(self, group_k, group_l, site_k, site_l,
                             epsAB_kl=0., kAB_kl=0, overwrite=False):
        """
        new_interaction_mie method

        Method that adds the unlike mie interactions between group k and
        group l

        Parameters
        ----------
        group_k: string
            name of the group k
        group_l: string
            name of the group l
        site_k: string
            name of the site k, available optiones are 'H', 'e1' and 'e2'
        site_l: string
            name of the site l, available optiones are 'H', 'e1' and 'e2'
        epsAB_kl: float
            unlike association energy between the the groups [K]
        kAB_kl: float
            unlike association volume between the groups [Amstrong^3]
        overwrite: bool, optional
            whether to overwrite or not current parameters of the database
        """
        bool_kk = self.df_asso_kl.group_k == group_k
        bool_ll = self.df_asso_kl.group_l == group_l
        bool_kl = self.df_asso_kl.group_k == group_l
        bool_lk = self.df_asso_kl.group_l == group_k
        boolsite_kk = self.df_asso_kl.iloc[:, 1] == site_k
        boolsite_ll = self.df_asso_kl.iloc[:, 3] == site_l

        index1 = self.df_asso_kl.index[bool_kk & bool_ll & boolsite_kk & boolsite_ll]
        len1 = index1.shape[0]

        index2 = self.df_asso_kl.index[bool_kl & bool_lk & boolsite_kk & boolsite_ll]
        len2 = index2.shape[0]

        if len1 == 1:
            index = index1
            n = len1
        elif len2 == 1:
            index = index2
            n = len2
        else:
            n = 0

        already_in = n > 0
        if already_in and not overwrite:
            raise Exception('Interaction parameters between site{} of'.format(site_k) +
                            ' of group {} '.format(group_k) + 'and site {}'.format(site_l) +
                            ' of group {}'.format(group_l) + 'are already in the database' +
                            ' set overwrite=True to overwrite the current parameters')

        elif already_in and overwrite:
            self.df_asso_kl.iloc[index, [4, 5]] = [epsAB_kl, kAB_kl]
        else:
            new_interaction = {'group_k': group_k,
                               'site\xa0a\xa0of group\xa0k': site_k,
                               'group_l': group_l,
                               'site\xa0b\xa0of group\xa0l': site_l,
                               'epsAB_kl': epsAB_kl, 'KAB_kl': kAB_kl}
            df_asso_new = self.df_asso_kl.append(new_interaction, ignore_index=True)
            self.df_asso_kl = df_asso_new

    def get_interactions(self, group_k, group_l):
        """
        get_interactions method

        Method that outputs the available interactions between group k and
        group l

        Parameters
        ----------
        group_k: string
            name of the group k
        group_l: string
            name of the group l

        Returns
        -------
        df_group : DataFrame
             parameters of each group
        df_mie : DataFrame
             unlike Mie Parameters
        df_asso : DataFrame
             unlike association parameters
        """
        df_group = self.df_groups.loc[[group_l, group_k]]

        bool_kk = self.df_mie_kl.group_k == group_k
        bool_ll = self.df_mie_kl.group_l == group_l
        bool_kl = self.df_mie_kl.group_k == group_l
        bool_lk = self.df_mie_kl.group_l == group_k

        df1 = self.df_mie_kl[bool_kk & bool_ll]
        len1 = df1.shape[0]

        df2 = self.df_mie_kl[bool_kl & bool_lk]
        len2 = df2.shape[0]

        if len1 == 1:
            df_mie = df1
        elif len2 == 1:
            df_mie = df2
        else:
            df_mie = 'There are no custom interaction parameters set for group {}'.format(group_k)
            df_mie += 'and group {}'.format(group_l)

        bool_kk = self.df_asso_kl.group_k == group_k
        bool_ll = self.df_asso_kl.group_l == group_l
        bool_kl = self.df_asso_kl.group_k == group_l
        bool_lk = self.df_asso_kl.group_l == group_k

        df1 = self.df_asso_kl[bool_kk & bool_ll]
        len1 = df1.shape[0]

        df2 = self.df_asso_kl[bool_kl & bool_lk]
        len2 = df2.shape[0]

        if len1 >= 1:
            df_asso = df1
        elif len2 >= 1:
            df_asso = df2
        else:
            df_asso = 'There are no association parameters set for group {}'.format(group_k)
            df_asso += ' and group {}'.format(group_l)

        return df_group, df_mie, df_asso

    def restore_database(self):
        """
        restore_database method

        Method that restores the database of its initial state.
        This method will erase any custom groups or interactions
        set up by the user.
        """
        self.df_groups = copy(self.df_groups_backup)
        self.df_mie_kl = copy(self.df_mie_kl_backup)
        self.df_asso_kl = copy(self.df_asso_kl_backup)
        self.df_secondorder = copy(self.df_secondorder_backup)
        self.df_secondasso = copy(self.df_secondasso_backup)

        self.group_list = list(self.df_groups.index)

database = GCdatabase()


#!/usr/bin/env python
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import time
import ctypes
import tempfile
import numpy
import h5py
import pyscf.lib as lib
from pyscf.lib import logger
import pyscf.ao2mo
from pyscf.cc import _ccsd

# t2,l2 as ijab

def kernel(mycc, eris, t1=None, t2=None, l1=None, l2=None,
           max_cycle=50, tol=1e-8, verbose=logger.INFO):
    cput0 = (time.clock(), time.time())
    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(mycc.stdout, verbose)

    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2
    if l1 is None: l1 = t1
    if l2 is None: l2 = t2

    nocc, nvir = t1.shape
    saved = make_intermediates(mycc, t1, t2, eris)

    if mycc.diis:
        adiis = lib.diis.DIIS(mycc, mycc.diis_file)
        adiis.space = mycc.diis_space
    else:
        adiis = lambda t1,t2,*args: (t1, t2)
    cput0 = log.timer('CCSD lambda initialization', *cput0)

    conv = False
    for istep in range(max_cycle):
        l1new, l2new = update_amps(mycc, t1, t2, l1, l2, eris, saved)
        normt = numpy.linalg.norm(l1new-l1) + numpy.linalg.norm(l2new-l2)
        l1, l2 = l1new, l2new
        l1new = l2new = None
        if mycc.diis:
            l1, l2 = mycc.diis(l1, l2, istep, normt, 0, adiis)
        log.info('cycle = %d  norm(lambda1,lambda2) = %.6g', istep+1, normt)
        cput0 = log.timer('CCSD iter', *cput0)
        if normt < tol:
            conv = True
            break
    return conv, l1, l2

# l2, t2 as ijab
def make_intermediates(mycc, t1, t2, eris):
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    nov = nocc * nvir
    foo = eris.fock[:nocc,:nocc]
    fov = eris.fock[:nocc,nocc:]
    fvv = eris.fock[nocc:,nocc:]

    class _Saved(object):
        pass
    saved = _Saved()

# As we don't have l2 in memory, hold tau temporarily in memory
    w1 = fvv - numpy.einsum('ja,jb->ba', fov, t1)
    w2 = foo + numpy.einsum('ib,jb->ij', fov, t1)
    w3 = _cp(numpy.einsum('kc,jkbc->bj', fov, t2) * 2 + fov.T)
    w3 -= numpy.einsum('kc,kjbc->bj', fov, t2)
    w3 += reduce(numpy.dot, (t1.T, fov, t1.T))
    w4 = fov.copy()

    eris_ovvv = _cp(eris.ovvv)
    eris_ovvv = lib.unpack_tril(eris_ovvv.reshape(nov,-1))
    eris_ovvv = eris_ovvv.reshape(nocc,nvir,nvir,nvir)

    wovvv = numpy.empty((nocc,nvir,nvir,nvir))
    t2tmp = numpy.empty((nocc,nvir,nocc,nvir))
    for i in range(nocc):
        wovvv[i] = eris_ovvv[i].transpose(1,0,2) * 2
        t2tmp[i] = t2[i].transpose(2,0,1)
    #:wovvv += numpy.einsum('jabd,kjdc->kabc', eris_ovvv, t2) * -1.5
    tmp = lib.dot(t2tmp.reshape(nov,-1), wovvv.reshape(-1,nvir**2),
                  -1.5/2).reshape(-1,nvir,nvir,nvir)
    g2ovvv = tmp
    for i in range(nocc):
        wovvv[i] -= eris_ovvv[i].transpose(1,2,0)
        wovvv[i] += tmp[i].transpose(1,2,0)
        g2ovvv[i] = eris_ovvv[i]*2
        g2ovvv[i] -= eris_ovvv[i].transpose(1,2,0)
    tmp = t2tmp = None

    w1 += numpy.einsum('jcba,jc->ba', eris_ovvv, t1*2)
    w1 -= numpy.einsum('jabc,jc->ba', eris_ovvv, t1)
    #:w3 += numpy.einsum('kdcb,kjdc->bj', eris_ovvv, theta)
    theta = numpy.empty(t2.shape)
    for i in range(nocc):
        theta[i] = t2[i] * 2
        theta[i] -= t2[i].transpose(0,2,1)
        lib.dot(eris_ovvv[i].reshape(-1,nvir).T,
                _cp(theta[i].reshape(nocc,-1)).T, 1, w3, 1)

    theta = _cp(theta.transpose(0,2,1,3))
    #:vkbca = numpy.einsum('jdca,kbjd->kbca', g2ovvv, theta)
    vkbca = lib.dot(_cp(theta.reshape(nov,-1)),
                    g2ovvv.reshape(-1,nvir*nvir)).reshape(-1,nvir,nvir,nvir)
    for i in range(nocc):
        wovvv[i] += vkbca[i].transpose(2,0,1)
        wovvv[i] -= vkbca[i].transpose(2,1,0) * .5
    vkabc = None

    #:wOVov = numpy.einsum('jbcd,kd->jbkc', eris_ovvv, t1)
    #:wOvOv = numpy.einsum('jdcb,kd->jbkc', eris_ovvv, -t1)
    wOVov = lib.dot(eris_ovvv.reshape(-1,nvir),
                    t1.T).reshape(-1,nvir,nvir,nocc).transpose(0,1,3,2).copy()
    for i in range(nocc):
        g2ovvv[i] = eris_ovvv[i].transpose(1,2,0) * 2
    wOvOv = lib.dot(g2ovvv.reshape(-1,nvir),
                    -t1.T, .5).reshape(-1,nvir,nvir,nocc).transpose(0,1,3,2).copy()
    for i in range(nocc):
        g2ovvv[i] -= eris_ovvv[i].transpose(1,0,2)
    eris_ovov = _cp(_cp(eris.ovov).transpose(0,2,1,3))
    tau = _ccsd.make_tau(t2, t1, t1)
    #:wooov[:,j0:j1] = numpy.einsum('icbd,jkbd->ijkc', g2ovvv, tau)
    #:woooo[:,:,j0:j1] = numpy.einsum('icjd,klcd->ijkl', eris_ovov, tau)
    tmp = lib.dot(g2ovvv.reshape(-1,nvir**2), tau.reshape(-1,nvir**2).T)
    wooov = _cp(tmp.reshape(-1,nvir,nocc,nocc).transpose(0,2,3,1))
    woooo = lib.dot(eris_ovov.reshape(-1,nvir**2),
                    tau.reshape(-1,nvir**2).T).reshape(-1,nocc,nocc,nocc)
    eris_ovov = eris_ovvv = g2ovvv = tau = tmp = None

    eris_ooov = _cp(eris.ooov)
    eris_ovoo = _cp(eris.ovoo)
    #:woooo += numpy.einsum('icjl,kc->ijkl', eris_ovoo, t1)
    #:wOVov += numpy.einsum('jblk,lc->jbkc', eris_ovoo, -t1)
    for i in range(nocc):
        woooo[i] += lib.dot(t1, eris_ovoo[i].reshape(nvir,-1)).reshape((nocc,)*3).transpose(1,0,2)
    lib.dot(eris_ovoo.reshape(-1,nocc), t1, -1, wOVov.reshape(-1,nvir), 1)
    #:wooov -= numpy.einsum('ibjl,lkcb->ijkc', eris_ovoo*1.5, t2)
    t2tmp = numpy.empty((nocc,nvir,nocc,nvir))
    for i in range(nocc):
        t2tmp[i] = t2[i].transpose(2,0,1)
    tmp_ooov = _cp(-eris_ooov.transpose(2,0,1,3)).reshape(-1,nov)
    lib.dot(tmp_ooov, t2tmp.reshape(nov,-1), 1.5, wooov.reshape(-1,nov), 1)
    t2tmp = None

    g2ooov, tmp_ooov = tmp_ooov.reshape(nocc,nocc,nocc,nvir), None
    g2ooov += eris_ooov * 2
    #:vikjc = numpy.einsum('iklb,jlcb->ikjc', g2ooov, theta)
    vikjc = lib.dot(g2ooov.reshape(-1,nov), theta.reshape(-1,nov).T)
    vikjc = vikjc.reshape(nocc,nocc,nocc,nvir)
    wooov += vikjc.transpose(0,2,1,3)
    wooov -= vikjc*.5
    g2ooov = vikjc = eris_ovoo = None

    w2 += numpy.einsum('ijkb,kb->ij', eris_ooov, t1) * 2
    w2 -= numpy.einsum('kjib,kb->ij', eris_ooov, t1)
    #:w3 -= numpy.einsum('kjlc,klbc->bj', eris_ooov, theta)
    for i in range(nocc):
        lib.dot(_cp(theta[i].transpose(1,2,0)).reshape(-1,nvir).T,
                eris_ooov[i].reshape(nocc,-1).T, -1, w3, 1)
    #:woooo += numpy.einsum('ikjc,lc->ijkl', eris_ooov, t1)
    #:wOvOv += numpy.einsum('jklb,lc->jbkc', eris_ooov, t1)
    woooo += lib.dot(eris_ooov.reshape(-1,nvir),
                     t1.T).reshape((-1,nocc,nocc,nocc)).transpose(0,2,1,3)
    for i in range(nocc):
        lib.dot(_cp(eris_ooov[i].transpose(2,0,1)).reshape(-1,nocc),
                t1, 1, wOvOv[i].reshape(-1,nvir), 1)
        wooov[i] += eris_ooov[i].transpose(1,0,2)*2
        wooov[i] -= eris_ooov[i]
    eris_ooov = theta = None

    eris_ovov = _cp(eris.ovov)
    g2ovov = numpy.empty((nocc,nocc,nvir,nvir))
    for i in range(nocc):
        g2ovov[i] = eris_ovov[i].transpose(1,0,2)*2
        g2ovov[i] -= eris_ovov[i].transpose(1,2,0)
    tmpw4 = numpy.einsum('klcd,ld->kc', g2ovov, t1)
    #:w1 -= numpy.einsum('kcja,kjcb->ba', g2ovov, t2)
    w1 -= lib.dot(t2.reshape(-1,nvir).T, g2ovov.reshape(-1,nvir))
    w1 -= numpy.einsum('ja,jb->ba', tmpw4, t1)
    #:w2 += numpy.einsum('ibkc,jkbc->ij', g2ovov, t2)
    w2 += lib.dot(g2ovov.reshape(nocc,-1), t2.reshape(nocc,-1).T)
    w2 += numpy.einsum('ib,jb->ij', tmpw4, t1)
    w3 += reduce(numpy.dot, (t1.T, tmpw4, t1.T))
    w4 += tmpw4
    vOVov  = eris_ovov.copy()
    #:vOVov += numpy.einsum('jbld,klcd->jbkc', g2ovov, t2)
    #:vOVov -= numpy.einsum('jbld,kldc->jbkc', eris_ovov, t2)
    lib.dot(_cp(g2ovov.transpose(0,2,1,3)).reshape(-1,nov),
            _cp(t2.transpose(0,2,1,3).reshape(nov,-1).T), 1,
            vOVov.reshape(nov,-1), 1)
    lib.dot(eris_ovov.reshape(-1,nov),
            _cp(t2.transpose(0,3,1,2).reshape(nov,-1).T), -1,
            vOVov.reshape(nov,-1), 1)
    g2ovov = None

    #:tmp = numpy.einsum('jbld,kd->ljbk', eris_ovov, t1)
    #:wOVov -= numpy.einsum('ljbk,lc->jbkc', tmp, t1)
    #:tmp = numpy.einsum('jdlb,kd->ljbk', eris_ovov, t1)
    #:wOvOv += numpy.einsum('ljbk,lc->jbkc', tmp, t1)
    tmp = numpy.empty((nocc,nvir,nocc))
    for j in range(nocc):
        lib.dot(_cp(eris_ovov[j].transpose(1,0,2)).reshape(-1,nvir),
                t1.T, 1, tmp.reshape(-1,nocc))
        lib.dot(tmp.reshape(nocc,-1).T, t1, -1, wOVov[j].reshape(-1,nvir), 1)
        lib.dot(eris_ovov[j].reshape(nvir,-1).T, t1.T, 1,
                tmp.reshape(-1,nocc))
        lib.dot(tmp.reshape(nocc,-1).T, t1, 1, wOvOv[j].reshape(-1,nvir), 1)
    tmp = None

    #:vOvOv = numpy.einsum('jdlb,kldc->jbkc', eris_ovov, t2)
    ovovtmp = _cp(eris_ovov.transpose(0,3,2,1).reshape(-1,nov))
    vOvOv = numpy.empty((nocc,nvir,nocc,nvir))
    for j in range(nocc):
        lib.dot(t2[j].reshape(-1,nvir).T, ovovtmp.T, 1,
                vOvOv[j].reshape(nvir,-1))
        vOvOv[j] -= eris.oovv[j].transpose(2,0,1)
    ovovtmp = eris_ovov = None
    vOvOv = lib.transpose(vOvOv.reshape(nov,-1)).reshape(nocc,nvir,nocc,nvir)
    wOVov += vOVov
    wOvOv += vOvOv
    saved.wOVov = wOVov
    saved.wOvOv = wOvOv
    ovovtmp = wOVov = wOvOv = eris_ovov = None

    ov2 = vOVov*2 + vOvOv
    w3 += numpy.einsum('kcjb,kc->bj', ov2, t1)
    #:wooov += numpy.einsum('ibjc,kb->ijkc', ov2, t1)
    #:wovvv -= numpy.einsum('jakb,jc->kabc', ov2, t1)
    for i in range(nocc):
        wooov[i] += lib.dot(t1, ov2[i].reshape(nvir,-1)).reshape(nocc,nocc,nvir).transpose(1,0,2)
    lib.dot(_cp(ov2.transpose(0,2,1,3).reshape(nocc,-1)).T,
            t1, -1, wovvv.reshape(-1,nvir), 1)
    ov2 = None
    ov1 = vOvOv*2 + vOVov
    #:wooov -= numpy.einsum('ibkc,jb->ijkc', ov1, t1)
    #:wovvv += numpy.einsum('jakc,jb->kabc', ov1, t1)
    for i in range(nocc):
        lib.dot(t1, ov1[i].reshape(nvir,-1), -1, wooov[i].reshape(nocc,-1), 1)
    wovvv += lib.dot(_cp(ov1.reshape(nocc,-1)).T,
                     t1).reshape(nvir,-1,nvir,nvir).transpose(1,0,3,2)
    ov1 = None

    woooo += _cp(eris.oooo).transpose(0,2,1,3)
    saved.woooo = woooo
    saved.wooov = wooov
    woooo = wooov = None

    w3 += numpy.einsum('bc,jc->bj', w1, t1)
    w3 -= numpy.einsum('kj,kb->bj', w2, t1)

    eris_ooov = _cp(eris.ooov)
    g2ooov = eris_ooov * 2
    g2ooov -= eris_ooov.transpose(2,0,1,3)
    #:tmp = numpy.einsum('kjla,jb->kabl', g2ooov, t1)
    #:wovvv = numpy.einsum('kabl,lc->kabc', tmp, t1)
    #:wovvv += numpy.einsum('kjla,jlbc->kabc', g2ooov, t2)
    tmp = lib.dot(g2ooov.reshape(nocc,-1).T, t1).reshape(-1,nocc,nvir,nvir).transpose(0,2,3,1)
    lib.dot(_cp(tmp.reshape(-1,nocc)), t1, 1, wovvv.reshape(-1,nvir), 1)
    tmp = None
    lib.dot(_cp(g2ooov.transpose(0,2,1,3).reshape(nocc**2,-1)).T,
            t2.reshape(nocc**2,-1), 1, wovvv.reshape(nov,-1), 1)
    g2ooov = eris_ooov = vOVov = vOvOv = None

    saved.wovvv = wovvv
    saved.w1 = w1
    saved.w2 = w2
    saved.w3 = w3
    saved.w4 = w4
    return saved

# update L1, L2
def update_amps(mycc, t1, t2, l1, l2, eris=None, saved=None):
    if saved is None:
        saved = make_intermediates(mycc, t1, t2, eris)
    time1 = time0 = time.clock(), time.time()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    nov = nocc * nvir
    foo = eris.fock[:nocc,:nocc]
    fov = eris.fock[:nocc,nocc:]
    fvv = eris.fock[:nocc,:nocc]

    #:mba = numpy.einsum('klca,klcb->ba', l2, t2*2-t2.transpose(0,1,3,2))
    #:mij = numpy.einsum('ikcd,jkcd->ij', l2, t2*2-t2.transpose(0,1,3,2))
    #:theta = t2*2 - t2.transpose(0,1,3,2)
    theta = _ccsd.make_0132(t2, t2, 2, -1)
    mba = lib.dot(theta.reshape(-1,nvir).T, l2.reshape(-1,nvir))
    mij = lib.dot(l2.reshape(nocc,-1), theta.reshape(nocc,-1).T)
    theta = None
    mba1 = numpy.einsum('jc,jb->bc', l1, t1) + mba
    mij1 = numpy.einsum('kb,jb->kj', l1, t1) + mij
    mia1 =(t1 + numpy.einsum('kc,jkbc->jb', l1, t2) * 2
         - numpy.einsum('kc,jkcb->jb', l1, t2)
         - reduce(numpy.dot, (t1, l1.T, t1))
         - numpy.einsum('bd,jd->jb', mba, t1)
         - numpy.einsum('lj,lb->jb', mij, t1))

    tmp = mycc.add_wvvVV(numpy.zeros_like(l1), l2, eris)
    l2new = numpy.empty((nocc,nocc,nvir,nvir))
    ij = 0
    for i in range(nocc):
        for j in range(i):
            tmp1 = tmp[ij] * .5  # *.5 because of l2+l2.transpose(1,0,3,2) later
            l2new[i,j] = tmp1
            l2new[j,i] = tmp1.T
            ij += 1
        l2new[i,i] = tmp[ij] * .5
        ij += 1
    l1new =(numpy.einsum('ijab,jb->ia', l2new, t1) * 4
          - numpy.einsum('jiab,jb->ia', l2new, t1) * 2)
    tmp = tmp1 = None

    l1new += eris.fock[:nocc,nocc:]
    l1new += numpy.einsum('ib,ba->ia', l1, saved.w1)
    l1new -= numpy.einsum('ja,ij->ia', l1, saved.w2)
    l1new -= numpy.einsum('ik,ka->ia', mij, saved.w4)
    l1new -= numpy.einsum('ca,ic->ia', mba, saved.w4)
    l1new += numpy.einsum('ijab,bj->ia', l2, saved.w3) * 2
    l1new -= numpy.einsum('ijba,bj->ia', l2, saved.w3)

    l2new += numpy.einsum('ia,jb->ijab', l1, saved.w4)
    #:l2new += numpy.einsum('jibc,ca->jiba', l2, saved.w1)
    #:l2new -= numpy.einsum('kiba,jk->jiba', l2, saved.w2)
    lib.dot(l2.reshape(-1,nvir), saved.w1, 1, l2new.reshape(-1,nvir), 1)
    lib.dot(saved.w2, l2.reshape(nocc,-1),-1, l2new.reshape(nocc,-1), 1)

    eris_ooov = _cp(eris.ooov)
    l1new -= numpy.einsum('jkia,kj->ia', eris_ooov, mij1) * 2
    l1new += numpy.einsum('ikja,kj->ia', eris_ooov, mij1)
    #:l2new -= numpy.einsum('ka,kijb->jiba', l1, eris_ooov)
    lib.dot(_cp(eris_ooov.transpose(0,2,1,3).reshape(nocc,-1)).T,
            l1, -1, l2new.reshape(-1,nvir), 1)
    eris_ooov = None

    tau = _ccsd.make_tau(t2, t1, t1)
    #:l2tau = numpy.einsum('ijcd,klcd->ijkl', l2, tau)
    l2tau = lib.dot(l2.reshape(nocc**2,-1),
                    tau.reshape(nocc**2,-1).T).reshape((nocc,)*4)
    tau = None
    l2t1 = numpy.einsum('ijcd,kc->ijkd', l2, t1)

    eris_ovvv = _cp(eris.ovvv)
    eris_ovvv = lib.unpack_tril(eris_ovvv.reshape(nov,-1))
    eris_ovvv = eris_ovvv.reshape(nocc,nvir,nvir,nvir)

    l1new += numpy.einsum('iabc,bc->ia', eris_ovvv, mba1) * 2
    l1new -= numpy.einsum('ibca,bc->ia', eris_ovvv, mba1)
    #:l2new += numpy.einsum('ic,jbac->jiba', l1, eris_ovvv)
    tmp = lib.dot(l1, eris_ovvv.reshape(-1,nvir).T).reshape(nocc,-1,nvir,nvir)
    for i in range(nocc):
        l2new[i] += tmp[i].transpose(0,2,1)
    #:m4 = numpy.einsum('ijkd,kadb->ijab', l2t1, eris_ovvv)
    m4 = tmp
    lib.dot(_cp(l2t1.reshape(nocc*nocc,-1)),
            _cp(eris_ovvv.transpose(0,2,1,3).reshape(-1,nvir**2)),
            1, m4.reshape(nocc*nocc,-1))
    l2new -= m4
    l1new -= numpy.einsum('ijab,jb->ia', m4, t1) * 2
    l1new -= numpy.einsum('ijab,ia->jb', m4, t1) * 2
    l1new += numpy.einsum('jiab,jb->ia', m4, t1)
    l1new += numpy.einsum('jiab,ia->jb', m4, t1)
    eris_ovvv = tmp = None

    eris_ovov = _cp(eris.ovov)
    l1new += numpy.einsum('jb,iajb->ia', l1, eris_ovov) * 2
    #:l2new -= numpy.einsum('jbic,ca->jiba', eris_ovov, mba1)
    #:l2new -= numpy.einsum('kajb,ik->ijab', eris_ovov, mij1)
    tmp = lib.dot(eris_ovov.reshape(-1,nvir), mba1).reshape(nocc,nvir,nocc,nvir)
    lib.dot(mij1, eris_ovov.reshape(nocc,-1), 1, tmp.reshape(nocc,-1), 1)
    tmp_oovv = numpy.empty((nocc,nocc,nvir,nvir))
    for i in range(nocc):
        tmp_oovv[i] = eris_ovov[i].transpose(1,0,2) * .5
        l2new[i] += tmp_oovv[i]
        l2new[i] -= tmp[i].transpose(1,0,2)
    tmp = None
    l1new += numpy.einsum('iajb,jb->ia', eris_ovov, mia1) * 2
    l1new -= numpy.einsum('ibja,jb->ia', eris_ovov, mia1)
    #:m4 = numpy.einsum('kalb,ijkl->ijab', eris_ovov, l2tau)
    lib.dot(l2tau.reshape(nocc*nocc,-1), tmp_oovv.reshape(-1,nvir**2),
            1, m4.reshape(nocc**2,-1))
    l2new += m4
    l1new += numpy.einsum('ijab,jb->ia', m4, t1) * 4
    l1new -= numpy.einsum('ijba,jb->ia', m4, t1) * 2
    eris_ovov = m4 = tmp_oovv = None

    eris_oovv = _cp(eris.oovv)
    l1new -= numpy.einsum('jb,ijba->ia', l1, eris_oovv)
    eris_oovv = None

    saved_wooov = _cp(saved.wooov)
    #:l1new -= numpy.einsum('jkca,ijkc->ia', l2, saved_wooov)
    l1new -= lib.dot(saved_wooov.reshape(nocc,-1), l2.reshape(-1,nvir))
    saved_wovvv = _cp(saved.wovvv)
    #:l1new += numpy.einsum('kibc,kabc->ia', l2, saved_wovvv)
    for j in range(nocc):
        l1new += lib.dot(l2[j].reshape(nocc,-1),
                         saved_wovvv[j].reshape(nvir,-1).T)
    saved_wooov = saved_wovvv = None

    saved_wOvOv = _cp(saved.wOvOv)
    tmp_ovov = _cp(saved.wOVov) * 2
    tmp_ovov += saved_wOvOv
    #:tmp = l2.transpose(0,2,1,3) - l2.transpose(0,3,1,2)*.5
    #:l2new += numpy.einsum('kcia,kcjb->jiba', tmp, tmp_ovov)
    tmp = numpy.empty((nocc,nvir,nocc,nvir))
    for i in range(nocc):
        tmp[i] = l2[i].transpose(2,0,1)*-.5
        tmp[i] += l2[i].transpose(1,0,2)
    tmp = lib.dot(tmp_ovov.reshape(-1,nov),
                  tmp.reshape(nov,-1)).reshape(-1,nvir,nocc,nvir)
    #:tmp = numpy.einsum('jkca,ibkc->ijab', l2, saved_wOvOv)
    for i in range(nocc):
        l2new[i] += tmp[i].transpose(1,0,2)
        tmp_ovov[i] = l2[i].transpose(2,0,1)
    lib.dot(saved_wOvOv.reshape(-1,nov), tmp_ovov.reshape(nov,-1),
            1, tmp.reshape(nov,-1))
    for i in range(nocc):
        l2new[i] += tmp[i].transpose(1,2,0)
        l2new[i] += tmp[i].transpose(1,0,2) * .5
    saved_wOvOv = tmp = tmp_ovov = None

    saved_woooo = _cp(saved.woooo)
    #:m3 = numpy.einsum('klab,ijkl->ijab', l2, saved_woooo)
    m3 = lib.dot(saved_woooo.reshape(-1,nocc**2),
                 l2.reshape(nocc**2,-1), .5).reshape(-1,nocc,nvir,nvir)
    l2new += m3
    l1new += numpy.einsum('ijab,jb->ia', m3, t1) * 4
    l1new -= numpy.einsum('ijba,jb->ia', m3, t1) * 2
    saved_woooo = m3 = None

    mo_e = eris.fock.diagonal()
    eia = lib.direct_sum('i-j->ij', mo_e[:nocc], mo_e[nocc:])
    l1new /= eia
    l1new += l1

#    l2new = l2new + l2new.transpose(1,0,3,2)
#    l2new /= lib.direct_sum('ia+jb->ijab', eia, eia)
#    l2new += l2
    ij = 0
    for i in range(nocc):
        for j in range(i):
            dab = lib.direct_sum('a+b->ab', eia[i], eia[j])
            tmp = (l2new[i,j]+l2new[j,i].T) / dab + l2[i,j]
            l2new[i,j] = tmp
            l2new[j,i] = tmp.T
            ij += 1
        dab = lib.direct_sum('a+b->ab', eia[i], eia[i])
        l2new[i,i] = (l2new[i,i]+l2new[i,i].T)/dab + l2[i,i]
        ij += 1

    time0 = log.timer_debug1('update l1 l2', *time0)
    return l1new, l2new

def _cp(a):
    return numpy.array(a, copy=False, order='C')

if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf.cc import ccsd
    from pyscf import ao2mo

    mol = gto.M()
    mf = scf.RHF(mol)

    mcc = ccsd.CCSD(mf)

    numpy.random.seed(12)
    nocc = 5
    nmo = 12
    nvir = nmo - nocc
    eri0 = numpy.random.random((nmo,nmo,nmo,nmo))
    eri0 = ao2mo.restore(1, ao2mo.restore(8, eri0, nmo), nmo)
    fock0 = numpy.random.random((nmo,nmo))
    fock0 = fock0 + fock0.T + numpy.diag(range(nmo))*2
    t1 = numpy.random.random((nocc,nvir))
    t2 = numpy.random.random((nocc,nocc,nvir,nvir))
    t2 = t2 + t2.transpose(1,0,3,2)
    l1 = numpy.random.random((nocc,nvir))
    l2 = numpy.random.random((nocc,nocc,nvir,nvir))
    l2 = l2 + l2.transpose(1,0,3,2)

    eris = lambda:None
    eris.oooo = eri0[:nocc,:nocc,:nocc,:nocc].copy()
    eris.ooov = eri0[:nocc,:nocc,:nocc,nocc:].copy()
    eris.ovoo = eri0[:nocc,nocc:,:nocc,:nocc].copy()
    eris.oovv = eri0[:nocc,:nocc,nocc:,nocc:].copy()
    eris.ovov = eri0[:nocc,nocc:,:nocc,nocc:].copy()
    idx = numpy.tril_indices(nvir)
    eris.ovvv = eri0[:nocc,nocc:,nocc:,nocc:][:,:,idx[0],idx[1]].copy()
    eris.vvvv = pyscf.ao2mo.restore(4,eri0[nocc:,nocc:,nocc:,nocc:],nvir)
    eris.fock = fock0

    saved = make_intermediates(mcc, t1, t2, eris)
    l1new, l2new = update_amps(mcc, t1, t2, l1, l2, eris, saved)
    print(abs(l1new).sum()-38172.7896467303)
    print(numpy.dot(l1new.flatten(), numpy.arange(35)) - 739312.005491083)
    print(numpy.dot(l1new.flatten(), numpy.sin(numpy.arange(35)))-7019.50937051188)
    print(numpy.dot(numpy.sin(l1new.flatten()), numpy.arange(35))-69.6652346635955)

    print(abs(l2new).sum()-72035.4931071527)
    print(abs(l2new-l2new.transpose(1,0,3,2)).sum())
    print(numpy.dot(l2new.flatten(), numpy.arange(35**2)) - 48427109.5409886)
    print(numpy.dot(l2new.flatten(), numpy.sin(numpy.arange(35**2)))-137.758016736487)
    print(numpy.dot(numpy.sin(l2new.flatten()), numpy.arange(35**2))-507.656936701192)

    mol = gto.Mole()
    mol.verbose = 0
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]

    mol.basis = 'cc-pvdz'
    mol.build()
    rhf = scf.RHF(mol)
    rhf.conv_tol = 1e-16
    rhf.scf()

    mcc = ccsd.CCSD(rhf)
    mcc.conv_tol = 1e-12
    ecc, t1, t2 = mcc.kernel()

    nmo = rhf.mo_energy.size
    fock0 = numpy.diag(rhf.mo_energy)
    nocc = mol.nelectron // 2
    nvir = nmo - nocc

    eris = mcc.ao2mo()
    conv, l1, l2 = kernel(mcc, eris, t1, t2, tol=1e-8)
    print(numpy.linalg.norm(l1)-0.0132626841292)
    print(numpy.linalg.norm(l2)-0.212575609057)

    import ccsd_rdm
    dm1 = ccsd_rdm.make_rdm1(mcc, t1, t2, l1, l2)
    dm2 = ccsd_rdm.make_rdm2(mcc, t1, t2, l1, l2)
    h1 = reduce(numpy.dot, (rhf.mo_coeff.T, rhf.get_hcore(), rhf.mo_coeff))
    eri = pyscf.ao2mo.full(rhf._eri, rhf.mo_coeff)
    eri = pyscf.ao2mo.restore(1, eri, nmo).reshape((nmo,)*4)
    e1 = numpy.einsum('pq,pq', h1, dm1)
    e2 = numpy.einsum('pqrs,pqrs', eri, dm2) * .5
    print e1+e2+mol.energy_nuc() - rhf.e_tot - ecc


import collections
import functools
import numpy as np
import scipy.signal as signal
from pathlib import Path
try:
    import backstaff.units as units
    import backstaff.plotting as plotting
    import backstaff.field_lines as field_lines
    import backstaff.beam_heating as beam_heating
except ModuleNotFoundError:
    import units
    import plotting
    import field_lines
    import beam_heating

class ElectronBeamSwarm(field_lines.FieldLineSet3):

    VALUE_DESCRIPTIONS = {
        'x': r'$x$ [Mm]',
        'y': r'$y$ [Mm]',
        'z': 'Height [Mm]',
        'z0': 'Initial height [Mm]',
        's': r'$s$ [Mm]',
        'sz0': r'$s(z=0) - s$ [Mm]',
        'initial_pitch_angle_cosine': r'$\mu_0$',
        'collisional_pitch_angle_cosine': r'$\mu$',
        'adiabatic_pitch_angle_cosine': r'$\mu$',
        'initial_pitch_angle': r'$\beta_0$ [deg]',
        'electric_field_angle_cosine': 'Electric field angle cosine',
        'total_power': 'Total power [erg/s]',
        'total_power_density': r'Total power density [erg/s/cm$^3$]',
        'total_energy_density': r'Total energy density [erg/cm$^3$]',
        'lower_cutoff_energy': r'$E_\mathrm{{c}}$ [keV]',
        'acceleration_volume': r'Acceleration site volume [cm$^3$]',
        'estimated_depletion_distance': r'$\tilde{{s}}_\mathrm{{dep}}$ [Mm]',
        'total_propagation_distance': r'$s_\mathrm{{dep}}$ [Mm]',
        'residual_factor': r'$r$',
        'acceleration_height': 'Acceleration site height [Mm]',
        'depletion_height': 'Depletion height [Mm]',
        'beam_electron_fraction': 'Beam electrons relative to total electrons',
        'return_current_speed_fraction': 'Speed relative to speed of light',
        'estimated_electron_density': r'Electron density [electrons/cm$^3$]',
        'deposited_power': 'Deposited power [erg/s]',
        'deposited_power_per_dist':
        r'$\mathrm{{d}}\mathcal{{E}}/\mathrm{{d}}s$ [erg/s/cm]',
        'deposited_power_density': r'$Q$ [erg/s/cm$^3$]',
        'power_change': 'Power change [erg/s]',
        'power_density_change': r'Power density change [erg/s/cm$^3$]',
        'beam_flux': r'Energy flux [erg/s/cm$^2$]',
        'conduction_flux': r'Energy flux [erg/s/cm$^2$]',
        'remaining_power': 'Remaining power [erg/s]',
        'relative_cumulative_power':
        r'$\mathcal{{E}}/P_\mathrm{{beam}}$ [$\%$]',
        'r': r'$\rho$ [g/cm$^3$]',
        'tg': r'$T$ [K]',
        'nel': r'$n_\mathrm{{e}}$ [electrons/cm$^3$]',
        'krec': r'$K$ [Bifrost units]',
        'qspitz': r'Power density change [erg/s/cm$^3$]',
        'r0': r'$\rho$ [g/cm$^3$]',
        'tg0': r'$T$ [K]',
        'p': r'$P$ [dyn/cm$^2$]',
        'b': r'$|B|$ [G]',
        'beta': r'$\beta$',
        'ux': r'$u_x$ [cm/s]',
        'uy': r'$u_y$ [cm/s]',
        'uz': r'$u_z$ [cm/s]',
        'us': r'$u_s$ [cm/s]',
        'uhor': r'$u_\mathrm{h}$ [cm/s]',
    }

    VALUE_UNIT_CONVERTERS = {
        'r': lambda f: f*units.U_R,
        'qspitz': lambda f: f*(units.U_E/units.U_T),
        'qjoule': lambda f: f*(units.U_E/units.U_T),
        'dedt': lambda f: f*(units.U_E/units.U_T),
        'r0': lambda f: f*units.U_R,
        'z': lambda f: -f,
        'z0': lambda f: -f,
        'bx': lambda f: f*units.U_B,
        'by': lambda f: f*units.U_B,
        'bz': lambda f: f*units.U_B,
        'b': lambda f: f*units.U_B,
        'p': lambda f: f*units.U_E,
        'ux': lambda f: f*units.U_U,
        'uy': lambda f: f*units.U_U,
        'uz': lambda f: f*units.U_U,
        'us': lambda f: f*units.U_U,
        'uhor': lambda f: f*units.U_U,
    }

    @staticmethod
    def from_file(file_path,
                  acceleration_data_type=None,
                  params={},
                  derived_quantities=[],
                  verbose=False):
        import backstaff.reading as reading
        file_path = Path(file_path)
        extension = file_path.suffix
        if extension == '.pickle':
            electron_beam_swarm = reading.read_electron_beam_swarm_from_combined_pickles(
                file_path,
                acceleration_data_type=acceleration_data_type,
                params=params,
                derived_quantities=derived_quantities,
                verbose=verbose)
        elif extension == '.fl':
            electron_beam_swarm = reading.read_electron_beam_swarm_from_custom_binary_file(
                file_path,
                acceleration_data_type=acceleration_data_type,
                params=params,
                derived_quantities=derived_quantities,
                verbose=verbose)
        else:
            raise ValueError(
                'Invalid file extension {} for electron beam data.'.format(
                    extension))

        return electron_beam_swarm

    @staticmethod
    def dummy(domain_bounds):
        return ElectronBeamSwarm(domain_bounds, 0, {}, {}, {}, {}, {})

    def __init__(self,
                 domain_bounds,
                 number_of_beams,
                 fixed_scalar_values,
                 fixed_vector_values,
                 varying_scalar_values,
                 varying_vector_values,
                 acceleration_data,
                 params={},
                 derived_quantities=[],
                 verbose=False):
        assert isinstance(acceleration_data, dict)
        self.number_of_beams = number_of_beams
        self.acceleration_data = acceleration_data
        super().__init__(domain_bounds,
                         number_of_beams,
                         fixed_scalar_values,
                         fixed_vector_values,
                         varying_scalar_values,
                         varying_vector_values,
                         params=params,
                         derived_quantities=derived_quantities,
                         verbose=verbose)

        if self.verbose:
            print('Acceleration data:\n    {}'.format('\n    '.join(
                self.acceleration_data.keys())))

    def get_subset(self,
                   only_quantities=None,
                   included_field_lines_finder=None,
                   **kwargs):
        return super().get_subset(
            self.acceleration_data,
            only_quantities=only_quantities,
            included_field_lines_finder=included_field_lines_finder,
            **kwargs)

    def get_number_of_beams(self):
        return self.number_of_beams

    def compute_number_of_sites(self):
        return np.unique(np.stack(
            (self.fixed_scalar_values['x0'], self.fixed_scalar_values['y0'],
             self.fixed_scalar_values['z0']),
            axis=1),
                         axis=0).shape[0]

    def get_acceleration_data(self, acceleration_data_type):
        return self.acceleration_data[acceleration_data_type]

    def get_acceleration_sites(self):
        return self.get_acceleration_data('acceleration_sites')

    def _derive_quantities(self, derived_quantities):

        super()._derive_quantities(derived_quantities)

        if 'initial_pitch_angle' in derived_quantities:
            self.fixed_scalar_values['initial_pitch_angle'] = np.arccos(
                self.get_fixed_scalar_values(
                    'initial_pitch_angle_cosine'))*180.0/np.pi

        if 'total_power_density' in derived_quantities:
            self.fixed_scalar_values[
                'total_power_density'] = self.get_fixed_scalar_values(
                    'total_power')/self.get_fixed_scalar_values(
                        'acceleration_volume')

        if 'total_energy_density' in derived_quantities:
            self._obtain_total_energy_densities()

        if 'non_thermal_energy_per_thermal_electron' in derived_quantities:
            self.fixed_scalar_values[
                'non_thermal_energy_per_thermal_electron'] = self._obtain_total_energy_densities(
                )/self.get_fixed_scalar_values(
                    'nel0' if self.has_fixed_scalar_values('nel0') else 'nel')

        if 'mean_electron_energy' in derived_quantities:
            self._obtain_mean_electron_energies()

        if 'acceleration_height' in derived_quantities:
            self.fixed_scalar_values['acceleration_height'] = np.asfarray(
                [-z[0] for z in self.get_varying_scalar_values('z')])

        if 'depletion_height' in derived_quantities:
            self.fixed_scalar_values['depletion_height'] = np.asfarray(
                [-z[-1] for z in self.get_varying_scalar_values('z')])

        if 'acceleration_site_electron_density' in derived_quantities:
            self._obtain_acceleration_site_electron_densities()

        if 'beam_electron_fraction' in derived_quantities:
            self._obtain_beam_electron_fractions()

        if 'return_current_speed_fraction' in derived_quantities:
            mean_electron_energies = self._obtain_mean_electron_energies(
            )*units.KEV_TO_ERG
            mean_electron_speed_fractions = np.sqrt(
                1.0 - 1.0/(1.0 + mean_electron_energies/units.MC2_ELECTRON)**2)
            beam_electron_fractions = self._obtain_beam_electron_fractions()
            self.fixed_scalar_values[
                'return_current_speed_fraction'] = beam_electron_fractions*mean_electron_speed_fractions

        if 'acceleration_current' in derived_quantities:
            self._obtain_acceleration_currents()

        if 'acceleration_induced_magnetic_field' in derived_quantities:
            self._obtain_acceleration_induced_magnetic_fields()

        if 'acceleration_ambient_magnetic_field' in derived_quantities:
            self._obtain_acceleration_ambient_magnetic_field()

        if 'relative_acceleration_induced_magnetic_field' in derived_quantities:
            B_induced = self._obtain_acceleration_induced_magnetic_fields()
            B = self._obtain_acceleration_ambient_magnetic_field()
            self.fixed_scalar_values[
                'relative_acceleration_induced_magnetic_field'] = B_induced/B

        if 'acceleration_induced_electric_field' in derived_quantities:
            self._obtain_acceleration_induced_electric_fields()

        if 'parallel_electric_field' in derived_quantities:
            self._obtain_parallel_electric_fields()

        if 'relative_acceleration_induced_electric_field' in derived_quantities:
            E_induced = self._obtain_acceleration_induced_electric_fields()
            E = np.abs(self._obtain_parallel_electric_fields())
            self.fixed_scalar_values[
                'relative_acceleration_induced_electric_field'] = E_induced/E

        if 'return_current_heating_ratio' in derived_quantities:
            self._obtain_return_current_heating_ratio()

        if 'estimated_electron_density' in derived_quantities:
            assert self.has_varying_scalar_values('r')
            self.varying_scalar_values['estimated_electron_density'] = [
                self.varying_scalar_values['r'][i]*units.U_R*
                units.MASS_DENSITY_TO_ELECTRON_DENSITY
                for i in range(self.get_number_of_beams())
            ]

        if 'deposited_power_per_dist' in derived_quantities:
            scale = 1.0/(self.get_param('dense_step_length')*units.U_L)
            self.varying_scalar_values['deposited_power_per_dist'] = [
                arr*scale
                for arr in self.varying_scalar_values['deposited_power']
            ]

        if 'power_change' in derived_quantities:
            self.varying_scalar_values['power_change'] = [
                arr.copy()
                for arr in self.varying_scalar_values['deposited_power']
            ]
            for i in range(self.get_number_of_beams()):
                self.varying_scalar_values['power_change'][i][
                    0] -= self.fixed_scalar_values['total_power'][i]

        if 'power_density_change' in derived_quantities:
            self.varying_scalar_values['power_density_change'] = [
                arr.copy() for arr in
                self.varying_scalar_values['deposited_power_density']
            ]
            for i in range(self.get_number_of_beams()):
                self.varying_scalar_values['power_density_change'][i][
                    0] -= self.get_fixed_scalar_values(
                        'total_power')[i]/self.get_fixed_scalar_values(
                            'acceleration_volume')[i]

        if 'padded_total_power_density' in derived_quantities:
            self.varying_scalar_values['padded_total_power_density'] = [
                np.zeros_like(arr) for arr in self.varying_scalar_values['x']
            ]
            for i in range(self.get_number_of_beams()):
                self.varying_scalar_values['padded_total_power_density'][i][
                    0] += self.get_fixed_scalar_values(
                        'total_power')[i]/self.get_fixed_scalar_values(
                            'acceleration_volume')[i]

        if 'beam_flux' in derived_quantities:
            self.varying_scalar_values['beam_flux'] = [
                arr.copy() for arr in
                self.varying_scalar_values['deposited_power_density']
            ]
            for i in range(self.get_number_of_beams()):
                self.varying_scalar_values['beam_flux'][i][
                    0] -= self.get_fixed_scalar_values(
                        'total_power')[i]/self.get_fixed_scalar_values(
                            'acceleration_volume')[i]
                self.varying_scalar_values['beam_flux'][i] *= self.get_param(
                    'dense_step_length')*units.U_L

        if 'conduction_flux' in derived_quantities:
            self.varying_scalar_values['conduction_flux'] = [
                arr.copy() for arr in self.varying_scalar_values['qspitz']
            ]
            for i in range(self.get_number_of_beams()):
                self.varying_scalar_values['conduction_flux'][
                    i] *= self.get_param(
                        'dense_step_length')*units.U_L*units.U_E/units.U_T

        if 'cumulative_power' in derived_quantities:
            self.varying_scalar_values['cumulative_power'] = [
                np.cumsum(self.varying_scalar_values['deposited_power'][i])
                for i in range(self.get_number_of_beams())
            ]

        if 'relative_cumulative_power' in derived_quantities:
            self.varying_scalar_values['relative_cumulative_power'] = [
                100*
                np.cumsum(self.varying_scalar_values['deposited_power'][i])/
                self.fixed_scalar_values['total_power'][i]
                for i in range(self.get_number_of_beams())
            ]

        if 'remaining_power' in derived_quantities:
            self.varying_scalar_values['remaining_power'] = [
                self.fixed_scalar_values['total_power'][i] -
                np.cumsum(self.varying_scalar_values['deposited_power'][i])
                for i in range(self.get_number_of_beams())
            ]

        if 'pitch_angle_intersection_energy' in derived_quantities:
            assert self.has_param('power_law_delta')
            delta = self.get_param('power_law_delta')
            self.varying_scalar_values['pitch_angle_intersection_energy'] = [
                Ec/np.sqrt(
                    (1.0 -
                     (np.sqrt(np.maximum(0.0, 1.0 - b*
                                         (1.0 - mu0**2)/b[0]))/mu0)**3)/
                    r**(-2/delta))
                for Ec, mu0, b, r in zip(
                    self.get_fixed_scalar_values('lower_cutoff_energy'),
                    self.get_fixed_scalar_values('initial_pitch_angle_cosine'),
                    self.get_varying_scalar_values('b'),
                    self.get_varying_scalar_values('residual_factor'))
            ]

        # if 'collisional_pitch_angle_cosine' in derived_quantities:
        #     assert self.has_param('power_law_delta')
        #     delta = self.get_param('power_law_delta')
        #     self.varying_scalar_values['collisional_pitch_angle_cosine'] = [
        #         mu0*np.cbrt(np.maximum(0.0, 1.0 - r**(-2/delta)))
        #         for mu0, r in zip(
        #             self.get_fixed_scalar_values('initial_pitch_angle_cosine'),
        #             self.get_varying_scalar_values('residual_factor'))
        #     ]

        if 'adiabatic_pitch_angle_cosine' in derived_quantities:
            self.varying_scalar_values['adiabatic_pitch_angle_cosine'] = [
                np.sqrt(np.maximum(0.0, 1.0 - b*(1.0 - mu0**2)/b[0]))
                for mu0, b in zip(
                    self.get_fixed_scalar_values('initial_pitch_angle_cosine'),
                    self.get_varying_scalar_values('b'))
            ]

    def _obtain_mean_electron_energies(self):
        if not self.has_fixed_scalar_values('mean_electron_energy'):
            assert self.has_param('power_law_delta')
            delta = self.get_param('power_law_delta')
            self.fixed_scalar_values['mean_electron_energy'] = (
                (delta - 0.5)/(delta - 1.5)
            )*self.get_fixed_scalar_values('lower_cutoff_energy')

        return self.get_fixed_scalar_values('mean_electron_energy')

    def _obtain_mean_electron_speeds(self):
        if not self.has_fixed_scalar_values('mean_electron_speed'):
            assert self.has_param('power_law_delta')
            delta = self.get_param('power_law_delta')
            self.fixed_scalar_values['mean_electron_speed'] = (
                (delta - 0.5)/(delta - 1))*np.sqrt(
                    2*self.get_fixed_scalar_values('lower_cutoff_energy')*
                    units.KEV_TO_ERG/units.M_ELECTRON)  # [cm/s]

        return self.get_fixed_scalar_values('mean_electron_speed')

    def _obtain_acceleration_site_electron_densities(self):
        if not self.has_fixed_scalar_values(
                'acceleration_site_electron_density'):
            assert self.has_fixed_scalar_values(
                'r0') or self.has_fixed_scalar_values('r')
            self.fixed_scalar_values[
                'acceleration_site_electron_density'] = self.get_fixed_scalar_values(
                    'r0' if self.has_fixed_scalar_values('r0') else 'r'
                )*units.U_R*units.MASS_DENSITY_TO_ELECTRON_DENSITY

        return self.get_fixed_scalar_values(
            'acceleration_site_electron_density')

    def _obtain_beam_electron_fractions(self):
        if not self.has_fixed_scalar_values('beam_electron_fraction'):
            assert self.has_param('particle_energy_fraction')
            assert self.has_fixed_scalar_values(
                'bx') and self.has_fixed_scalar_values(
                    'by') and self.has_fixed_scalar_values('bz')
            assert self.has_fixed_scalar_values(
                'ix') and self.has_fixed_scalar_values(
                    'iy') and self.has_fixed_scalar_values('iz')

            bx = self.get_fixed_scalar_values('bx')*units.U_B
            by = self.get_fixed_scalar_values('by')*units.U_B
            bz = self.get_fixed_scalar_values('bz')*units.U_B
            ix = self.get_fixed_scalar_values('ix')
            iy = self.get_fixed_scalar_values('iy')
            iz = self.get_fixed_scalar_values('iz')
            free_energy = (bx*bx + by*by + bz*bz - (bx*ix + by*iy + bz*iz)**2/
                           (ix*ix + iy*iy + iz*iz))/(8.0*np.pi)
            mean_electron_energies = self._obtain_mean_electron_energies(
            )*units.KEV_TO_ERG
            electron_densities = self._obtain_acceleration_site_electron_densities(
            )
            self.fixed_scalar_values[
                'beam_electron_fraction'] = self.get_param(
                    'particle_energy_fraction')*free_energy/(
                        mean_electron_energies*electron_densities)

        return self.get_fixed_scalar_values('beam_electron_fraction')

    def _obtain_total_energy_densities(self):
        if not self.has_fixed_scalar_values('total_energy_density'):
            assert self.has_param('acceleration_duration')
            self.fixed_scalar_values[
                'total_energy_density'] = self.get_fixed_scalar_values(
                    'total_power')*self.get_param(
                        'acceleration_duration')/self.get_fixed_scalar_values(
                            'acceleration_volume')

        return self.get_fixed_scalar_values('total_energy_density')

    def _obtain_acceleration_currents(self):
        if not self.has_fixed_scalar_values('acceleration_current'):
            u_acc = self._obtain_total_energy_densities()
            E_mean = self._obtain_mean_electron_energies()*units.KEV_TO_ERG
            v_mean = self._obtain_mean_electron_speeds()
            self.fixed_scalar_values[
                'acceleration_current'] = units.Q_ELECTRON*u_acc*v_mean/E_mean

        return self.get_fixed_scalar_values('acceleration_current')

    def _obtain_acceleration_induced_magnetic_fields(self):
        if not self.has_fixed_scalar_values(
                'acceleration_induced_magnetic_field'):
            j = self._obtain_acceleration_currents()
            L = np.cbrt(self.get_fixed_scalar_values('acceleration_volume'))
            self.fixed_scalar_values[
                'acceleration_induced_magnetic_field'] = np.pi*j*L/units.CLIGHT

        return self.get_fixed_scalar_values(
            'acceleration_induced_magnetic_field')

    def _obtain_resistivities(self):
        if not self.has_fixed_scalar_values('resistivity'):
            nel = self.get_fixed_scalar_values('nel0')
            tg = self.get_fixed_scalar_values('tg0')
            r = self.get_fixed_scalar_values('r0')*units.U_R
            x = beam_heating.compute_equilibrium_hydrogen_ionization_fraction(
                tg, nel)
            nH = beam_heating.compute_total_hydrogen_density(r)

            self.fixed_scalar_values['resistivity'] = (
                7.26e-9*x/tg**(3/2))*np.log(
                    3*np.sqrt((units.KBOLTZMANN*tg)**3/(np.pi*nH))/
                    (2*units.Q_ELECTRON**3)) + 7.6e-18*(1 - x)*np.sqrt(tg)/x

        return self.get_fixed_scalar_values('resistivity')

    def _obtain_acceleration_induced_electric_fields(self):
        if not self.has_fixed_scalar_values(
                'acceleration_induced_electric_field'):
            j = self._obtain_acceleration_currents()
            eta = self._obtain_resistivities()
            self.fixed_scalar_values[
                'acceleration_induced_electric_field'] = eta*j*units.STATV_TO_V*1e2

        return self.get_fixed_scalar_values(
            'acceleration_induced_electric_field')

    def _obtain_parallel_electric_fields(self):
        if not self.has_fixed_scalar_values('parallel_electric_field'):
            bx = self.get_fixed_scalar_values('bx0')*units.U_B
            by = self.get_fixed_scalar_values('by0')*units.U_B
            bz = self.get_fixed_scalar_values('bz0')*units.U_B
            ex = self.get_fixed_scalar_values('ex0')*units.U_EL
            ey = self.get_fixed_scalar_values('ey0')*units.U_EL
            ez = self.get_fixed_scalar_values('ez0')*units.U_EL

            self.fixed_scalar_values['parallel_electric_field'] = (
                bx*ex + by*ey + bz*ez)/np.sqrt(bx**2 + by**2 + bz**2)

        return self.get_fixed_scalar_values('parallel_electric_field')

    def _obtain_return_current_heating_ratio(self):
        if not self.has_fixed_scalar_values('return_current_heating_ratio'):
            E_mean = self._obtain_mean_electron_energies()  # [keV]
            mu = self.get_fixed_scalar_values('initial_pitch_angle_cosine')

            nel = self.get_fixed_scalar_values('nel0')
            tg = self.get_fixed_scalar_values('tg0')
            r = self.get_fixed_scalar_values('r0')*units.U_R
            x = beam_heating.compute_equilibrium_hydrogen_ionization_fraction(
                tg, nel)
            r = self.get_fixed_scalar_values('r0')*units.U_R
            nH = beam_heating.compute_total_hydrogen_density(r)  # [1/cm^3]

            electron_coulomb_logarithm = beam_heating.compute_electron_coulomb_logarithm(
                nel, E_mean)
            neutral_hydrogen_coulomb_logarithm = beam_heating.compute_neutral_hydrogen_coulomb_logarithm(
                E_mean)
            gamma = beam_heating.compute_effective_coulomb_logarithm(
                x,
                electron_coulomb_logarithm,
                neutral_hydrogen_coulomb_logarithm,
            )

            E = self._obtain_acceleration_induced_electric_fields()/(
                units.STATV_TO_V*1e2)  # [statV/cm]
            e = units.Q_ELECTRON  # [statC]

            self.fixed_scalar_values['return_current_heating_ratio'] = (
                E*e/nH)/(2*np.pi*e**4*gamma/(mu*E_mean*units.KEV_TO_ERG))

        return self.get_fixed_scalar_values('return_current_heating_ratio')

    def _obtain_acceleration_ambient_magnetic_field(self):
        if not self.has_fixed_scalar_values(
                'acceleration_ambient_magnetic_field'):
            bx = self.get_fixed_scalar_values('bx0')
            by = self.get_fixed_scalar_values('by0')
            bz = self.get_fixed_scalar_values('bz0')

            B = np.sqrt(bx**2 + by**2 + bz**2)*units.U_B

            self.fixed_scalar_values['acceleration_ambient_magnetic_field'] = B

        return self.get_fixed_scalar_values(
            'acceleration_ambient_magnetic_field')

class AccelerationSites(field_lines.FieldLineSet3):

    VALUE_DESCRIPTIONS = ElectronBeamSwarm.VALUE_DESCRIPTIONS
    VALUE_UNIT_CONVERTERS = ElectronBeamSwarm.VALUE_UNIT_CONVERTERS

    def __init__(self,
                 domain_bounds,
                 number_of_sites,
                 fixed_scalar_values,
                 fixed_vector_values,
                 varying_scalar_values,
                 varying_vector_values,
                 params={},
                 derived_quantities=[],
                 verbose=False):
        self.number_of_sites = number_of_sites
        super().__init__(domain_bounds,
                         number_of_sites,
                         fixed_scalar_values,
                         fixed_vector_values,
                         varying_scalar_values,
                         varying_vector_values,
                         params=params,
                         derived_quantities=derived_quantities,
                         verbose=verbose)

    def get_number_of_sites(self):
        return self.number_of_sites

def find_beams_starting_in_coords(x_coords,
                                  y_coords,
                                  z_coords,
                                  propagation_senses,
                                  fixed_scalar_values,
                                  max_propagation_sense_diff=1e-6,
                                  max_distance=1e-5):
    return [
        i for i, (x, y, z, s) in enumerate(
            zip(fixed_scalar_values['x0'], fixed_scalar_values['y0'],
                fixed_scalar_values['z0'],
                fixed_scalar_values['propagation_sense']))
        if np.any(
            np.logical_and(
                (x - x_coords)**2 + (y - y_coords)**2 +
                (z - z_coords)**2 <= max_distance**2,
                np.abs(s - propagation_senses) < max_propagation_sense_diff))
    ]

def find_beams_propagating_longer_than_distance(min_distance,
                                                fixed_scalar_values):
    return list(
        np.nonzero(
            fixed_scalar_values['total_propagation_distance'] > min_distance)
        [0])

def find_beams_in_temperature_height_region(height_lims, tg_lims,
                                            varying_scalar_values):
    return [
        i for i, (z, tg) in enumerate(
            zip(varying_scalar_values['z'], varying_scalar_values['tg']))
        if np.any((z > -height_lims[1])*(z < -height_lims[0])*
                  (tg > tg_lims[0])*(tg < tg_lims[1]))
    ]

def find_peak_deposition_point(varying_scalar_values, field_line_idx):
    deposited_power = varying_scalar_values['deposited_power'][field_line_idx]
    indices, _ = signal.find_peaks(deposited_power/np.mean(deposited_power),
                                   prominence=2)
    return slice(np.max(indices) if indices.size > 0 else -1, None, None)

def plot_electron_beams(*args, **kwargs):
    return field_lines.plot_field_lines(*args, **kwargs)

def plot_electron_beam_properties(*args, **kwargs):
    return field_lines.plot_field_line_properties(*args, **kwargs)

def plot_beam_value_histogram(*args, **kwargs):
    return field_lines.plot_field_line_value_histogram(*args, **kwargs)

def plot_beam_value_histogram_difference(*args, **kwargs):
    return field_lines.plot_field_line_value_histogram_difference(
        *args, **kwargs)

def plot_beam_value_2d_histogram(*args, **kwargs):
    return field_lines.plot_field_line_value_2d_histogram(*args, **kwargs)

def plot_beam_value_2d_histogram_difference(*args, **kwargs):
    return field_lines.plot_field_line_value_2d_histogram_difference(
        *args, **kwargs)

def plot_beam_value_2d_histogram_comparison(*args, **kwargs):
    return field_lines.plot_field_line_value_2d_histogram_comparison(
        *args, **kwargs)


""" General framework for parallel tempering MCMC. """
import copy
from logging import Logger
import math
import numpy as np
import os
import random
import time
from typing import List

from rdkit import Chem
from rdkit.Chem import AllChem, rdForceFieldHelpers
from rdkit.Chem.Lipinski import RotatableBondSmarts
# noinspection PyPackageRequirements
from tap import Tap
from tqdm import tqdm

from conformation.rdkit_hmc import calc_energy_grad, hmc_step

class Args(Tap):
    """
    System arguments.
    """
    bin_path: str  # Path to RDKit binary file containing molecule
    max_attempts: int = 10000  # Max number of embedding attempts
    num_minimization_iters: int = 0  # Number of minimization steps
    temperatures: List[float] = [300, 500]  # Temperature ladder, with the first being the target temperature
    epsilon: float = 1  # Leapfrog step size in femtoseconds
    L: int = 10  # Number of leapfrog steps
    num_steps: int = 1  # Number of parallel tempering steps
    swap_prob: float = 0.2  # Probability of performing a swap operation on a given step
    subsample_frequency: int = 1  # Frequency at which configurations are saved from MH steps
    log_frequency: int = 100  # Log frequency
    save_dir: str  # Path to directory containing output files

def parallel_tempering(args: Args, logger: Logger) -> None:
    """
    Parallel tempering scheme.
    :param args: System arguments.
    :param logger: System logger.
    :return: None.
    """
    # Set up logger
    debug, info = logger.debug, logger.info

    # Define constants
    k_b = 3.297e-24  # Boltzmann constant in cal/K
    avogadro = 6.022e23
    args.epsilon *= 1e-15  # Convert to femtoseconds

    # Load molecule
    # noinspection PyUnresolvedReferences
    mol = Chem.Mol(open(args.bin_path, "rb").read())
    mol.RemoveAllConformers()

    debug(f'Starting search...')

    # Generate initial conformations
    # Here, we consider the variables of interest, q, to effectively be the atomic coordinates
    initial_q = copy.deepcopy(mol)
    num_atoms = initial_q.GetNumAtoms()
    AllChem.EmbedMultipleConfs(initial_q, maxAttempts=args.max_attempts, numConfs=len(args.temperatures), numThreads=0)
    AllChem.MMFFOptimizeMoleculeConfs(initial_q, maxIters=args.num_minimization_iters, numThreads=0)
    current_q_list = []
    for i in range(len(args.temperatures)):
        current_q = copy.deepcopy(mol)
        c = initial_q.GetConformers()[i]
        c.SetId(0)
        current_q.AddConformer(c)
        current_q_list.append(current_q)

    # Compute force field information
    mmff_p = rdForceFieldHelpers.MMFFGetMoleculeProperties(initial_q)
    force_field = rdForceFieldHelpers.MMFFGetMoleculeForceField(initial_q, mmff_p)

    # Masses in kg
    mass = np.array([mol.GetAtomWithIdx(i).GetMass() / (1000. * avogadro) for i in range(num_atoms)])

    # Add the first conformation to the list
    energy, _ = calc_energy_grad(current_q_list[0].GetConformer().GetPositions(), force_field)
    all_conformation_molecules = [current_q_list[0]]

    debug(f'Running HMC steps...')
    start_time = time.time()
    num_internal_accepted = [0]*len(args.temperatures)
    num_swap_accepted = [0]*(len(args.temperatures) - 1)
    num_swap_attempted = [0]*(len(args.temperatures) - 1)
    total_num_swap_accepted = 0
    total_swap_attempted = 0
    for step in tqdm(range(args.num_steps)):
        alpha = random.uniform(0, 1)
        if alpha > args.swap_prob:
            results = []
            for i in range(len(args.temperatures)):
                accepted, current_q, current_energy = hmc_step(current_q_list[i], force_field, args.temperatures[i],
                                                               k_b, avogadro, mass, num_atoms, args.epsilon, args.L)
                results.append([accepted, current_q, current_energy])

            for i in range(len(args.temperatures)):
                accepted, current_q, current_energy = results[i]

                if accepted:
                    current_q_list[i] = current_q  # Necessary because Python is pass by object reference!
                    num_internal_accepted[i] += 1

                if i == 0:
                    if step % args.subsample_frequency == 0:
                        all_conformation_molecules.append(current_q)

        elif len(args.temperatures) > 1:
            swap_index = random.randint(0, len(args.temperatures) - 2)
            energy_k0, _ = calc_energy_grad(current_q_list[swap_index].GetConformer().GetPositions(), force_field)
            energy_k0 *= (1000.0 * 4.184 / avogadro)
            energy_k1, _ = calc_energy_grad(current_q_list[swap_index + 1].GetConformer().GetPositions(), force_field)
            energy_k1 *= (1000.0 * 4.184 / avogadro)
            delta_beta = (1./(k_b * args.temperatures[swap_index] * 4.184) -
                          1./(k_b * args.temperatures[swap_index + 1] * 4.184))
            delta_energy = energy_k0 - energy_k1
            delta = delta_beta * delta_energy

            prob_ratio = math.exp(-delta)
            mu = random.uniform(0, 1)
            if swap_index == 0:
                energy, _ = calc_energy_grad(current_q_list[0].GetConformer().GetPositions(), force_field)
            if mu <= prob_ratio:
                tmp = copy.deepcopy(current_q_list[swap_index])
                current_q_list[swap_index] = copy.deepcopy(current_q_list[swap_index + 1])
                current_q_list[swap_index + 1] = tmp
                num_swap_accepted[swap_index] += 1
                total_num_swap_accepted += 1
            if swap_index == 0:
                if step % args.subsample_frequency == 0:
                    all_conformation_molecules.append(current_q_list[0])
            num_swap_attempted[swap_index] += 1
            total_swap_attempted += 1

        if step % args.log_frequency == 0:
            for i in range(len(args.temperatures)):
                debug(f'% Moves accepted for temperature {args.temperatures[i]}: '
                      f'{float(num_internal_accepted[i]) / float(step + 1) * 100.0}')

            for i in range(len(args.temperatures) - 1):
                if num_swap_attempted[i] == 0:
                    debug(f'% Moves accepted for swap at base temperature {args.temperatures[i]}: NA')
                else:
                    debug(f'% Moves accepted for swap at base temperature {args.temperatures[i]}: '
                          f'{float(num_swap_accepted[i]) / num_swap_attempted[i] * 100.0}')

            debug(f'# Swap moves attempted: {total_swap_attempted}')
            if total_swap_attempted == 0:
                debug(f'% Moves accepted for swap: {0.0}')
                debug(f'% Moves accepted for swap: NA')
            else:
                debug(f'% Moves accepted for swap: '
                      f'{float(total_num_swap_accepted) / float(total_swap_attempted) * 100.0}')
    end_time = time.time()
    debug(f'Total Time (s): {end_time - start_time}')
    for i in range(len(args.temperatures)):
        debug(f'% Moves accepted for temperature {args.temperatures[i]}: '
              f'{float(num_internal_accepted[i]) / float(args.num_steps) * 100.0}')
    debug(f'# Swap moves attempted: {total_swap_attempted}')
    if total_swap_attempted > 0:
        debug(f'% Moves accepted for swap: {float(total_num_swap_accepted) / float(total_swap_attempted) * 100.0}')

    # Discover the rotatable bonds
    rotatable_bonds = mol.GetSubstructMatches(RotatableBondSmarts)
    debug(f'Num rotatable bonds: {len(rotatable_bonds)}')

    # Save all sub sampled conformations in molecule object
    # noinspection PyUnresolvedReferences
    all_mol = Chem.Mol(open(args.bin_path, "rb").read())
    all_mol.RemoveAllConformers()
    for i in range(len(all_conformation_molecules)):
        c = all_conformation_molecules[i].GetConformer()
        c.SetId(i)
        all_mol.AddConformer(c)

    # Save molecule to binary file
    bin_str = all_mol.ToBinary()
    with open(os.path.join(args.save_dir, "conformations.bin"), "wb") as b:
        b.write(bin_str)