# Encoding: UTF-8
# Copyright (c) Marnik Bercx, University of Antwerp
# Distributed under the terms of the MIT License

import itertools
import math
import json
import os
import pdb

import numpy as np

from monty.io import zopen
from monty.json import jsanitize, MSONable
from pymatgen.core import Structure, Composition, Molecule, Site, Element
from pymatgen.analysis.chemenv.coordination_environments.voronoi \
    import DetailedVoronoiContainer
from pymatgen.io.vasp.outputs import Outcar
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.analysis.transition_state import NEBAnalysis
from pymatgen.util.plotting import pretty_plot

scipy_old_piecewisepolynomial = True
try:
    from scipy.interpolate import PiecewisePolynomial
except ImportError:
    from scipy.interpolate import CubicSpline

    scipy_old_piecewisepolynomial = False

"""
Module that contains tools to represent and calculate the properties of
battery cathodes.

"""

__author__ = "Marnik Bercx"
__copyright__ = "Copyright 2018, Marnik Bercx, University of Antwerp"
__version__ = "0.1"
__maintainer__ = "Marnik Bercx"
__email__ = "marnik.bercx@uantwerpen.be"
__date__ = "May 2018"

# TODO Currently, the dimers are defined by their indices.
# This is a consequence of the fact that the DetailedVoronoiContainer
# expects indices for its neighbor method. Frankly, I would prefer sites as
# the basis of the dimer definition as well as it's environment.

# Values for determining the neighbors of a site in a voronoi decomposition
VORONOI_DIST_FACTOR = 1.4
VORONOI_ANG_FACTOR = 0.6

# Tuple of possible cations. This idea should work fine, considering the
# fact that these cation elements rarely serve another purpose than being
# a cation.
CATIONS = ("Li", "Na", "Mg")

# Tolerance for determining whether two oxygens are on opposite sides of an
# octahedron. If the angle between the two vectors connecting the site and
# the corresponding oxygens is larger than this value, the oxygens are
# considered to be opposites.
OXYGEN_ANGLE_TOL = math.pi * 9 / 10

# Tolerance for the representation determination. This can be pretty big, since
# the dimer environment structure is known.
REPRESENTATION_DIST_TOL = 5e-1
REPRESENTATION_ANGLE_TOL = 2e-1

# Dimer representation symmetry permutations
SYMMETRY_PERMUTATIONS = [[1, 2, 4, 3, 6, 5, 8, 7, 9, 10, 11, 12],
                         [2, 1, 3, 4, 7, 8, 5, 6, 11, 12, 9, 10],
                         [1, 2, 3, 4, 5, 6, 7, 8, 10, 9, 12, 11],
                         [2, 1, 4, 3, 8, 7, 6, 5, 11, 12, 9, 10],
                         [1, 2, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11],
                         [2, 1, 4, 3, 8, 7, 6, 5, 12, 11, 10, 9],
                         [2, 1, 3, 4, 7, 8, 5, 6, 12, 11, 10, 9],
                         [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]]

class Cathode(Structure):
    """
    A class representing a cathode material in a battery.

    The main idea of this class is to keep track of the original sites by
    using sites with empty Compositions. This is important to make sure that
    the voronoi decomposition is successful, and hence interesting if we
    want to look at coordinations and neighbors. Another advantage is that
    we can consider the empty cation sites for final positions of transition
    metal migrations.

    """

    def __init__(self, lattice, species, coords, charge=None,
                 validate_proximity=False,
                 to_unit_cell=False, coords_are_cartesian=False,
                 site_properties=None):

        super(Cathode, self).__init__(
            lattice=lattice, species=species, coords=coords,
            validate_proximity=validate_proximity, to_unit_cell=to_unit_cell,
            coords_are_cartesian=coords_are_cartesian,
            site_properties=site_properties
        )

        self._voronoi = None

    @property
    def cation_configuration(self):
        """
        A list of all sites which correspond to cations in the Cathode.

        Returns:

        """
        return [site for site in self.sites if site.species_string in CATIONS]

    @cation_configuration.setter
    def cation_configuration(self, configuration):

        # TODO Add checks

        # Remove all cations
        for cation in [cat for cat in CATIONS if Element[cat] in set(
                self.composition.keys())]:
            self.replace_species({cation: {cation: 0}})

        # Add the cation sites
        if isinstance(configuration, dict):
            for cation in configuration.keys():
                for index in configuration[cation]:
                    self.replace(index, cation, properties={"magmom": 0})

        elif all([isinstance(item, Site) for item in configuration]):
            for catsite in configuration:
                for i, site in enumerate(self):
                    if np.linalg.norm(site.distance(catsite)) < 0.05:
                        self.replace(i, catsite.specie,
                                     properties={"magmom": 0})

        else:
            raise TypeError("Cation configurations should be a dictionary "
                            "mapping cations to site indices or a list of "
                            "sites.")

    def __str__(self):
        """
        Overwritten string representation, in order to provide information
        about the cation configuration, as well as the VESTA index, which
        is useful when defining structural changes.

        """
        outs = ["Full Formula ({s})".format(s=self.composition.formula),
                "Reduced Formula: {}".format(self.composition.reduced_formula)]
        to_s = lambda x: "%0.6f" % x
        outs.append("abc   : " + " ".join([to_s(i).rjust(10)
                                           for i in self.lattice.abc]))
        outs.append("angles: " + " ".join([to_s(i).rjust(10)
                                           for i in self.lattice.angles]))
        if self._charge:
            if self._charge >= 0:
                outs.append("Overall Charge: +{}".format(self._charge))
            else:
                outs.append("Overall Charge: -{}".format(self._charge))
        outs.append("Sites ({i})".format(i=len(self)))
        data = []
        props = self.site_properties
        keys = sorted(props.keys())
        vesta_index = 1
        for i, site in enumerate(self):
            if site.species_and_occu.num_atoms == 0:
                row = [str(i), "-", "Vac"]

            else:
                row = [str(i), vesta_index, site.species_string]
                vesta_index += 1

            row.extend([to_s(j) for j in site.frac_coords])
            for k in keys:
                row.append(props[k][i])
            data.append(row)

        from tabulate import tabulate
        outs.append(
            tabulate(data,
                     headers=["#", "#VESTA", "SP", "a", "b", "c"] + keys,
                     ))
        return "\n".join(outs)

    @property
    def voronoi(self):
        """
        ChemEnv voronoi decomposition of the cathode structure.

        Returns:
            pymatgen.analysis.chemenv.coordination_environments.voronoi.\
            DetailedVoronoiContainer

        """
        if self._voronoi is None:
            self._voronoi = DetailedVoronoiContainer(self)

        return self._voronoi

    @voronoi.setter
    def voronoi(self, voronoi_container):
        self._voronoi = voronoi_container

    def add_cations(self, sites=None):
        """
        Args:
            sites:

        Returns:

        """

        # Add the cation sites
        if isinstance(sites, dict):
            for cation in sites.keys():
                for index in sites[cation]:
                    self.replace(index, cation, properties={"magmom": 0})

        elif all([isinstance(item, Site) for item in sites]):
            for catsite in sites:
                for i, site in enumerate(self):
                    if np.linalg.norm(site.distance(catsite)) < 0.05:
                        self.replace(i, catsite.specie,
                                     properties={"magmom": 0})

        else:
            raise TypeError("Cation configurations should be a dictionary "
                            "mapping cations to site indices or a list of "
                            "sites.")

    def remove_cations(self, sites=None):
        """
        Remove the cations from the cathode, i.e. delithiate the structure in
        case Li is the cation of the cathode.

        Note that this does not remove the sites from the pymatgen Structure.
        The occupancy is simply adjusted to an empty Composition object.

        Args:
            sites: List of indices
            List of pymatgen.core.Sites which are to be removed.

        """

        # TODO add checks

        # If no sites are given
        if sites is None:
            # Remove all the cations
            self.cation_configuration = []

        # If a List of integers is given
        elif all([isinstance(item, int) for item in sites]):
            for index in sites:
                self.replace(index, Composition(), properties={"magmom": 0})

        # If a List of sites is given
        elif all([isinstance(item, Site) for item in sites]):
            for site in sites:
                # Check if the provided site corresponds to a cation site
                if site in self.cation_configuration:
                    cat_conf = self.cation_configuration.copy()
                    cat_conf.remove(site)
                    self.cation_configuration = cat_conf
                else:
                    raise Warning("Requested site not found in cation "
                                  "configuration.")
        else:
            raise IOError("Incorrect site input.")

    def change_site_distance(self, site_indices, distance):
        """
        Change the coordinates of two sites in a structure in order to adjust
        their distance.

        Args:
            site_indices:
            distance:
        """

        # TODO Add possibility of site_indices simply being the sites

        site_a = self.sites[site_indices[0]]
        site_b = self.sites[site_indices[1]]

        # Find the distance between the sites, as well as the image of site B
        # closest to site A
        (original_distance, closest_image_b) = site_a.distance_and_image(
            site_b)

        image_cart_coords = self.lattice.get_cartesian_coords(
            site_b.frac_coords + closest_image_b
        )

        # Calculate the vector that connects site A with site B
        connection_vector = image_cart_coords - site_a.coords

        # Make it a unit vector
        connection_vector /= np.linalg.norm(connection_vector)

        # Calculate the distance the sites need to be moved.
        site_move_distance = (original_distance - distance) / 2

        # Calculate the new cartesian coordinates of the sites
        new_site_a_coords = site_a.coords \
                            + site_move_distance * connection_vector
        new_site_b_coords = site_b.coords \
                            - site_move_distance * connection_vector

        # Change the sites in the structure
        self.replace(i=site_indices[0], species=site_a.species_string,
                     coords=new_site_a_coords,
                     coords_are_cartesian=True,
                     properties=site_a.properties)

        self.replace(i=site_indices[1], species=site_b.species_string,
                     coords=new_site_b_coords,
                     coords_are_cartesian=True,
                     properties=site_b.properties)

    def update_sites(self, directory, ignore_magmom=False):
        """
        Based on the CONTCAR and OUTCAR of a geometry optimization, update the
        site coordinates and magnetic moments that were optimized. Note that
        this method relies on the cation configuration of the cathode not
        having changed.

        Args:
            directory (str): Directory in which the geometry optimization
                output files (i.e. CONTCAR and OUTCAR) are stores.
            ignore_magmom (bool): Flag that indicates that the final magnetic
                moments of the optimized structure should be ignored. This means
                that the magnetic moments of the Cathode structure will
                remain the same.
        """

        new_cathode = Cathode.from_file(os.path.join(directory, "CONTCAR"))

        out = Outcar(os.path.join(directory, "OUTCAR"))

        if ignore_magmom:
            magmom = [site.properties["magmom"] for site in self.sites
                      if site.species_and_occu != Composition()]
            new_cathode.add_site_property("magmom", magmom)
        else:
            magmom = [site["tot"] for site in out.magnetization]
            new_cathode.add_site_property("magmom", magmom)

        # Update the lattice
        self.modify_lattice(new_cathode.lattice)

        # Update the coordinates of the occupied sites.
        new_index = 0
        for i, site in enumerate(self):

            # If the site is not empty
            if site.species_and_occu != Composition():
                new_site = new_cathode.sites[new_index]
                # Update the site coordinates
                self.replace(i, species=new_site.species_and_occu,
                             coords=new_site.frac_coords,
                             properties=new_site.properties)
                new_index += 1

    def find_noneq_cations(self):
        """
        Find a list of the site indices of all non-equivalent cations.

        Returns:
            List of site indices

        """
        symmops = SpacegroupAnalyzer(self).get_space_group_operations()

        cation_indices = [
            index for index in range(len(self.sites))
            if not self.sites[index].species_string == "O"
        ]

        # Start with adding the first cation
        inequiv_cations = [cation_indices[0], ]

        for index in cation_indices[1:]:

            s1 = [self.sites[index], ]

            # Check if the site is equivalent with one of the sites in the
            # inequivalent list.
            inequivalent = True

            for inequive_index in inequiv_cations:

                s2 = [self.sites[inequive_index], ]

                if symmops.are_symmetrically_equivalent(s1, s2):
                    inequivalent = False

            if inequivalent:
                inequiv_cations.append(index)

        return inequiv_cations

    def find_cation_configurations(self):
        """
        Plan is to find all non-equivalent cation configurations. Is probably
        already implemented elsewhere.

        Returns:

        """
        raise NotImplementedError

    def set_to_high_spin(self):
        """

        :return:
        """
        raise NotImplementedError

    def set_to_low_spin(self):
        """

        :return:
        """
        raise NotImplementedError

    def as_ordered_structure(self):
        """
        Return the structure as a pymatgen.core.Structure, removing the
        unoccupied sites. This is because many of the IO methods of pymatgen
        run into issues when empty occupancies are present.

        Returns:
            pymatgen.core.Structure

        """

        return Structure.from_sites(
            [site for site in self.sites
             if site.species_and_occu != Composition()]
        )

    def to(self, fmt=None, filename=None, **kwargs):
        """
        Structure method override to solve issue with writing the Cathode to a
        POSCAR file

        # TODO Figure out what exactly was the problem here again... Should
        have written this down immediately! I think it had something to do
        with the order of the Sites changing...

        Args:
            fmt:
            filename:
            **kwargs:

        Returns:

        """

        if fmt == "poscar":
            structure = self.as_ordered_structure()
            return structure.to(fmt, filename, **kwargs)

        else:
            return super(Cathode, self).to(fmt, filename, **kwargs)

    @classmethod
    def from_structure(cls, structure):
        """
        Initializes a Cathode from a pymatgen.core.Structure.

        Args:
            structure (pymatgen.core.Structure): Structure from which to
            initialize the Cathode.

        Returns:
            pybat.core.Structure

        """

        return cls.from_sites(structure.sites)

