#! /usr/bin/env python

import mdtraj as md
import argparse
import numpy as np
from LLC_Membranes.llclib import physical, topology
from LLC_Membranes.analysis import p2p
import matplotlib.pyplot as plt
import tqdm
import matplotlib as mpl
import os.path as path

def initialize():

    parser = argparse.ArgumentParser(description='Calculate the number density of selected groups along axis')

    parser.add_argument('-g', '--gro', default='wiggle.gro', help='Name of coordinate file')
    parser.add_argument('-t', '--traj', default='traj_whole.xtc', help='Trajectory file (.trr, .xtc should work)')
    parser.add_argument('-b', '--begin', default=0, type=int, help='Frame to begin calculations')
    parser.add_argument('-c', '--components', nargs='+', help="Region(s) or atoms to include in calculation. Special"
                                                              "groups include 'head groups' and 'tails'")
    parser.add_argument('-m', '--monomer', default='NaGA3C11', help='Name of monomer (no file extension)')
    parser.add_argument('-e', '--end', default=-1, type=int, help='Frame to stop doing calculations')
    parser.add_argument('-bins', default=100, type=int, help='Number of bins to use')
    parser.add_argument('-s', '--solvate', action="store_true",
                        help='If the system is solvated, plot the number density of water as well')
    parser.add_argument('-l', '--load', help='Name of compressed .npz to load')

    args = parser.parse_args()

    return args

def grps(name):
    """ Return names of atoms making up specialized groups. This is specifically for NaGA3C11.

    TODO: incorporate these groups as annotations (maybe)

    :param name: specialized group names. Valid options are 'head groups' and 'tails'

    :type name: str

    :return: atom names that constitute specialized groups
    """

    if name == 'head groups':

        return ['C', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'O3', 'O4']

    elif name == 'tails':

        return ['C7', 'C8', 'C9', 'C10', 'C11', 'C12', 'C13', 'C14', 'C15', 'C16', 'C17', 'C18', 'C19', 'C20', 'C21',
                'C22', 'C23', 'C24', 'C25', 'C26', 'C27', 'C28', 'C29', 'C30', 'C31', 'C32', 'C33', 'C34', 'C35',
                'C36', 'C37', 'C38', 'C39', 'C40', 'C41', 'C42', 'C43', 'C44', 'C45', 'C46', 'C47', 'C48']

class Region(object):

    def __init__(self, traj, gro, components, monomer, begin=0, end=-1):
        """ Determine the density of components. Similar to a radial density, but instead of center of mass, the radial
        density of each atom in the region is averaged together.

        :param traj: name of trajectory to analyze (.xtc or .trr)
        :param gro: name of GROMACS coordinate file with same topology as traj
        :param component: list of atoms or name of groups of atoms
        :param begin: first frame to analyze
        :param end: last grame to analyze

        :type traj: str
        :type gro: str
        :type component: list
        :type begin: int
        :type end: int
        """

        special_groups = ['head groups', 'tails']

        if components[0].lower() in special_groups:
            self.components = grps(components[0].lower())
        else:
            self.components = components

        print('Loading trajectory...', end="")
        self.t = md.load('%s' % traj, top='%s' % gro)[begin:end]
        print('done')

        self.monomer = topology.LC('%s.gro' % monomer)

        self.r = None
        self.density = None

    def component_density(self, bins=50, spline=True, progress=True, npts_spline=10):

        self.r = np.zeros([bins])
        self.density = np.zeros([self.t.n_frames, bins])

        pore_defining_atoms = [a.index for a in self.t.topology.atoms if a.name in self.monomer.pore_defining_atoms
                               and a.residue.name in self.monomer.residues]

        pore_centers = physical.avg_pore_loc(4, self.t.xyz[:, pore_defining_atoms, :], self.t.unitcell_vectors,
                                             spline=spline, progress=progress, npts=npts_spline)

        print('Calculating component density')
        self.r, self.density = physical.compdensity(self.com, pore_centers, self.t.unitcell_vectors,
                                                    nbins=bins, spline=spline, cut=cut)

def duplicate(pos, box):
    """
    Duplicate a set of positions periodically once in the +/- xy directions
    :param pos: xyz positions of a set of coordinates to be duplicated periodically
    :param box: box vectors in mdtraj format (t.unitcell_vectors : [nframes, 3, 3]) for every frame
    :return: Periodically duplicated system
    """

    n = pos.shape[1]  # number atoms in original unit cell

    p = np.zeros([pos.shape[0], n*9, 3])  # will hold periodically duplicated system
    p[:, :n, :] = pos

    # x-direction
    for t in range(pos.shape[0]):
        p[t, n:2*n, :] = pos[t, :, :] + box[t, 0, :]
        p[t, 2*n:3*n, :] = pos[t, :, :] - box[t, 0, :]

    # y-direction
    n *= 3
    for t in range(pos.shape[0]):
        p[t, n:2*n, :] = p[t, :n, :] + box[t, 1, :]
        p[t, 2*n:3*n, :] = p[t, :n, :] - box[t, 1, :]

    return p

def compdensity(component, pore_centers, start, box, cut=1.5, pores=4, nbins=50, rmax=3.5, buffer=0.0):
    """ Measure the density of a component as a function of the distance from the pore centers

    :param component: the coordinates of the component(s) which you want a radial distribution of at each frame
    :param pore_centers: a numpy array of the locations of each pore center at each trajectory frame
    :param start: the frame number at which to start calculations (should be after equilibration)
    :param cut: cutoff distance for distance calculations. Will not count anything further than cut from the pore center
    :param pores: number of pores (int) default=4
    :param rmax: maximum distance from pore center to calculate density for, default = 3.5 nm
    :param buffer: percentage used to define the location of z planes between which component density will be computed,
           float, default = 0 (i.e. no buffer). Should be between 0 and 1. e.g. for 1 percent, use 0.01 as the buffer

    :type component: numpy.ndarray
    :type pore_centers: numpy.ndarray
    :type start: int
    :type cut: float
    :type pores: int
    :type rmax: float
    :type buffer: float

    :return: the density of "component" as a function the distance from the pore center. Also
             returns the calculated bin width for plotting
    """

    # Extract basic system information. It's important to follow the format of the component array to get it right

    tot_atoms = component.shape[1]  # the total number of components in a single frame
    n_ppore = tot_atoms // pores  # the total number of components in each pore

    nT = component.shape[0]
    zbox = np.mean(box[:, 2, 2])
    # Find the approximate max and minimum z values of the components based on the last frame

    zmax = np.max(component[-1, :, 2])
    zmin = np.min(component[-1, :, 2])
    thickness = zmax - zmin  # approximate membrane thickness

    # now find the maximum and minimum permissible z dimensions based on the buffer
    zmax_buff = zmax - thickness * buffer  # Could use buffer/2 depending on how you interpret what buffer % means
    zmin_buff = zmin + thickness * buffer

    density = np.zeros([nbins])  # number / nm^3
    for t in tqdm.tqdm(range(start, nT)):
        for p in range(pores):
            # narrow down the positions to those that are within 'cut' of at least one pore
            distances = np.linalg.norm(component[t, :, :2] - pore_centers[t, p, :], axis=1)
            d_sorted = np.sort(distances)
            # find where the distances exceed the cutoff
            stop = 0
            while d_sorted[stop] < cut:
                stop += 1

            hist, bin_edges = np.histogram(d_sorted[:stop], bins=nbins, range=(0, cut))  # the range option is necessary
            #  to make sure we have equal sized bins on every iteration

            density += hist

    density /= zbox * (nT - start)  # take average
    bin_width = cut / nbins

    # normalize based on area of anulus where bin is located
    r = np.zeros([nbins])
    normalization = []
    for i in range(nbins):
        normalization.append(np.pi * (bin_edges[i + 1] ** 2 - bin_edges[i] ** 2))
        density[i] /= np.pi * (bin_edges[i + 1] ** 2 - bin_edges[i] ** 2)
        r[i] = (bin_edges[i + 1] + bin_edges[i]) / 2

    return density, r, bin_width

