"""
Module containing a container which can hold an arbitrary number of
`distpy.distribution.Distribution.Distribution` objects, each of which can have
any number of parameters which it describes (as long as the specific
`distpy.distribution.Distribution.Distribution` supports that number of
parameters). `distpy.distribution.Distribution.Distribution` objects can be
added through `DistributionList.add_distribution`. Once all the distributions
are added, points can be drawn using the `DistributionList.draw` method and the
log value of the entire set of distributions can be evaluated at a point using
the `DistributionList.log_value` method. See documentation of individual
methods for further details. This class represents a list-like container of
`distpy.distribution.Distribution.Distribution` objects; see
`distpy.distribution.DistributionSet.DistributionSet` for a dictionary- or set-
like container of `distpy.distribution.Distribution.Distribution` objects.
Unlike the `distpy.distribution.DistributionSet.DistributionSet` class,
`DistributionList` is a subclass of
`distpy.distribution.Distribution.Distribution` and implements all of its
methods and properties.

**File**: $DISTPY/distpy/distribution/DistributionList.py
**Author**: Keith Tauscher
**Date**: 31 May 2021
"""
import numpy as np
import numpy.random as rand
import scipy.linalg as scila
from ..util import int_types, numerical_types, sequence_types
from ..transform import TransformList, NullTransform
from .Distribution import Distribution