class LiRichCathode(Cathode):
    """
    A class representing a Li-rich cathode material.

    """

    def __init__(self, lattice, species, coords, charge=None,
                 validate_proximity=False,
                 to_unit_cell=False, coords_are_cartesian=False,
                 site_properties=None):

        super(LiRichCathode, self).__init__(
            lattice=lattice, species=species, coords=coords,
            validate_proximity=validate_proximity, to_unit_cell=to_unit_cell,
            coords_are_cartesian=coords_are_cartesian,
            site_properties=site_properties
        )

    def find_oxygen_dimers(self, site_index=None,
                           oxygen_angle_tol=OXYGEN_ANGLE_TOL):
        """
        Returns a list of index pairs corresponding to the oxygen dimers that
        can be formed around the site provided by the user, i.e. with oxygens
        that are neighbours of the provided site.

        Args:
            site_index:

        Returns:

        """

        if site_index is None:

            # TODO This is a lazy way of dealing with this...
            # Just using the method for single indices and then filtering
            # duplicates is probably not the fastest implementation. But it
            # sure is the fastest to implement :P

            oxygen_dimers = set()

            cation_indices = [
                index for index in range(len(self.sites))
                if not self.sites[index].species_string == "O"
            ]

            for index in cation_indices:
                for dimer in self.find_oxygen_dimers(index):
                    oxygen_dimers.add(dimer)

            return list(oxygen_dimers)

        else:

            # Determine the oxygen neighbors for the provided site
            oxygen_neighbors_indices = [
                neighbor["index"] for neighbor
                in self.voronoi.neighbors(site_index, VORONOI_DIST_FACTOR,
                                          VORONOI_ANG_FACTOR)
                if self.sites[neighbor["index"]].species_string == "O"
            ]

            # pdb.set_trace()

            if len(oxygen_neighbors_indices) <= 1:
                raise ValueError("Provided site does not have two oxygen "
                                 "neighbours.\n")

            # Find all oxygen neighbour combinations that can form dimers. This
            # means they are not opposites on the site's octahedron environment.
            oxygen_dimers = []
            oxygen_combinations = itertools.combinations(
                oxygen_neighbors_indices, 2
            )

            for oxygen_pair in oxygen_combinations:

                site = self.sites[site_index]
                oxygen_site_A = self.sites[oxygen_pair[0]]
                oxygen_site_B = self.sites[oxygen_pair[1]]

                oxygen_image_A = site.distance_and_image(oxygen_site_A)[1]
                oxygen_image_B = site.distance_and_image(oxygen_site_B)[1]

                image_A_cart_coords = oxygen_site_A.coords \
                                      + np.dot(oxygen_image_A,
                                               self.lattice.matrix)
                image_B_cart_coords = oxygen_site_B.coords \
                                      + np.dot(oxygen_image_B,
                                               self.lattice.matrix)

                oxygen_vector_A = image_A_cart_coords - site.coords
                oxygen_vector_B = image_B_cart_coords - site.coords

                if angle_between(oxygen_vector_A, oxygen_vector_B) < \
                        oxygen_angle_tol:
                    oxygen_dimers.append(oxygen_pair)

            if len(oxygen_dimers) > 12:
                raise Warning(
                    "Found more than 12 oxygen pairs around a single "
                    "site. This can be caused by the use of a small "
                    "unit cell. Results may not be useful.")

            return oxygen_dimers

    def remove_dimer_cations(self, dimer_indices):
        """

        Args:
            dimer_indices:

        Returns:

        """

        # Find the sites of the dimer environment
        dimer_environment = Dimer(self, dimer_indices).sites

        remove_sites = [site for site in dimer_environment
                        if site.species_string in CATIONS]

        self.remove_cations(remove_sites)

    def find_noneq_dimers(self, site_index=None, method="symmops"):
        """
        A script that distills the non-equivalent oxygen dimers around a site
        index.

        In case no site index is provided, the method will loop over all sites
        which do not correspond to an oxygen.

        Args:
            site_index (int): Index of the site around which the dimers
            should be considered.
            method (str): Method for determining equivalency:

            "symmops" - Two dimers are considered when they are
            symmetrically equivalent, using the symmetry operations of the
            cathode unit cell.

            "representation" - Two dimers are equivalent if their
            environments are the same, i.e. when they have the same
            representation.

        Returns:
            List of Tuples with the dimer indices

        """

        ineq_dimers = []

        if method == "symmops":

            symmops = SpacegroupAnalyzer(self).get_space_group_operations()

            # If no site is provided, consider all inequivalent cation sites
            if site_index is None:

                ineq_cations = self.find_noneq_cations()

                for index in ineq_cations:

                    site_dimers = self.find_oxygen_dimers(index)

                    for dimer in site_dimers:

                        d1 = [self.sites[dimer[0]], self.sites[dimer[1]]]

                        inequivalent = True

                        for ineq_dimer in ineq_dimers:

                            d2 = [self.sites[ineq_dimer[0]],
                                  self.sites[ineq_dimer[1]]]

                            if symmops.are_symmetrically_equivalent(d1, d2):
                                inequivalent = False

                        if inequivalent:
                            ineq_dimers.append(dimer)

            # Else only consider the site provided
            else:
                site_dimers = self.find_oxygen_dimers(site_index)

                for dimer in site_dimers:

                    d1 = [self.sites[dimer[0]], self.sites[dimer[1]]]

                    inequivalent = True

                    for ineq_dimer in ineq_dimers:

                        d2 = [self.sites[ineq_dimer[0]],
                              self.sites[ineq_dimer[1]]]

                        if symmops.are_symmetrically_equivalent(d1, d2):
                            inequivalent = False

                    if inequivalent:
                        ineq_dimers.append(dimer)

            return ineq_dimers

        elif method == "representation":
            # TODO update this method

            raise NotImplementedError()

            # if site_index is None:
            #
            #     ignore = set(CATIONS).union(("O",))
            #
            #     transition_metal_indices = [index for index in
            #                                 range(len(self.sites))
            #                                 if self.sites[index].species_string
            #                                 not in ignore]
            #
            #     for index in transition_metal_indices:
            #
            #         dimers = [Dimer(self, dimer_indices) for dimer_indices
            #                   in self.find_oxygen_dimers(index)]
            #
            #         for dimer in dimers:
            #             if dimer not in noneq_dimers:
            #                 noneq_dimers.append(dimer)
            #
            # else:
            #
            #     dimers = [Dimer(self, dimer_indices) for dimer_indices
            #               in self.find_oxygen_dimers(site_index)]
            #
            #     for dimer in dimers:
            #         if dimer not in noneq_dimers:
            #             noneq_dimers.append(dimer)
            #
            # return noneq_dimers

        else:
            raise IOError("Method for finding non-equivalent dimers is not "
                          "recognized.")

    def list_noneq_dimers(self):
        """
        Create a list of lists of equivalent dimers of the various
        non-equivalent dimers, i.e. group all dimers in the structure in
        lists of dimers that are equivalent to each other.

        Returns:

        """

        symmops = SpacegroupAnalyzer(self).get_space_group_operations()

        dimers = self.find_oxygen_dimers()

        noneq_dimer_lists = [[dimer, ] for dimer in self.find_noneq_dimers()]

        for dimer in dimers:

            d1 = [self.sites[dimer[0]], self.sites[dimer[1]]]

            for noneq_dimer_list in noneq_dimer_lists:

                d2 = [self.sites[noneq_dimer_list[0][0]],
                      self.sites[noneq_dimer_list[0][1]]]

                if symmops.are_symmetrically_equivalent(d1, d2):
                    noneq_dimer_list.append(dimer)

        return noneq_dimer_lists