if __name__ == "__main__":

    args = initialize()  # parse the args

    # run : regional_density.py - m offset_disordered_regional_density.npz offset_regional_density.npz
    # layered_disordered_regional_density.npz layered_regional_density.npz solvated_regional_density.npz

    # mpl.style.use('seaborn')

    # regions = ['Tails', 'Head Groups', 'Sodium']
    regions = ['Tails', 'Head Groups', 'Sodium']

    # if args.multi:
    if False:

        # colors = ['blue', 'red', 'green', 'xkcd:orange']
        # colors = ['xkcd:red', 'xkcd:green', 'blue', 'xkcd:yellow']
        colors = ['xkcd:blue', 'xkcd:olive', 'xkcd:orangered', 'xkcd:magenta', 'xkcd:gold']
        names = ['Ordered Parallel Displaced', 'Ordered Sandwiched', 'Disordered Sandwiched',
                 'Disordered Parallel Displaced', 'Solvated Parallel Displaced']

        # import pylab
        #
        # fig = pylab.figure()
        # figlegend = pylab.figure(figsize=(7.75, 0.6))
        # ax = fig.add_subplot(111)
        # for i in range(5):
        #     ax.plot(range(10), pylab.randn(10), color=colors[i], label=names[i])
        #
        # #lines = ax.plot(range(10), pylab.randn(10), range(10), pylab.randn(10),range(10), pylab.randn(10), range(10), pylab.randn(10), range(10), pylab.randn(10))
        # figlegend.legend(*ax.get_legend_handles_labels(), 'best', ncol=3)
        # fig.show()
        # figlegend.show()
        # fig.tight_layout()
        # figlegend.savefig('legend.pdf')
        #
        # plt.show()
        # exit()

        # names = ['Parallel Displaced (d=3.7)', 'Sandwiched (d=3.7)', 'Sandwiched (d=5.0)', 'Parallel Displaced (d=5.0)',
        #          'Solvated Parallel Displaced']
        # names = ['Dry', 'Solvated']
        # colors = ['xkcd:orange', 'xkcd:blue', 'xkcd:orange']

        n = len(args.multi)
        system = np.load(args.multi[0])

        # It is assumed that all of the data uses the same number of bins with the same bin width
        r = system['r']
        bin_width = system['bw']

        results = np.zeros([n, len(regions), len(r)])
        # results = np.zeros([n, 4, len(r)])

        # results[0, :3, :] = system['results']

        for i in range(n - 1):
            system = np.load(args.multi[i])
            results[i, :, :] = system['results']

        system = np.load(args.multi[-1])
        results[-1, :, :] = system['results'][:3, :]

        outline = np.zeros([4, n, r.shape[0]*2 + 2, 2])
        for i in range(len(regions)):
            print(i)
            fig = plt.figure(i)
            for j in range(n):
                #plt.bar(r, results[j, i, :], bin_width, color=colors[j], alpha=1, label=names[j])

                half_width = bin_width / 2
                # outline = np.zeros([r.shape[0]*2 + 2, 2])
                outline[i, j, 0, 0] = r[0] - half_width
                outline[i, j, -1, 0] = r[-1] + half_width
                outline[i, j, 1:-1:2, 0] = r - half_width
                outline[i, j, 2:-1:2, 0] = r + half_width
                outline[i, j, 1:-1:2, 1] = results[j, i, :]
                outline[i, j, 2:-1:2, 1] = results[j, i, :]
                if i == 0:
                    plt.plot(outline[i, j, :-1, 0], outline[i, j, :-1, 1], color=colors[j], linewidth=2, label=names[j])
                else:
                    plt.plot(outline[i, j, 1:, 0], outline[i, j, 1:, 1], color=colors[j], linewidth=2,
                             label=names[j])
            # plt.title(regions[i], fontsize=14)

            #plt.legend()#fontsize=12)

            plt.ylabel('Component Number Density \n (number/nm$^3$)', fontsize=18)
            plt.xlabel('Distance from pore center, r (nm)', fontsize=18)
            plt.axes().tick_params(labelsize=14)
            # plt.ylim([0, 0.6])
            plt.tight_layout()
            plt.savefig("%s_density.pdf" % regions[i])

        plt.show()

    if args.solvate:
        colors = ['xkcd:red', 'xkcd:green', 'blue', 'xkcd:yellow']
    else:
        colors = ['red', 'green', 'blue']

    if not args.load:

        R = Region(args.traj, args.gro, args.components, begin=args.begin, end=args.end)

        print('Loading trajectory...', end="")
        t = md.load('%s' % args.traj, top='%s' % args.gro)[args.begin:args.end]
        print('done')

        box = t.unitcell_vectors
        nT = t.n_frames
        npores = 4
        r_max = 0

        if args.solvate:
            results = np.zeros([len(regions) + 1, args.bins])
        else:
            results = np.zeros([len(regions), args.bins])

        #keep = [a.index for a in t.topology.atoms if a.residue.name != 'HOH']  # everything kept if system not solvated
        components = ['C', 'C1', 'C2', 'C3', 'C4', 'C5']
        comp = [a.index for a in t.topology.atoms if a.name in components]

        p_centers = physical.avg_pore_loc(npores, t.xyz[:, comp, :], box)

        for i, reg in enumerate(regions):

            print('Calculating number density of %s region' % reg)

            pos = p2p.restrict_atoms(t, reg)  # restrict trajectory to region

            p = duplicate(pos, t.unitcell_vectors)  # duplicate things periodically

            equil = 0
            density, r, bin_width = compdensity(pos, p_centers, equil, box, pores=npores, nbins=args.bins)

            results[i, :] = density

            plt.bar(r, density, bin_width, color=colors[i], alpha=0.6, label=reg)

        np.savez_compressed("regional_density", results=results, r=r, bw=bin_width, box=t.unitcell_vectors)
        print('Arrays saved as regional_density.npz')

        if args.solvate:

            print('Calculating number density of solvent')
            keep = [a.index for a in t.topology.atoms if a.residue.name == 'HOH' and a.name == 'O']
            pos = t.xyz[:, keep, :]
            p = duplicate(pos, t.unitcell_vectors)
            equil = 0
            density, r, bin_width = Structure_char.compdensity(pos, p_centers, equil, t.unitcell_vectors, pores=npores, buffer=0, nbins=args.bins)
            results[-1, :] = density
            plt.bar(r, density, bin_width, color='xkcd:gold', alpha=0.75, label='Water')

    else:

        d = np.load(args.load)

        bw = d['bw']
        r = d['r']
        results = d['results']
        box = d['box']

        if args.solvate:
            regions.append('Water')

        for i in range(results.shape[0]):
            plt.bar(r, results[i, :], bw, color=colors[i], alpha=0.75, label=regions[i])

    un_normalize = np.zeros_like(results)

    annulus_area = []
    dr = r[-1] - r[-2]  # width of annulus
    # r from above is calculated as the average of the edges of the bins. So I need to reverse that here
    for i in r:
        inner_r = i - (dr/2)
        outer_r = i + (dr/2)
        annulus_area.append(np.pi*(outer_r**2 - inner_r**2))

    for i in range(results.shape[0]):
        un_normalize[i, :] = results[i, :] * annulus_area

    r_cut = 0.6
    cut_index = 0
    while r[cut_index] < r_cut:
        cut_index += 1

    print('Percentage of sodium within %s nm of pore center: %2.2f %%' % (r_cut, 100*(sum(un_normalize[2, :cut_index])/sum(un_normalize[2, :]))))

    r_cut = 0.6
    cut_index = len(r) - 1
    while r[cut_index] > r_cut:
        cut_index -= 1

    print('Percentage of tails within %s nm of pore center: %2.2f %%' % (r_cut, 100*(sum(un_normalize[0, :cut_index])/sum(un_normalize[0, :]))))

    r_cut_pore = r_cut
    pore_cut_index = cut_index

    r_cut = 1
    cut_index = len(r) - 1
    while r[cut_index] > r_cut:
        cut_index -= 1

    print('Percentage of tails between %s and %s nm of pore center: %2.2f %%' % (r_cut_pore, r_cut,
        100*(sum(un_normalize[0, pore_cut_index:cut_index])/sum(un_normalize[0, :]))))

    print('Maximum of Head group region: r = %s' % r[np.argmax(results[1, :])])

    plt.legend(prop={'size':14}, loc=1)
    plt.ylabel('Component number density (count/nm$^3$)', fontsize=14)
    plt.xlabel('Distance from pore center (nm)', fontsize=14)
    plt.axes().tick_params(labelsize=14)
    plt.xlim(0, 1.5)
    plt.tight_layout()
    plt.savefig("regional_density.png")
    plt.show()