class DistributionList(Distribution):
    """
    An object which keeps track of many distributions which can be univariate
    or multivariate. It provides methods like log_value, which calls log_value
    on all of its constituent distributions, and draw, which draws from all of
    its constituent distributions.
    """
    def __init__(self, distribution_tuples=[]):
        """
        Creates a new `DistributionList` with the given distributions inside.

        Parameters
        ----------
        distribution_tuples : sequence
            a list of lists/tuples of the form `(distribution,)` or
            `(distribution, transforms)` where:

            - `distribution` is a
            `distpy.distribution.Distribution.Distribution` object
            - `transforms` is either a
            `distpy.transform.TransformList.TransformList` or something that
            can be cast to one (see the
            `distpy.transform.TransformList.TransformList.cast` method). It
            should describe the space in which the distribution applies. For
            example, to draw a variable from a normal distribution in log10
            space, `distribution` should be a
            `distpy.distribution.GaussianDistribution.GaussianDistribution` and
            `transforms` should be a
            `distpy.transform.Log10Transform.Log10Transform`. The result will
            contain only positive numbers.
        """
        self._data = []
        if type(distribution_tuples) in sequence_types:
            for idistribution in range(len(distribution_tuples)):
                this_tup = distribution_tuples[idistribution]
                if (type(this_tup) in sequence_types):
                    self.add_distribution(*this_tup)
                else:
                    raise ValueError("One of the distribution tuples " +\
                        "provided to the initializer of a DistributionList " +\
                        "was not a sequence of length 2 like " +\
                        "(distribution, transforms).")
        else:
            raise ValueError("The distribution_tuples argument given to " +\
                "the initializer was not list-like. It should be a list of " +\
                "tuples of the form (distribution, transformations) " +\
                "where distribution is a Distribution object and " +\
                "transformations are lists of strings.")

    @property
    def empty(self):
        """
        Boolean describing whether this `DistributionList` is empty.
        """
        return (len(self._data) == 0)

    @property
    def numparams(self):
        """
        The total number of parameters described by in this `DistributionList`.
        """
        return len(self.transform_list)

    @property
    def mean(self):
        """
        The approximate mean of this distribution. If the transform has a large
        second derivative at the mean, then this approximation is poor.
        """
        if not hasattr(self, '_mean'):
            mean = []
            for (distribution, transforms) in self._data:
                if distribution.numparams == 1:
                    mean.append(transforms[0].apply_inverse(distribution.mean))
                else:
                    this_mean = distribution.mean
                    for (itransform, transform) in enumerate(transforms):
                        mean.append(\
                            transform.apply_inverse(this_mean[itransform]))
            self._mean = np.array(mean)
        return self._mean

    @property
    def variance(self):
        """
        The covariance of this distribution.
        """
        if not hasattr(self, '_variance'):
            variances = []
            for (distribution, transforms) in self._data:
                if distribution.numparams == 1:
                    this_mean = np.array([distribution.mean])
                    this_covariance = np.array([[distribution.variance]])
                    variances.append(transforms.inverse.transform_covariance(\
                        this_covariance, this_mean))
                else:
                    this_mean = distribution.mean
                    this_covariance = distribution.variance
                    variances.append(transforms.inverse.transform_covariance(\
                        this_covariance, this_mean))
            self._variance = scila.block_diag(*variances)
        return self._variance

    def add_distribution(self, distribution, transforms=None):
        """
        Adds a `distpy.distribution.Distribution.Distribution` and the
        parameters it describes to the `DistributionList`.

        Parameters
        ----------
        distribution : `distpy.distribution.Distribution.Distribution`
            the distribution to add
        transforms : `distpy.transform.TransformList.TransformList` or\
        `distpy.transform.Transform.Transform` or sequence or str or None
            a `distpy.transform.TransformList.TransformList` object (or
            something castable to one, see
            `distpy.transform.TransformList.TransformList.cast`) which apply to
            the parameters (can be a single
            `distpy.transform.Transform.Transform` or something castable to
            one, see `distpy.transform.CastTransform.cast_to_transform`, if the
            distribution is univariate). If `transforms` is None, then the
            transforms are assumed to be
            `distpy.transform.NullTransform.NullTransform`
        """
        if isinstance(distribution, Distribution):
            transforms = TransformList.cast(transforms,\
                num_transforms=distribution.numparams)
            self._data.append((distribution, transforms))
        else:
            raise ValueError("The distribution given to a DistributionList " +\
                "was not recognized as a distribution.")

    @property
    def num_distributions(self):
        """
        The number of distributions stored in this `DistributionList` object.
        """
        return len(self._data)

    def __add__(self, other):
        """
        Adds this `DistributionList` to another to create a combined set.

        Parameters
        ----------
        other : `DistributionList`
            another `DistributionList` with parameters distinct from this one

        Returns
        -------
        sum : `DistributionList`
            the combination of the two `DistributionList` objects being added
        """
        if isinstance(other, DistributionList):
            return DistributionList(distribution_tuples=self._data+other._data)
        else:
            raise TypeError("Can only add DistributionList objects to " +\
                "other DistributionList objects.")

    def __iadd__(self, other):
        """
        Adds all distributions from `other` to this `DistributionList`.

        Parameters
        ----------
        other : `DistributionList`
            set of distributions to add into this one

        Returns
        -------
        enlarged : `DistributionList`
            this `DistributionList` after the distributions from `other` have
            been added in
        """
        if isinstance(other, DistributionList):
            for distribution_tuple in other._data:
                self.add_distribution(*distribution_tuple)
            return self
        else:
            raise TypeError("DistributionList objects can only have other " +\
                "DistributionList objects added to them.")

    def modify_transforms(self, new_transform_list):
        """
        Creates a `DistributionList` with the same distribution and parameters
        but different transforms. Draws from this `DistributionList` and the
        returned `DistributionList` will differ by the given transforms.

        Parameters
        ----------
        new_transform_list : `distpy.transform.TransformList.TransformList` or\
        `distpy.transform.Transform.Transform` or str or None
            a `distpy.transform.TransformList.TransformList` containing the new
            transforms (or something that can be cast to one with a number of
            transforms equal to the `DistributionList.numparams`; see
            `distpy.transform.TransformList.TransformList.cast` for details on
            what can be cast successfully)

        Returns
        -------
        modified : `DistributionList`
            new `DistributionList` object with the same distribution and
            parameters but different transforms
        """
        new_transform_list =\
            TransformList.cast(transforms, num_transforms=self.numparams)
        (new_data, running_index) = ([], 0)
        for (distribution, transforms) in self._data:
            new_running_index = running_index + distribution.numparams
            new_transforms =\
                new_transform_list[running_index:new_running_index]
            running_index = new_running_index
            new_data.append((distribution, new_transforms))
        return DistributionList(distribution_tuples=new_data)

    def draw(self, shape=None, random=rand):
        """
        Draws a point from all distributions.

        Parameters
        ----------
        shape : int or tuple or None
            shape of arrays which are values of return value
        random : `numpy.random.RandomState`
            the random number generator to use (default: numpy.random)

        Returns
        -------
        drawn_points : `numpy.ndarray`
            random variates drawn from the distribution (in the following `p`
            is `DistributionList.numparams`):
            - if `shape` is None, then `drawn_points` is a 1D `numpy.ndarray`
            of length `p`
            - if `shape` is an integer, then `drawn_points` is a 2D
            `numpy.ndarray` of shape `(shape,p)`
            - if `shape` is a tuple, then drawn points is a `numpy.ndarray` of
            shape `shape + (p,)`
        """
        none_shape = (type(shape) is type(None))
        if none_shape:
            shape = 1
        if type(shape) in int_types:
            shape = (shape,)
        point = np.ndarray(shape+(self.numparams,))
        params_included = 0
        for (distribution, transforms) in self._data:
            numparams = distribution.numparams
            if numparams == 1:
                transform = transforms[0]
                point[...,params_included] = transforms[0].apply_inverse(\
                    distribution.draw(shape=shape, random=random))
            else:
                this_draw = distribution.draw(shape=shape, random=random)
                for (itransform, transform) in enumerate(transforms):
                    point[...,params_included+itransform] =\
                        transform.apply_inverse(this_draw[...,itransform])
            params_included += numparams
        if self.numparams == 1:
            point = point[...,0]
        if none_shape:
            return point[0]
        else:
            return point

    def log_value(self, point):
        """
        Evaluates the log of the product of the values of the
        `distpy.distribution.Distribution.Distribution` objects contained in
        this `DistributionList`, which is the sum of their log values.

        Parameters
        ----------
        point : `numpy.ndarray`
            array of parameter values

        Returns
        -------
        total_log_value : float
            total log_value coming from contributions from all distributions
        """
        if type(point) in numerical_types:
            point = [point]
        if type(point) in sequence_types:
            point = np.array(point)
            if point.shape == (self.numparams,):
                result = 0.
                params_included = 0
                for (distribution, transforms) in self._data:
                    numparams = distribution.numparams
                    if numparams == 1:
                        subpoint = transforms[0].apply(point[params_included])
                    else:
                        subpoint = [\
                            transform.apply(point[params_included+itransform])\
                            for (itransform, transform) in\
                            enumerate(transforms)]
                    result += distribution.log_value(subpoint)
                    for (itransform, transform) in enumerate(transforms):
                        result += transform.log_derivative(\
                            point[params_included+itransform])
                    params_included += numparams
                    if not np.isfinite(result):
                        return -np.inf
                return result
            else:
                raise ValueError("point given to log_value function of a " +\
                    "DistributionList did not have the correct length.")
        else:
            raise ValueError("point given to log_value function of a " +\
                "DistributionList was not an array.")

    def __getitem__(self, which):
        """
        Gets a `DistributionList` with only the specified distributions.

        Parameters
        ----------
        which : int or slice or sequence
            the index or indices of the distributions to include in the
            returned value

        Returns
        -------
        sublist : `DistributionList`
            a `DistributionList` object with only the specified distribution(s)
        """
        if type(which) in int_types:
            distribution_list =\
                DistributionList(distribution_tuples=[self._data[which]])
        elif isinstance(which, slice):
            distribution_list =\
                DistributionList(distribution_tuples=self._data[which])
        elif type(which) in sequence_types:
            distribution_list = DistributionList()
            for which_element in which:
                distribution_list.add_distribution(*self._data[which_element])
        else:
            raise ValueError("Only integers, sequences of integers, and " +\
                "slices are allowed as arguments to __getitem__.")
        return distribution_list

    def delete_distribution(self, index):
        """
        Deletes a distribution from this `DistributionList`.

        Parameters
        ----------
        index : int
            the index of the distribution to delete
        """
        self._data = self._data[:index] + self._data[index+1:]

    def __delitem__(self, which):
        """
        Deletes a distribution from this `DistributionList`. Alias of
        `DistributionList.delete_distribution` that allows for usage of the
        `del` keyword.

        Parameters
        ----------
        index : int
            the index of the distribution to delete
        """
        if type(which) in int_types:
            self.delete_distribution(which)
        elif isinstance(which, slice) or (type(which) in sequence_types):
            if isinstance(which, slice):
                which = list(range(*which.indices(self.num_distributions)))
            which = sorted(which)[-1::-1]
            for index in which:
                self.delete_distribution(index)
        else:
            raise ValueError("Only integers, sequences of integers, and " +\
                "slices are allowed as arguments to __delitem__.")

    @property
    def summary_string(self):
        """
        A string with the dimenstionality of the distribution.
        """
        return '{:d}D DistributionList'.format(self.numparams)

    @property
    def minimum(self):
        """
        The minimum allowable value(s) in this distribution.
        """
        if not hasattr(self, '_minimum'):
            self._minimum = []
            for (distribution, transforms) in self._data:
                if distribution.numparams == 1:
                    self._minimum.append(transforms[0].untransform_minimum(\
                        distribution.minimum))
                else:
                    self._minimum.extend([transform.untransform_minimum(\
                        minimum) for (minimum, transform) in\
                        zip(distribution.minimum, transforms)])
        return self._minimum

    @property
    def maximum(self):
        """
        The maximum allowable value(s) in this distribution.
        """
        if not hasattr(self, '_maximum'):
            self._maximum = []
            for (distribution, transforms) in self._data:
                if distribution.numparams == 1:
                    self._maximum.append(transforms[0].untransform_maximum(\
                        distribution.maximum))
                else:
                    self._maximum.extend([transform.untransform_maximum(\
                        maximum) for (maximum, transform) in\
                        zip(distribution.maximum, transforms)])
        return self._maximum

    @property
    def is_discrete(self):
        """
        Boolean describing whether all of the distributions in this list are
        discrete.
        """
        return all([distribution.is_discrete\
            for (distribution, transforms) in self._data])

    @staticmethod
    def _distribution_tuples_equal(first, second):
        """
        Checks whether two distribution tuples are equal.

        Parameters
        ----------
        first : tuple
            tuple of form `(distribution, parameters)` as internally
            represented in a `DistributionList`
        second : tuple
            tuple of form `(distribution, transforms)` as internally
            represented in a `DistributionList`

        Returns
        -------
        result : bool
            True if and only if the distribution and transformations stored in
            `first` are the same as those stored in `second`.
        """
        (first_distribution, first_transforms) = first
        (second_distribution, second_transforms) = second
        numparams = first_distribution.numparams
        if second_distribution.numparams == numparams:
            for (first_transform, second_transform) in\
                zip(first_transforms, second_transforms):
                if first_transform != second_transform:
                    return False
            return (first_distribution == second_distribution)
        else:
            return False

    def __eq__(self, other):
        """
        Checks for equality of this `DistributionList` with `other`.

        Parameters
        ----------
        other : object
            object to check for equality

        Returns
        -------
        result : bool
            True if and only if `other` is a `DistributionList` with the same
            distribution tuples (though they need not be internally stored in
            the same order).
        """
        if isinstance(other, DistributionList):
            if len(self._data) == len(other._data):
                return all([DistributionList._distribution_tuples_equal(\
                    *tuples) for tuples in zip(self._data, other._data)])
            else:
                return False
        else:
            return False

    @property
    def transform_list(self):
        """
        The `distpy.transform.TransformList.TransformList` object describing
        the transforms in this `DistributionList`.
        """
        answer = TransformList()
        for (distribution, transforms) in self._data:
            answer += transforms
        return answer

    def discrete_sublist(self):
        """
        Compiles the subset of the `Distribution` objects in this
        `DistributionList` that represent discrete variables.

        Returns
        -------
        subset : `DistributionList`
            a `DistributionList` object containing all
            `distpy.distribution.Distribution.Distribution` objects in this
            `DistributionList` which describe discrete variables
        """
        answer = DistributionList()
        for (distribution, transforms) in self._data:
            if distribution.is_discrete:
                answer.add_distribution(distribution, transforms)
        return answer

    def continuous_sublist(self):
        """
        Compiles the subset of the `Distribution` objects in this
        `DistributionList` that represent continuous variables.

        Returns
        -------
        subset : `DistributionList`
            a `DistributionList` object containing all
            `distpy.distribution.Distribution.Distribution` objects in this
            `DistributionList` which describe continuous variables
        """
        answer = DistributionList()
        for (distribution, transforms) in self._data:
            if not distribution.is_discrete:
                answer.add_distribution(distribution, transforms)
        return answer

    def transformed_version(self):
        """
        Compiles a version of this `DistributionList` where the parameters
        exist in transformed space (instead of transforms being carried through
        this object).

        Returns
        -------
        transformless : `DistributionList`
            a `DistributionList` with the same distributions and parameter
            names but without transforms
        """
        answer = DistributionList()
        for (distribution, transforms) in self._data:
            answer.add_distribution(distribution)
        return answer

    def fill_hdf5_group(self, group, save_metadata=True):
        """
        Fills the given hdf5 file group with data about this
        `DistributionList`.

        Parameters
        ----------
        group : h5py.Group
            the hdf5 file group to fill
        save_metadata : bool
            - if True, attempts to save metadata alongside distribution list
            and throws error if it fails
            - if False, metadata is ignored in saving process
        """
        group.attrs['class'] = 'DistributionList'
        for (ituple, distribution_tuple) in enumerate(self._data):
            (distribution, transforms) = distribution_tuple
            subgroup = group.create_group('distribution_{}'.format(ituple))
            distribution.fill_hdf5_group(subgroup, save_metadata=save_metadata)
            transforms.fill_hdf5_group(subgroup)

    @staticmethod
    def load_from_hdf5_group(group, *distribution_classes):
        """
        Loads a `DistributionList` object from the given group.

        Parameters
        ----------
        group : h5py.Group
            the hdf5 file group in which a `DistributionList` was saved
        distribution_classes : sequence
            sequence of Distribution subclasses with which to load
            subdistributions (if the sub-distributions are defined in `distpy`,
            then these are not necessary because all distributions save their
            class as an attribute when they run
            `distpy.distribution.Distribution.Distribution.fill_hdf5_group`)

        Returns
        -------
        loaded : `DistributionList`
            a `DistributionList` that was saved in `group`
        """
        ituple = 0
        distribution_tuples = []
        while ('distribution_{}'.format(ituple)) in group:
            subgroup = group['distribution_{}'.format(ituple)]
            distribution_class_name = subgroup.attrs['class']
            if ituple >= len(distribution_classes):
                module = __import__('distpy')
                distribution_class = getattr(module, distribution_class_name)
            else:
                distribution_class =  distribution_classes[ituple]
            distribution = distribution_class.load_from_hdf5_group(subgroup)
            transform_list = TransformList.load_from_hdf5_group(subgroup)
            distribution_tuples.append((distribution, transform_list))
            ituple += 1
        return DistributionList(distribution_tuples=distribution_tuples)

    @property
    def gradient_computable(self):
        """
        Boolean describing whether the gradient of the distributions inside
        this `DistributionList` have been implemented.
        """
        answer = True
        for (distribution, transforms) in self._data:
            answer = (answer and distribution.gradient_computable)
        return answer

    def gradient_of_log_value(self, point):
        """
        Computes the derivatives of the log value with respect to the
        parameters.

        Parameters
        ----------
        point : `numpy.ndarray`
            array of parameter values

        Returns
        -------
        gradient : `numpy.ndarray`
            1D array of length `DistributionSet.numparams` of derivative values
            corresponding to the parameters
        """
        if type(point) in sequence_types:
            point = np.array(point)
            result = np.zeros((self.numparams,))
            params_included = 0
            for (idistribution, distribution_tuple) in enumerate(self._data):
                (distribution, transforms) = distribution_tuple
                numparams = distribution.numparams
                if numparams == 1:
                    result[params_included] +=\
                        distribution.gradient_of_log_value(\
                        transforms[0].apply(point[params_included]))
                else:
                    subpoint = np.array([\
                        transform.apply(point[params_included+itransform])\
                        for (itransform, transform) in enumerate(transforms)])
                    result[params_included:params_included+numparams] +=\
                        distribution.gradient_of_log_value(subpoint)
                for index in range(numparams):
                    result[params_included+index] +=\
                        transforms[index].derivative_of_log_derivative(\
                        point[params_included+index])
                params_included += numparams
            return result
        else:
            raise ValueError("point given to gradient_of_log_value " +\
                "function of a DistributionList was not a sequence of values.")

    @property
    def hessian_computable(self):
        """
        Boolean describing whether the hessian of the distributions inside this
        `DistributionList` have been implemented.
        """
        answer = True
        for (distribution, transforms) in self._data:
            answer = (answer and distribution.hessian_computable)
        return answer

    def hessian_of_log_value(self, point):
        """
        Computes the second derivatives of the log value with respect to the
        parameters.

        Parameters
        ----------
        point : `numpy.ndarray`
            array of parameter values

        Returns
        -------
        gradient : `numpy.ndarray`
            1D array of length `DistributionSet.numparams` of second derivative
            values corresponding to the parameters
        """
        if type(point) in sequence_types:
            point = np.array(point)
            result = np.zeros((self.numparams,) * 2)
            params_included = 0
            for (idistribution, distribution_tuple) in enumerate(self._data):
                (distribution, transforms) = distribution_tuple
                numparams = distribution.numparams
                if numparams == 1:
                    subpoint = transforms[0].apply(point[params_included])
                    result[params_included,params_included] +=\
                        distribution.hessian_of_log_value(subpoint)
                else:
                    subpoint = np.array([\
                        transform.apply(point[params_included+itransform])\
                        for (itransform, transform) in enumerate(transforms)])
                    result_slice = 2 *\
                        (slice(params_included, params_included + numparams),)
                    result[result_slice] +=\
                        distribution.hessian_of_log_value(subpoint)
                for i in range(numparams):
                    result[params_included+index,params_included+index] +=\
                        transforms[index].second_derivative_of_log_derivative(\
                        point[params_included+index])
                params_included += numparams
            return result
        else:
            raise ValueError("point given to hessian_of_log_value " +\
                "function of a DistributionList was not a sequence of values.")

    def copy(self):
        """
        Finds a deep copy of this `DistributionList`.

        Returns
        -------
        copied : `DistributionList`
            deep copy of this `DistributionList`
        """
        copied = DistributionList()
        for (distribution, transforms) in self._data:
            copied_distribution = distribution.copy()
            copied_transforms = [transform for transform in transforms]
            copied.add_distribution(copied_distribution, copied_transforms)
        return copied

    def reset(self):
        """
        Resets this distribution. This allows ideal distributions to live
        alongside samples as the same kind of object.
        """
        for (distribution, transforms) in self._data:
            distribution.reset()