# TODO Currently the whole dimer representation only works for the O-O
# dimers in the O3 stacking. Allowing for different oxygen frameworks will
# require some more possible representations. One way is to figure out the
# structure of the molecule and set up the representation for each
# structure, but I think it would pay off to think if there isn't some
# better way of checking the equivalency of dimer environments. Moreover,
# we should probably first test to see if considering only the immediate
# environment is sufficient.

class Dimer(MSONable):
    """
    Class definition of an oxygen dimer in a Li-rich cathode structure.

    """

    # TODO Give the definition of this class a good hard thinking over.

    def __init__(self, cathode, dimer_indices):
        """
        Initialize

        Args:
            cathode
            dimer_indices

        Returns:
            pybat.core.Dimer

        """
        self._cathode = cathode
        self._indices = tuple(dimer_indices)
        self._sites = list
        self._center = None
        self._representation = list

    def __eq__(self, other):
        """
        Checks if the dimer environments of two dimers are the same.

        Args:
            other:

        Returns:

        """
        is_equal = False

        for permutation in SYMMETRY_PERMUTATIONS:

            if [self.representation[index] for index in range(1, 13)] == \
                    [other.representation[key] for key in permutation]:
                is_equal = True

        return is_equal

    @property
    def cathode(self):
        return self._cathode

    @property
    def indices(self):
        return self._indices

    @property
    def sites(self):

        if self._sites is list:
            # Find the oxygen neighbours
            oxygen_a_neighbors = [
                neighbor["index"] for neighbor
                in self.cathode.voronoi.neighbors(self.indices[0],
                                                  VORONOI_DIST_FACTOR,
                                                  VORONOI_ANG_FACTOR)
            ]
            oxygen_b_neighbors = [
                neighbor["index"] for neighbor
                in self.cathode.voronoi.neighbors(self.indices[1],
                                                  VORONOI_DIST_FACTOR,
                                                  VORONOI_ANG_FACTOR)
            ]

            # Determine the indices of the oxygen environment. The indices are
            # sorted in such a way that the oxygen indices come first,
            # followed by the indices of the shared neighbors.
            # TODO This can be done better. Really.
            # TODO Fix issue for small unit cells
            # The issue for small unit cells is that the oxygen atoms have
            # more shared neighbors than two according to the voronoi
            # decomposition, which messes up the assignment of the
            # environment atoms in the representation. This needs to be fixed.
            shared_neighbors = tuple(
                set(oxygen_a_neighbors).intersection(oxygen_b_neighbors)
            )
            other_neighbors = set(oxygen_a_neighbors).union(oxygen_b_neighbors)
            other_neighbors.remove(shared_neighbors[0])
            other_neighbors.remove(shared_neighbors[1])
            other_neighbors = tuple(other_neighbors)

            dimer_environment_indices = self._indices + shared_neighbors \
                                        + other_neighbors

            # Recover the corresponding sites
            self._sites = [self.cathode.sites[index] for index
                           in dimer_environment_indices]

        return self._sites

    @property
    def center(self):

        if self._center is None:

            # Find the center of the oxygen sites
            oxygen_sites = []
            for site in self.sites:
                if site.species_string == "O":
                    oxygen_sites.append(site)

            (distance, oxygen_image) = oxygen_sites[0].distance_and_image(
                oxygen_sites[1])

            image_cart_coords = oxygen_sites[1].coords \
                                + np.dot(oxygen_image,
                                         self.cathode.lattice.matrix)

            self._center = (oxygen_sites[0].coords + image_cart_coords) / 2

        return self._center

    @property
    def representation(self):
        if self._representation is list:

            # A representation # TODO Add definition

            # The representation idea seems like the fastest way of being able
            # to compare two dimers. By reducing them to a representation, we
            # are able compare dimers by applying permutations that represent
            # symmetry transformations. This is not a very general approach,
            # clearly.

            # Note that the sites property is built in such a way that the
            # first two sites correspond to the oxygen atoms and the next two
            # correspond to the shared neighbours. This convention is made to
            # save us some work here.

            dimer_molecule = self.get_dimer_molecule()

            oxy_1 = dimer_molecule.sites[0]
            oxy_2 = dimer_molecule.sites[1]

            shared_neighbor_3 = dimer_molecule.sites[2]
            shared_neighbor_4 = dimer_molecule.sites[3]

            # The representation is defined as a dictionary between site
            # numbers and dimer environment sites
            representation = {1: oxy_1.species_and_occu,
                              2: oxy_2.species_and_occu,
                              3: shared_neighbor_3.species_and_occu,
                              4: shared_neighbor_4.species_and_occu}

            # TODO Find a cleaner way of assigning representation positions

            # Loop over the remaining sites to find their representation
            # positions
            for site in dimer_molecule.sites[4:]:

                # Find the sites which are in the plane of the oxygens and
                # their shared neighbors.
                if np.linalg.norm(oxy_1.coords
                                  - (
                                          shared_neighbor_4.coords - oxy_1.coords)
                                  - site.coords) < REPRESENTATION_DIST_TOL:
                    representation[5] = site.species_and_occu

                if np.linalg.norm(oxy_1.coords
                                  - (
                                          shared_neighbor_3.coords - oxy_1.coords)
                                  - site.coords) < REPRESENTATION_DIST_TOL:
                    representation[6] = site.species_and_occu

                if np.linalg.norm(oxy_2.coords
                                  - (
                                          shared_neighbor_4.coords - oxy_2.coords)
                                  - site.coords) < REPRESENTATION_DIST_TOL:
                    representation[7] = site.species_and_occu

                if np.linalg.norm(oxy_2.coords
                                  - (
                                          shared_neighbor_3.coords - oxy_2.coords) \
                                  - site.coords) < REPRESENTATION_DIST_TOL:
                    representation[8] = site.species_and_occu

                # Find the sites which are out of plane
                oxy_1_oop = np.cross(
                    shared_neighbor_4.coords - oxy_1.coords,
                    shared_neighbor_3.coords - oxy_1.coords
                )
                oxy_2_oop = np.cross(
                    shared_neighbor_3.coords - oxy_2.coords,
                    shared_neighbor_4.coords - oxy_2.coords
                )

                if angle_between(oxy_1_oop, site.coords - oxy_1.coords) \
                        < REPRESENTATION_ANGLE_TOL:
                    representation[9] = site.species_and_occu

                if angle_between(oxy_1_oop, site.coords - oxy_1.coords) \
                        > math.pi - REPRESENTATION_ANGLE_TOL:
                    representation[10] = site.species_and_occu

                if angle_between(oxy_2_oop, site.coords - oxy_2.coords) \
                        < REPRESENTATION_ANGLE_TOL:
                    representation[11] = site.species_and_occu

                if angle_between(oxy_2_oop, site.coords - oxy_2.coords) \
                        > math.pi - REPRESENTATION_ANGLE_TOL:
                    representation[12] = site.species_and_occu

            self._representation = representation

            if len(representation) != len(self.sites):
                raise ValueError("Failed to create dimer representation.")

        return self._representation

    def get_dimer_molecule(self):

        molecule_sites = []

        for site in self.sites:
            (distance, jimage) = site.distance_and_image_from_frac_coords(
                self.cathode.lattice.get_fractional_coords(self.center))

            image_cart_coords = \
                site.coords - np.dot(jimage, self.cathode.lattice.matrix)

            molecule_sites.append(Site(site.species_and_occu,
                                       image_cart_coords))

        return Molecule.from_sites(molecule_sites)

    def visualize_dimer_environment(self, filename=None):
        """
        Creates a .xyz file of the oxygen dimer environment in order to
        visualize it in e.g. VESTA.

        Returns:

        """

        # Turn the structure into a molecule, ignoring the sites with zero
        # occupancy
        dimer_environment_molecule = Molecule.from_sites(
            [site for site in self.get_dimer_molecule().sites
             if not site.species_and_occu == Composition({"": 0})]
        )

        if filename is None:
            filename = str(self.cathode.composition.reduced_composition
                           ).replace(" ", "") + "_" \
                       + str(self.indices[0]) + "_" + str(self.indices[1]) \
                       + ".xyz"

        dimer_environment_molecule.to("xyz", filename)

    # TODO Check if MSONable methods need to be implemented

    def to(self, fmt="json", filename=None):

        if fmt == "json":
            if filename:
                with zopen(filename, "wt", encoding='utf8') as file:
                    return json.dump(self.as_dict(), file)
            else:
                return json.dumps(self.as_dict())
        else:
            raise NotImplementedError("Currently only json format is "
                                      "supported.")

    @classmethod
    def from_str(cls, input_string, fmt="json"):
        """
        Initialize a Facet from a string.

        Currently only supports 'json' formats.

        Args:
            input_string (str): String from which the Facet is initialized.
            fmt (str): Format of the string representation.

        Returns:
            (*cage.Facet*)
        """
        if fmt == "json":
            d = json.loads(input_string)
            return cls.from_dict(d)
        else:
            raise NotImplementedError('Only json format has been '
                                      'implemented.')

    @classmethod
    def from_file(cls, filename):

        with zopen(filename) as file:
            contents = file.read()

        return cls.from_str(contents)

    def as_dict(self):

        d = {"cathode": self.cathode.as_dict(),
             "dimer_indices": self.indices}

        return d

    @classmethod
    def from_dict(cls, d):

        return cls(cathode=LiRichCathode.from_dict(d["cathode"]),
                   dimer_indices=d["dimer_indices"])