"""Sequential Monte Carlo sampler also known as
Adaptive Transitional Marcov Chain Monte Carlo sampler.

Runs on any pymc3 model.

Created on March, 2016

Various significant updates July, August 2016

Made pymc3 compatible November 2016
Renamed to SMC and further improvements March 2017

@author: Hannes Vasyura-Bathke
"""
import numpy as np
import pymc3 as pm
from tqdm import tqdm

import theano
import copy
import warnings

from ..model import modelcontext
from ..vartypes import discrete_types
from ..theanof import inputvars, make_shared_replacements, join_nonshared_inputs
import numpy.random as nr

from .arraystep import metrop_select
from ..backends import smc_text as atext

__all__ = ['SMC', 'ATMIP_sample']

EXPERIMENTAL_WARNING = "Warning: SMC is an experimental step method, and not yet"\
    " recommended for use in PyMC3!"

class Proposal(object):
    """Proposal distributions modified from pymc3 to initially create all the
    Proposal steps without repeated execution of the RNG - significant speedup!

    Parameters
    ----------
    s : :class:`numpy.ndarray`
    """
    def __init__(self, s):
        self.s = np.atleast_1d(s)

class MultivariateNormalProposal(Proposal):
    def __call__(self, num_draws=None):
        return np.random.multivariate_normal(
                mean=np.zeros(self.s.shape[0]), cov=self.s, size=num_draws)

proposal_dists = {
    'MultivariateNormal': MultivariateNormalProposal,
        }

def choose_proposal(proposal_name, scale=1.):
    """Initialise and select proposal distribution.

    Parameters
    ----------
    proposal_name : string
        Name of the proposal distribution to initialise
    scale : float or :class:`numpy.ndarray`

    Returns
    -------
    class:`pymc3.Proposal` Object
    """
    return proposal_dists[proposal_name](scale)

class SMC(atext.ArrayStepSharedLLK):
    """Adaptive Transitional Markov-Chain Monte-Carlo sampler class.

    Creates initial samples and framework around the (C)ATMIP parameters

    Parameters
    ----------
    vars : list
        List of variables for sampler
    out_vars : list
        List of output variables for trace recording. If empty unobserved_RVs
        are taken.
    n_chains : int
        Number of chains per stage has to be a large number
        of number of n_jobs (processors to be used) on the machine.
    scaling : float
        Factor applied to the proposal distribution i.e. the step size of the
        Markov Chain
    covariance : :class:`numpy.ndarray`
        (n_chains x n_chains)
        Initial Covariance matrix for proposal distribution,
        if None - identity matrix taken
    likelihood_name : string
        name of the :class:`pymc3.determinsitic` variable that contains the
        model likelihood - defaults to 'like'
    proposal_name :
        Type of proposal distribution, see
        smc.proposal_dists.keys() for options
    tune : boolean
        Flag for adaptive scaling based on the acceptance rate
    tune_interval : int
        Number of steps to tune for
    coef_variation : scalar, float
        Coefficient of variation, determines the change of beta
        from stage to stage, i.e.indirectly the number of stages,
        low coef_variation --> slow beta change,
        results in many stages and vice verca (default: 1.)
    check_bound : boolean
        Check if current sample lies outside of variable definition
        speeds up computation as the forward model wont be executed
        default: True
    model : :class:`pymc3.Model`
        Optional model for sampling step.
        Defaults to None (taken from context).
    random_seed : int
        Optional to set the random seed.  Necessary for initial population.

    References
    ----------
    .. [Ching2007] Ching, J. and Chen, Y. (2007).
        Transitional Markov Chain Monte Carlo Method for Bayesian Model
        Updating, Model Class Selection, and Model Averaging.
        J. Eng. Mech., 10.1061/(ASCE)0733-9399(2007)133:7(816), 816-832.
        `link <http://ascelibrary.org/doi/abs/10.1061/%28ASCE%290733-9399
        %282007%29133:7%28816%29>`__
    """
    default_blocked = True

    def __init__(self, vars=None, out_vars=None, n_chains=100, scaling=1., covariance=None,
                 likelihood_name='like', proposal_name='MultivariateNormal', tune=True,
                 tune_interval=100, coef_variation=1., check_bound=True, model=None,
                 random_seed=-1):
        warnings.warn(EXPERIMENTAL_WARNING)
        if random_seed != -1:
            nr.seed(random_seed)

        model = modelcontext(model)

        if vars is None:
            vars = model.vars

        vars = inputvars(vars)

        if out_vars is None:
            out_vars = model.unobserved_RVs

        out_varnames = [out_var.name for out_var in out_vars]

        if covariance is None and proposal_name == 'MultivariateNormal':
            self.covariance = np.eye(sum(v.dsize for v in vars))
            scale = self.covariance
        elif covariance is None:
            scale = np.ones(sum(v.dsize for v in vars))
        else:
            scale = covariance

        self.proposal_name = proposal_name
        self.proposal_dist = choose_proposal(self.proposal_name, scale=scale)

        self.scaling = np.atleast_1d(scaling)
        self.tune = tune
        self.check_bnd = check_bound
        self.tune_interval = tune_interval
        self.steps_until_tune = tune_interval

        self.proposal_samples_array = self.proposal_dist(n_chains)

        self.stage_sample = 0
        self.accepted = 0

        self.beta = 0
        self.stage = 0
        self.chain_index = 0
        self.resampling_indexes = np.arange(n_chains)

        self.coef_variation = coef_variation
        self.n_chains = n_chains
        self.likelihoods = np.zeros(n_chains)

        self.likelihood_name = likelihood_name
        self._llk_index = out_varnames.index(likelihood_name)
        self.discrete = np.concatenate([[v.dtype in discrete_types] * (v.dsize or 1) for v in vars])
        self.any_discrete = self.discrete.any()
        self.all_discrete = self.discrete.all()

        # create initial population
        self.population = []
        self.array_population = np.zeros(n_chains)
        for _ in range(self.n_chains):
            self.population.append(pm.Point({v.name: v.random() for v in vars}, model=model))

        self.chain_previous_lpoint = copy.deepcopy(self.population)

        shared = make_shared_replacements(vars, model)
        self.logp_forw = logp_forw(out_vars, vars, shared)
        self.check_bnd = logp_forw([model.varlogpt], vars, shared)

        super(SMC, self).__init__(vars, out_vars, shared)

    def astep(self, q0):
        if self.stage == 0:
            l_new = self.logp_forw(q0)
            q_new = q0

        else:
            if not self.stage_sample:
                self.proposal_samples_array = self.proposal_dist(self.n_steps)

            if not self.steps_until_tune and self.tune:
                # Tune scaling parameter
                self.scaling = tune(self.accepted / float(self.tune_interval))
                # Reset counter
                self.steps_until_tune = self.tune_interval
                self.accepted = 0

            delta = self.proposal_samples_array[self.stage_sample, :] * self.scaling

            if self.any_discrete:
                if self.all_discrete:
                    delta = np.round(delta, 0)
                    q0 = q0.astype(int)
                    q = (q0 + delta).astype(int)
                else:
                    delta[self.discrete] = np.round(delta[self.discrete], 0).astype(int)
                    q = q0 + delta
                    q = q[self.discrete].astype(int)
            else:
                q = q0 + delta

            l0 = self.chain_previous_lpoint[self.chain_index]

            if self.check_bnd:
                varlogp = self.check_bnd(q)

                if np.isfinite(varlogp):
                    logp = self.logp_forw(q)
                    q_new, accepted = metrop_select(
                        self.beta * (logp[self._llk_index] - l0[self._llk_index]), q, q0)

                    if accepted:
                        self.accepted += 1
                        l_new = logp
                        self.chain_previous_lpoint[self.chain_index] = l_new
                    else:
                        l_new = l0
                else:
                    q_new = q0
                    l_new = l0

            else:
                logp = self.logp_forw(q)
                q_new, accepted = metrop_select(
                    self.beta * (logp[self._llk_index] - l0[self._llk_index]), q, q0)

                if accepted:
                    self.accepted += 1
                    l_new = logp
                    self.chain_previous_lpoint[self.chain_index] = l_new
                else:
                    l_new = l0

            self.steps_until_tune -= 1
            self.stage_sample += 1

            # reset sample counter
            if self.stage_sample == self.n_steps:
                self.stage_sample = 0

        return q_new, l_new

    def calc_beta(self):
        """Calculate next tempering beta and importance weights based on
        current beta and sample likelihoods.

        Returns
        -------
        beta(m+1) : scalar, float
            tempering parameter of the next stage
        beta(m) : scalar, float
            tempering parameter of the current stage
        weights : :class:`numpy.ndarray`
            Importance weights (floats)
        """
        low_beta = self.beta
        up_beta = 2.
        old_beta = self.beta

        while up_beta - low_beta > 1e-6:
            current_beta = (low_beta + up_beta) / 2.
            temp = np.exp((current_beta - self.beta) * (self.likelihoods - self.likelihoods.max()))
            cov_temp = np.std(temp) / np.mean(temp)
            if cov_temp > self.coef_variation:
                up_beta = current_beta
            else:
                low_beta = current_beta

        beta = current_beta
        weights = temp / np.sum(temp)
        return beta, old_beta, weights

    def calc_covariance(self):
        """Calculate trace covariance matrix based on importance weights.

        Returns
        -------
        cov : :class:`numpy.ndarray`
            weighted covariances (NumPy > 1.10. required)
        """
        cov = np.cov(self.array_population, aweights=self.weights.ravel(), bias=False, rowvar=0)
        if np.isnan(cov).any() or np.isinf(cov).any():
            raise ValueError('Sample covariances not valid! Likely "n_chains" is too small!')
        return np.atleast_2d(cov)

    def select_end_points(self, mtrace):
        """Read trace results (variables and model likelihood) and take end points
        for each chain and set as start population for the next stage.

        Parameters
        ----------
        mtrace : :class:`.base.MultiTrace`

        Returns
        -------
        population : list
            of :func:`pymc3.Point` dictionaries
        array_population : :class:`numpy.ndarray`
            Array of trace end-points
        likelihoods : :class:`numpy.ndarray`
            Array of likelihoods of the trace end-points
        """
        array_population = np.zeros((self.n_chains, self.ordering.dimensions))
        n_steps = len(mtrace)

        # collect end points of each chain and put into array
        for var, slc, shp, _ in self.ordering.vmap:
            slc_population = mtrace.get_values(varname=var, burn=n_steps - 1, combine=True)
            if len(shp) == 0:
                array_population[:, slc] = np.atleast_2d(slc_population).T
            else:
                array_population[:, slc] = slc_population
        # get likelihoods
        likelihoods = mtrace.get_values(varname=self.likelihood_name,
                                        burn=n_steps - 1,
                                        combine=True)

        # map end array_endpoints to dict points
        population = [self.bij.rmap(row) for row in array_population]

        return population, array_population, likelihoods

    def get_chain_previous_lpoint(self, mtrace):
        """Read trace results and take end points for each chain and set as
        previous chain result for comparison of metropolis select.

        Parameters
        ----------
        mtrace : :class:`.base.MultiTrace`

        Returns
        -------
        chain_previous_lpoint : list
            all unobservedRV values, including dataset likelihoods
        """
        array_population = np.zeros((self.n_chains, self.lordering.dimensions))
        n_steps = len(mtrace)
        for var, (_, slc, shp, _) in zip(mtrace.varnames, self.lordering.vmap):
            slc_population = mtrace.get_values(varname=var, burn=n_steps - 1, combine=True)
            if len(shp) == 0:
                array_population[:, slc] = np.atleast_2d(slc_population).T
            else:
                array_population[:, slc] = slc_population

        return [self.lij.rmap(row) for row in array_population[self.resampling_indexes, :]]

    def mean_end_points(self):
        """Calculate mean of the end-points and return point.

        Returns
        -------
        Dictionary of trace variables
        """
        return self.bij.rmap(self.array_population.mean(axis=0))

    def resample(self):
        """Resample pdf based on importance weights.
        based on Kitagawas deterministic resampling algorithm.

        Returns
        -------
        outindex : :class:`numpy.ndarray`
            Array of resampled trace indexes
        """
        parents = np.arange(self.n_chains)
        N_childs = np.zeros(self.n_chains, dtype=int)

        cum_dist = np.cumsum(self.weights)
        u = (parents + np.random.rand()) / self.n_chains
        j = 0
        for i in parents:
            while u[i] > cum_dist[j]:
                j += 1

            N_childs[j] += 1

        indx = 0
        outindx = np.zeros(self.n_chains, dtype=int)
        for i in parents:
            if N_childs[i] > 0:
                for j in range(indx, (indx + N_childs[i])):
                    outindx[j] = parents[i]

            indx += N_childs[i]

        return outindx