#   Copyright 2020 Google LLC

#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
"""Wavefunction class for accessing and manipulation of the state of interest
"""
#zeros_like is incorrectly flagged by pylint
#pylint: disable=unsupported-assignment-operation
#there are many instances where access to protected members/methods simplify the
#code structure. They are not exposed to the users, and therefore, harmless
#pylint: disable=protected-access

import copy
import os
import math
from typing import (Any, Callable, cast, Dict, KeysView, List, Optional, Tuple,
                    Union)

import pickle
import numpy
from scipy import linalg
from scipy.special import factorial, jv

from fqe.fqe_decorators import wrap_apply, wrap_apply_generated_unitary
from fqe.fqe_decorators import wrap_time_evolve, wrap_rdm
from fqe.fqe_data import FqeData
from fqe.fqe_data_set import FqeDataSet
from fqe.util import alpha_beta_electrons
from fqe.util import map_broken_symmetry
from fqe.util import sort_configuration_keys
from fqe.util import vdot
from fqe.hamiltonians import hamiltonian, sparse_hamiltonian, \
                             diagonal_hamiltonian, diagonal_coulomb, \
                             restricted_hamiltonian
from fqe.bitstring import count_bits
from fqe.fqe_ops import fqe_operator, fqe_ops_utils
from fqe.wick import wick