class DimerNEBAnalysis(NEBAnalysis):
    """
    Subclass of the NEBAnalysis class in order to change the plotting of the
    barriers, as well as allowing for saving the NEB analysis to a json file.
    """

    def __init__(self, r, energies, forces, structures, spline_options=None,
                 dimer_indices=None):
        super().__init__(
            r, energies, forces, structures, spline_options
        )
        self._dimer_indices = tuple(dimer_indices)

    @property
    def dimer_indices(self):
        return self._dimer_indices

    @dimer_indices.setter
    def dimer_indices(self, indices):
        self._dimer_indices = indices

    @property
    def dimer_distances(self):
        return np.array([s.distance_matrix[self.dimer_indices]
                         for s in self.structures])

    def as_dict(self):
        """
        Dict representation of NEBAnalysis.

        Returns:
            JSON serializable dict representation.
        """
        return {"@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                'r': jsanitize(self.r),
                'energies': jsanitize(self.energies),
                'forces': jsanitize(self.forces),
                'structures': [s.as_dict() for s in self.structures],
                "dimer_indices": self.dimer_indices}

    def to(self, fmt="json", filename=None):

        if fmt == "json":
            if filename:
                with zopen(filename, "wt", encoding='utf8') as file:
                    return json.dump(self.as_dict(), file)
            else:
                return json.dumps(self.as_dict())
        else:
            raise NotImplementedError("Currently only json format is "
                                      "supported.")

    @classmethod
    def from_dir(cls, root_dir, relaxation_dirs=None, **kwargs):
        """

        Args:
            root_dir:
            relaxation_dirs:
            **kwargs:

        Returns:

        """
        if relaxation_dirs is not None:
            raise NotImplementedError

        indices = tuple(
            [int(el) for el in
             os.path.abspath(root_dir).split('/')[-1].split('_')
             if all([is_number(c) for c in el])]
        )

        neb = super().from_dir(root_dir, relaxation_dirs, **kwargs)

        # Because the dimer indices are based on the internal indices of the
        # Cathode object, we need to load the cathode json files to
        # determine the distance between the dimers properly.
        image_dirs = [file for file in os.listdir(root_dir)
                      if len(file) == 2 and os.path.isdir(file)]

        structures = [
            Cathode.from_file(os.path.join(image_dir, "final_cathode.json"))
            for image_dir in image_dirs
        ]

        # Sort the data according to the directory numbers
        structure_data = sorted(
            zip(image_dirs, structures),
            key=lambda z: int(z[0])
        )

        dimer_neb = DimerNEBAnalysis(
            r=neb.r,
            energies=neb.energies,
            forces=neb.forces,
            structures=[el[1] for el in structure_data],
            spline_options=neb.spline_options,
            dimer_indices=indices,
        )

        return dimer_neb

    @classmethod
    def from_str(cls, input_string, fmt="json"):
        """
        Initialize a DimerNEBAnalysis from a string.

        Currently only supports 'json' formats.

        Args:
            input_string (str): String from which the object is initialized.
            fmt (str): Format of the string representation.

        Returns:
            (*pybat.core.DimerNEBAnalysis*)
        """
        if fmt == "json":
            d = json.loads(input_string)
            return cls.from_dict(d)
        else:
            raise NotImplementedError('Only json format has been '
                                      'implemented.')

    @classmethod
    def from_file(cls, filename):

        with zopen(filename) as file:
            contents = file.read()

        return cls.from_str(contents)

    @classmethod
    def from_dict(cls, d):

        return cls(r=d['r'], energies=d["energies"], forces=d["forces"],
                   structures=[LiRichCathode.from_dict(structure) for
                               structure in d["structures"]],
                   dimer_indices=d["dimer_indices"])

    def get_plot(self, normalize_rnx_coodinate=True, label_barrier=True):
        """
        Returns the NEB plot. Uses Henkelman's approach of spline fitting
        each section of the reaction path based on tangent force and energies.

        Args:
            label_barrier (bool): Whether to label the maximum barrier.

        Returns:
            matplotlib.pyplot object.
        """
        plt = pretty_plot(12, 8)

        spline_x = np.arange(0, np.max(self.r), 0.01)
        spline_y = self.spline(spline_x) * 1000

        relative_energies = self.energies - self.energies[0]

        plt.plot(spline_x, spline_y, 'k--',
                 self.r, relative_energies * 1000, 'ro',
                 linewidth=2,
                 markersize=10)

        plt.xlabel("O-O Distance ($\mathrm{\AA}$)")
        plt.xticks(self.r[::2], [str(round(d, 2)) for d in
                                 self.dimer_distances[::2]])
        plt.ylabel("Energy (meV)")
        plt.ylim((np.min(spline_y) - 10, np.max(spline_y) * 1.02 + 20))

        if label_barrier:
            data = zip(spline_x, spline_y)
            barrier = max(data, key=lambda d: d[1])
            plt.plot([0, barrier[0]], [barrier[1], barrier[1]], 'k--')
            plt.annotate('%.0f meV' % (np.max(spline_y) - np.min(spline_y)),
                         xy=(barrier[0] / 2, barrier[1] * 1.02),
                         xytext=(barrier[0] / 2, barrier[1] * 1.02),
                         horizontalalignment='center')

        plt.tight_layout()
        return plt

# SO Plagiarism
def unit_vector(vector):
    """ Returns the unit vector of the vector.  """
    return vector / np.linalg.norm(vector)

def angle_between(v1, v2):
    """
    Returns the angle in radians between vectors 'v1' and 'v2'::
    """
    v1_u = unit_vector(v1)
    v2_u = unit_vector(v2)
    return np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))

# def rotation_matrix(axis, theta):
#     """
#     Return the rotation matrix associated with clockwise rotation about
#     the given axis by theta radians.
#     """
#     axis = np.asarray(axis)
#     axis = axis/math.sqrt(np.dot(axis, axis))
#     a = math.cos(theta/2.0)
#     b, c, d = axis*math.sin(theta/2.0)
#     aa, bb, cc, dd = a*a, b*b, c*c, d*d
#     bc, ad, ac, ab, bd, cd = b*c, a*d, a*c, a*b, b*d, c*d
#     return np.array([[aa+bb-cc-dd, 2*(bc+ad), 2*(bd-ac)],
#                      [2*(bc-ad), aa+cc-bb-dd, 2*(cd+ab)],
#                      [2*(bd+ac), 2*(cd-ab), aa+dd-bb-cc]])

def permute(iterable, permutation):
    if len(iterable) != len(permutation):
        raise ValueError("Length of list does not match permutation length!")

    permutation_numbers = set([i for i in range(1, len(iterable) + 1)])

    if len(permutation_numbers.intersection(permutation)) != len(permutation):
        raise ValueError("Permutation is ill-defined.")

    return [iterable[index - 1] for index in permutation]

def is_number(s):
    try:
        float(s)
        return True
    except ValueError:
        return False

def new_function(arg1, arg2):
    """
    This describes the function.

    Args:
        arg1 (list):
        arg2:

    Returns:

    """


#!/usr/bin/env python
import time
import pandas as pd
import collections
from collections import defaultdict
import re
import numpy as np
from Bio.SeqRecord import SeqRecord
from Bio.Seq import Seq
import numpy as np
import re

def peptide_mass(peptide, fixed_modifications=["Carbamidomethylation of C"], variable_modifications=[], nterm=True, cterm=True):
    if 'X' in peptide:
        return [np.inf]
    if 'Z' in peptide:
        return [np.inf]
    if 'B' in peptide:
        return [np.inf]

    masses = []

    mass_h20=18.01056
    nterm_mass =    1.007825035
    cterm_mass =    1.007825035 + 15.99491463

    monoisotopic = {'A':   71.037114,
                    'C':   103.009185,
                    'D':   115.026943,
                    'E':   129.042593,
                    'F':   147.068414,
                    'G':   57.021464,
                    'H':   137.058912,
                    'U':   150.95363,
                    'O':   237.147735,
                    'I':   113.084064,
                    'K':   128.09496,
                    'L':   113.084064,
                    'J':   113.084064,
                    'M':   131.04049,
                    'N':   114.042927,
                    'P':   97.05276,
                    'Q':   128.05858,
                    'R':   156.101111,
                    'S':   87.03203,
                    'T':   101.04768,
                    'V':   99.06841,
                    'W':   186.07931,
                    'Y':   163.06333  }
    ##############
    # Fixed mods #
    ##############
    options = ["Carbamidomethylation of C"]
    for i in fixed_modifications:
        assert i in options
    if "Carbamidomethylation of C" in fixed_modifications:
        monoisotopic['C'] = monoisotopic['C'] + 57.021464

    for aa in peptide:
        masses.append(monoisotopic[aa])

    # < insert directed graph here > #
    masslist = [sum(masses)]
    # < insert directed graph here > #

    if nterm == True:
        masslist = [i + nterm_mass for i in masslist]
    if cterm == True:
        masslist = [i + cterm_mass for i in masslist]

    return masslist

def mz2mw(mz, charge):
    proton_mass=1.007276
    mw = (mz * charge) - (proton_mass * charge)
    return np.round(mw, 4)