def ATMIP_sample(n_steps, step=None, start=None, homepath=None, chain=0, stage=0, n_jobs=1,
                 tune=None, progressbar=False, model=None, random_seed=-1, rm_flag=False):
    """(C)ATMIP sampling algorithm (Cascading - (C) not always relevant)

    Samples the solution space with n_chains of Metropolis chains, where each
    chain has n_steps iterations. Once finished, the sampled traces are
    evaluated:

    (1) Based on the likelihoods of the final samples, chains are weighted
    (2) the weighted covariance of the ensemble is calculated and set as new
        proposal distribution
    (3) the variation in the ensemble is calculated and the next tempering
        parameter (beta) calculated
    (4) New n_chains Metropolis chains are seeded on the traces with high
        weight for n_steps iterations
    (5) Repeat until beta > 1.

    Parameters
    ----------
    n_steps : int
        The number of samples to draw for each Markov-chain per stage
    step : :class:`SMC`
        SMC initialisation object
    start : List of dictionaries
        with length of (n_chains)
        Starting points in parameter space (or partial point)
        Defaults to random draws from variables (defaults to empty dict)
    homepath : string
        Result_folder for storing stages, will be created if not existing.
    chain : int
        Chain number used to store sample in backend. If `n_jobs` is
        greater than one, chain numbers will start here.
    stage : int
        Stage where to start or continue the calculation. It is possible to
        continue after completed stages (stage should be the number of the
        completed stage + 1). If None the start will be at stage = 0.
    n_jobs : int
        The number of cores to be used in parallel. Be aware that theano has
        internal parallelisation. Sometimes this is more efficient especially
        for simple models.
        step.n_chains / n_jobs has to be an integer number!
    tune : int
        Number of iterations to tune, if applicable (defaults to None)
    progressbar : bool
        Flag for displaying a progress bar
    model : :class:`pymc3.Model`
        (optional if in `with` context) has to contain deterministic
        variable name defined under step.likelihood_name' that contains the
        model likelihood
    random_seed : int or list of ints
        A list is accepted, more if `n_jobs` is greater than one.
    rm_flag : bool
        If True existing stage result folders are being deleted prior to
        sampling.

    References
    ----------
    .. [Minson2013] Minson, S. E. and Simons, M. and Beck, J. L., (2013),
        Bayesian inversion for finite fault earthquake source models
        I- Theory and algorithm. Geophysical Journal International, 2013,
        194(3), pp.1701-1726,
        `link <https://gji.oxfordjournals.org/content/194/3/1701.full>`__
    """
    warnings.warn(EXPERIMENTAL_WARNING)

    model = modelcontext(model)
    step.n_steps = int(n_steps)
    if random_seed != -1:
        nr.seed(random_seed)

    if step is None:
        raise TypeError('Argument `step` has to be a SMC step object.')

    if homepath is None:
        raise TypeError('Argument `homepath` should be path to result_directory.')

    if n_jobs > 1:
        if not (step.n_chains / float(n_jobs)).is_integer():
            raise TypeError('n_chains / n_jobs has to be a whole number!')

    if start is not None:
        if len(start) != step.n_chains:
            raise TypeError('Argument `start` should have dicts equal the '
                            'number of chains (`step.n_chains`)')
        else:
            step.population = start

    if not any(step.likelihood_name in var.name for var in model.deterministics):
        raise TypeError('Model (deterministic) variables need to contain a variable %s '
                        'as defined in `step`.' % step.likelihood_name)

    stage_handler = atext.TextStage(homepath)

    if progressbar and n_jobs > 1:
        progressbar = False

    if stage == 0:
        # continue or start initial stage
        step.stage = stage
        draws = 1
    else:
        step = stage_handler.load_atmip_params(stage)
        draws = step.n_steps

    stage_handler.clean_directory(stage, None, rm_flag)
    with model:
        chains = stage_handler.recover_existing_results(stage, draws, step, n_jobs)
        if chains is not None:
            rest = len(chains) % n_jobs
            if rest > 0:
                pm._log.info('Fixing %i chains ...' % rest)
                chains, rest_chains = chains[:-rest], chains[-rest:]
                # process traces that are not a multiple of n_jobs
                sample_args = {
                    'draws': draws,
                    'step': step,
                    'stage_path': stage_handler.stage_path(stage),
                    'progressbar': progressbar,
                    'model': model,
                    'n_jobs': rest,
                    'chains': rest_chains}

                _iter_parallel_chains(**sample_args)
                pm._log.info('Back to normal!')

    with model:
        while step.beta < 1:
            if step.stage == 0:
                # Initial stage
                pm._log.info('Sample initial stage: ...')
                draws = 1
            else:
                draws = n_steps

            pm._log.info('Beta: %f Stage: %i' % (step.beta, step.stage))

            # Metropolis sampling intermediate stages
            chains = stage_handler.clean_directory(stage, chains, rm_flag)
            sample_args = {
                    'draws': draws,
                    'step': step,
                    'stage_path': stage_handler.stage_path(step.stage),
                    'progressbar': progressbar,
                    'model': model,
                    'n_jobs': n_jobs,
                    'chains': chains}

            _iter_parallel_chains(**sample_args)

            mtrace = stage_handler.load_multitrace(step.stage)

            step.population, step.array_population, step.likelihoods = step.select_end_points(
                mtrace)
            step.beta, step.old_beta, step.weights = step.calc_beta()

            if step.beta > 1.:
                pm._log.info('Beta > 1.: %f' % step.beta)
                step.beta = 1.
                stage_handler.dump_atmip_params(step)
                if stage == -1:
                    chains = []
                else:
                    chains = None
            else:
                step.covariance = step.calc_covariance()
                step.proposal_dist = choose_proposal(step.proposal_name, scale=step.covariance)
                step.resampling_indexes = step.resample()
                step.chain_previous_lpoint = step.get_chain_previous_lpoint(mtrace)

                stage_handler.dump_atmip_params(step)

                step.stage += 1
                del(mtrace)

        # Metropolis sampling final stage
        pm._log.info('Sample final stage')
        step.stage = -1
        temp = np.exp((1 - step.old_beta) * (step.likelihoods - step.likelihoods.max()))
        step.weights = temp / np.sum(temp)
        step.covariance = step.calc_covariance()
        step.proposal_dist = choose_proposal(step.proposal_name, scale=step.covariance)
        step.resampling_indexes = step.resample()
        step.chain_previous_lpoint = step.get_chain_previous_lpoint(mtrace)

        sample_args['step'] = step
        sample_args['stage_path'] = stage_handler.stage_path(step.stage)
        sample_args['chains'] = chains
        _iter_parallel_chains(**sample_args)

        stage_handler.dump_atmip_params(step)
        return stage_handler.load_multitrace(step.stage, model=model)