class Wavefunction:
    """Wavefunction is the central object for manipulaion in the
    OpenFermion-FQE.
    """

    def __init__(self,
                 param: Optional[List[List[int]]] = None,
                 broken: Optional[Union[List[str], str]] = None) -> None:
        """
        Args:
            param (list[list[n, ms, norb]]): the constructor accepts a list of \
              parameter lists.  The parameter lists are comprised of

              p[0] (integer) - number of particles;

              p[1] (integer) - z component of spin angular momentum;

              p[2] (integer) - number of spatial orbitals

            broken (str): pass in the symmetries that should be preserved by \
                the wavefunction.

        Member Variables:
            _conserve_spin (bool): When this flag is true, the wavefunction \
                will maintain a constant m_s

            _conserve_number (bool): When this flag is true, the wavefunction \
                will maintain a constant nele

            _civec (dict[(int, int)] -> FqeData): This is a dictionary for \
                FqeData objects.  The key is a tuple defined by the number of \
                electrons and the spin projection of the system.
        """
        self._symmetry_map: Dict[Tuple[int, int], Tuple[int, int]] = {}
        self._conserved: Dict[str, int] = {}

        self._conserve_spin: bool = False
        if broken is None or 'spin' not in broken:
            self._conserve_spin = True

        self._conserve_number: bool = False
        if broken is None or 'number' not in broken:
            self._conserve_number = True

        self._norb: int = 0
        self._civec: Dict[Tuple[int, int], 'FqeData'] = {}

        if not self._conserve_spin and not self._conserve_number:
            raise TypeError('Number and spin non-conserving waveunfction is' \
                            ' the complete Fock space.')

        if param:
            user_input_norbs = set([x[2] for x in param])
            if len(user_input_norbs) != 1:
                raise ValueError('Number of orbitals is not consistent')

            self._norb = list(user_input_norbs)[0]
            for i in param:
                nalpha, nbeta = alpha_beta_electrons(i[0], i[1])
                self._civec[(i[0], i[1])] = FqeData(nalpha, nbeta, self._norb)

            if self._conserve_number:
                self._conserved['n'] = param[0][0]

            if self._conserve_spin:
                self._conserved['s_z'] = param[0][1]

                if not self._conserve_number:
                    self._symmetry_map = map_broken_symmetry(
                        param[0][1], param[0][2])

    def __add__(self, other: 'Wavefunction') -> 'Wavefunction':
        """Intrinsic addition function to combine two wavefunctions.  This acts \
        to iterate through the wavefunctions, combine coefficients of \
        configurations they have in common and add configurations that are \
        unique to each one.  The values are all combined into a new \
        wavefunction object

        Args:
            other (wavefunction.Wavefunction): the second wavefunction to \
                add with the local wavefunction

        Returns:
            wfn (wavefunction.Wavefunction): a new wavefunction with the \
                values set by adding together values
        """
        out = copy.deepcopy(self)
        out.ax_plus_y(1.0, other)
        return out

    def __iadd__(self, wfn: 'Wavefunction') -> 'Wavefunction':
        """Same is __add___ but performed in-place

        Args:
            wfn (wavefunction.Wavefunction): the second wavefunction to \
                add with the local wavefunction

        Returns:
            Wavefunction: self
        """
        self.ax_plus_y(1.0, wfn)
        return self

    def __sub__(self, other: 'Wavefunction') -> 'Wavefunction':
        """Intrinsic subtraction function to combine two wavefunctions.  This
        acts to iterate through the wavefunctions, combine coefficients of
        configurations they have in common and include configurations that are
        unique to each one. The values are all combined into a new
        wavefunction object

        Args:
            other (wavefunction.Wavefunction): the second wavefunction that \
                will be subtracted from the first wavefunction

        Returns:
            wfn (wavefunction.Wavefunction): a new wavefunction with the
                values set by subtracting the wfn from the first
        """
        out = copy.deepcopy(self)
        out.ax_plus_y(-1.0, other)
        return out

    def __getitem__(self, key: Tuple[int, int]) -> complex:
        """Element read access to the wave function.
        Args:
            key (Tuple[int, int]): a pair of strings for alpha and beta

        Returns:
            (complex): the value of the wave function
        """
        astr, bstr = key[0], key[1]
        sector = (count_bits(astr) + count_bits(bstr),
                  count_bits(astr) - count_bits(bstr))
        return self._civec[sector][key]

    def __setitem__(self, key: Tuple[int, int], value: complex) -> None:
        """Element write access to the wave function.
        Args:
            key (Tuple[int, int]): a pair of strings for alpha and beta

            value (complex): the value to be set to the wave function
        """
        astr, bstr = key[0], key[1]
        sector = (count_bits(astr) + count_bits(bstr),
                  count_bits(astr) - count_bits(bstr))
        self._civec[sector][key] = value

    def empty_copy(self) -> 'Wavefunction':
        """create a copy of self with zero coefficients

        Returns:
            Wavefunction: a new object with zero coefficients
        """
        out = Wavefunction()
        out._norb = self._norb
        out._conserved = self._conserved
        out._symmetry_map = self._symmetry_map
        for key, civec in self._civec.items():
            out._civec[key] = civec.empty_copy()
        return out

    def _copy_beta_inversion(self) -> 'Wavefunction':
        """Return a copy of the wavefunction with the beta particle and hole
        inverted.

        Returns:
            Wavefunction: wavefuction with beta particle/hole conjugation
        """
        norb = self._norb
        m_s = self._conserved['s_z']

        nele = norb + m_s
        param = []
        maxb = min(norb, nele)
        minb = nele - maxb
        param = [
            [nele, nele - nbeta * 2, norb] for nbeta in range(minb, maxb + 1)
        ]

        inverted = Wavefunction(param, broken=['spin'])

        data = {}
        for key, sector in self._civec.items():
            work = ((key[0] + key[1]) // 2, (key[0] - key[1]) // 2)
            nkey = self._symmetry_map[work]
            data[(nkey[0] + nkey[1],
                  nkey[0] - nkey[1])] = sector.beta_inversion()

        inverted.set_wfn(strategy='from_data', raw_data=data)

        return inverted

    def ax_plus_y(self, sval: complex, wfn: 'Wavefunction') -> None:
        """Perform scale and add of the wavefunction. The result will be stored
        in self.

        Args:
            sval (complex): a factor to be multiplied to wfn

            wfn (Wavefunction): a wavefunction to be added to self
        """
        if self._civec.keys() != wfn._civec.keys():
            raise ValueError('inconsistent sectors in Wavefunction.ax_plus_y')

        for sector in self._civec:
            self._civec[sector].ax_plus_y(sval, wfn._civec[sector])

    def sector(self, key: Tuple[int, int]) -> 'FqeData':
        """Return a specific sector of the wavefunction using a key.

        Args:
            key (Tuple[int, int]): key for ci vector
        Returns:
            FqeData: corresponding sector as an FqeData object
        """
        return self._civec[key]

    def sectors(self) -> KeysView[Tuple[int, int]]:
        """
        Return:
            KeysView[Tuple[int, int]]: a list of the configuration keys \
                in the wavefunction
        """
        return self._civec.keys()

    def conserve_number(self) -> bool:
        """
        Returns:
            (bool): True if this wave function conserves the number symmetry
        """
        return self._conserve_number

    def conserve_spin(self) -> bool:
        """
        Returns:
            (bool): True if this wave function conserves the spin (Sz) \
                symmetry. Otherwise False
        """
        return self._conserve_spin

    def norb(self) -> int:
        """
        Returns:
            (int): the number of orbitals
        """
        return self._norb

    def norm(self) -> float:
        """Calculate the norm of the wavefuntion

        Returns:
            (float): the norm
        """
        normall = 0.0
        for sector in self._civec.values():
            normall += sector.norm()**2
        normall = math.sqrt(normall)
        return normall

    def normalize(self) -> None:
        """Generte the wavefunction norm and then scale each element by that
        value.
        """
        self.scale(1.0 / self.norm())

    @wrap_apply
    def apply(self, hamil: 'hamiltonian.Hamiltonian') -> 'Wavefunction':
        """ Returns a wavefunction subject to application of the Hamiltonian
        (or more generally, the operator).

        Args:
            hamil (Hamiltonian): Hamiltonian to be applied

        Returns:
            (Wavefunction): resulting wave function array
        """
        if not self._conserve_number or not hamil.conserve_number():
            if self._conserve_number:
                raise TypeError('Number non-conserving hamiltonian passed to'
                                ' number conserving wavefunction')
            if hamil.conserve_number():
                raise TypeError('Number conserving hamiltonian passed to'
                                ' number non-conserving wavefunction')

        if isinstance(hamil, sparse_hamiltonian.SparseHamiltonian):

            transformed = self._apply_few_nbody(hamil)

        else:
            if self._conserve_spin and not self._conserve_number:
                out = self._copy_beta_inversion()
            else:
                out = self

            if isinstance(hamil, diagonal_hamiltonian.Diagonal):
                transformed = out._apply_diagonal(hamil)
            elif isinstance(hamil, diagonal_coulomb.DiagonalCoulomb):
                transformed = out._apply_diagonal_coulomb(hamil)
            else:
                if isinstance(hamil,
                              restricted_hamiltonian.RestrictedHamiltonian):
                    expected = self._norb
                else:
                    expected = self._norb * 2
                if hamil.dim() != expected:
                    raise ValueError('Hamiltonian has incorrect size:' \
                                     + ' expected {}'.format(expected) \
                                     + ' provided {}'.format(hamil.dim()))

                transformed = out._apply_array(hamil.tensors(), hamil.e_0())

            if self._conserve_spin and not self._conserve_number:
                transformed = transformed._copy_beta_restore(
                    self._conserved['s_z'], self._norb, self._symmetry_map)

        return transformed

    def _apply_array(self, array: Tuple[numpy.ndarray, ...],
                     e_0: complex) -> 'Wavefunction':
        """Return a wavefunction subject to application of the numpy array as

        .. math::
            h[i, j]a_i^+ a_j|Psi>
            h[i, j, k, l]a_i^+(rho) a_j^+(eta) a_k(rho) a_l(eta)|Psi>

        Arg:
            array (numpy.array): numpy array

            e_0 (complex): scalar part of the Hamiltonian

        Returns:
            newwfn (Wavefunction): a new intialized wavefunction object

        """
        if self._conserve_spin:
            assert array[0].shape[0] == self._norb or array[0].shape[
                0] == self._norb * 2
            out = copy.deepcopy(self)
            for _, sector in out._civec.items():
                sector.apply_inplace(array)
        else:
            assert array[0].shape[0] == self._norb * 2
            out = copy.deepcopy(self)
            nsectors = out._number_sectors()
            for _, nsector in nsectors.items():
                nsector.apply_inplace(array)

        if numpy.abs(e_0) > 1.e-15:
            out.ax_plus_y(e_0, self)

        return out

    def _apply_diagonal(self, hamil: 'diagonal_hamiltonian.Diagonal'
                       ) -> 'Wavefunction':
        """Applies the diagonal operator to the wavefunction

        Args:
            hamil (Diagonal): diagonal Hamiltonian to be applied

        Returns:
            (Wavefunction): resulting wave function
        """
        out = copy.deepcopy(self)

        for _, sector in out._civec.items():
            sector.apply_diagonal_inplace(hamil.diag_values())

        if numpy.abs(hamil.e_0()) > 1.e-15:
            out.ax_plus_y(hamil.e_0(), self)

        return out

    def _apply_diagonal_coulomb(self, hamil: 'diagonal_coulomb.DiagonalCoulomb'
                               ) -> 'Wavefunction':
        """Applies the diagonal coulomb operator to the wavefunction

        Args:
            hamil (DiagonalCoulomb): diagonal coulomb Hamiltonian to be applied

        Returns:
            (Wavefunction): resulting wave function
        """
        out = copy.deepcopy(self)

        for _, sector in out._civec.items():
            diag, array = hamil._tensor[1], hamil._tensor[2]
            sector.apply_diagonal_coulomb(diag, array, inplace=True)

        if numpy.abs(hamil.e_0()) > 1.e-15:
            out.ax_plus_y(hamil.e_0(), self)

        return out

    def _number_sectors(self) -> Dict[int, FqeDataSet]:
        """An internal utility function that groups FqeData into a set of
        FqeDataSet that corresponds to the same number of electrons.  It checks
        spin completeness and raises an exception if the wave function space is
        not spin complete

        Returns:
            Dict[int, FqeDataSet]: stores FqeDataSet for each number of
                particles. Keys are the number of particles.
        """
        norb = self.norb()
        numbers = set(key[0] for key in self._civec)
        numbersectors = {}
        for nele in numbers:
            # generate all possible sz
            maxalpha = min(norb, nele)
            minalpha = nele - maxalpha
            sectors = {}
            sp_compl = set(((nele, 2 * nalpha - nele)
                            for nalpha in range(minalpha, maxalpha + 1)))

            if set(self._civec.keys()).intersection(sp_compl) != sp_compl:
                raise ValueError('Wave function space is not spin complete.')

            for nalpha in range(minalpha, maxalpha + 1):
                nbeta = nele - nalpha
                sectors[(nalpha, nbeta)] = self._civec[(nele, nalpha - nbeta)]
            dataset = FqeDataSet(nele, norb, sectors)
            numbersectors[nele] = dataset
        return numbersectors

    def _copy_beta_restore(self, s_z: int, norb: int,
                           map_symm: Dict[Tuple[int, int], Tuple[int, int]]
                          ) -> 'Wavefunction':
        """Return a copy of the wavefunction with beta restored back to number
        breaking/spin conserving.

        Args:
            s_z (int): the value of Sz

            norb (int): the number of orbitals in the system

            map_symm (Dict[Tuple[int,int], Tuple[int,int]]): dictionary that maps \
                between number-conserved and spin-conserved wave function sectors

        Returns:
            Wavefunction: restored wavefunction
        """
        max_alpha = min(norb, norb + s_z)
        min_alpha = max(s_z, 0)

        param = [[2 * nalpha - s_z, s_z, norb]
                 for nalpha in range(min_alpha, max_alpha + 1)]

        restored = Wavefunction(param, broken=['number'])

        data = {}
        for key in param:
            work = ((key[0] + key[1]) // 2, (key[0] - key[1]) // 2)
            nkey = map_symm[work]
            other = (nkey[0] + nkey[1], nkey[0] - nkey[1])
            data[(key[0], key[1])] = self._civec[other].beta_inversion()

        restored.set_wfn(strategy='from_data', raw_data=data)

        return restored

    @wrap_apply_generated_unitary
    def apply_generated_unitary(self,
                                time: float,
                                algo: str,
                                hamil: 'hamiltonian.Hamiltonian',
                                accuracy: float = 1.0E-14,
                                expansion: int = 30,
                                spec_lim: Optional[List[float]] = None
                               ) -> 'Wavefunction':
        """Perform the exponentiation of fermionic algebras to the
        wavefunction according the method and accuracy.

        Args:
            time (float): the final time value to evolve to

            algo (string): polynomial expansion algorithm to be used

            hamil (Hamiltonian): the Hamiltonian used to generate the unitary

            accuracy (float): the accuracy to which the system should be evolved

            expansion (int): the maximum number of terms in the polynomial expansion

            spec_lim (List[float]): spectral range of the Hamiltonian, the length of \
                the list should be 2. Optional.

        Returns:
            newwfn (Wavefunction): a new intialized wavefunction object
        """

        assert isinstance(hamil, hamiltonian.Hamiltonian)

        algo_avail = ['taylor', 'chebyshev']

        assert algo in algo_avail

        if not isinstance(hamil, sparse_hamiltonian.SparseHamiltonian) \
            and self._conserve_spin and not self._conserve_number:
            base = self._copy_beta_inversion()
        else:
            base = self

        max_expansion = max(30, expansion)

        if algo == 'taylor':
            ham_arrays = hamil.iht(time)

            time_evol = copy.deepcopy(base)
            work = copy.deepcopy(base)

            for order in range(1, max_expansion):
                work = work.apply(ham_arrays)
                coeff = 1.0 / factorial(order)
                time_evol.ax_plus_y(coeff, work)
                if work.norm() * numpy.abs(coeff) < accuracy:
                    break

        elif algo == 'chebyshev':

            assert spec_lim, 'Spectral range was not provided.' + \
                                 ' Provide upper and lower limits.'

            wprime = 0.9875
            ascale = (spec_lim[1] - spec_lim[0]) / (2.0 * wprime)
            eshift = -(spec_lim[0] + ascale * wprime)

            time_evol = copy.deepcopy(base)
            time_evol.scale(jv(0, ascale * time))
            minus = copy.deepcopy(base)

            current = minus.apply(hamil)
            current.ax_plus_y(eshift, minus)
            current.scale(1.0 / ascale)
            time_evol.ax_plus_y(2.0 * jv(1, ascale * time) * (-1.j), current)

            for order in range(2, max_expansion):
                minus.scale(-1.0)
                minus.ax_plus_y(2.0 / ascale, current.apply(hamil))
                minus.ax_plus_y(2.0 * eshift / ascale, current)
                current, minus = minus, current

                coeff = 2.0 * jv(order, ascale * time) * (-1.j)**order
                time_evol.ax_plus_y(coeff, current)

                if current.norm() * numpy.abs(coeff) < accuracy:
                    break

            time_evol.scale(numpy.exp(eshift * time * 1.j))

        if numpy.abs(hamil.e_0() * time) > 1.e-15:
            time_evol.scale(numpy.exp(-1.j * time * hamil.e_0()))

        if self._conserve_spin and not self._conserve_number:
            time_evol = time_evol._copy_beta_restore(self._conserved['s_z'],
                                                     self._norb,
                                                     self._symmetry_map)
        return time_evol

    def get_coeff(self, key: Tuple[int, int]) -> numpy.ndarray:
        """Retrieve a vector from a configuration in the wavefunction

        key indicates wavefunction sector by [num_alpha, num_beta]

        Args:
            key (Tuple[int, int]): a key identifying the configuration to access

        Returns:
            numpy.array(dtype=numpy.complex128): coeff that corresponds to key
        """
        return self._civec[key].coeff

    def max_element(self) -> complex:
        """
        Return:
            (complex): the largest magnitude value in the wavefunction
        """
        maxval = 0.0
        for config in self._civec.values():
            configmax = max(config.coeff.max(),
                            config.coeff.min(),
                            key=numpy.abs)
            maxval = max(configmax, maxval)

        return maxval

    def print_wfn(self, threshold: float = 0.001, fmt: str = 'str') -> None:
        """Print occupations and coefficients to the screen.

        Args:
            threshhold (float): only print CI vector values such that \
              :math:`|c|` > threshold.

            fmt (string): formats print according to argument
        """

        def _print_format(fmt: str) -> Callable[[int, int], str]:
            """ Select the function which will perform formatted printing
            """
            if fmt == 'occ':

                def _occupation_format(iastring: int, ibstring: int):
                    """occ - Prints a string indicating occupation state of each spatial
                    orbital.  A doubly occupied orbital will be indicated with "2",
                    a singly occupied orbital will get "a" or "b" depending on the
                    spin state.  An empy orbital will be ".".
                    """
                    astring = int(iastring)
                    bstring = int(ibstring)
                    occstr = [
                        '.' for _ in range(
                            max(astring.bit_length(), bstring.bit_length()))
                    ]
                    docc = astring & bstring

                    def build_occ_value(bstr: int, char: str,
                                        occstr: List[str]):
                        """Fill a list with a character corresponding to the
                        location of '1' bits in bitstring

                        Args:
                            bstr (int): a bitstring to examine

                            char (str): a string to put into the list

                            occstr (list): a list to store the value
                                corresponding to flipped bits
                        """
                        ind = 1
                        while bstr:
                            if bstr & 1:
                                occstr[-ind] = char
                            ind += 1
                            bstr = bstr >> 1

                    build_occ_value(astring, 'a', occstr)
                    build_occ_value(bstring, 'b', occstr)
                    build_occ_value(docc, '2', occstr)

                    return ''.join(occstr).rjust(self._norb, '.')

                return _occupation_format

            def _string_format(iastring: int, ibstring: int) -> str:
                """Prints a binary string indicating which orbital creation
                operators are acting on the vacuum with a 1. The position in
                the string indicates the index of the orbital. The beta string
                is shown acting on the vacuum first.

                Args:
                    iastring (int): alpha string

                    ibstring (int): alpha string

                Returns:
                    str: representation of binary string
                """
                astring = int(iastring)
                bstring = int(ibstring)
                fmt_a = bin(1 << self._norb | astring)
                fmt_b = bin(1 << self._norb | bstring)
                return "a'{}'b'{}'".format(fmt_a[3:], fmt_b[3:])

            return _string_format

        print_format = _print_format(fmt)

        config_in_order = sort_configuration_keys(self.sectors())
        for key in config_in_order:
            self._civec[key].print_sector(pformat=print_format,
                                          threshold=threshold)

    def read(self, filename: str, path: str = os.getcwd()) -> None:
        """Initialize a wavefunction from a binary file.

        Args:
            filename (str): the name of the file to write the wavefunction to.

            path (str): the path to save the file.  If no path is given then \
              it is saved in the current working directory.
        """
        with open(os.path.join(path, filename), 'r+b') as wfnfile:
            wfn_data = pickle.load(wfnfile)

        self._symmetry_map = wfn_data[0]
        self._conserved = wfn_data[1]
        self._conserve_spin = wfn_data[2]
        self._conserve_number = wfn_data[3]
        self._norb = wfn_data[4]

        for sector in wfn_data[5:]:
            self._civec[(sector[0][0], sector[0][1])] = sector[1]

    def save(self, filename: str, path: str = os.getcwd()) -> None:
        """Save the wavefunction into path/filename.

        Args:
            filename (str): the name of the file to write the wavefunction to.

            path (str): the path to save the file.  If no path is given, then \
              it is saved in the current working directory.
        """
        wfn_data = [
            self._symmetry_map, self._conserved, self._conserve_spin,
            self._conserve_number, self._norb
        ]

        for key in self._civec:
            wfn_data.append([key, self._civec[key]])

        with open(os.path.join(path, filename), 'w+b') as wfnfile:
            pickle.dump(wfn_data, wfnfile)

    def scale(self, sval: complex) -> None:
        """ Scale each configuration space by the value sval

        Args:
            sval (complex): value to scale by
        """

        sval = complex(sval)  # type: ignore

        for sector in self._civec.values():
            sector.scale(sval)

    def set_wfn(self,
                strategy: str = 'ones',
                raw_data: Optional[Dict[Tuple[int, int], numpy.ndarray]] = None
               ) -> None:
        """Set the values of the wavefunction inplace based on a strategy.

        Args:
            strategy (string): The procedure to follow to set the wavefunction \
                coefficients. One of 'random', 'hartree-fock', 'ones', 'zeros', \
                or 'from_data'. If 'from_data', raw_data must be provided.

            raw_data (Dict[Tuple[int, int], numpy.ndarray]): The values to use \
                if setting from data. Optional.
        """
        if strategy == 'from_data' and not raw_data:
            raise ValueError('No data provided for set_wfn')

        if strategy == 'from_data':
            for key, data in raw_data.items():  # type: ignore
                self._civec[key].set_wfn(strategy='from_data', raw_data=data)
        elif strategy == 'hartree-fock':
            # make sure we only have 1 sector
            if len(self.sectors()) != 1:
                raise ValueError(("Hartree-Fock wf initialization only works "
                                  "with single sector wavefunctions"))
            for sector in self._civec.values():
                sector.set_wfn(strategy=strategy)
        else:
            for sector in self._civec.values():
                sector.set_wfn(strategy=strategy)

        if strategy == 'random':
            self.normalize()

    def transform(
            self,
            rotation: numpy.ndarray,
            low: Optional[numpy.ndarray] = None,
            upp: Optional[numpy.ndarray] = None
    ) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, 'Wavefunction']:
        """Transform the wavefunction using the orbtial rotation matrix and
        return the new wavefunction and the permutation matrix for the unitary
        transformation. This is an internal code, so performs minimal checking

        Args:
            rotation (numpy.ndarray): MO rotation matrix, which is unitary

            low (numpy.ndarray): L in the LU decomposition (optional)

            upp (numpy.ndarray): U in the LU decomposition (optional)

        Returns:
            (numpy.ndarray, numpy.ndarray, numpy.ndarray, 'Wavefunction'): \
                permutation, L, U, and transformed wavefunction
        """
        norb = self._norb
        external = low is not None
        assert external == (upp is not None)
        if external:
            assert numpy.allclose(rotation, low @ upp)  # type: ignore

        def ludecomp(rotmat: numpy.ndarray
                    ) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
            """Returns permutation, lower triangular, and upper triangular
            matrices from the LU decomposition. Note that the rotation matrix
            is Hermitian conjugated.

            Args:
                rotmat (numpy.ndarray): MO rotation matrix, which is unitary

                (numpy.ndarray, numpy.ndarray, numpy.ndarray): permutation, \
                    L, and U from the LU decomposition.
            """
            tmat = rotmat.transpose().conjugate()
            return linalg.lu(tmat)

        def transpose_matrix(low: numpy.ndarray, upp: numpy.ndarray
                            ) -> Tuple[numpy.ndarray, numpy.ndarray]:
            """Returns Hermitian conjugate of the L and U factors. The diagonal
            elements are in the upper triagular matrix. This transposition seeks to
            compensate the transposition in ludecomp above.

            Args:
                low (numpy.ndarray): L in the LU decomposition

                upp (numpy.ndarray): U in the LU decomposition

            Returns:
                (numpy.ndarray, numpy.ndarray): L and U after transposition \
                    where L and U are lower- and upper-triangular, respectively
            """
            ndim = low.shape[0]
            assert low.shape[1] == ndim and upp.shape == (ndim, ndim)
            lowt = copy.deepcopy(low)
            uppt = copy.deepcopy(upp)
            for irow in range(ndim):
                for icol in range(irow + 1, ndim):
                    uppt[irow, icol] /= uppt[irow, irow]
                lowt[irow, irow], uppt[irow, irow] = uppt[irow, irow], lowt[
                    irow, irow]
                for icol in range(irow):
                    lowt[irow, icol] *= lowt[icol, icol]
            return uppt.T.conj(), lowt.T.conj()

        def process_matrix(low: numpy.ndarray,
                           upp: numpy.ndarray) -> numpy.ndarray:
            """Returns an operator using which the wavefuction will be transformed.

            Args:
                low (numpy.ndarray): L in the LU decomposition

                upp (numpy.ndarray): U in the LU decomposition

            Returns:
                (numpy.ndarray): matrix elements of the transformation operator T
            """
            ndim = low.shape[0]
            assert low.shape[1] == ndim and upp.shape == (ndim, ndim)
            unitmat = numpy.identity(ndim)
            output = linalg.solve_triangular(upp, unitmat)

            for icol in range(ndim):
                for irow in range(icol + 1, ndim):
                    output[irow, icol] -= low[irow, icol]
                output[icol, icol] -= 1.0
            return output

        current = self
        perm = None

        if not self._conserve_spin:
            if not external:
                perm, low, upp = ludecomp(rotation)
                lowt, uppt = transpose_matrix(low, upp)
            else:
                lowt, uppt = low, upp
            output = process_matrix(lowt, uppt)
            for icol in range(norb * 2):
                work = numpy.zeros_like(rotation)
                work[:, icol] = output[:, icol]
                onefwfn = current.apply((work,))
                current.ax_plus_y(1.0, onefwfn)
        else:
            if rotation.shape[0] == norb:
                if not external:
                    perm, low, upp = ludecomp(rotation)
                    lowt, uppt = transpose_matrix(low, upp)
                else:
                    lowt, uppt = low, upp
                output = process_matrix(lowt, uppt)
                for _, civec in current._civec.items():
                    civec.apply_columns_recursive_inplace(output, output)
            elif rotation.shape[0] == norb * 2:
                assert numpy.std(rotation[:norb, norb:]) \
                       + numpy.std(rotation[norb:, :norb]) < 1.0e-8
                if not external:
                    perm1, low1, upp1 = ludecomp(rotation[:norb, :norb])
                    perm2, low2, upp2 = ludecomp(rotation[norb:, norb:])
                    lowt1, uppt1 = transpose_matrix(low1, upp1)
                    lowt2, uppt2 = transpose_matrix(low2, upp2)
                else:
                    lowt1 = low[:norb, :norb]  # type: ignore
                    lowt2 = low[norb:, norb:]  # type: ignore
                    uppt1 = upp[:norb, :norb]  # type: ignore
                    uppt2 = upp[norb:, norb:]  # type: ignore
                output1 = process_matrix(lowt1, uppt1)
                output2 = process_matrix(lowt2, uppt2)
                for _, civec in current._civec.items():
                    civec.apply_columns_recursive_inplace(output1, output2)
                if not external:
                    perm = numpy.zeros_like(rotation)
                    perm[:norb, :norb] = perm1[:, :]
                    perm[norb:, norb:] = perm2[:, :]
                    upp = numpy.zeros_like(rotation)
                    upp[:norb, :norb] = upp1[:, :]
                    upp[norb:, norb:] = upp2[:, :]
                    low = numpy.zeros_like(rotation)
                    low[:norb, :norb] = low1[:, :]
                    low[norb:, norb:] = low2[:, :]

        return perm, low, upp, current

    @wrap_time_evolve
    def time_evolve(
            self,
            time: float,
            hamil: Union['fqe_operator.FqeOperator', 'hamiltonian.Hamiltonian'],
            inplace: bool = False) -> 'Wavefunction':
        """Perform time evolution of the wavefunction given Fermion Operators
        either as raw operations or wrapped up in a Hamiltonian.

        Args:
            time (float): the duration by which to evolve the operators

            hamil (Hamiltoninan or FermionOperator): Hamiltonian to be used for \
                time evolution

            inplace (bool): whether the result will be stored in place

        Returns:
            (Wavefunction): a wavefunction object that has been time evolved.
        """
        assert isinstance(hamil, hamiltonian.Hamiltonian)

        if not self._conserve_number or not hamil.conserve_number():
            if self._conserve_number:
                raise TypeError('Number non-conserving hamiltonian passed to'
                                ' number conserving wavefunction')
            if hamil.conserve_number():
                raise TypeError('Number conserving hamiltonian passed to'
                                ' number non-conserving wavefunction')

        if isinstance(
                hamil,
                sparse_hamiltonian.SparseHamiltonian) and hamil.is_individual():

            final_wfn = self._evolve_individual_nbody(time, hamil, inplace)

        else:
            is_diag = ((hamil.quadratic() and hamil.diagonal()) or
                       hamil.diagonal_coulomb())
            if inplace and (not is_diag and not hamil.quadratic()):
                raise ValueError("Inplace is not implemented for this case")

            if self._conserve_spin and not self._conserve_number:
                work_wfn = self._copy_beta_inversion()
            elif inplace:
                work_wfn = self
            else:
                work_wfn = copy.deepcopy(self)

            if hamil.quadratic():
                if hamil.diagonal():

                    ihtdiag = -1.j * time * hamil.diag_values()
                    final_wfn = work_wfn._evolve_diagonal(ihtdiag, inplace)

                else:
                    transformation = hamil.calc_diag_transform()

                    permu, low, upp, work_wfn = work_wfn.transform(
                        transformation)

                    ci_trans = transformation @ permu

                    h1e = hamil.transform(ci_trans)

                    ihtdiag = -1.j * time * h1e.diagonal()
                    evolved_wfn = work_wfn._evolve_diagonal(ihtdiag,
                                                            inplace=True)

                    _, _, _, final_wfn = evolved_wfn.transform(
                        ci_trans.T.conj(), low, upp)

            elif hamil.diagonal_coulomb():

                diag, vij = hamil.iht(time)

                final_wfn = work_wfn._evolve_diagonal_coulomb_inplace(diag, vij)

            else:

                final_wfn = work_wfn.apply_generated_unitary(
                    time, 'taylor', hamil)

            if self._conserve_spin and not self._conserve_number:
                final_wfn = final_wfn._copy_beta_restore(
                    self._conserved['s_z'], self._norb, self._symmetry_map)

        if numpy.abs(hamil.e_0()) > 1.0e-15:
            final_wfn.scale(numpy.exp(-1.j * time * hamil.e_0()))

        return final_wfn

    def _evolve_diagonal(self, ithdiag: numpy.ndarray,
                         inplace: bool = False) -> 'Wavefunction':
        """Evolve a diagonal Hamiltonian on the wavefunction

        Args:
            ithdiag (numpy.ndarray): preprocessed diagonal array

            inplace (bool): whether the result will be stored in-place

        Returns:
            (Wavefunction): resulting wave function. If in-place is True,
                self is returned.
        """
        if inplace:
            wfn = self
        else:
            wfn = copy.deepcopy(self)

        for key, sector in self._civec.items():
            wfn._civec[key].coeff = sector.evolve_diagonal(ithdiag, inplace)

        return wfn

    def _evolve_diagonal_coulomb_inplace(self, diag: numpy.ndarray,
                                         vij: numpy.ndarray) -> 'Wavefunction':
        """Evolve a diagonal coulomb Hamiltonian on the wavefunction and store the
        result inplace.  (Not in-place version is no longer used and thus has been
        removed).

        Args:
            diag (numpy.ndarray): 1-body part of the Hamiltonian

            vij (numpy.ndarray): 2-body part of the Hamiltonian

        Returns:
            (Wavefunction): wavefunction after 1 time step of the time evolution.
                since this function is in-place, what is returned is always self.
        """

        for key, sector in self._civec.items():
            self._civec[key].coeff = sector.evolve_diagonal_coulomb(
                diag, vij, inplace=True)
        return self

    def expectationValue(
            self,
            ops: Union['fqe_operator.FqeOperator', 'hamiltonian.Hamiltonian'],
            brawfn: 'Wavefunction' = None) -> Union[complex, numpy.ndarray]:
        """Calculates expectation values given operators

        Args:
            ops (FqeOperator or Hamiltonian): operator for which the expectation value is \
                computed

            brawfn (Wavefunction): bra-side wave function for transition quantity (optional)

        Returns:
            (complex or numpy.ndarray): resulting expectation value or RDM
        """
        if isinstance(ops, fqe_operator.FqeOperator):
            if brawfn:
                return ops.contract(brawfn, self)
            return ops.contract(self, self)

        if isinstance(ops, str):
            if any(char.isdigit() for char in ops):
                ops = sparse_hamiltonian.SparseHamiltonian(ops)
            else:
                return self.rdm(ops, brawfn=brawfn)

        if not isinstance(ops, hamiltonian.Hamiltonian):
            raise TypeError('Expected an Fqe Hamiltonian or Operator' \
                            ' but recieved {}'.format(type(ops)))
        workwfn = self.apply(ops)

        if brawfn:
            return vdot(brawfn, workwfn)
        return vdot(self, workwfn)

    def _apply_individual_nbody(self,
                                hamil: 'sparse_hamiltonian.SparseHamiltonian',
                                base: 'Wavefunction' = None) -> 'Wavefunction':
        """
        Applies an individual n-body operator to the wave function self.

        Args:
            hamil (SparseHamiltonian): Sparse Hamiltonian to be applied to the wavefunction
`
            base (Wavefunction): the result will be accumulated to base

        Returns:
            (Wavefunction): resulting wavefunciton. When base is provided, base is returned
        """
        assert isinstance(hamil, sparse_hamiltonian.SparseHamiltonian)

        if hamil.nterms() > 1:
            raise ValueError(
                'Indivisual n-body code is called with multiple terms')

        [(coeff, alpha, beta)] = hamil.terms()
        daga = []
        dagb = []
        undaga = []
        undagb = []
        for oper in alpha:
            assert oper[0] < self._norb
            if oper[1] == 1:
                daga.append(oper[0])
            else:
                undaga.append(oper[0])
        for oper in beta:
            assert oper[0] < self._norb
            if oper[1] == 1:
                dagb.append(oper[0])
            else:
                undagb.append(oper[0])

        if len(daga) + len(dagb) != len(undaga) + len(undagb):
            raise ValueError('Number non-conserving operators specified')

        if base is None:
            out = self.empty_copy()
        else:
            out = base
        if len(daga) == len(undaga) and len(dagb) == len(undagb):
            for key in self._civec.keys():
                out._civec[key].apply_individual_nbody_accumulate(
                    coeff, self._civec[key], daga, undaga, dagb, undagb)
        else:
            ssectors = self._number_sectors()
            nsectors = out._number_sectors()
            for skey, nsector in nsectors.items():
                nsector.apply_individual_nbody_accumulate(
                    coeff, ssectors[skey], daga, undaga, dagb, undagb)
        return out

    def _evolve_individual_nbody(self,
                                 time: float,
                                 hamil: 'sparse_hamiltonian.SparseHamiltonian',
                                 inplace: bool = False) -> 'Wavefunction':
        """Apply up to 4-body individual operator.

        This routine assumes the Hamiltonian is normal ordered.

        Args:
            time (float): time for evolution

            hamil (SparseHamiltonian): Sparse Hamiltonian using which \
                the wavefunction is evolved

            inplace (bool): whether to store the results in-place

        Returns:
            (Wavefunction): resulting wavefunciton. If inplace is True, \
                self is returned.
        """
        if not isinstance(hamil, sparse_hamiltonian.SparseHamiltonian):
            raise TypeError('Expected a Hamiltonian Object but received' \
                            ' {}'.format(hamil))

        if hamil.nterms() > 2:
            raise ValueError(
                'Individual n-body code is called with multiple terms')

        # check if everything is paired
        if hamil.nterms() == 2:
            [(coeff0, alpha0, beta0), (coeff1, alpha1, beta1)] = hamil.terms()
            check = True
            for aop in alpha0:
                check &= (aop[0], aop[1] ^ 1) in alpha1
            for bop in beta0:
                check &= (bop[0], bop[1] ^ 1) in beta1

            if self._conserve_number:
                if not check:
                    raise ValueError(
                        'Operators in _evolve_individual_nbody is not Hermitian'
                    )
        else:
            [
                (coeff0, alpha0, beta0),
            ] = hamil.terms()
            check = True
            for aop in alpha0:
                check &= (aop[0], aop[1] ^ 1) in alpha0
            for bop in beta0:
                check &= (bop[0], bop[1] ^ 1) in beta0

            if self._conserve_number:
                if not check:
                    raise ValueError(
                        'Operators in _evolve_individual_nbody is not Hermitian'
                    )
            coeff0 *= 0.5

        daga = []
        dagb = []
        undaga = []
        undagb = []

        for oper in alpha0:
            assert oper[0] < self._norb
            if oper[1] == 1:
                daga.append(oper[0])
            else:
                undaga.append(oper[0])

        for oper in beta0:
            assert oper[0] < self._norb
            if oper[1] == 1:
                dagb.append(oper[0])
            else:
                undagb.append(oper[0])

        if hamil.nterms() == 2:
            parity = (-1)**(len(alpha0) * len(beta0) + len(daga) * (len(daga)-1)//2 \
                            + len(dagb) * (len(dagb) - 1) // 2 \
                            + len(undaga) * (len(undaga) - 1) // 2 \
                            + len(undagb) * (len(undagb) - 1) // 2)
            if not numpy.abs(coeff0 - numpy.conj(coeff1) * parity) < 1.0e-8:
                raise ValueError(
                    'Coefficients in _evolve_individual_nbody is not Hermitian')

        if daga == undaga and dagb == undagb:
            out = self if inplace else copy.deepcopy(self)
            for _, sector in out._civec.items():
                sector.evolve_inplace_individual_nbody_trivial(
                    time, coeff0, daga, dagb)
        else:
            out = Wavefunction()
            out._norb = self._norb
            out._conserved = self._conserved
            out._symmetry_map = self._symmetry_map
            if len(daga) == len(undaga) and len(dagb) == len(undagb):
                for label, isec in self._civec.items():
                    osec = isec.evolve_individual_nbody_nontrivial(time, coeff0, daga, \
                                                                   undaga, dagb, undagb)
                    out._civec[label] = osec
            else:
                nsectors = self._number_sectors()
                for _, nsector in nsectors.items():
                    osector = nsector.evolve_individual_nbody(time, coeff0, daga, \
                                                              undaga, dagb, undagb)
                    for (nalpha, nbeta), osec in osector.sectors().items():
                        out._civec[(nalpha + nbeta, nalpha - nbeta)] = osec
        return out

    def _apply_few_nbody(self, hamil: 'sparse_hamiltonian.SparseHamiltonian'
                        ) -> 'Wavefunction':
        """ Applies SparseHamiltonian by looping over all of the operators.
        Useful when the operator is extremely sparse

        Args:
            hamil (SparseHamiltonian): Sparse Hamiltonian to be applied to the wavefunction

        Returns:
            (Wavefunction): resulting wavefunciton
        """
        out = None
        for oper in hamil.terms_hamiltonian():
            if out is None:
                out = self._apply_individual_nbody(oper)
            else:
                out = self._apply_individual_nbody(oper, base=out)

        if out is None:
            out = copy.deepcopy(self)

        if numpy.abs(hamil.e_0()) > 1.e-15:
            out.ax_plus_y(hamil.e_0(), self)

        return out

    @wrap_rdm
    def rdm(self, string: str, brawfn: Optional['Wavefunction'] = None
           ) -> Union[complex, numpy.ndarray]:
        """ Returns rank-1 RDM. The operator order is specified by string.
        Note that, if the entire RDM is requested for N-broken wave function,
        this returns a packed format.

        Args:
            string (str): character strings that specify the quantity to be computed

            brawfn (Wavefunction): bra-side wave function for transition RDM (optional)

        Returns:
            Union[complex, numpy.ndarray]: Resulting RDM in numpy.ndarray or \
                an RDM element in complex
        """
        rank = len(string.split()) // 2
        if any(char.isdigit() for char in string):
            result = self.apply(sparse_hamiltonian.SparseHamiltonian(string))
            if brawfn is None:
                return vdot(self, result)
            return vdot(brawfn, result)

        fqe_ops_utils.validate_rdm_string(string)
        rdm = list(self._compute_rdm(rank, brawfn))
        return wick(string, rdm, self._conserve_spin)

    def _compute_rdm(self, rank: int, brawfn: 'Wavefunction' = None
                    ) -> Tuple[numpy.ndarray, ...]:
        """Internal function that computes RDM up to rank = rank

        Args:
            rank (int): the rank up to which RDMs are computed

            brawfn (Wavefunction): bra wavefunction for transition RDMs (optional)

        Returns:
            Tuple[numpy.ndarray, ...]: tuple of RDMs
        """
        assert rank > 0
        assert rank < 5

        out: List[Any] = [None, None, None, None]
        tmp: List[Any] = [None, None, None, None]

        if self._conserve_spin:
            for key, sector in self._civec.items():
                assert brawfn is None or key in brawfn.sectors()
                bra = None if brawfn is None else brawfn._civec[key]
                if rank == 1:
                    (tmp[0],) = sector.rdm1(bra)
                elif rank == 2:
                    (tmp[0], tmp[1]) = sector.rdm12(bra)
                elif rank == 3:
                    (tmp[0], tmp[1], tmp[2]) = sector.rdm123(bra)
                elif rank == 4:
                    (tmp[0], tmp[1], tmp[2], tmp[3]) = sector.rdm1234(bra)
                for i in range(4):
                    if tmp[i] is not None:
                        out[i] = tmp[i] if out[i] is None else out[i] + tmp[i]
            out2: List[numpy.ndarray] = []
            for i in range(4):
                if out[i] is not None:
                    out2.append(out[i])
            return tuple(out2)

        numbersectors = self._number_sectors()
        for nkey, dataset in numbersectors.items():
            assert brawfn is None or nkey in brawfn._number_sectors().keys()
            nbra = None if brawfn is None else brawfn._number_sectors()[nkey]
            if rank == 1:
                (tmp[0],) = dataset.rdm1(nbra)
            elif rank == 2:
                (tmp[0], tmp[1]) = dataset.rdm12(nbra)
            elif rank == 3:
                (tmp[0], tmp[1], tmp[2]) = dataset.rdm123(nbra)
            elif rank == 4:
                (tmp[0], tmp[1], tmp[2], tmp[3]) = dataset.rdm1234(nbra)
            for i in range(4):
                if tmp[i] is not None:
                    out[i] = tmp[i] if out[i] is None else out[i] + tmp[i]
        out2 = []
        for i in range(4):
            if out[i] is not None:
                out2.append(out[i])
        return tuple(out2)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Ed Mountjoy (March 2018)
#
"""
Wrapper for Verena Zuber's colocalisation method

Assumptions:
    - GWAS and cov matrix are referring to the same strand
"""

import sys
import argparse
import pandas as pd
import numpy as np
import Finemap
import pickle

def main():

    # Get args
    args = parse_args()
    chrom, start, end = parse_range(args.range)

    #
    # Load and run 1-D finemapping on left dataset -----------------------------
    #

    # Load left sumstats
    print("Loading left sumstats...")
    left_z, left_efal, left_n = parse_sumstats(args.left_sumstats, chrom, start,
        end, args.left_rsidcol, args.left_chromcol, args.left_poscol,
        args.left_betacol, args.left_secol, args.left_effalcol, args.left_ncol,
        args.sep)
    print("  {0} variants loaded...".format(left_z.shape[0]))

    # Load left cov matrix
    print("Loading left cov matrix...")
    left_cov, left_cov_efal = parse_cov_matrix(args.left_cov, args.left_covmeta)
    print("  {0} variants loaded...".format(left_cov.shape[0]))

    # Harmonise left z-scores and cov matrix
    print("Harmonising left sumstat and cov matrix...")
    left_z, left_cov = harmonise_z_covmat(left_z, left_cov, left_efal,
                                          left_cov_efal)
    print("  {0} overlapping variants...".format(left_z.shape[0]))

    # Run finemap on left
    print("Running finemap on left...")
    left_res = Finemap.finemap(left_z.as_matrix(),
                               left_cov.as_matrix(),
                               left_n,
                               left_z.index.tolist(),
                               kmax=args.left_kmax,
                               kstart=args.left_kstart,
                               max_iter=args.maxiter,
                               prior=args.prior,
                               v_scale=args.v_scale,
                               g=args.g)
    # Write output as a pickled object and pandas df
    save_object(left_res, "{0}.left_configurations.pkl".format(args.outprefix))
    left_res_df = finemap_obj_to_df(left_res)
    left_res_df.to_csv("{0}.left_configurations.tsv".format(args.outprefix),
                       sep="\t", index=None)

    #
    # Load and run 1-D finemapping on 2nd (right) dataset ----------------------
    #

    # Skip if no 2nd dataset is provided
    if args.right_sumstats is None:
        print("No right dataset provided (--right_sumstats). Finishing here!")
        return 0

    # Load right sumstats
    print("Loading right sumstats...")
    right_z, right_efal, right_n = parse_sumstats(args.right_sumstats, chrom,
        start, end, args.right_rsidcol, args.right_chromcol, args.right_poscol,
        args.right_betacol, args.right_secol, args.right_effalcol,
        args.right_ncol, args.sep)
    print("  {0} variants loaded...".format(right_z.shape[0]))

    # Load right cov matrix
    print("Loading right cov matrix...")
    right_cov, right_cov_efal = parse_cov_matrix(args.right_cov,
                                                 args.right_covmeta)
    print("  {0} variants loaded...".format(right_cov.shape[0]))

    # Harmonise right z-scores and cov matrix
    print("Harmonising right sumstat and cov matrix...")
    right_z, right_cov = harmonise_z_covmat(right_z, right_cov, right_efal,
                                            right_cov_efal)
    print("  {0} overlapping variants...".format(right_z.shape[0]))

    # Run finemap on right
    print("Running finemap on right...")
    right_res = Finemap.finemap(right_z.as_matrix(),
                                right_cov.as_matrix(),
                                right_n,
                                right_z.index.tolist(),
                                kmax=args.right_kmax,
                                kstart=args.right_kstart,
                                max_iter=args.maxiter,
                                prior=args.prior,
                                v_scale=args.v_scale,
                                g=args.g)
    # Write output as a pickled object and pandas df
    save_object(right_res, "{0}.right_configurations.pkl".format(args.outprefix))
    right_res_df = finemap_obj_to_df(right_res)
    right_res_df.to_csv("{0}.right_configurations.tsv".format(args.outprefix),
                       sep="\t", index=None)

    #
    # Caluclate 2D joint posteriors (colocalisation)
    #

    print("Running colocalisation between left and right...")

    coloc_evidence, joint_res = left_res.joint_posterior(right_res)

    # Write table of joint posterior configurations
    joint_res_df = joint_posterior_to_df(joint_res)
    joint_res_df.to_csv("{0}.joint_configurations.tsv".format(args.outprefix),
                       sep="\t", index=None)
    save_object(joint_res, "{0}.joint_configurations.pkl".format(args.outprefix))
    # Write file containing single colocalisation evidence score
    outf = "{0}.joint_evidence.tsv".format(args.outprefix)
    with open(outf, "wb") as out_h:
        out_h.write("{0}\n".format(coloc_evidence).encode("utf8"))

    print("Finished!")

    return 0

def joint_posterior_to_df(obj):
    """ Converts POSTGAP joint posterior (colocalisation) in a pandas df
    Args:
        obj (TwoDConfigurationSample)
    Returns:
        Dataframe
    """
    # Add each configuration as a row
    rows = []
    for configuration in obj.configurations:
        snp_labels = ";".join([obj.labels[position] for position in configuration])
        index = obj.configurations[configuration]
        posterior  = obj.posterior[index]
        posterior1 = obj.posterior1[index]
        posterior2 = obj.posterior2[index]
        rows.append([snp_labels, posterior1, posterior2, posterior])
    df = pd.DataFrame(rows, columns=["snps", "left_posterior", "right_posterior", "joint_posterior"])
    df = df.sort_values("joint_posterior", ascending=False)
    return df

def finemap_obj_to_df(obj):
    """ Converts POSTGAP finemap object to a pandas df
    Args:
        obj (OneDConfigurationSample)
    Returns:
        Dataframe
    """
    # Add each configuration as a row
    rows = []
    for configuration in obj.configurations:
        index_of_configuration = obj.configurations[configuration]
        posterior = obj.posterior[index_of_configuration]
        prior = np.exp(obj.log_prior[index_of_configuration])
        logBF = obj.log_BF[index_of_configuration]
        snps = ';'.join([obj.labels[position] for position in configuration])
        rows.append([snps, prior, logBF, posterior])
    df = pd.DataFrame(rows, columns=["snps", "prior", "logBF", "posterior"])
    df = df.sort_values("posterior", ascending=False)

    return df

def save_object(obj, filename):
    """ Write object as a pickle
    Args:
        obj: any object
        filename: output file
    """
    with open(filename, 'wb') as output:
        pickle.dump(obj, output, pickle.HIGHEST_PROTOCOL)
    return 0

def harmonise_z_covmat(z, cov, z_efal, cov_efal):
    """ Harmonise z-scores and covariance matrix. This includes:
        1. Take intersection of rsids between the two
        2. sort into the same orders
        3. if effect alleles are not the same, flip in the cov matrix
    Args:
        z (Series)
        cov (Dataframe)
        z_efal (Series)
        cov_efal (Series)
    Returns:
        z (Series), cov (Dataframe)
    """
    intersect = z.index[z.index.isin(cov.index)]
    z = z.loc[intersect]
    cov = cov.loc[intersect, intersect]

    # Make harmonisation vector to flip cov if effect alleles don't match
    vec = [1 if z_efal[snp] == cov_efal[snp] else -1 for snp in intersect]
    vec = pd.Series(vec, index=intersect)

    # Multiply the cov matrix by the harmonisation vector across both axes
    cov = cov.multiply(vec, axis=0)
    cov = cov.multiply(vec, axis=1)

    return z, cov

def parse_cov_matrix(in_cov, in_meta):
    """ Parses ldstore cov matrix and associated meta data file
    Args:
        in_cov (str): file containing matrix of correlations
        in_meta (str): file containing variant meta-data
    Returns:
        cov matrix (dataframe), effect alleles (dict)
    """
    cov = pd.read_csv(in_cov, sep=" ", header=None)
    meta = pd.read_csv(in_meta, sep=" ", header=0)
    cov.index = meta.RSID
    meta.index = meta.RSID
    cov.columns = meta.RSID

    # Remove duplicates
    isdupe = meta.index.duplicated()
    cov = cov.loc[~isdupe, ~isdupe]
    meta = meta[~isdupe]

    return cov, meta["A_allele"].to_dict()

def parse_sumstats(inf, chrom, start, end, rsid_col, chrom_col, pos_col,
                   beta_col, se_col, effal_col, n_col, sep):
    """ Parse a summary stat file and return zscores
    Args:
        inf (str):   input summary stat file
        chrom (str): chrom to extract
        start (int): position to extract from
        end (int):   position to extract to
        ...column names
    Returns:
        zscores (Series), effect alleles (dict), n (int)
    """
    # Load iteratively
    iter_csv = pd.read_csv(inf, sep=sep, header=0, iterator=True,
               chunksize=100000)
    iter_chunks = []
    for chunk in iter_csv:
        chunk[chrom_col] = chunk[chrom_col].astype(str)
        iter_chunks.append(chunk[((chunk[chrom_col] == chrom) &
                               (chunk[pos_col] >= start) &
                               (chunk[pos_col] <= end))])
    df = pd.concat(iter_chunks)
    # df.to_csv("temp/test_sumstat_slice.tsv", index=None, sep="\t") # DEBUG

    # Only keep rows with a valid rsid
    df = df[df[rsid_col].str.startswith("rs")]
    df.index = df[rsid_col]

    # Remove duplicates
    df = df.loc[~df.index.duplicated(), :]

    # Calculate z-scores
    zscores = df[beta_col] / df[se_col]

    # Get mean N
    n_mean = int(np.mean(df[n_col]))

    return zscores, df[effal_col].to_dict(), n_mean

def parse_range(range_str):
    """ Parse the chrom, start and end from the range string
    Args:
        range_str (str): In format chrom:start-end
    Returns:
        chrom (str), start (int), end (int)
    """
    chrom, start_end = range_str.split(":")
    start, end = start_end.split("-")
    return str(chrom), int(start), int(end)

def parse_args():
    """ Load command line args.
    """
    parser = argparse.ArgumentParser()
    # Input args
    parser.add_argument('--left_sumstats', metavar="<file>", help=('Summary statistics file'), type=str, required=True)
    parser.add_argument('--left_cov', metavar="<file>", help=("Covariance matrix - correlation structure between variants"), type=str, required=True)
    parser.add_argument('--left_covmeta', metavar="<file>", help=("Covariance matrix SNP info (from LDstore)"), type=str, required=True)
    parser.add_argument('--range', metavar="<str>", help=("Genomic range in format chrom:start-end"), type=str, required=True)
    parser.add_argument('--right_sumstats', metavar="<file>", help=('(Optional) Summary statistics file'), type=str, required=False)
    parser.add_argument('--right_cov', metavar="<file>", help=("Covariance matrix - leave blank if same as left_cov"), type=str, required=False)
    parser.add_argument('--right_covmeta', metavar="<file>", help=("Covariance matrix SNP info (from LDstore)"), type=str, required=False)
    # Output args
    parser.add_argument('--outprefix', metavar="<str>", help=("Output prefix"), type=str, required=True)
    # Finemapping args
    parser.add_argument('--left_kmax', metavar="<int>", help=('Maximum number of causal variants (default: 5)'), default=5, type=int)
    parser.add_argument('--left_kstart', metavar="<int>", help=('Full exploration of sets with #kstart causal variants (default: 1)'), default=1, type=int)
    parser.add_argument('--right_kmax', metavar="<int>", help=('Maximum number of causal variants (default: left_kmax)'), type=int)
    parser.add_argument('--right_kstart', metavar="<int>", help=('Full exploration of sets with #kstart causal variants (default: left_kstart)'), type=int)
    parser.add_argument('--v_scale', metavar="<float>", help=('Prior variance of the independence prior (default: 0.0025)'), default=0.0025, type=float)
    parser.add_argument('--maxiter', metavar="<int>", help=('Max iterations of stochastic search (default: 100000)'), default=100000, type=int)
    parser.add_argument('--prior', metavar="<str>", help=('Choice of "independence" or "gprior" (default: independence)'), default="independence", type=str, choices=["independence", "gprior"])
    parser.add_argument('--g', metavar="<str>", help=('g-parameter of the g-prior (default: BRIC). see Mixtures of g Priors for Bayesian Variable Selection Liang et al 2008'), default="BRIC", type=str, choices=["BRIC", "BIC", "RIC"])
    # File parsing args
    parser.add_argument('--sep', metavar="<str>", help=('Column sep (default: tab)'), type=str, default="\t")
    parser.add_argument('--left_rsidcol', metavar="<str>", help=('RSID column (default: rsid)'), default="rsid", type=str)
    parser.add_argument('--left_chromcol', metavar="<str>", help=('Chromosome column (default: chrom)'), default="chrom", type=str)
    parser.add_argument('--left_poscol', metavar="<str>", help=('Position column (default: pos)'), default="pos", type=str)
    parser.add_argument('--left_betacol', metavar="<str>", help=('Beta column (default: beta)'), default="beta", type=str)
    parser.add_argument('--left_secol', metavar="<str>", help=('Standard error column (default: se)'), default="se", type=str)
    parser.add_argument('--left_effalcol', metavar="<str>", help=('Effect allele column (default: effect_allele)'), default="effect_allele", type=str)
    parser.add_argument('--left_ncol', metavar="<str>", help=('Sample size column (default: n)'), default="n", type=str)
    parser.add_argument('--right_rsidcol', metavar="<str>", help=('RSID column (default: left_rsidcol)'), type=str)
    parser.add_argument('--right_chromcol', metavar="<str>", help=('Chromosome column (default: left_chromcol)'), type=str)
    parser.add_argument('--right_poscol', metavar="<str>", help=('Position column (default: left_poscol)'), type=str)
    parser.add_argument('--right_betacol', metavar="<str>", help=('Beta column (default: left_betacol)'), type=str)
    parser.add_argument('--right_secol', metavar="<str>", help=('Standard error column (default: left_secol)'), type=str)
    parser.add_argument('--right_effalcol', metavar="<str>", help=('Effect allele columns (default: left_effalcol)'), type=str)
    parser.add_argument('--right_ncol', metavar="<str>", help=('Sample size column (default: n)'), default="n", type=str)

    args = parser.parse_args()

    # If right argument is not set, set to same as left argument
    for rarg in ["right_cov", "right_covmeta", "right_kmax", "right_kstart",
                 "right_rsidcol", "right_chromcol", "right_poscol",
                 "right_betacol", "right_secol", "right_effalcol",
                 "right_ncol"]:
        if getattr(args, rarg) is None:
            larg_val = getattr(args, rarg.replace("right_", "left_"))
            setattr(args, rarg, larg_val)

    return args

if __name__ == '__main__':

    main()