class TagMatch:
    def __init__(self, query, tag_mass_list, target, fixed_modifications=['Carbamidomethylation of C'], variable_modifications=[], enzymes=['Trypsin'], specificity='specific', prec_tol=0.02, gap_tol=0.5, max_missed_cleavages=3):
        self.possible_enzymes = ['Trypsin', 'Trypsin, no P rule','Whole protein']
        self.enzymes = enzymes
        self.max_missed_cleavages = max_missed_cleavages
        for enzyme in self.enzymes:
            assert enzyme in self.possible_enzymes

        self.possible_specificity = ['specific', 'semi-specific', 'unspecific']
        self.specificity = specificity
        assert self.specificity in self.possible_specificity

        self.target = target
        self.target_length = len(target)

        self.tag_mass_list = tag_mass_list

        self.gap_mass_dict = defaultdict(set)

        if not ((self.specificity=='specific') and (self.enzymes==['Trypsin'])):
            self.gaps = self.gaps()
            self.max_gaps = (max(self.gaps[0]), max(self.gaps[1]))
            self.maxngap = self.max_gaps[0]
            self.maxcgap = self.max_gaps[1]

        self.query=query
        self.query_length = len(query)
        self.fixed_modifications=fixed_modifications
        self.variable_modifications=variable_modifications
        self.query_positions = self.query_positions()
        self.prec_tol = prec_tol
        self.gap_tol  = gap_tol

        self.validated=False
        self.specificity = specificity

        self.positions = self.match()
        self.peptides  = self.peptides()
        self.masses = self.masses()

        self.amino_gaps = self.amino_gaps()
        self.validate_masses(self.masses)

    def query_positions(self):
        positions = [(m.start(), m.start()+self.query_length ) for m in re.finditer('(?={})'.format(self.query), self.target)]
        return positions

    def validate_masses(self, mass_dict):
        validated_precursor = []
        precursor_peptides = defaultdict(list)
        validated_peptides = set()

        for position_key in mass_dict:
            target_masses = mass_dict[position_key]
            target_peptide = self.peptides[position_key]

            for target_mass in target_masses:
                for tag_mass_set in self.tag_mass_list:
                    assert len(tag_mass_set) == 3
                    query_mass = tag_mass_set[1]
                    diff = abs(target_mass - query_mass)
                    if diff < self.prec_tol:
                        nterm_mass_gap = tag_mass_set[0]
                        cterm_mass_gap = tag_mass_set[2]
                        for amino_gaps in self.amino_gaps[target_peptide]:
                            nterm_gap_match=False
                            cterm_gap_match=False
                            if not amino_gaps[0] in self.gap_mass_dict:
                                nterm_amino_gaps = peptide_mass(amino_gaps[0], nterm=False, cterm=False, fixed_modifications=self.fixed_modifications, variable_modifications = self.variable_modifications)
                            else:
                                nterm_amino_gaps = self.gap_mass_dict[amino_gaps[0]]
                            for nterm_amino_gap in nterm_amino_gaps:
                                if abs(nterm_amino_gap - nterm_mass_gap) < self.gap_tol:
                                    nterm_gap_match=True
                            if not amino_gaps[1] in self.gap_mass_dict:
                                cterm_amino_gaps = peptide_mass(amino_gaps[1], nterm=False, cterm=False, fixed_modifications = self.fixed_modifications, variable_modifications=self.variable_modifications)
                            else:
                                cterm_amino_gaps = self.gap_mass_dict[amino_gaps[1]]
                            for cterm_amino_gap in cterm_amino_gaps:
                                if abs(cterm_amino_gap - cterm_mass_gap) < self.gap_tol:
                                    cterm_gap_match=True

                            if (cterm_gap_match==True) and (nterm_gap_match==True):
                                validated_precursor.append(tag_mass_set)
                                precursor_peptides[tag_mass_set].append(self.peptides[position_key])
                                validated_peptides.add(self.peptides[position_key])
        self.validated_precursor = validated_precursor
        self.precursor_peptides  = precursor_peptides
        self.validated_peptides = list(validated_peptides)
        if len(validated_precursor) > 0:
            self.validated = True

    def masses(self):
        masses = {}
        for key in self.peptides:
            peptide = self.peptides[key]
            masslist = peptide_mass(peptide, fixed_modifications=self.fixed_modifications, variable_modifications=self.variable_modifications, nterm=True, cterm=True)
            masses[key] = masslist
        return masses

    def peptides(self):
        peptides = {}
        for coord in self.positions:
            p = self.target[coord[0]:coord[1]]
            key='{}:{}'.format(str(coord[0]), str(coord[1]))
            peptides[key] = p
        return peptides

    def match(self):
        coords = []
        for coord in self.query_positions:
            nterms = self.nterms(coord[0])
            cterms = self.cterms(coord[1]-1)

            for nterm in nterms:
                for cterm in cterms:
                    amino_acid_before = self.target[nterm -1: nterm]
                    first_amino_acid  = self.target[nterm]
                    last_amino_acid   = self.target[cterm -1]
                    amino_acid_after   = self.target[cterm : cterm + 1]
                    valid = valid_cleavage( amino_acid_before, first_amino_acid, last_amino_acid, amino_acid_after, self.enzymes, self.specificity)
                    if valid == True:
                        mc = missed_cleavages(self.target[nterm:cterm], self.enzymes)
                        if (mc <= self.max_missed_cleavages):
                            coords.append((nterm, cterm))
        return coords

    def nterms(self, pos):
        nterm = pos
        nterms=[]

        if ((self.specificity=='specific') and (self.enzymes==['Trypsin'])):
            valid_gaps=True
            missed=0
            while valid_gaps==True:
                if nterm == 0:
                    nterms.append(nterm)
                    return nterms
                else:
                    current_amino = self.target[nterm]
                    previous_amino = self.target[nterm-1]
                    if (current_amino != 'P') and (previous_amino in {'R', 'K'}):
                        nterms.append(nterm)
                        missed += 1
                        if missed > self.max_missed_cleavages:
                            return nterms
                    nterm -= 1
                    assert nterm >= 0
        else:
            valid_gaps=True
            limit = self.maxngap + self.gap_tol
            while valid_gaps==True:
                amino_gap = self.target[nterm:pos]
                validated=False
                gaps = peptide_mass(amino_gap, fixed_modifications=self.fixed_modifications, variable_modifications=self.variable_modifications, nterm=False, cterm=False)
                self.gap_mass_dict[amino_gap].update(gaps)
                for gap in gaps:
                    diff_gaps = set([(abs(i-gap) < self.gap_tol) for i in self.gaps[0]])
                    if True in diff_gaps:
                        validated=True

                if nterm == 0:
                    if validated == True:
                        nterms.append(nterm)
                    return nterms
                else:
                    if validated == True:
                        nterms.append(nterm)
                    nterm -= 1
                    assert nterm >= 0
                if min(gaps) > limit:
                    return nterms

    def cterms(self, pos):
        cterm = pos
        cterms=[]

        if ((self.specificity=='specific') and (self.enzymes==['Trypsin'])):
            valid_gaps =True
            missed=0
            while valid_gaps ==True:
                amino_gap = self.target[pos +1:cterm+1]

                if cterm == self.target_length-1:
                    cterms.append(cterm + 1)
                    return cterms
                else:
                    current_amino = self.target[cterm]
                    next_amino = self.target[cterm +1]
                    if (next_amino != 'P') and (current_amino in {'R', 'K'}):
                        cterms.append(cterm + 1)
                        missed += 1
                        if missed > self.max_missed_cleavages:
                            return cterms
                    cterm += 1
                    assert cterm <= self.target_length -1

        else:
            valid_gaps =True
            limit = self.maxcgap + self.gap_tol
            while valid_gaps ==True:
                amino_gap = self.target[pos +1:cterm+1]
                gaps = peptide_mass(amino_gap, fixed_modifications=self.fixed_modifications, variable_modifications=self.variable_modifications, nterm=False, cterm=False)
                self.gap_mass_dict[amino_gap].update(gaps)
                validated=False

                for gap in gaps:
                    diff_gaps = set([(abs(i-gap) < self.gap_tol) for i in self.gaps[1]])
                    if True in diff_gaps:
                        validated=True

                if cterm == self.target_length-1:
                    if validated == True:
                        cterms.append(cterm + 1)
                    return cterms
                else:
                    if validated == True:
                        cterms.append(cterm + 1)
                    cterm += 1
                    assert cterm <= self.target_length -1
                if min(gaps) > limit:
                    return cterms

    def gaps(self):
        ngaps = []
        cgaps = []
        for tag in self.tag_mass_list:
            n = tag[0]
            c = tag[2]
            ngaps.append(n)
            cgaps.append(c)
        return (list(set(ngaps)), list(set(cgaps)))

    def amino_gaps(self):
        amino_gap_dict = defaultdict(list)
        for key in self.peptides:
            peptide = self.peptides[key]
            positions = [(m.start(), m.start()+self.query_length ) for m in re.finditer('(?={})'.format(self.query), peptide)]
            for coord in positions:
                tag_start = coord[0]
                tag_end = tag_start + self.query_length
                nterm_gap = peptide[:tag_start]
                cterm_gap = peptide[tag_end:]
                amino_gap_dict[peptide].append( (nterm_gap, cterm_gap) )
        return amino_gap_dict

def valid_cleavage(amino_acid_before, first_amino_acid, last_amino_acid, amino_acid_after, enzymes, specificity):
    valid = False
    nterm_valid=False
    cterm_valid=False
    for enzyme in enzymes:
        if enzyme=='Trypsin, no P rule':
            cleavage_aminos=['K','R']
            if amino_acid_before == '':
                nterm_valid=True
            elif amino_acid_before in cleavage_aminos:
                nterm_valid=True
            if amino_acid_after =='':
                cterm_valid=True
            elif last_amino_acid in cleavage_aminos:
                cterm_valid=True

        elif enzyme == 'Trypsin':
            cleavage_aminos=['K','R']

            if amino_acid_before == '':
                nterm_valid=True
            elif (amino_acid_before in cleavage_aminos) and (first_amino_acid != 'P'):
                nterm_valid=True

            if amino_acid_after =='':
                cterm_valid=True
            elif last_amino_acid in cleavage_aminos:
                cterm_valid=True

        elif enzyme == 'Whole protein':

            if amino_acid_before == '':
                nterm_valid=True

            if amino_acid_after =='':
                cterm_valid=True
    if ((nterm_valid and cterm_valid) == True) and (specificity == 'specific'):
        valid = True
    elif ((nterm_valid or cterm_valid) == True) and (specificity == 'semi-specific'):
        valid = True
    elif specificity == 'unspecific':
        valid=True
    return valid

def missed_cleavages(peptide, enzymes):
    new_peptides = [peptide]
    for enzyme in enzymes:
        holder=[]
        for datum in new_peptides:
            if enzyme=='Trypsin, no P rule':
                peptides = re.sub(r'(?<=[RK])','\n', datum).split('\n')
                holder += peptides
            elif enzyme == 'Trypsin':
                peptides = re.sub(r'(?<=[RK])(?=[^P])','\n', datum).split('\n')
                holder += peptides
            elif enzyme == 'Whole protein':
                peptides = [datum]
                holder += peptides
        new_peptides = holder
    return len(new_peptides) - 1

def character_strip(sequence, character):
    started=False
    start_pos=0
    while (started==False) and (start_pos < len(sequence)):
        if sequence[start_pos]==character:
            start_pos += 1
        else:
            started = True

    end_pos=len(sequence)-1
    ended=False
    while (ended==False) and (end_pos > 0):
        if sequence[end_pos]==character:
            end_pos -= 1
        else:
            ended = True
    return sequence[start_pos: end_pos + 1]

class gap_sequence:
    def __init__(self, gap, fixed_modifications=['Carbamidomethylation of C'], variable_modifications=[], prec_tol=0.02, gap_tol=0.5):
        self.gap = gap
        self.prec_tol = prec_tol
        self.gap_tol  = gap_tol
        self.fixed_modifications = fixed_modifications
        self.variable_modifications = variable_modifications
        self.gap_sequences=[]

        self.possible = set([0])
        self.valid = False

        self.aminos     =           ['A',
                                     'C',
                                     'D',
                                     'E',
                                     'F' ,
                                     'G',
                                     'H',
                                     'U',
                                     'O',
                                     'I',
                                     'K',
                                     'L',
                                     'J',
                                     'M',
                                     'N',
                                     'P' ,
                                     'Q',
                                     'R',
                                     'S',
                                     'T' ,
                                     'V' ,
                                     'W',
                                     'Y']

    def sequences(self, mass=0, sequence=''):
        diff = np.absolute(mass - self.gap)

        if diff < self.prec_tol:
            self.gap_sequences.append(sequence)

        elif (mass < self.gap):
            for amino in self.aminos:
                masses=peptide_mass(amino, variable_modifications = self.variable_modifications, fixed_modifications=self.fixed_modifications, nterm=False, cterm=False)
                for m in masses:
                    newmass=mass + m
                    newsequence = sequence + amino
                    self.sequences(newmass, newsequence)

    def validate(self):
        newpossible=set()
        for mass in self.possible:
            diff = np.absolute(mass - self.gap)

            if diff < self.prec_tol:
                self.valid = True

            elif (mass < self.gap) and (self.valid==False):
                for amino in self.aminos:
                    masses=peptide_mass(amino, variable_modifications = self.variable_modifications, fixed_modifications=self.fixed_modifications, nterm=False, cterm=False)
                    for m in masses:
                        newmass = mass + m
                        newpossible.add(np.round(newmass,4))
        if (len(newpossible) > 0) and (self.valid ==False):
            self.possible=newpossible
            self.validate()