def _sample(draws, step=None, start=None, trace=None, chain=0, tune=None,
            progressbar=True, model=None, random_seed=-1):

    sampling = _iter_sample(draws, step, start, trace, chain,
                            tune, model, random_seed)

    if progressbar:
        sampling = tqdm(sampling, total=draws)

    try:
        for strace in sampling:
            pass

    except KeyboardInterrupt:
        pass

    return chain

def _iter_sample(draws, step, start=None, trace=None, chain=0, tune=None,
                 model=None, random_seed=-1):
    """Modified from :func:`pymc3.sampling._iter_sample` to be more efficient with SMC algorithm."""
    model = modelcontext(model)
    draws = int(draws)
    if draws < 1:
        raise ValueError('Argument `draws` should be above 0.')

    if start is None:
        start = {}

    if random_seed != -1:
        nr.seed(random_seed)

    try:
        step = pm.step_methods.CompoundStep(step)
    except TypeError:
        pass

    point = pm.Point(start, model=model)
    step.chain_index = chain
    trace.setup(draws, chain)
    for i in range(draws):
        if i == tune:
            step = pm.sampling.stop_tuning(step)

        point, out_list = step.step(point)
        trace.record(out_list)
        yield trace

def _work_chain(work):
    """Wrapper function for parallel execution of _sample i.e. the Markov Chains.

    Parameters
    ----------
    work : List
        Containing all the information that is unique for each Markov Chain
        i.e. [:class:'SMC', chain_number(int),
        sampling index(int), start_point(dictionary)]

    Returns
    -------
    chain : int
        Index of chain that has been sampled
    """
    return _sample(*work)

def _iter_parallel_chains(draws, step, stage_path, progressbar, model, n_jobs, chains=None):
    """Do Metropolis sampling over all the chains with each chain being
    sampled 'draws' times. Parallel execution according to n_jobs.
    """
    if chains is None:
        chains = range(step.n_chains)

    pm._log.info('Initialising chain traces ...')

    max_int = np.iinfo(np.int32).max
    random_seeds = nr.randint(1, max_int, size=len(chains))

    pm._log.info('Sampling ...')
    work = [(draws,
             step,
             step.population[step.resampling_indexes[chain]],
             atext.TextChain(stage_path, model=model),
             chain,
             None,
             False,
             model,
             rseed) for chain, rseed in
            zip(chains, random_seeds)]

    if draws < 10:
        chunksize = n_jobs
    else:
        chunksize = 1

    p = atext.paripool(_work_chain, work, chunksize=chunksize, nprocs=n_jobs)

    if n_jobs == 1 and progressbar:
        p = tqdm(p, total=len(chains))

    for _ in p:
        pass

def tune(acc_rate):
    """Tune adaptively based on the acceptance rate.

    Parameters
    ----------
    acc_rate: scalar, float
        Acceptance rate of the Metropolis sampling

    Returns
    -------
    scaling: scalar float
    """
    # a and b after Muto & Beck 2008 .
    a = 1. / 9
    b = 8. / 9
    return np.power((a + (b * acc_rate)), 2)

def logp_forw(out_vars, vars, shared):
    """Compile Theano function of the model and the input and output variables.

    Parameters
    ----------
    out_vars : List
        containing :class:`pymc3.Distribution` for the output variables
    vars : List
        containing :class:`pymc3.Distribution` for the input variables
    shared : List
        containing :class:`theano.tensor.Tensor` for dependend shared data
    """
    out_list, inarray0 = join_nonshared_inputs(out_vars, vars, shared)
    f = theano.function([inarray0], out_list)
    f.trust_input = True
    return f

import numpy,import theano,import pymc3,from pymc3


'''Copyright (c) 2021 James Gayvert

    Permission is hereby granted, free of charge, to any person obtaining a copy
    of this software and associated documentation files (the "Software"), to deal
    in the Software without restriction, including without limitation the rights
    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    copies of the Software, and to permit persons to whom the Software is
    furnished to do so, subject to the following conditions:

    The above copyright notice and this permission notice shall be included in all
    copies or substantial portions of the Software.

    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    SOFTWARE.'''

import numpy as np
from numpy import linalg as LA
import warnings

# https://physics.nist.gov/cgi-bin/cuu/Value?hrev
au2eV = 27.211386245988

def _delete_indices(H0, W, exclude_states):
    exclude_states = np.unique(exclude_states)
    exclude_states = sorted(exclude_states, reverse=True)
    for i in exclude_states:
        H0 = np.delete(H0, i, axis=0)
        H0 = np.delete(H0, i, axis=1)
        W = np.delete(W, i, axis=0)
        W = np.delete(W, i, axis=1)
    return H0, W

class Root():
    '''
    Root obtained from diagonalizing CAP Hamiltonian at given value of eta.

    Attributes
    ----------
    eta: float
        CAP strength parameter
    energy: float
        Total energy of state
    eigvc: list of float
        Eigenvector for state
    '''
    def __init__(self, energy, eta, eigvc):
        '''
        Initializes root object.

        Parameters
        ----------
        eta: float
            CAP strength parameter
        energy: float
            Total energy of state
        eigvc: list of float
            Eigenvector for state
        '''
        self.eta = eta
        self.energy = energy
        self.eigvc = eigvc

