'''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 scipy 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


def _biorthogonalize(Leigvc, Reigvc):
    M = Leigvc.T @ Reigvc
    P, L, U = LA.lu(M)
    Linv = LA.inv(L)
    Uinv = LA.inv(U)
    Leigvc = np.dot(Linv, Leigvc.T)
    Reigvc = np.dot(Reigvc, Uinv)
    Leigvc = Leigvc.T
    return Leigvc, Reigvc


def _sort_eigenvectors(eigv, eigvc):
    idx = eigv.argsort()
    eigv = eigv[idx]
    eigvc = eigvc[:, idx]
    return eigv, eigvc


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
    Reigvc: list of float
        Bi-orthogonalized right eigenvector for state
    Leigvc: list of float
        Bi-orthogonalized left eigenvector for state
    '''
    def __init__(self, energy, eta, Reigvc, Leigvc):
        '''
        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.Reigvc = Reigvc
        self.Leigvc = Leigvc


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 are obtained using density matrix 
    or first order correction.
    
    Attributes
    ----------
    roots: 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.
    W: np.ndarray of float: default=None
        CAP matrix in state basis. -1 prefactor is assumed.
    etas: list of float
        List of CAP strengths in trajectory.
    tracking: str, default="overlap"
        Method to use to track the state
    correction: str, default="density"
        Choice of correction scheme. Either "density" or "derivative".
    '''
    def __init__(self,
                 state_idx,
                 init_roots,
                 W,
                 tracking="overlap",
                 correction="density"):
        '''
        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
        W: np.ndarray of float: default=None
            CAP matrix in state basis. -1 prefactor is assumed.
        tracking: str, default="overlap"
            Method to use to track the state
        correction: str, default="density"
            Choice of correction scheme. Either "density" or "derivative".
        '''
        self._last = init_roots[state_idx]
        self.roots = [self._last]
        self.uncorrected_energies = [self._last.energy]
        self.corrected_energies = []
        self.W = W
        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.")
        if correction == "derivative":
            self._calculate_corrected_energies = self._derivative_correction
        elif correction == "density":
            self._calculate_corrected_energies = self._density_matrix_correction
        else:
            raise RuntimeError(
                "Invalid choice of correction. Choose either derivative or density."
            )

    def _overlap_tracking(self, states):
        maxo = 0.0
        state_idx = -1
        for st in states:
            ov = np.abs(np.dot(st.Reigvc, self._last.Reigvc))
            if ov > maxo:
                maxo = ov
                cur = st
        self._last = cur
        self.roots.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.roots.append(cur)
        self.uncorrected_energies.append(cur.energy)
        self.etas.append(cur.eta)

    def _density_matrix_correction(self):
        for i in range(0, len(self.roots)):
            eigvc = self.roots[i].Reigvc
            Leigvc = self.roots[i].Leigvc
            total = 0
            for k in range(0, len(self.W)):
                for l in range(0, len(self.W)):
                    total += Leigvc[k] * eigvc[l] * self.W[k][l]
            total *= 1.0j
            self.corrected_energies.append(self.uncorrected_energies[i] -
                                           self.etas[i] * total)

    def _derivative_correction(self):
        derivs = list(
            np.gradient(self.uncorrected_energies) / np.gradient(self.etas))
        for i in range(0, len(self.roots)):
            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",
                     return_root=False):
        '''
        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"
        return_root: bool, default=False
            Whether to return :class:`~pyopencap.analysis.Root` object associated with optimal value of eta
        
        Returns
        -------
        E_res: complex float
            Complex energy at optimal value of eta
        eta_opt: float
            Optimal value of eta
        root: :class:`~pyopencap.analysis.Root`
            Only returned when return_root is set to true

        '''
        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
            if return_root:
                return E, self.etas[opt_idx], self.roots[opt_idx]
            else:
                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
            if return_root:
                return E, self.etas[opt_idx], self.roots[opt_idx]
            else:
                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",
                   return_root=False):
        '''
        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"
        return_root: bool, default=False
            Whether to return :class:`~pyopencap.analysis.Root` object associated with optimal value of eta
        
        Returns
        -------
        E: float
            Energy at given value of eta
        root: :class:`~pyopencap.analysis.Root`
            Only returned when return_root is set to true
        '''
        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:
            if return_root:
                return (self.corrected_energies[idx] -
                        ref_energy) * scaling_factor, self.roots[idx]
            else:
                return (self.corrected_energies[idx] -
                        ref_energy) * scaling_factor
        else:
            if return_root:
                return (self.uncorrected_energies[idx] -
                        ref_energy) * scaling_factor, self.roots[idx]
            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, Reigvc = _sort_eigenvectors(*LA.eig(CAPH))
            Leigv, Leigvc = _sort_eigenvectors(*LA.eig(CAPH.T))
            Leigvc, Reigvc = _biorthogonalize(Leigvc, Reigvc)
            for j, eig in enumerate(eigv):
                roots.append(Root(eigv[j], eta, Reigvc[:, j], Leigvc[:, j]))
                self.total_energies.append(eigv[j])
            self._all_roots.append(roots)

    def track_state(self, state_idx, tracking="overlap", correction="density"):
        '''
        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.
        correction: str, default="density"
            Choice of correction scheme. Either "density" or "derivative".

        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],
                                    self.W,
                                    tracking=tracking,
                                    correction=correction)
        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