class blast_tags:
    def __init__(self, subject, tags, fixed_modifications=['Carbamidomethylation of C'], variable_modifications=[], prec_tol=0.02, gap_tol=0.5, min_identities = 4, IL_equivalence=True):

        self.subject=subject
        self.tags=tags
        self.fixed_modifications = fixed_modifications
        self.variable_modifications = variable_modifications
        self.prec_tol  = prec_tol
        self.gap_tol = gap_tol

        self.min_identities = min_identities
        self.IL_equivalence = IL_equivalence

        self.mweights = set()
        self.nmasses  = set()
        self.cmasses  = set()

        self.samples  = set()
        self.scans    = set()

        record_holder = set()
        self.newrecords = []

        for qtag in self.tags:
            qseq = str(qtag.seq)
            qsubstrings = self.substrings(qseq, self.subject)

            modified_tags = self.segment_matching(qseq, self.subject, min_identities=self.min_identities, IL_equivalence=self.IL_equivalence)
            idlist = qtag.id.split(';')
            file = idlist[0]
            scan = idlist[1]
            mw = float(idlist[2].split('mw=')[1])
            ngap = float(idlist[3].split('ngap=')[1])
            cgap = float(idlist[4].split('cgap=')[1])

            for i in modified_tags:
                newseq=i[0]
                newngap = np.round(ngap + i[1],6)
                newcgap = np.round(cgap + i[2],6)
                newseq, newngap, newcgap = self.fix_negative_gaps(newseq, newngap, newcgap)
                valid_tag=True
                if newngap <= -(self.gap_tol):
                    valid_tag=False
                elif newcgap <= -(self.gap_tol):
                    valid_tag=False
                if valid_tag == True:
                    newid ='{};{};mw={};ngap={};cgap={}'.format(file, scan, mw, newngap, newcgap)
                    newrec=SeqRecord(seq=Seq(newseq),id=newid,description='Segment correction of {}-{}-{}'.format(ngap,qseq,cgap))
                    temp = newrec.format('fasta')
                    if not temp in record_holder:
                        record_holder.add(temp)
                        self.newrecords.append(newrec)

    def valid_tag(self, tag):
        sequence_masses = peptide_mass(tag[0], fixed_modifications=self.fixed_modifications, variable_modifications = self.variable_modifications, cterm=True, nterm=True)
        for mass in sequence_masses:
            total = mass + tag[2] + tag[3]
            diff = np.absolute(total - tag[1])
            if diff < self.prec_tol:
                return True

    def substrings(self, reference, match_sequence):
        substrings=[]
        for charpos in range(len(match_sequence)):
            mismatch=False
            pos= charpos
            while (mismatch == False) & (pos < len(match_sequence)):
                if match_sequence[charpos:pos +1]  in reference:
                    pos += 1
                else:
                    mismatch=True
                subs = match_sequence[charpos:pos ]
            if subs != '':
                substrings.append(subs)
        return substrings

    def longest_substring(self, substrings):
        maxlen = len(max(substrings, key=len))
        maxes = [i for i in substrings if len(i) == maxlen]
        return maxes

    def substring_coordinates(self, reference, substr):
        coordinates=[(m.start(), m.start() + len(substr)) for m in re.finditer('(?={})'.format(substr), reference)]
        return coordinates

    def cmass_gap(self, reference, coordinates):
        mass_gaps=[]
        for coordinate in coordinates:
            amino_gap = reference[coordinate[1]:]
            amino_masses = peptide_mass(amino_gap, fixed_modifications = self.fixed_modifications, variable_modifications = self.variable_modifications, nterm=False, cterm=False)
            mass_gaps += amino_masses
        return mass_gaps

    def nmass_gap(self, reference, coordinates):
        mass_gaps=[]
        for coordinate in coordinates:
            amino_gap = reference[:coordinate[0]]
            amino_masses = peptide_mass(amino_gap, fixed_modifications = self.fixed_modifications, variable_modifications = self.variable_modifications, nterm=False, cterm=False)
            mass_gaps += amino_masses
        return mass_gaps

    def new_records(self):
        newrecords = []
        test = []
        for newtag in self.newtags:
            rec = SeqRecord(seq=Seq(newtag[0]), id = '{};{};mw={};ngap={};cgap={}'.format(list(self.samples)[0], list(self.scans)[0], str(newtag[1]), str(newtag[2]), str(newtag[3])))
            trec = rec.format('fasta')
            if not trec in test:
                test.append(trec)
                newrecords.append(rec)
        return newrecords

    def match_segments(self, tagsequence, blastsequence, substrings, min_identities, matchsequence='', matchngap=None, matchcgap=None, identities=[], matches = set()):
        if (len(substrings) == 0):
            if len(identities) > 0:
                if (len(''.join(identities)) >= min_identities):
                    for cg in peptide_mass(blastsequence, cterm=False, nterm=False, fixed_modifications = self.fixed_modifications, variable_modifications = self.variable_modifications):
                        for cg2 in peptide_mass(tagsequence, cterm=False, nterm=False, fixed_modifications = self.fixed_modifications, variable_modifications = self.variable_modifications ) :
                            matches.add((matchsequence + blastsequence, matchngap, cg2-cg, '-'.join(identities)))
            return matches

        subst = substrings[0]
        newsubstrings = substrings[1:]
        refpos = [(m.start(), m.start()+len(subst)) for m in re.finditer('(?={})'.format(subst), blastsequence)]
        tagpos = [(m.start(), m.start()+len(subst)) for m in re.finditer('(?={})'.format(subst), tagsequence)]

        newtags = self.match_segments(tagsequence, blastsequence, newsubstrings, min_identities, matchsequence=matchsequence, matchngap=matchngap, matchcgap=matchcgap, identities=identities, matches=matches)
        matches.update(newtags)

        for r in refpos:
            for t in tagpos:
                refgap = blastsequence[:r[0]]
                matchgap = tagsequence[:t[0]]
                newblastsequence = blastsequence[r[1] :]
                newtagsequence = tagsequence[t[1] :]

                refgapmasses=peptide_mass(refgap, nterm=False, cterm=False, fixed_modifications=self.fixed_modifications, variable_modifications=self.variable_modifications)
                matchgapmasses = peptide_mass(matchgap, nterm=False, cterm=False, fixed_modifications=self.fixed_modifications, variable_modifications = self.variable_modifications)

                for rm in refgapmasses:
                    for tm in matchgapmasses:
                        if matchngap==None:
                            newmatchngap = tm - rm
                            newmatchngap = np.round(newmatchngap, 6)
                            newmatchsequence = refgap + subst
                            newtags = self.match_segments(newtagsequence, newblastsequence, newsubstrings, min_identities, matchsequence=newmatchsequence, matchngap=newmatchngap, matchcgap = matchcgap, identities=identities + [subst] , matches=matches)
                            matches.update(newtags)
                        elif rm==tm==0:
                            return matches
                        else:
                            diff = abs(rm - tm)
                            if diff < self.gap_tol:
                                newmatchsequence = matchsequence + refgap + subst
                                newtags = self.match_segments(newtagsequence, newblastsequence, newsubstrings, min_identities, matchsequence=newmatchsequence, matchngap=matchngap, matchcgap=matchcgap, identities=identities + [subst], matches=matches)
                                matches.update(newtags)
        return matches

    def segment_matching(self, tagsequence, blastsequence, min_identities=4, IL_equivalence=True):
        if IL_equivalence==True:
            newtagsequence = 'L'.join(tagsequence.split('I'))
            newblastsequence = 'L'.join(blastsequence.split('I'))
        else:
            newtagsequence = tagsequence
            newblastsequence = blastsequence

        substrings = self.substrings(newblastsequence, newtagsequence)
        newtags = self.match_segments(newtagsequence, newblastsequence, substrings, min_identities=min_identities, matchsequence='', matchngap=None, matchcgap=None, identities=[], matches=set())
        besttags=[]
        bestlen=0

        for tag in newtags:
            identities = ''.join(tag[3].split('-'))
            if len(identities) > bestlen:
                bestlen = len(identities)

        for tag in newtags:
            identities = ''.join(tag[3].split('-'))
            if len(identities) == bestlen:
                assert tag[0] == newblastsequence
                tag = (blastsequence, tag[1], tag[2], tag[3])
                besttags.append(tag)
        return set(besttags)

    def fix_negative_gaps(self, sequence, ngap, cgap):

        if ngap < 0:
            ngap_fixed=False
            for position in range(1, len(sequence)):
                newsequence = sequence[position:]
                newgapsequence = sequence[:position]
                ng = peptide_mass(newgapsequence, nterm=False, cterm=False,  fixed_modifications=self.fixed_modifications, variable_modifications= self.variable_modifications)
                mng = min(ng)

                if mng + ngap >= self.gap_tol:
                    break

                for _ in ng:
                    newg = ngap + _
                    if np.absolute(newg) < self.gap_tol:
                        ngap = newg
                        sequence = newsequence
                        ngap_fixed=True
                        break
                if ngap_fixed==True:
                    break

        if cgap < 0:
            cgap_fixed=False
            for position in range(1, len(sequence)):
                newsequence = sequence[:-position]
                newgapsequence = sequence[-position:]
                cg = peptide_mass(newgapsequence, nterm=False, cterm=False,  fixed_modifications=self.fixed_modifications, variable_modifications= self.variable_modifications)

                mcg = min(cg)

                if mcg + cgap >= self.gap_tol:
                    break
    #
                for _ in cg:
                    newg = cgap + _
                    if np.absolute(newg) < self.gap_tol:
                        cgap = newg
                        sequence = newsequence
                        cgap_fixed=True
                        break
                if cgap_fixed==True:
                    break

        return sequence, np.round(ngap,4), np.round(cgap,4)


"""Implementation of HSD-scripted.

High-level Q-functions Q(s,\zbf) are trained either with QMIX (with decentralized execution)
or Q-learning (centralized execution) using the global reward
Low-level Q-functions Q(o^n,z^n,a^n) are trained with independent Q-learning using local rewards
"""

import tensorflow as tf
import numpy as np
import sys
import networks