class EigenvalueTrajectory():
    '''
    Eigenvalue trajectory generated by repeated diagonalizations of CAP Hamiltonian over range of eta values.

    States are tracked using either eigenvector overlap or energy criterion. Corrected energies(U) are obtained as:
    U = E - eta*dE/deta.

    Attributes
    ----------
    states: list of :class:`~pyopencap.analysis.Root`
        List of states in trajectory.
    uncorrected_energies: list of float
        List of uncorrected energies in trajectory.
    corrected_energies: list of float
        List of corrected energies in trajectory.
    etas: list of float
        List of CAP strengths in trajectory.
    '''
    def __init__(self, state_idx, init_roots, tracking="overlap"):
        '''
        Initializes EigenvalueTrajectory object, which tracks a state starting from the first diagonalization at eta = 0.

        Parameters
        ----------
        state_idx: int
            Index of state to track
        init_roots: list of :class:`~pyopencap.analysis.Root`
            Initial set of roots generated by diagonalization at eta = 0
        tracking: str, default="overlap"
            Method to use to track the state
        '''
        self._last = init_roots[state_idx]
        self.states = [self._last]
        self.uncorrected_energies = [self._last.energy]
        self.corrected_energies = []
        self.etas = [0.0]
        if tracking == "overlap":
            self._add_state = self._overlap_tracking
        elif tracking == "energy":
            self._add_state = self._energy_tracking
        else:
            raise RuntimeError(
                "Invalid choice of tracking. Choose either overlap or energy.")

    def _overlap_tracking(self, states):
        maxo = 0.0
        state_idx = -1
        for st in states:
            ov = np.abs(np.dot(st.eigvc, self._last.eigvc))
            if ov > maxo:
                maxo = ov
                cur = st
        self._last = cur
        self.states.append(cur)
        self.uncorrected_energies.append(cur.energy)
        self.etas.append(cur.eta)

    def _energy_tracking(self, states):
        min = float("inf")
        cur = -1
        for st in states:
            if np.absolute(st.energy - self._last.energy) < min or cur == -1:
                cur = st
                min = np.absolute(st.energy - self._last.energy)
        self._last = cur
        self.states.append(cur)
        self.uncorrected_energies.append(cur.energy)
        self.etas.append(cur.eta)

    def _calculate_corrected_energies(self):
        derivs = list(
            np.gradient(self.uncorrected_energies) / np.gradient(self.etas))
        # E_cor = E_uc - eta * dE/deta
        for i in range(0, len(self.states)):
            self.corrected_energies.append(self.uncorrected_energies[i] -
                                           self.etas[i] * derivs[i])

    def find_eta_opt(self,
                     corrected=False,
                     start_idx=1,
                     end_idx=-1,
                     ref_energy=0.0,
                     units="au"):
        '''
        Finds optimal cap strength parameter for eigenvalue trajectory, as defined by eta_opt = min|eta*dE/deta|.

        The range of self.etas[start_idx:end_idx] (in python slice notation) is searched for the optimal value of CAP strength parameter.

        Parameters
        ----------
        corrected: bool, default=False
            Set to true if searching for stationary point on corrected trajectory
        start_idx: int, default=1
            Starting slice index
        end_idx: int, default=1
            Ending slice index
        ref_energy: float, default=0.0
            Reference energy to define excitation energy.
        units: str, default="au"
            Options are "au" or "eV"

        Returns
        -------
        E_res: complex float
            Complex energy at optimal value of eta
        eta_opt: float
            Optimal value of eta

        '''
        if units == "au":
            scaling_factor = 1.0
        elif units == "eV":
            scaling_factor = au2eV
        else:
            raise RuntimeError("Units should be either au or eV.")
        if corrected:
            derivs = np.array(self.etas) * np.absolute(
                np.gradient(self.corrected_energies) / np.gradient(self.etas))
            min_val = np.min(derivs[start_idx:end_idx])
            opt_idx = list(derivs).index(min_val)
            E = self.corrected_energies[opt_idx]
            E = (E - ref_energy) * scaling_factor
            return E, self.etas[opt_idx]
        else:
            derivs = np.array(self.etas) * np.absolute(
                np.gradient(self.uncorrected_energies) /
                np.gradient(self.etas))
            min_val = np.min(derivs[start_idx:end_idx])
            opt_idx = list(derivs).index(min_val)
            E = self.uncorrected_energies[opt_idx]
            E = (E - ref_energy) * scaling_factor
            return E, self.etas[opt_idx]

    def energies_ev(self, ref_energy, corrected=False):
        '''
        Returns excitation energies of all states in trajectory in eV with respect to specified reference energy.

        Parameters
        ----------
        ref_energy: float
            Reference energy
        corrected: bool, default=False
            Set to true if analyzing corrected trajectory

        Returns
        -------
        E_eV: list of floats
            Excitation energies in eV with respect to specified reference energy.
        '''
        E_eV = []
        if corrected:
            E_hartree = self.corrected_energies
        else:
            E_hartree = self.uncorrected_energies
        for E in E_hartree:
            E_eV.append((E - ref_energy) * au2eV)
        return E_eV

    def get_energy(self, eta, corrected=False, ref_energy=0.0, units="au"):
        '''
        Returns total energy at given value of eta.

        Note that if the eta provided is not in self.etas, the nearest value will be used.

        Parameters
        ----------
        eta: float
            Value of CAP strength parameter
        corrected: bool, default=False
            Set to true if analyzing corrected trajectory
        ref_energy: float, default=0.0
            Reference energy to define excitation energy.
        units: str, default="au"
            Options are "au" or "eV"

        Returns
        -------
        E: float
            Energy at given value of eta
        '''
        if units == "au":
            scaling_factor = 1.0
        elif units == "eV":
            scaling_factor = au2eV
        else:
            raise RuntimeError("Invalid unit. Options are au or eV.")
        if eta in self.etas:
            idx = self.etas.index(eta)
        else:
            idx = np.nanargmin(np.abs(np.asarray(self.etas) - eta))
            warnings.warn("Warning: "+str(eta) + " is not in the list of values for this trajectory." \
                +" Defaulting to nearest value of " + str(self.etas[idx]))
        if corrected:
            return (self.corrected_energies[idx] - ref_energy) * scaling_factor
        else:
            return (self.uncorrected_energies[idx] -
                    ref_energy) * scaling_factor

    def get_logarithmic_velocities(self, corrected=False):
        '''
        Returns eta*dE/deta for each point on eigenvalue trajectory.

        Useful for plotting when dealing with multiple potential stationary points.

        Parameters
        ----------
        corrected: bool, default=False
            Set to true if analyzing corrected trajectory

        Returns
        -------
        derivs: np.array of float
            eta*dE/deta for each point on eigenvalue trajectory
        '''
        if corrected:
            energies = self.corrected_energies
        else:
            energies = self.uncorrected_energies
        return np.array(self.etas) * np.absolute(
            np.gradient(energies) / np.gradient(self.etas))

class CAPHamiltonian():
    '''
    Projected CAP Hamiltonian handler for generating eigenvalue trajectories.

    The instance variables H0,W etc. are only set after `run_trajectory` is executed. The original matrices
    passed/parsed when the object is constructed are stored in _H0, _W, etc. This makes it easy to run multiple
    trajectories with different states included in the projection scheme without having to construct a new object.

    Attributes
    ----------
    H0: np.ndarray of float: default=None
        Zeroth order Hamiltonian in state basis
    W: np.ndarray of float: default=None
        CAP matrix in state basis. -1 prefactor is assumed.
    nstates: int
        Number of states
    total_energies: list of float
        Energies of all states found by repeated diagonalization of CAP Hamiltonian
    etas: list of float
        List of CAP strengths in trajectory.
    cap_lambda: float
        Real CAP strength used for continuum remover CAP. Set to 0.0 by default.
    '''
    def _init_from_matrices(self, H0, W):
        H0 = np.array(H0)
        W = np.array(W)
        assert H0.shape == W.shape
        self._H0 = H0
        self._W = W
        self._nstates = len(H0)

    def __init__(self,
                 pc=None,
                 H0=None,
                 W=None,
                 output=None,
                 irrep="",
                 onset=""):
        '''
        Initializes CAPHamiltonian object from H0 and W matrix in state basis.

        Object can be initialized in one of two ways. The user can pass H0 and W directly
        as numpy matrices, or they can specify a path to a properly formatted output file
        (either OpenCAP or Q-Chem) which contains these two matrices.

        W matrix is assumed to already have a -1 prefactor, as that is how OpenCAP output is formatted.

        Parameters
        ----------
        pc: :class:`~pyopencap.CAP`: default=None
            PyOpenCAP CAP object
        H0: np.ndarray of float: default=None
            Zeroth order Hamiltonian in state basis
        W: np.ndarray of float: default=None
            CAP matrix in state basis. -1 prefactor is assumed.
        output: str: default=None
            Path to Q-Chem or OpenCAP output file.
        irrep: str: default=None
            Title of irreducible representation of state of interest. Only compatible with Q-Chem projected CAP-EOM-CC outputs. Set to 'all' to include all symmetries in CAP projection.
        onset: str: default=None
            Title of CAP onset. Only compatible with Q-Chem projected CAP-ADC outputs.

        '''
        if pc is not None:
            self._init_from_matrices(pc.get_H(), pc.get_projected_cap())
        elif H0 is not None and W is not None:
            self._init_from_matrices(H0, W)
        elif output is not None:
            with open(output, 'r') as file:
                filedata = file.readlines()
                for l in filedata:
                    if "Welcome to OpenCAP" in l:
                        self._init_from_opencap_output(output)
                        return
                    elif "Welcome to Q-Chem" in l:
                        self._init_from_qchem_output(output, irrep, onset)
                        return
            raise RuntimeError(
                "Only Q-Chem and OpenCAP outputs are supported.")
        else:
            raise RuntimeError(
                "Error: Either pass a CAP object, H0 and W as matrices, or specify a path to an OpenCAP/Q-Chem output file."
            )

    def _init_from_qchem_eomcc_old(self, output_file, irrep):
        with open(output_file, 'r') as file:
            filedata = file.readlines()
        cur_idx = -1
        nstates = 0
        for i in range(0, len(filedata)):
            if "Performing Projected CAP-EOM calculation for " + str(
                    irrep) in filedata[i]:
                cur_idx = i + 1
                break
        if cur_idx == -1:
            if not irrep == "":
                raise RuntimeError("Error: could not find matrices for " +
                                   str(irrep) + " states in " + output_file)
            else:
                raise RuntimeError("Error: could not find matrices in " +
                                   output_file)
        nstates = int(filedata[cur_idx].split()[-2])
        while "zeroth order hamiltonian" not in filedata[cur_idx].lower(
        ) and cur_idx < len(filedata):
            cur_idx += 1
        cur_idx += 1
        H0 = []
        for i in range(0, nstates):
            l1 = filedata[cur_idx].split()
            l1 = [float(x) for x in l1]
            H0 += l1
            cur_idx += 1
        H0 = np.reshape(H0, (nstates, nstates))
        cur_idx += 1
        W = []
        for i in range(0, nstates):
            l1 = filedata[cur_idx].split()
            l1 = [float(x) for x in l1]
            W += l1
            cur_idx += 1
        W = np.reshape(W, (nstates, nstates))
        assert H0.shape == W.shape
        self._H0 = H0
        self._W = W
        self._nstates = len(H0)

    def _init_from_qchem_eomcc(self, output_file, irrep):
        try:
            self._init_from_qchem_eomcc_new(output_file, irrep)
        except ValueError as e:
            self._init_from_qchem_eomcc_old(output_file, irrep)

    def _init_from_qchem_eomcc_new(self, output_file, irrep):
        with open(output_file, 'r') as file:
            filedata = file.readlines()
        cur_idx = -1
        nstates = 0
        if irrep=="all":
            for i in range(0, len(filedata)):
                if "Total Projected CAP Hamiltonian" in filedata[i]:
                    cur_idx = i + 1
                    break
        else:
            for i in range(0, len(filedata)):
                if "Performing Projected CAP-EOM calculation for " + str(
                        irrep) in filedata[i]:
                    cur_idx = i + 1
                    break
        if cur_idx == -1:
            if not irrep == "":
                raise RuntimeError("Error: could not find matrices for " +
                                   str(irrep) + " states in " + output_file)
            else:
                raise RuntimeError("Error: could not find matrices in " +
                                   output_file)
        nstates = int(filedata[cur_idx].split()[-1])
        cur_idx +=2
        H0 = []
        for i in range(0, nstates):
            l1 = filedata[cur_idx].split()
            l1 = [float(x) for x in l1]
            H0 += l1
            cur_idx += 1
        H0 = np.reshape(H0, (nstates, nstates))
        cur_idx += 1
        W = []
        for i in range(0, nstates):
            l1 = filedata[cur_idx].split()
            l1 = [float(x) for x in l1]
            W += l1
            cur_idx += 1
        W = np.reshape(W, (nstates, nstates))
        assert H0.shape == W.shape
        self._H0 = H0
        self._W = W
        self._nstates = len(H0)

    def _init_from_qchem_adc(self, output_file, onset):
        with open(output_file, 'r') as file:
            filedata = file.readlines()
        cur_idx = -1
        found = False
        for i in range(0, len(filedata)):
            if "Hamiltonian subspace matrix" in filedata[i]:
                cur_idx = i + 2
                found = True
                break
        if not found:
            raise RuntimeError("Error: could not find matrices in " +
                               output_file)
        H0 = []
        done = False
        nstates = 0
        while cur_idx < len(filedata) and not done:
            l1 = filedata[cur_idx].split()
            try:
                l1 = [float(x) for x in l1]
                if len(l1) > 0:
                    H0 += l1
                    cur_idx += 1
                    nstates += 1
                else:
                    done = True
            except:
                done = True
        H0 = np.reshape(H0, (nstates, nstates))
        found = False
        for i in range(0, len(filedata)):
            if "The imaginary part of the CAP subspace matrix, onset=" + str(
                    onset) in filedata[i]:
                cur_idx = i + 2
                found = True
                break
        if not found:
            raise RuntimeError("Error: could not find matrices in " +
                               output_file)
        W = []
        done = False
        nstates = 0
        while cur_idx < len(filedata) and not done:
            l1 = filedata[cur_idx].split()
            try:
                l1 = [float(x) for x in l1]
                if len(l1) > 0:
                    W += l1
                    cur_idx += 1
                    nstates += 1
                else:
                    done = True
            except:
                done = True
        W = np.reshape(W, (nstates, nstates))
        assert H0.shape == W.shape
        self._H0 = H0
        self._W = W
        self._nstates = len(H0)

    def _init_from_qchem_output(self, output_file, irrep, onset):
        with open(output_file, 'r') as file:
            filedata = file.readlines()
        method = ""
        for i in range(0, len(filedata)):
            if "Performing Projected CAP-EOM calculation for " in filedata[i]:
                method = "eomcc"
                break
            elif "The imaginary part of the CAP subspace matrix, onset=" in filedata[
                    i]:
                method = "adc"
                break
        if method == "eomcc":
            self._init_from_qchem_eomcc(output_file, irrep)
        elif method == "adc":
            self._init_from_qchem_adc(output_file, onset)
        else:
            raise RuntimeError("Error: Incompatible Q-Chem output.")

    def _init_from_opencap_output(self, output_file):
        with open(output_file, 'r') as file:
            filedata = file.readlines()
        cur_idx = -1
        nstates = 0
        for i in range(0, len(filedata)):
            if "Number of states" in filedata[i]:
                cur_idx = i
                nstates = int(filedata[cur_idx].split()[-1])
                break
        if nstates == 0:
            raise RuntimeError("Error: incompatible output. 0 states found.")
        cur_idx += 2
        # zeroth order hamiltonian first
        H0 = []
        for i in range(0, nstates):
            l1 = filedata[cur_idx].split()
            l1 = [float(x) for x in l1]
            H0 += l1
            cur_idx += 1
        H0 = np.reshape(H0, (nstates, nstates))
        cur_idx += 1
        # now CAP matrix
        W = []
        for i in range(0, nstates):
            l1 = filedata[cur_idx].split()
            l1 = [float(x) for x in l1]
            W += l1
            cur_idx += 1
        W = np.reshape(W, (nstates, nstates))
        assert H0.shape == W.shape
        self._H0 = H0
        self._W = W
        self._nstates = len(H0)

    def run_trajectory(self,
                       eta_list,
                       cap_lambda=0.0,
                       exclude_states=None,
                       include_states=None):
        '''
        Diagonalizes CAP Hamiltonian over range of eta values.

        CAP Hamiltonian is defined as H^CAP = H0 + i*eta*W - cap_lambda * W.
        W matrix is assumed to already have a -1 prefactor, as that is how OpenCAP output is formatted.
        Recommended range for eta_list is between 1E-5 and 1E-2, though this can vary widely based on system and CAP shape.

        Parameters
        ----------
        eta_list: iterable object
            List of eta values to use
        cap_lambda: float, default=0.0
            Real CAP strength to use for continuum remover CAP
        exclude_states: list of int, default=None
            List of states to exclude from subspace projection. Not compatible with include_states parameter.
        include_states: list of int, default=None
            List of states to include in subspace projection. Not compatible with exclude_states parameter.

        '''
        if exclude_states is not None and include_states is not None:
            raise RuntimeError(
                "Error: exclude_states and include_states keyword arguments are incompatible."
            )
        if exclude_states is not None:
            self.H0, self.W = _delete_indices(self._H0, self._W,
                                              exclude_states)
            self.nstates = len(self.H0)
        elif include_states is not None:
            all_states = [i for i in range(0, self._nstates)]
            exclude_states = [
                item for item in all_states if item not in include_states
            ]
            self.H0, self.W = _delete_indices(self._H0, self._W,
                                              exclude_states)
            self.nstates = len(self.H0)
        else:
            self.H0 = self._H0
            self.W = self._W
            self.nstates = self._nstates
        self._all_roots = []
        self.total_energies = []
        self.cap_lambda = cap_lambda
        self.etas = []
        for i in range(0, len(eta_list)):
            eta = eta_list[i]
            self.etas.append(eta)
            roots = []
            CAPH = self.H0 + 1.0j * eta * self.W - cap_lambda * self.W
            eigv, eigvc = LA.eig(CAPH)
            for j, eig in enumerate(eigv):
                roots.append(Root(eigv[j], eta, eigvc[:, j]))
                self.total_energies.append(eigv[j])
            self._all_roots.append(roots)

    def track_state(self, state_idx, tracking="overlap"):
        '''
        Diagonalizes CAP Hamiltonian over range of eta values.

        CAP Hamiltonian is defined as H^CAP = H0 + i*eta*W - cap_lambda * W.
        CAP matrix W is assumed to already have a -1 prefactor.

        Parameters
        ----------
        state_idx: int
            Index of state to track
        tracking: str, default="overlap"
            Method to use to track the state. Options are "overlap", which tracks based on eigenvector overlap, and "energy" which
            tracks based on energy.

        Returns
        -------
        traj: :class:`~pyopencap.analysis.EigenvalueTrajectory`
            Eigenvalue trajectory for further analysis.
        '''
        if len(self._all_roots) == 0:
            raise RuntimeError(
                "Nothing to track. Execute `run_trajectory` first.")
        traj = EigenvalueTrajectory(state_idx,
                                    self._all_roots[0],
                                    tracking=tracking)
        for i in range(1, len(self._all_roots)):
            traj._add_state(self._all_roots[i])
        traj._calculate_corrected_energies()
        return traj

    def energies_ev(self, ref_energy):
        '''
        Returns excitation energies of all calculated states in eV with respect to specified reference energy.

        Parameters
        ----------
        ref_energy: float
            Reference energy

        Returns
        -------
        E_eV: list of floats
            Excitation energies in eV with respect to specified reference energy.
        '''
        E_eV = []
        for E in self.total_energies:
            E_eV.append((E - ref_energy) * au2eV)
        return E_eV

    def export(self, finame):
        '''
        Exports Zeroth order Hamiltonian and CAP matrix to an OpenCAP formatted output file for further analysis.
        Useful for saving the results of an expensive electronic structure calculation performed in a python environment.

        Parameters
        ----------
        finame: str
            File handle to export data.
        '''
        from datetime import datetime
        from pandas import DataFrame
        now = datetime.now()
        dt_string = now.strftime("%m/%d/%Y %H:%M:%S")
        with open(finame, 'w') as f:
            f.write(
                "Welcome to OpenCAP: An open-source program for studying resonances in molecules.\n"
            )
            f.write("OpenCAP output file generated on: " + dt_string + "\n")
            f.write(
                "Printing out matrices required for Projected CAP calculation.\n"
            )
            f.write("Number of states: " + str(self._nstates) + "\n")
            f.write("Zeroth order Hamiltonian\n")
            f.write(DataFrame(self._H0).to_string(index=False, header=False,float_format='%.15g'))
            f.write("\nCAP Matrix\n")
            f.write(DataFrame(self._W).to_string(index=False, header=False,float_format='%.15g'))

    def __str__(self):
        ''' Returns formatted matrices.
        '''
        from pandas import DataFrame
        import pandas
        pandas.set_option("display.precision", 15)
        my_str = "Zeroth order Hamiltonian\n"
        my_str += DataFrame(self._H0).to_string(index=False, header=False,float_format='%.15g')
        my_str+= "\nCAP Matrix\n"
        my_str+= DataFrame(self._W).to_string(index=False, header=False,float_format='%.15g')
        return my_str