class Alg(object):

    def __init__(self, alg_name, config_alg, n_agents, l_state, l_obs, l_action, N_roles, nn):
        """
        Args:
            alg_name: currently supports 'hsd-scripted' for QMIX at high level or 'mara-c' for Q-learning at high level. Low level is always independent Q-learning
            config_alg: dictionary of general RL params
            n_agents: number of agents on the team controlled by this alg
            l_state, l_obs, l_action, N_roles: int
            nn: dictionary with neural net sizes
        """
        self.alg_name = alg_name

        self.l_state = l_state
        self.l_obs = l_obs
        self.l_action = l_action
        self.N_roles = N_roles
        self.nn = nn

        self.n_agents = n_agents
        self.tau = config_alg['tau']
        self.lr_Q = config_alg['lr_Q']
        self.gamma = config_alg['gamma']

        self.agent_labels = np.eye(self.n_agents)

        if self.alg_name == 'mara-c':
            # Combinatorial action space, but we don't allow duplicate roles
            assert(self.N_roles >= self.n_agents)
            self.dim_role_space = int(np.math.factorial(self.N_roles) / np.math.factorial(self.N_roles - self.n_agents))
            self.list_list_indices = []
            self.populate_list_list_roles(0, [])

        # Initialize computational graph
        self.create_networks()
        self.list_initialize_target_ops, self.list_update_target_ops, self.list_update_target_ops_low = self.get_assign_target_ops()
        self.create_train_op()
        self.create_train_op_IQL()

        # TF summaries
        self.create_summary()

    def populate_list_list_roles(self, agent_idx, list_indices):

        if len(list_indices) == self.n_agents:
            self.list_list_indices.append(list_indices)
        else:
            for idx_role in range(self.N_roles):
                if idx_role in list_indices:
                    # Skip over duplicate roles
                    continue
                else:
                    l = list(list_indices)
                    l.append(idx_role)
                    self.populate_list_list_roles(agent_idx+1, l)

    def create_networks(self):

        # Placeholders
        self.state = tf.placeholder(tf.float32, [None, self.l_state], 'state')
        self.obs = tf.placeholder(tf.float32, [None, self.l_obs], 'obs')
        self.role = tf.placeholder(tf.float32, [None, self.N_roles], 'role')

        # Low-level Q-functions
        with tf.variable_scope("Qlow_main"):
            self.Q_low = networks.Q_low(self.obs, self.role, self.nn['n_h1_low'], self.nn['n_h2_low'], self.l_action)
        with tf.variable_scope("Qlow_target"):
            self.Q_low_target = networks.Q_low(self.obs, self.role, self.nn['n_h1_low'], self.nn['n_h2_low'], self.l_action)

        self.argmax_Q_low = tf.argmax(self.Q_low, axis=1)
        self.argmax_Q_low_target = tf.argmax(self.Q_low_target, axis=1)

        # Low level action
        self.actions_low_1hot = tf.placeholder(tf.float32, [None, self.l_action], 'actions_low_1hot')

        # High-level Q-functions
        if self.alg_name == 'hsd-scripted':
            # Individual agent networks
            # output dimension is [time * n_agents, q-values]
            with tf.variable_scope("Agent_main"):
                self.agent_qs = networks.Qmix_single(self.obs, self.nn['n_h1'], self.nn['n_h2'], self.N_roles)
            with tf.variable_scope("Agent_target"):
                self.agent_qs_target = networks.Qmix_single(self.obs, self.nn['n_h1'], self.nn['n_h2'], self.N_roles)

            self.argmax_Q = tf.argmax(self.agent_qs, axis=1)
            self.argmax_Q_target = tf.argmax(self.agent_qs_target, axis=1)

            # To extract Q-value from agent_qs and agent_qs_target
            # [batch*n_agents, N_roles]
            self.actions_1hot = tf.placeholder(tf.float32, [None, self.N_roles], 'actions_1hot')
            # [batch*n_agents, 1]
            self.q_selected = tf.reduce_sum(tf.multiply(self.agent_qs, self.actions_1hot), axis=1)
            # [batch, n_agents]
            self.mixer_q_input = tf.reshape( self.q_selected, [-1, self.n_agents] )

            self.q_target_selected = tf.reduce_sum(tf.multiply(self.agent_qs_target, self.actions_1hot), axis=1)
            self.mixer_target_q_input = tf.reshape( self.q_target_selected, [-1, self.n_agents] )

            # Mixing network
            with tf.variable_scope("Mixer_main"):
                self.mixer = networks.Qmix_mixer(self.mixer_q_input, self.state, self.l_state, self.n_agents, self.nn['n_h_mixer'])
            with tf.variable_scope("Mixer_target"):
                self.mixer_target = networks.Qmix_mixer(self.mixer_target_q_input, self.state, self.l_state, self.n_agents, self.nn['n_h_mixer'])
        elif self.alg_name == 'mara-c':
            # Standard Q-learning for role assignment
            with tf.variable_scope("Qhigh_main"):
                self.Q_high = networks.Q_high(self.state, self.nn['n_h1'], self.nn['n_h2'], self.dim_role_space)
            with tf.variable_scope("Qhigh_target"):
                self.Q_high_target = networks.Q_high(self.state, self.nn['n_h1'], self.nn['n_h2'], self.dim_role_space)
            self.argmax_Q_high = tf.argmax(self.Q_high, axis=1)
            self.argmax_Q_high_target = tf.argmax(self.Q_high_target, axis=1)
            self.actions_high_1hot = tf.placeholder(tf.float32, [None, self.dim_role_space], 'actions_high_1hot')

    def get_assign_target_ops(self):

        # ops for equating main and target
        list_initial_ops = []
        # ops for slow update of target toward main
        list_update_ops = []
        # ops for slow update of low-level target toward low-level main
        list_update_ops_low = []

        if self.alg_name == 'hsd-scripted':
            list_Agent_main = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'Agent_main')
            map_name_Agent_main = {v.name.split('main')[1] : v for v in list_Agent_main}
            list_Agent_target = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'Agent_target')
            map_name_Agent_target = {v.name.split('target')[1] : v for v in list_Agent_target}

            if len(list_Agent_main) != len(list_Agent_target):
                raise ValueError("get_initialize_target_ops : lengths of Agent_main and Agent_target do not match")

            for name, var in map_name_Agent_main.items():
                # create op that assigns value of main variable to
                # target variable of the same name
                list_initial_ops.append( map_name_Agent_target[name].assign(var) )

            for name, var in map_name_Agent_main.items():
                # incremental update of target towards main
                list_update_ops.append( map_name_Agent_target[name].assign( self.tau*var + (1-self.tau)*map_name_Agent_target[name] ) )

            list_Mixer_main = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'Mixer_main')
            map_name_Mixer_main = {v.name.split('main')[1] : v for v in list_Mixer_main}
            list_Mixer_target = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'Mixer_target')
            map_name_Mixer_target = {v.name.split('target')[1] : v for v in list_Mixer_target}

            if len(list_Mixer_main) != len(list_Mixer_target):
                raise ValueError("get_initialize_target_ops : lengths of Mixer_main and Mixer_target do not match")

            # ops for equating main and target
            for name, var in map_name_Mixer_main.items():
                # create op that assigns value of main variable to
                # target variable of the same name
                list_initial_ops.append( map_name_Mixer_target[name].assign(var) )

            # ops for slow update of target toward main
            for name, var in map_name_Mixer_main.items():
                # incremental update of target towards main
                list_update_ops.append( map_name_Mixer_target[name].assign( self.tau*var + (1-self.tau)*map_name_Mixer_target[name] ) )
        elif self.alg_name == 'mara-c':
            list_Qhigh_main = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'Qhigh_main')
            map_name_Qhigh_main = {v.name.split('main')[1] : v for v in list_Qhigh_main}
            list_Qhigh_target = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'Qhigh_target')
            map_name_Qhigh_target = {v.name.split('target')[1] : v for v in list_Qhigh_target}

            if len(list_Qhigh_main) != len(list_Qhigh_target):
                raise ValueError("get_initialize_target_ops : lengths of Qhigh_main and Qhigh_target do not match")

            for name, var in map_name_Qhigh_main.items():
                # create op that assigns value of main variable to
                # target variable of the same name
                list_initial_ops.append( map_name_Qhigh_target[name].assign(var) )

            for name, var in map_name_Qhigh_main.items():
                # incremental update of target towards main
                list_update_ops.append( map_name_Qhigh_target[name].assign( self.tau*var + (1-self.tau)*map_name_Qhigh_target[name] ) )

        list_Qlow_main = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'Qlow_main')
        map_name_Qlow_main = {v.name.split('main')[1] : v for v in list_Qlow_main}
        list_Qlow_target = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'Qlow_target')
        map_name_Qlow_target = {v.name.split('target')[1] : v for v in list_Qlow_target}

        if len(list_Qlow_main) != len(list_Qlow_target):
            raise ValueError("get_initialize_target_ops : lengths of Qlow_main and Qlow_target do not match")

        for name, var in map_name_Qlow_main.items():
            # create op that assigns value of main variable to
            # target variable of the same name
            list_initial_ops.append( map_name_Qlow_target[name].assign(var) )

        for name, var in map_name_Qlow_main.items():
            # incremental update of target towards main
            list_update_ops_low.append( map_name_Qlow_target[name].assign( self.tau*var + (1-self.tau)*map_name_Qlow_target[name] ) )

        return list_initial_ops, list_update_ops, list_update_ops_low

    def run_actor(self, list_obs, roles, epsilon, sess):
        """Get low-level actions for all agents as a batch.

        Args:
            list_obs: list of vectors, one per agent
            roles: np.array where each row is a 1-hot vector
            epsilon: exploration parameter
            sess: TF session

        Returns: np.array of action integers
        """
        # convert to batch
        obs = np.array(list_obs)
        feed = {self.obs : obs, self.role : roles}
        actions_argmax = sess.run(self.argmax_Q_low, feed_dict=feed)

        actions = np.zeros(self.n_agents, dtype=int)
        for idx in range(self.n_agents):
            if np.random.rand() < epsilon:
                actions[idx] = np.random.randint(0, self.l_action)
            else:
                actions[idx] = actions_argmax[idx]

        return actions

    def assign_roles(self, list_obs, epsilon, sess):
        """Get high-level role assignment actions for all agents.

        Args:
            list_obs: list of vectors, one per agent
            epsilon: exploration parameter
            sess: TF session

        Returns: np.array of role indices
        """
        obs = np.array(list_obs)
        feed = {self.obs : obs}
        roles_argmax = sess.run(self.argmax_Q, feed_dict=feed)

        roles = np.zeros(self.n_agents, dtype=int)
        for idx in range(self.n_agents):
            if np.random.rand() < epsilon:
                roles[idx] = np.random.randint(0, self.N_roles)
            else:
                roles[idx] = roles_argmax[idx]

        return roles

    def assign_roles_centralized(self, state, epsilon, sess):
        """Centralized skill selection for all agents.

        Directly samples one single action index from high-level Q function
        Maps action index to roles for all agents

        Returns np.array of role indices
        """
        if np.random.rand() < epsilon:
            idx = np.random.randint(0, self.dim_role_space)
        else:
            feed = {self.state : np.array([state])}
            idx = sess.run(self.argmax_Q_high, feed_dict=feed)
        roles = np.array( self.list_list_indices[int(idx)] )

        return roles, np.squeeze(idx)

    def create_train_op(self):

        self.td_target = tf.placeholder(tf.float32, [None], 'td_target')
        if self.alg_name == 'hsd-scripted':
            # TD target calculated in train_step() using Mixer_target
            self.loss_Q_high = tf.reduce_mean(tf.square(self.td_target - tf.squeeze(self.mixer)))
        elif self.alg_name == 'mara-c':
            # Treat the role assignment as single-agent Q-learning
            self.td_error_Q = self.td_target - tf.reduce_sum(tf.multiply(self.Q_high, self.actions_high_1hot), axis=1)
            self.loss_Q_high = tf.reduce_mean(tf.square(self.td_error_Q))

        self.Q_opt = tf.train.AdamOptimizer(self.lr_Q)
        self.Q_op = self.Q_opt.minimize(self.loss_Q_high)

    def create_train_op_IQL(self):
        self.td_target_IQL = tf.placeholder(tf.float32, [None], 'td_target_IQL')
        self.td_error = self.td_target_IQL - tf.reduce_sum(tf.multiply(self.Q_low, self.actions_low_1hot), axis=1)
        self.loss_IQL = tf.reduce_mean(tf.square(self.td_error))

        self.IQL_opt = tf.train.AdamOptimizer(self.lr_Q)
        self.IQL_op = self.IQL_opt.minimize(self.loss_IQL)

    def create_summary(self):

        if self.alg_name == 'hsd-scripted':
            summaries = [tf.summary.scalar('loss_Q_high', self.loss_Q_high)]
            mixer_main_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'Mixer_main')
            for v in mixer_main_variables:
                summaries.append(tf.summary.histogram(v.op.name, v))
            grads = self.Q_opt.compute_gradients(self.loss_Q_high, mixer_main_variables)
            for grad, var in grads:
                if grad is not None:
                    summaries.append( tf.summary.histogram(var.op.name+'/gradient', grad) )

            agent_main_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'Agent_main')
            for v in agent_main_variables:
                summaries.append(tf.summary.histogram(v.op.name, v))
            grads = self.Q_opt.compute_gradients(self.loss_Q_high, agent_main_variables)
            for grad, var in grads:
                if grad is not None:
                    summaries.append( tf.summary.histogram(var.op.name+'/gradient', grad) )
        elif self.alg_name == 'mara-c':
            summaries = [tf.summary.scalar('loss_Q_high', self.loss_Q_high)]
            Q_main_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'Qhigh_main')
            for v in Q_main_variables:
                summaries.append(tf.summary.histogram(v.op.name, v))
            grads = self.Q_opt.compute_gradients(self.loss_Q_high, Q_main_variables)
            for grad, var in grads:
                if grad is not None:
                    summaries.append( tf.summary.histogram(var.op.name+'/gradient', grad) )

        self.summary_op = tf.summary.merge(summaries)

        summaries_low = [tf.summary.scalar('loss_IQL', self.loss_IQL)]
        Qlow_main_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'Qlow_main')
        for v in Qlow_main_variables:
            summaries_low.append(tf.summary.histogram(v.op.name, v))
        grads = self.IQL_opt.compute_gradients(self.loss_IQL, Qlow_main_variables)
        for grad, var in grads:
            if grad is not None:
                summaries_low.append( tf.summary.histogram(var.op.name+'/gradient', grad) )

        self.summary_op_low = tf.summary.merge(summaries_low)

    def process_actions(self, n_steps, actions, n_actions):
        """
        Args:
            n_steps: number of steps in trajectory
            actions: must have shape [time, n_agents], and values are action indices
            n_actions: dimension of action space

        Returns: 1-hot representation of actions
        """
        # Each row of actions is one time step,
        # row contains action indices for all agents
        # Convert to [time, agents, N_roles]
        # so each agent gets its own 1-hot row vector
        actions_1hot = np.zeros([n_steps, self.n_agents, n_actions], dtype=int)
        grid = np.indices((n_steps, self.n_agents))
        actions_1hot[grid[0], grid[1], actions] = 1

        # In-place reshape of actions to [time*n_agents, N_roles]
        actions_1hot.shape = (n_steps*self.n_agents, n_actions)

        return actions_1hot

    def process_batch(self, batch):
        """Used for high-level buffer

        Extract quantities of the same type from batch.
        Format batch so that each agent at each time step is one batch entry.
        """
        # shapes are [time, ...original dims...]
        if self.alg_name == 'hsd-scripted':
            state = np.stack(batch[:,0]) # [time, l_state]
            obs = np.stack(batch[:,1]) # [time, agents, l_obs]
            actions = np.stack(batch[:,2]) # [time, agents]
            reward = np.stack(batch[:,3]) # [time]
            state_next = np.stack(batch[:,4]) # [time, l_state]
            obs_next = np.stack(batch[:,5]) # [time, agents, l_obs]
            done = np.stack(batch[:,6]) # [time]
        elif self.alg_name == 'mara-c':
            state = np.stack(batch[:,0]) # [time, l_state]
            actions = np.stack(batch[:,1]) # [time]
            reward = np.stack(batch[:,2]) # [time]
            state_next = np.stack(batch[:,3]) # [time, l_state]
            done = np.stack(batch[:,4]) # [time]

        # Try to free memory
        batch = None

        n_steps = state.shape[0]

        if self.alg_name == 'hsd-scripted':
            # In-place reshape for obs, so that one time step
            # for one agent is considered one batch entry
            obs.shape = (n_steps * self.n_agents, self.l_obs)
            obs_next.shape = (n_steps * self.n_agents, self.l_obs)
            actions_1hot = self.process_actions(n_steps, actions, self.N_roles)
            return n_steps, state, obs, actions_1hot, reward, state_next, obs_next, done
        elif self.alg_name == 'mara-c':
            actions_1hot = np.zeros([n_steps, self.dim_role_space])
            actions_1hot[np.arange(n_steps), actions] = 1
            return n_steps, state, actions_1hot, reward, state_next, done

    def process_batch_low(self, batch):
        """
        Extract quantities of the same type from batch.
        Format batch so that each agent at each time step is one batch entry.
        """
        # shapes are [time, ...original dims...]
        # state = np.stack(batch[:,0]) # [time, l_state]
        obs = np.stack(batch[:,0]) # [time, agents, l_obs]
        actions = np.stack(batch[:,1]) # [time, agents]
        rewards = np.stack(batch[:,2]) # [time, agents]
        # state_next = np.stack(batch[:,4]) # [time, l_state]
        obs_next = np.stack(batch[:,3]) # [time, agents, l_obs]
        # done = np.stack(batch[:,6]) # [time]
        roles = np.stack(batch[:,4]) # [time, agents, N_roles]

        # Try to free memory
        batch = None

        n_steps = obs.shape[0]

        # In-place reshape for obs, so that one time step
        # for one agent is considered one batch entry
        obs.shape = (n_steps * self.n_agents, self.l_obs)
        obs_next.shape = (n_steps * self.n_agents, self.l_obs)
        rewards.shape = (n_steps * self.n_agents)
        roles.shape = (n_steps * self.n_agents, self.N_roles)

        actions_1hot = self.process_actions(n_steps, actions, self.l_action)

        return n_steps, obs, actions_1hot, rewards, obs_next, roles

    def train_step(self, sess, batch, step_train=0, summarize=False, writer=None):
        """Training step for role assignment policy via QMIX or Q-learning."""
        if self.alg_name == 'hsd-scripted':
            # Each agent for each time step is now a batch entry
            n_steps, state, obs, actions_1hot, reward, state_next, obs_next, done = self.process_batch(batch)

            # Get argmax actions from target networks
            feed = {self.obs : obs_next}
            argmax_actions = sess.run(self.argmax_Q_target, feed_dict=feed) # [batch*n_agents]
            # Convert to 1-hot
            actions_target_1hot = np.zeros([n_steps * self.n_agents, self.N_roles], dtype=int)
            actions_target_1hot[np.arange(n_steps*self.n_agents), argmax_actions] = 1

            # Get Q_tot target value
            feed = {self.state : state_next,
                    self.actions_1hot : actions_target_1hot,
                    self.obs : obs_next}
            Q_tot_target = sess.run(self.mixer_target, feed_dict=feed)

            done_multiplier = -(done - 1)
            target = reward + self.gamma * np.squeeze(Q_tot_target) * done_multiplier
        elif self.alg_name == 'mara-c':
            n_steps, state, actions_1hot, reward, state_next, done = self.process_batch(batch)

            feed = {self.state : state_next}
            Q_target = sess.run(self.Q_high_target, feed_dict=feed)
            target = reward + self.gamma * np.max(Q_target, axis=1)

        feed = {self.state : state, self.td_target : target}
        if self.alg_name == 'hsd-scripted':
            feed[self.obs] = obs
            feed[self.actions_1hot] = actions_1hot
        elif self.alg_name == 'mara-c':
            feed[self.actions_high_1hot] = actions_1hot

        if summarize:
            summary, _ = sess.run([self.summary_op, self.Q_op], feed_dict=feed)
            writer.add_summary(summary, step_train)
        else:
            _ = sess.run(self.Q_op, feed_dict=feed)

        sess.run(self.list_update_target_ops)

    def train_step_low(self, sess, batch, step_train=0, summarize=False, writer=None):
        """Training step for low-level action policy.

        Runs independent Q-learning on each agent's experiences using local rewards
        """
        n_steps, obs, actions_1hot, rewards, obs_next, roles = self.process_batch_low(batch)

        # Get target values
        feed = {self.obs : obs_next, self.role : roles}
        Q_target = sess.run(self.Q_low_target, feed_dict=feed)
        target = rewards + self.gamma * np.max(Q_target, axis=1)

        feed = {self.obs : obs,
                self.actions_low_1hot : actions_1hot,
                self.role : roles,
                self.td_target_IQL : target}
        if summarize:
            summary, _ = sess.run([self.summary_op_low, self.IQL_op], feed_dict=feed)
            writer.add_summary(summary, step_train)
        else:
            _ = sess.run(self.IQL_op, feed_dict=feed)

        sess.run(self.list_update_target_ops_low)