"""
Tool to calculate accuracy for loadgen accuracy output found in mlperf_log_accuracy.json
We assume that loadgen's query index is in the same order as
the images in coco's annotations/instances_val2017.json.
"""

from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import argparse
import json
import os

import numpy as np

from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval

# pylint: disable=missing-docstring

def get_args():
    """Parse commandline."""
    parser = argparse.ArgumentParser()
    parser.add_argument("--mlperf-accuracy-file", required=True, help="path to mlperf_log_accuracy.json")
    parser.add_argument("--coco-dir", required=True, help="coco directory")
    parser.add_argument("--verbose", action="store_true", help="verbose messages")
    args = parser.parse_args()
    return args

def main():
    args = get_args()

    cocoGt = COCO(os.path.join(args.coco_dir, "annotations/instances_val2017.json"))

    with open(args.mlperf_accuracy_file, "r") as f:
        results = json.load(f)

    detections = []
    image_ids = set()
    seen = set()
    image_map = cocoGt.dataset["images"]

    for j in results:
        idx = j['qsl_idx']
        # de-dupe in case loadgen sends the same image multiple times
        if idx in seen:
            continue
        seen.add(idx)

        # reconstruct from mlperf accuracy log
        # what is written by the benchmark is an array of float32's:
        # id, box[0], box[1], box[2], box[3], score, detection_class
        # note that id is a index into instances_val2017.json, not the actual image_id
        data = np.frombuffer(bytes.fromhex(j['data']), np.float32)
        for i in range(0, len(data), 7):
            image_idx, ymin, xmin, ymax, xmax, score, label = data[i:i + 7]
            image_idx = int(image_idx)
            image = image_map[image_idx]
            image_id = image["id"]
            height, width = image["height"], image["width"]
            ymin *= height
            xmin *= width
            ymax *= height
            xmax *= width
            loc = os.path.join(args.coco_dir, "val2017", image["file_name"])
            # pycoco wants {imageID,x1,y1,w,h,score,class}
            detections.append({
                "image_id": image_id,
                "image_loc": loc,
                "category_id": int(label),
                "bbox": [float(xmin), float(ymin), float(xmax - xmin), float(ymax - ymin)],
                "score": float(score)})
            image_ids.add(image_id)

    with open("coco-results.json", "w") as fp:
        json.dump(detections, fp, sort_keys=True, indent=4)

    cocoDt = cocoGt.loadRes(detections)
    cocoEval = COCOeval(cocoGt, cocoDt, iouType='bbox')
    cocoEval.params.imgIds = list(image_ids)
    cocoEval.evaluate()
    cocoEval.accumulate()
    cocoEval.summarize()

    print("mAP={:.3f}%".format(100. * cocoEval.stats[0]))
    if args.verbose:
        print("found and ignored {} dupes".format(len(results) - len(seen)))

if __name__ == "__main__":
    main()