class AtomInfo(object):
class PeriodicData(object):
    def __init__(self):
        def append_convertor(word):
    def _add_atom_info(self, atom_info):
    def __len__(self):
    def __getitem__(self, index):
    def iter_numbers(self): def parse_unit(expression): class CriteriaSet(object):
    """A set of criteria that can be associated with a custum pattern."""

    def __init__(self, vertex_criteria=None, edge_criteria=None, **kwargs):
        """
           Arguments:
            | ``vertex_criteria``  --  a dictionary with criteria for the
                                       vertices, ``key=vertex_index``,
                                       ``value=criterion``
            | ``edge_criteria``  --  a dictionary with criteria for the edges
                                     ``key=edge_index``, ``value=criterion``

           Any other keyword argument will be assigned as attribute to matches
           that fulfill the criteria of this set.
        """
        if vertex_criteria is None:
            self.vertex_criteria = {}
        else:
            self.vertex_criteria = vertex_criteria
        if edge_criteria is None:
            self.edge_criteria = {}
        else:
            self.edge_criteria = edge_criteria
        self.info = kwargs

    def test_match(self, match, pattern_graph, subject_graph):
        """Test if a match satisfies the criteria"""
        for vertex0, c in self.vertex_criteria.items():
            vertex1 = match.forward[vertex0]
            if not c(vertex1, subject_graph):
                return False
        for edge0_index, c in self.edge_criteria.items():
            vertex0a, vertex0b = pattern_graph.edges[edge0_index]
            edge1_index = subject_graph.edge_index[frozenset([
                match.forward[vertex0a],
                match.forward[vertex0b],
            ])]
            if not c(edge1_index, subject_graph):
                return False
        return True class GraphSearch(object):
    """An algorithm that searches for all matches of a pattern in a graph

       Usage:

         >>> gs = GraphSearch(pattern)
         >>> for match in gs(graph):
         ...     print match.forward
    """

    def __init__(self, pattern, debug=False):
        """
           Arguments:
            | ``pattern``  --  A Pattern instance, describing the pattern to
                               look for
            | ``debug``  --  When true, debugging info is printed on screen
                             [default=False]
        """
        self.pattern = pattern
        self.debug = debug

    def __call__(self, subject_graph, one_match=False):
        """Iterator over all matches of self.pattern in the given graph.

           Arguments:
            | subject_graph  --  The subject_graph in which the matches
                                 according to self.pattern have to be found.
            | one_match --  If True, only one match will be returned. This
                            allows certain optimizations.
        """
        # Matches are grown iteratively.
        for vertex0, vertex1 in self.pattern.iter_initial_relations(subject_graph):
            init_match = self.pattern.MatchClass.from_first_relation(vertex0, vertex1)
            # init_match cotains only one source -> dest relation. starting from
            # this initial match, the function iter_matches extends the match
            # in all possible ways and yields the completed matches
            for canonical_match in self._iter_matches(init_match, subject_graph, one_match):
                # Some patterns my exclude symmetrically equivalent matches as
                # to aviod dupplicates. with such a 'canonical' solution,
                # the pattern is allowed to generate just those symmatrical
                # duplicates of interest.
                ifm = self.pattern.iter_final_matches(canonical_match, subject_graph, one_match)
                for final_match in ifm:
                    self.print_debug("final_match: %s" % final_match)
                    yield final_match
                    if one_match: return

    def print_debug(self, text, indent=0):
        """Only prints debug info on screen when self.debug == True."""
        if self.debug:
            if indent > 0:
                print(" "*self.debug, text)
            self.debug += indent
            if indent <= 0:
                print(" "*self.debug, text)

    def _iter_candidate_groups(self, init_match, edges0, edges1):
        """Divide the edges into groups"""
        # collect all end vertices0 and end vertices1 that belong to the same
        # group.
        sources = {}
        for start_vertex0, end_vertex0 in edges0:
            l = sources.setdefault(start_vertex0, [])
            l.append(end_vertex0)
        dests = {}
        for start_vertex1, end_vertex1 in edges1:
            start_vertex0 = init_match.reverse[start_vertex1]
            l = dests.setdefault(start_vertex0, [])
            l.append(end_vertex1)
        for start_vertex0, end_vertices0 in sources.items():
            end_vertices1 = dests.get(start_vertex0, [])
            yield end_vertices0, end_vertices1


    def _iter_new_relations(self, init_match, subject_graph, edges0, constraints0, edges1):
        """Given an onset for a match, iterate over all possible new key-value pairs"""
        # Count the number of unique edges0[i][1] values. This is also
        # the number of new relations.
        num_new_relations = len(set(j for i, j in edges0))

        def combine_small(relations, num):
            """iterate over all compatible combinations within one set of relations"""
            if len(relations) == 0:
                return
            for i, pivot in enumerate(relations):
                if num == 1:
                    yield (pivot, )
                else:
                    compatible_relations = list(
                        item for item in relations[:i]
                        if pivot[0]!=item[0] and pivot[1]!=item[1]
                    )
                    for tail in combine_small(compatible_relations, num-1):
                        yield (pivot, ) + tail

        # generate candidate relations
        candidate_relations = []
        icg = self._iter_candidate_groups(init_match, edges0, edges1)
        for end_vertices0, end_vertices1 in icg:
            if len(end_vertices0) > len(end_vertices1):
                return # this can never work, the subject graph is 'too small'
            elif not self.pattern.sub and \
                 len(end_vertices0) != len(end_vertices1):
                return # an exact match is sought, this can never work
            l = []
            for end_vertex0 in end_vertices0:
                for end_vertex1 in end_vertices1:
                    if self.pattern.compare(end_vertex0, end_vertex1, subject_graph):
                        l.append((end_vertex0, end_vertex1))
            # len(end_vertices0) = the total number of relations that must be
            # made in this group
            if len(l) > 0:
                # turn l into a list of sets of internally compatible candidate
                # relations in this group
                l = list(combine_small(l, len(end_vertices0)))
                candidate_relations.append(l)
        if len(candidate_relations) == 0:
            return
        self.print_debug("candidate_relations: %s" % candidate_relations)

        def combine_big(pos=0):
            """Iterate over all possible sets of relations"""
            # pos is an index in candidate_relations
            crs = candidate_relations[pos]
            if pos == len(candidate_relations)-1:
                for relations in crs:
                    yield relations
            else:
                for tail in combine_big(pos+1):
                    for relations in crs:
                        yield relations + tail

        # final loop
        for new_relations in combine_big():
            new_relations = set(new_relations)
            self.print_debug("new_relations: %s" % (new_relations, ))
            # check the total number of new relations
            if len(new_relations) != num_new_relations:
                continue
            # check sanity of relations
            forward = dict(new_relations)
            if len(forward) != num_new_relations:
                continue
            reverse = dict((j, i) for i, j in new_relations)
            if len(reverse) != num_new_relations:
                continue
            # check the constraints
            for a0, b0 in constraints0:
                if forward[a0] not in subject_graph.neighbors[forward[b0]]:
                    forward = None
                    break
            if forward is None:
                continue
            yield forward

    def _iter_matches(self, input_match, subject_graph, one_match, level=0):
        """Given an onset for a match, iterate over all completions of that match

           This iterator works recursively. At each level the match is extended
           with a new set of relations based on vertices in the pattern graph
           that are at a distances 'level' from the starting vertex
        """
        self.print_debug("ENTERING _ITER_MATCHES", 1)
        self.print_debug("input_match: %s" % input_match)
        # A) collect the new edges in the pattern graph and the subject graph
        # to extend the match.
        #
        # Note that the edges are ordered. edge[0] is always in the match.
        # edge[1] is never in the match. The constraints contain information
        # about the end points of edges0. It is a list of two-tuples where
        # (a, b) means that a and b must be connected.
        #
        # Second note: suffix 0 indicates the pattern graph and suffix 1
        # is used for the subject graph.
        edges0, constraints0 = self.pattern.get_new_edges(level)
        edges1 = input_match.get_new_edges(subject_graph)
        self.print_debug("edges0: %s" % edges0)
        self.print_debug("constraints0: %s" % constraints0)
        self.print_debug("edges1: %s" % edges1)

        # B) iterate over the sets of new relations: [(vertex0[i], vertex1[j]),
        # ...] that contain all endpoints of edges0, that satisfy the
        # constraints0 and where (vertex0[i], vertex1[j]) only occurs if these
        # are end points of a edge0 and edge1 whose starting points are already
        # in init_match. These conditions are implemented in an iterator as to
        # separate concerns. This iterator also calls the routines that check
        # whether vertex1[j] also satisfies additional conditions inherent
        # vertex0[i].
        inr = self._iter_new_relations(input_match, subject_graph, edges0,
                                       constraints0, edges1)
        for new_relations in inr:
            # for each set of new_relations, construct a next_match and recurse
            next_match = input_match.copy_with_new_relations(new_relations)
            if not self.pattern.check_next_match(next_match, new_relations, subject_graph, one_match):
                continue
            if self.pattern.complete(next_match, subject_graph):
                yield next_match
            else:
                for match in self._iter_matches(next_match, subject_graph, one_match, level+1):
                    yield match
        self.print_debug("LEAVING_ITER_MATCHES", -1) class MolecularGraph(Graph):
    """Describes a molecular graph: connectivity, atom numbers and bond orders.

       This class inherits all features from the Graph class and adds methods
       and attributes that are specific from molecular graphs. Instances are
       immutable, so if you want to modify a molecular graph, just instantiate a
       new object with modified connectivity, numbers and orders. The advantage
       is that various graph analysis and properties can be cached.
    """
    def _check_numbers(self, numbers):
        """the number of vertices and atomic numbers must be the same"""
        if len(numbers) != self.num_vertices:
            raise TypeError("The number of vertices must match the length of "
                "the atomic numbers array.")

    def _check_orders(self, orders):
        """the number of edges and bond orders must be the same"""
        if len(orders) != self.num_edges:
            raise TypeError("The number of edges must match the length of "
                "the bond orders array.")

    def _check_symbols(self, symbols):
        """the size must be the same as the length of the array numbers and all elements must be strings"""
        if len(symbols) != self.size:
            raise TypeError("The number of symbols in the graph does not "
                "match the length of the atomic numbers array.")
        for symbol in symbols:
            if not isinstance(symbol, str):
                raise TypeError("All symbols must be strings.")

    numbers = ReadOnlyAttribute(np.ndarray, none=False, check=_check_numbers,
        npdim=1, npdtype=np.signedinteger, doc="the atomic numbers associated "
        "with the vertices")
    orders = ReadOnlyAttribute(np.ndarray, none=False, check=_check_orders,
        npdim=1, npdtype=np.floating, doc="the bond orders associated with the "
        "edges")
    symbols = ReadOnlyAttribute(tuple, _check_symbols, doc="symbols for the "
        "atoms, which can be element names for force-field atom types")

    @classmethod
    def from_geometry(cls, molecule, do_orders=False, scaling=1.0):
        """Construct a MolecularGraph object based on interatomic distances

           All short distances are computed with the binning module and compared
           with a database of bond lengths. Based on this comparison, bonded
           atoms are detected.

           Before marking a pair of atoms A and B as bonded, it is also checked
           that there is no third atom C somewhat between A and B.
           When an atom C exists that is closer to B (than A) and the angle
           A-B-C is less than 45 degrees, atoms A and B are not bonded.
           Similarly if C is closer to A (than B) and the angle B-A-C is less
           then 45 degrees, A and B are not connected.

           Argument:
            | ``molecule``  --  The molecule to derive the graph from

           Optional arguments:
            | ``do_orders``  --  set to True to estimate the bond order
            | ``scaling``  --  scale the threshold for the connectivity. increase
                               this to 1.5 in case of transition states when a
                               fully connected topology is required.
        """
        from molmod.bonds import bonds

        unit_cell = molecule.unit_cell
        pair_search = PairSearchIntra(
            molecule.coordinates,
            bonds.max_length*bonds.bond_tolerance*scaling,
            unit_cell
        )

        orders = []
        lengths = []
        edges = []

        for i0, i1, delta, distance in pair_search:
            bond_order = bonds.bonded(molecule.numbers[i0], molecule.numbers[i1], distance/scaling)
            if bond_order is not None:
                if do_orders:
                    orders.append(bond_order)
                lengths.append(distance)
                edges.append((i0,i1))

        if do_orders:
            result = cls(edges, molecule.numbers, orders, symbols=molecule.symbols)
        else:
            result = cls(edges, molecule.numbers, symbols=molecule.symbols)

        # run a check on all neighbors. if two bonds point in a direction that
        # differs only by 45 deg. the longest of the two is discarded. the
        # double loop over the neighbors is done such that the longest bonds
        # are eliminated first
        slated_for_removal = set([])
        threshold = 0.5**0.5
        for c, ns in result.neighbors.items():
            lengths_ns = []
            for n in ns:
                delta = molecule.coordinates[n] - molecule.coordinates[c]
                if unit_cell is not None:
                    delta = unit_cell.shortest_vector(delta)
                length = np.linalg.norm(delta)
                lengths_ns.append([length, delta, n])
            lengths_ns.sort(reverse=True, key=(lambda r: r[0]))
            for i0, (length0, delta0, n0) in enumerate(lengths_ns):
                for i1, (length1, delta1, n1) in enumerate(lengths_ns[:i0]):
                    if length1 == 0.0:
                        continue
                    cosine = np.dot(delta0, delta1)/length0/length1
                    if cosine > threshold:
                        # length1 > length0
                        slated_for_removal.add((c,n1))
                        lengths_ns[i1][0] = 0.0
        # construct a mask
        mask = np.ones(len(edges), bool)
        for i0, i1 in slated_for_removal:
            edge_index = result.edge_index.get(frozenset([i0,i1]))
            if edge_index is None:
                raise ValueError('Could not find edge that has to be removed: %i %i' % (i0, i1))
            mask[edge_index] = False
        # actual removal
        edges = [edges[i] for i in range(len(edges)) if mask[i]]
        if do_orders:
            orders = [orders[i] for i in range(len(orders)) if mask[i]]
            result = cls(edges, molecule.numbers, orders)
        else:
            result = cls(edges, molecule.numbers)

        lengths = [lengths[i] for i in range(len(lengths)) if mask[i]]
        result.bond_lengths = np.array(lengths)

        return result

    @classmethod
    def from_blob(cls, s):
        """Construct a molecular graph from the blob representation"""
        atom_str, edge_str = s.split()
        numbers = np.array([int(s) for s in atom_str.split(",")])
        edges = []
        orders = []
        for s in edge_str.split(","):
            i, j, o = (int(w) for w in s.split("_"))
            edges.append((i, j))
            orders.append(o)
        return cls(edges, numbers, np.array(orders))

    def __init__(self, edges, numbers, orders=None, symbols=None, num_vertices=None):
        """
           Arguments:
            | ``edges``  --  See base class (Graph) documentation
            | ``numbers``  --  consecutive atom numbers

           Optional arguments:
            | ``orders``  --  bond orders
            | ``symbols``  --  atomic symbols
            | ``num_vertices``  --  must be the same as the number of atoms or None

           When the nature of an atom or a bond is unclear ambiguous, set the
           corresponding integer to zero. This means the nature of the atom or
           bond is unspecified. When the bond orders are not given, they are all
           set to  zero.

           If you want to use 'special' atom types, use negative numbers. The
           same for bond orders. e.g. a nice choice for the bond order of a
           hybrid bond is -1.
        """
        if num_vertices is not None and num_vertices != len(numbers):
            raise ValueError('The number of vertices must be the same as the number of atoms.')
        if orders is None:
            orders = np.ones(len(edges), float)
        Graph.__init__(self, edges, len(numbers))
        self.numbers = numbers
        self.orders = orders
        self.symbols = symbols

    def __mul__(self, repeat):
        """Construct a graph that repeats this graph a number of times

           Arguments:
            | ``repeat`` -- The number of repetitions.
        """
        if not isinstance(repeat, int):
            raise TypeError("Can only multiply a graph with an integer")
        # copy edges
        new_edges = []
        for i in range(repeat):
            for vertex1, vertex2 in self.edges:
                new_edges.append(frozenset([vertex1+i*self.num_vertices, vertex2+i*self.num_vertices]))
        # copy numbers
        new_numbers = np.zeros((repeat, len(self.numbers)), int)
        new_numbers[:] = self.numbers
        new_numbers = new_numbers.ravel()
        # copy orders
        new_orders = np.zeros((repeat, len(self.orders)), int)
        new_orders[:] = self.orders
        new_orders = new_orders.ravel()
        # copy symbols
        if self.symbols is not None:
            new_symbols = self.symbols*repeat
        else:
            new_symbols = None
        return MolecularGraph(new_edges, new_numbers, new_orders, new_symbols)

    __rmul__ = __mul__

    @cached
    def blob(self):
        """A compact text representation of the graph"""
        atom_str = ",".join(str(number) for number in self.numbers)
        edge_str = ",".join("%i_%i_%i" % (i, j, o) for (i, j), o in zip(self.edges, self.orders))
        return "%s %s" % (atom_str, edge_str)

    def get_vertex_string(self, i):
        """Return a string based on the atom number"""
        number = self.numbers[i]
        if number == 0:
            return Graph.get_vertex_string(self, i)
        else:
            # pad with zeros to make sure that string sort is identical to number sort
            return "%03i" % number

    def get_edge_string(self, i):
        """Return a string based on the bond order"""
        order = self.orders[i]
        if order == 0:
            return Graph.get_edge_string(self, i)
        else:
            # pad with zeros to make sure that string sort is identical to number sort
            return "%03i" % order

    def get_subgraph(self, subvertices, normalize=False):
        """Creates a subgraph of the current graph

           See :meth:`molmod.graphs.Graph.get_subgraph` for more information.
        """
        graph = Graph.get_subgraph(self, subvertices, normalize)
        if normalize:
            new_numbers = self.numbers[graph._old_vertex_indexes] # vertices do change
        else:
            new_numbers = self.numbers # vertices don't change!
        if self.symbols is None:
            new_symbols = None
        elif normalize:
            new_symbols = tuple(self.symbols[i] for i in graph._old_vertex_indexes)
        else:
            new_symbols = self.symbols
        new_orders = self.orders[graph._old_edge_indexes]
        result = MolecularGraph(graph.edges, new_numbers, new_orders, new_symbols)
        if normalize:
            result._old_vertex_indexes = graph._old_vertex_indexes
        result._old_edge_indexes = graph._old_edge_indexes
        return result

    def add_hydrogens(self, formal_charges=None):
        """Returns a molecular graph where hydrogens are added explicitely

           When the bond order is unknown, it assumes bond order one. If the
           graph  has an attribute formal_charges, this routine will take it
           into account when counting the number of hydrogens to be added. The
           returned graph will also have a formal_charges attribute.

           This routine only adds hydrogen atoms for a limited set of atoms from
           the periodic system: B, C, N, O, F, Al, Si, P, S, Cl, Br.
        """

        new_edges = list(self.edges)
        counter = self.num_vertices
        for i in range(self.num_vertices):
            num_elec = self.numbers[i]
            if formal_charges is not None:
                num_elec -= int(formal_charges[i])
            if num_elec >= 5 and num_elec <= 9:
                num_hydrogen = num_elec - 10 + 8
            elif num_elec >= 13 and num_elec <= 17:
                num_hydrogen = num_elec - 18 + 8
            elif num_elec == 35:
                num_hydrogen = 1
            else:
                continue
            if num_hydrogen > 4:
                num_hydrogen = 8 - num_hydrogen
            for n in self.neighbors[i]:
                bo = self.orders[self.edge_index[frozenset([i, n])]]
                if bo <= 0:
                    bo = 1
                num_hydrogen -= int(bo)
            for j in range(num_hydrogen):
                new_edges.append((i, counter))
                counter += 1
        new_numbers = np.zeros(counter, int)
        new_numbers[:self.num_vertices] = self.numbers
        new_numbers[self.num_vertices:] = 1
        new_orders = np.zeros(len(new_edges), int)
        new_orders[:self.num_edges] = self.orders
        new_orders[self.num_edges:] = 1
        result = MolecularGraph(new_edges, new_numbers, new_orders)
        return result class BondPattern(CustomPattern):
    """Pattern for two consecutive vertices"""
    def __init__(self, criteria_sets=None, vertex_tags=None):
        """
           Arguments: see :class:`molmod.graphs.CustomPattern`
        """
        if vertex_tags is None:
            vertex_tags = {}
        pattern_graph = Graph([(0, 1)])
        CustomPattern.__init__(self, pattern_graph, criteria_sets, vertex_tags) class BendingAnglePattern(CustomPattern):
    """Pattern for three consecutive vertices"""
    def __init__(self, criteria_sets=None, vertex_tags=None):
        """
           Arguments: see :class:`molmod.graphs.CustomPattern`
        """
        if vertex_tags is None:
            vertex_tags = {}
        pattern_graph = Graph([(0, 1), (1, 2)])
        CustomPattern.__init__(self, pattern_graph, criteria_sets, vertex_tags) class DihedralAnglePattern(CustomPattern):
    """Pattern for four consecutive vertices"""
    def __init__(self, criteria_sets=None, vertex_tags=None):
        """
           Arguments: see :class:`molmod.graphs.CustomPattern`
        """
        if vertex_tags is None:
            vertex_tags = {}
        pattern_graph = Graph([(0, 1), (1, 2), (2, 3)])
        CustomPattern.__init__(self, pattern_graph, criteria_sets, vertex_tags) class OutOfPlanePattern(CustomPattern):
    """Pattern for a central vertex connected to three other vertices"""
    def __init__(self, criteria_sets=None, vertex_tags=None):
        """
           Arguments: see :class:`molmod.graphs.CustomPattern`
        """
        if vertex_tags is None:
            vertex_tags = {}
        pattern_graph = Graph([(0, 1), (0, 2), (0, 3)])
        CustomPattern.__init__(self, pattern_graph, criteria_sets, vertex_tags) class HasNumNeighbors(object):
    """Criterion for the number of neighboring vertexes"""

    def __init__(self, count):
        """
           Arguments:
            | ``count``  --  the expected number of neighbors
        """
        self.count = count

    def __call__(self, index, graph):
        """Return True only if the number of neighbors is correct

           Arguments:
            | ``index``  --  the index of the vertex/edge on which the criterion is
                             applied
            | ``graph``  --  the graph on which the criterion is tested
        """
        return len(graph.neighbors[index]) == self.count class Graph(ReadOnly):
    """An undirected graph, where edges have equal weight

       The edges attribute is the most elementary and is always available. All
       other attributes are optional.

       Graphs are meant to be immutable objects. Once created they are not
       supposed to be modified. If you want an adapted graph, create a new
       object. The main reason is that all graph analysis work is cached in this
       object. When the edges change, the cache becomes invalid and should be
       erased. The latter is not supported as it is easier to create just a new
       graph object with the modified edges.

       If you want to associate data with vertices or edges, add custom
       dictionary or list attributes to the graph object. This is the way to
       work attributes that are not applicable to all vertices or edges:

       >>> graph.vertex_property = {vertex1: blablabla, vertex2: blablabla, ...}
       >>> graph.edge_property = {frozenset([vertex1, vertex2]): blablabla, ..}

       If a certain property applies to all vertices or edges, it is sometimes
       more practical to work with lists or ``numpy`` arrays that have the same
       ordering as the vertices or the edges:

       >>> # atom numbers of ethene
       >>> graph.vertex_property = np.array([6, 6, 1, 1, 1, 1], int)
       >>> # bond orders of ethene
       >>> graph.edge_property = np.array([2, 1, 1, 1, 1], int)
    """
    edges = ReadOnlyAttribute(tuple, none=False, doc="the incidence list")
    num_vertices = ReadOnlyAttribute(int, none=False, doc="the number of vertices")

    def __init__(self, edges, num_vertices=None):
        """
           Arguments:
            | ``edges`` -- ``tuple(frozenset([vertex1, vertex2]), ...)``
            | ``num_vertices`` -- number of vertices

           vertex1 to vertexN must be integers from 0 to N-1. If vertices above
           the highest vertex value are not connected by edges, use the
           num_vertices argument to tell what the total number of vertices is.

           If the edges argument does not have the correct format, it will be
           converted.
        """

        tmp = []
        for edge in edges:
            if len(edge) != 2:
                raise TypeError("The edges must be a iterable with 2 elements")
            i, j = edge
            i = int(i)
            j = int(j)
            if i == j:
                raise ValueError("A edge must contain two different values.")
            if i < 0 or j < 0:
                raise TypeError("The edges must contain positive integers.")
            tmp.append(frozenset([i, j]))
        edges = tuple(tmp)

        if len(edges) == 0:
            real_num_vertices = 0
        else:
            real_num_vertices = max(max(a, b) for a, b in edges)+1
        if num_vertices is not None:
            if not isinstance(num_vertices, int):
                raise TypeError("The optional argument num_vertices must be an "
                    "integer when given.")
            if num_vertices < real_num_vertices:
                raise ValueError("num_vertices must be equal or larger to the "
                    "number of vertices deduced from the edge list.")
            real_num_vertices = num_vertices

        self.edges = edges
        self.num_vertices = real_num_vertices

    num_edges = property(lambda self: len(self.edges),
        doc="*Read-only attribute:* the number of edges in the graph.")

    def __mul__(self, repeat):
        """Construct a graph that repeats this graph a number of times

           Arguments:
            | ``repeat`` -- The number of repetitions.
        """
        if not isinstance(repeat, int):
            raise TypeError("Can only multiply a graph with an integer")
        new_edges = []
        for i in range(repeat):
            for vertex1, vertex2 in self.edges:
                new_edges.append(frozenset([
                    vertex1+i*self.num_vertices,
                    vertex2+i*self.num_vertices
                ]))
        return Graph(new_edges, self.num_vertices*repeat)

    __rmul__ = __mul__

    def __str__(self):
        return " ".join("%i-%i" % tuple(sorted([i, j])) for i, j in self.edges)

    # functions that should be implemented by derived classes

    def get_vertex_string(self, i):
        """Returns a string that fully characterizes the nature of the vertex

           In case of an ordinary graph, all vertices are the same.
        """
        return ""

    def get_edge_string(self, i):
        """Returns a string that fully characterizes the nature of an edge

           In case of an ordinary graph, all edges are the same.
        """
        return ""

    # cached attributes:

    @cached
    def edge_index(self):
        """A map to look up the index of a edge"""
        return dict((edge, index) for index, edge in enumerate(self.edges))

    @cached
    def neighbors(self):
        """A dictionary with neighbors

           The dictionary will have the following form:
           ``{vertexX: (vertexY1, vertexY2, ...), ...}``
           This means that vertexX and vertexY1 are connected etc. This also
           implies that the following elements are part of the dictionary:
           ``{vertexY1: (vertexX, ...), vertexY2: (vertexX, ...), ...}``.
        """
        neighbors = dict(
            (vertex, []) for vertex
            in range(self.num_vertices)
        )
        for a, b in self.edges:
            neighbors[a].append(b)
            neighbors[b].append(a)
        # turn lists into frozensets
        neighbors = dict((key, frozenset(val)) for key, val in neighbors.items())
        return neighbors

    @cached
    def distances(self):
        """The matrix with the all-pairs shortest path lenghts"""
        from molmod.ext import graphs_floyd_warshall
        distances = np.zeros((self.num_vertices,)*2, dtype=int)
        #distances[:] = -1 # set all -1, which is just a very big integer
        #distances.ravel()[::len(distances)+1] = 0 # set diagonal to zero
        for i, j in self.edges: # set edges to one
            distances[i, j] = 1
            distances[j, i] = 1
        graphs_floyd_warshall(distances)
        return distances

    @cached
    def max_distance(self):
        """The maximum value in the distances matrix."""
        if self.distances.shape[0] == 0:
            return 0
        else:
            return self.distances.max()

    @cached
    def central_vertices(self):
        """Vertices that have the lowest maximum distance to any other vertex"""
        max_distances = self.distances.max(0)
        max_distances_min = max_distances[max_distances > 0].min()
        return (max_distances == max_distances_min).nonzero()[0]

    @cached
    def central_vertex(self):
        """The vertex that has the lowest maximum distance to any other vertex

           This definition does not lead to a unique solution. One arbitrary
           solution is selected.
        """
        return self.central_vertices[0]

    @cached
    def independent_vertices(self):
        """Lists of vertices that are only interconnected within each list

           This means that there is no path from a vertex in one list to a
           vertex in another list. In case of a molecular graph, this would
           yield the atoms that belong to individual molecules.
        """
        candidates = set(range(self.num_vertices))

        result = []
        while len(candidates) > 0:
            pivot = candidates.pop()
            group = [
                vertex for vertex, distance
                in self.iter_breadth_first(pivot)
            ]
            candidates.difference_update(group)

            # this sort makes sure that the order of the vertices is respected
            group.sort()
            result.append(group)
        return result

    @cached
    def fingerprint(self):
        """A total graph fingerprint

           The result is invariant under permutation of the vertex indexes. The
           chance that two different (molecular) graphs yield the same
           fingerprint is small but not zero. (See unit tests.)"""
        if self.num_vertices == 0:
            return np.zeros(20, np.ubyte)
        else:
            return sum(self.vertex_fingerprints)

    @cached
    def vertex_fingerprints(self):
        """A fingerprint for each vertex

           The result is invariant under permutation of the vertex indexes.
           Vertices that are symmetrically equivalent will get the same
           fingerprint, e.g. the hydrogens in methane would get the same
           fingerprint.
        """
        return self.get_vertex_fingerprints(
            [self.get_vertex_string(i) for i in range(self.num_vertices)],
            [self.get_edge_string(i) for i in range(self.num_edges)],
        )

    @cached
    def equivalent_vertices(self):
        """A dictionary with symmetrically equivalent vertices."""
        level1 = {}
        for i, row in enumerate(self.vertex_fingerprints):
            key = row.tobytes()
            l = level1.get(key)
            if l is None:
                l = set([i])
                level1[key] = l
            else:
                l.add(i)
        level2 = {}
        for key, vertices in level1.items():
            for vertex in vertices:
                level2[vertex] = vertices
        return level2

    @cached
    def symmetries(self):
        """Graph symmetries (permutations) that map the graph onto itself."""

        symmetry_cycles = set([])
        symmetries = set([])
        for match in GraphSearch(EqualPattern(self))(self):
            match.cycles = match.get_closed_cycles()
            if match.cycles in symmetry_cycles:
                raise RuntimeError("Duplicates in EqualMatch")
            symmetry_cycles.add(match.cycles)
            symmetries.add(match)
        return symmetries

    @cached
    def symmetry_cycles(self):
        """The cycle representations of the graph symmetries"""
        result = set([])
        for symmetry in self.symmetries:
            result.add(symmetry.cycles)
        return result

    @cached
    def canonical_order(self):
        """The vertices in a canonical or normalized order.

           This routine will return a list of vertices in an order that does not
           depend on the initial order, but only depends on the connectivity and
           the return values of the function self.get_vertex_string.

           Only the vertices that are involved in edges will be included. The
           result can be given as first argument to self.get_subgraph, with
           reduce=True as second argument. This will return a complete canonical
           graph.

           The routine is designed not to use symmetry relations that are
           obtained with the GraphSearch routine. We also tried to create an
           ordering that feels like natural, i.e. starting in the center and
           pushing vertices with few equivalents to the front. If necessary, the
           nature of the vertices and  their bonds to atoms closer to the center
           will also play a role, but only as a last resort.
        """
        # A) find an appropriate starting vertex.
        # Here we take a central vertex that has a minimal number of symmetrical
        # equivalents, 'the highest atom number', and the highest fingerprint.
        # Note that the symmetrical equivalents are computed from the vertex
        # fingerprints, i.e. without the GraphSearch.
        starting_vertex = max(
            (
                -len(self.equivalent_vertices[vertex]),
                self.get_vertex_string(vertex),
                self.vertex_fingerprints[vertex].tobytes(),
                vertex
            ) for vertex in self.central_vertices
        )[-1]

        # B) sort all vertices based on
        #      1) distance from central vertex
        #      2) number of equivalent vertices
        #      3) vertex string, (higher atom numbers come first)
        #      4) fingerprint
        #      5) vertex index
        # The last field is only included to collect the result of the sort.
        # The fingerprint on itself would be sufficient, but the three first are
        # there to have a naturally appealing result.
        l = [
            [
                -distance,
                -len(self.equivalent_vertices[vertex]),
                self.get_vertex_string(vertex),
                self.vertex_fingerprints[vertex].tobytes(),
                vertex
            ] for vertex, distance in self.iter_breadth_first(starting_vertex)
            if len(self.neighbors[vertex]) > 0
        ]
        l.sort(reverse=True)

        # C) The order of some vertices is still not completely set. e.g.
        # consider the case of allene. The four hydrogen atoms are equivalent,
        # but one can have two different orders: make geminiles consecutive or
        # don't. It is more trikcy than one would think at first sight. In the
        # case of allene, geminility could easily solve the problem. Consider a
        # big flat rotationally symmetric molecule (order 2). The first five
        # shells are order 4 and one would just give a random order to four
        # segemnts in the first shell. Only when one reaches the outer part that
        # has order two, it turns out that the arbitrary choices in the inner
        # shell play a role. So it does not help to look at relations with
        # vertices at inner or current shells only. One has to consider the
        # whole picture. (unit testing reveals troubles like these)

        # I need some sleep now. The code below checks for potential fuzz and
        # will raise an error if the ordering is not fully determined yet. One
        # day, I'll need this code more than I do now, and I'll fix things up.
        # I know how to do this, but I don't care enough right now.
        # -- Toon
        for i in range(1, len(l)):
            if l[i][:-1] == l[i-1][:-1]:
                raise NotImplementedError

        # D) Return only the vertex indexes.
        return [record[-1] for record in l]

    # other usefull graph functions

    def iter_breadth_first(self, start=None, do_paths=False, do_duplicates=False):
        """Iterate over the vertices with the breadth first algorithm.

           See http://en.wikipedia.org/wiki/Breadth-first_search for more info.
           If not start vertex is given, the central vertex is taken.

           By default, the distance to the starting vertex is also computed. If
           the path to the starting vertex should be computed instead, set path
           to True.

           When duplicate is True, then vertices that can be reached through
           different  paths of equal length, will be iterated twice. This
           typically only makes sense when path==True.
        """
        if start is None:
            start = self.central_vertex
        else:
            try:
                start = int(start)
            except ValueError:
                raise TypeError("First argument (start) must be an integer.")
            if start < 0 or start >= self.num_vertices:
                raise ValueError("start must be in the range [0, %i[" %
                                 self.num_vertices)
        from collections import deque
        work = np.zeros(self.num_vertices, int)
        work[:] = -1
        work[start] = 0
        if do_paths:
            result = (start, 0, (start, ))
        else:
            result = (start, 0)
        yield result
        todo = deque([result])
        while len(todo) > 0:
            if do_paths:
                parent, parent_length, parent_path = todo.popleft()
            else:
                parent, parent_length = todo.popleft()
            current_length = parent_length + 1
            for current in self.neighbors[parent]:
                visited = work[current]
                if visited == -1 or (do_duplicates and visited == current_length):
                    work[current] = current_length
                    if do_paths:
                        current_path = parent_path + (current, )
                        result = (current, current_length, current_path)
                    else:
                        result = (current, current_length)
                    #print "iter_breadth_first", result
                    yield result
                    todo.append(result)

    def iter_shortest_paths(self, a, b):
        """Iterate over all the shortest paths between vertex a and b."""
        max_len = None
        ibf = self.iter_breadth_first(a, do_paths=True, do_duplicates=True)
        for vertex, length, path in ibf:
            if max_len is not None:
                if length > max_len:
                    return
            if vertex == b:
                max_len = length
                yield path

    def iter_breadth_first_edges(self, start=None):
        """Iterate over the edges with the breadth first convention.

           We need this for the pattern matching algorithms, but a quick look at
           Wikipedia did not result in a known and named algorithm.

           The edges are yielded one by one, together with the distance of the
           edge from the starting vertex and a flag that indicates whether the
           yielded edge connects two vertices that are at the same distance from
           the starting vertex. If that flag is False, the distance from the
           starting vertex to edge[0] is equal to the distance variable and the
           distance from edge[1] to the starting vertex is equal to distance+1.
           One item has the following format: ((i, j), distance, flag)
        """
        if start is None:
            start = self.central_vertex
        else:
            try:
                start = int(start)
            except ValueError:
                raise TypeError("First argument (start) must be an integer.")
            if start < 0 or start >= self.num_vertices:
                raise ValueError("start must be in the range [0, %i[" %
                                 self.num_vertices)
        from collections import deque
        work = np.zeros(self.num_vertices, int)
        work[:] = -1
        work[start] = 0
        todo = deque([start])
        while len(todo) > 0:
            parent = todo.popleft()
            distance = work[parent]
            for current in self.neighbors[parent]:
                if work[current] == -1:
                    yield (parent, current), distance, False
                    work[current] = distance+1
                    todo.append(current)
                elif work[current] == distance and current > parent:
                    # second equation in elif avoids duplicates
                    yield (parent, current), distance, True
                elif work[current] == distance+1:
                    yield (parent, current), distance, False

    def get_subgraph(self, subvertices, normalize=False):
        """Constructs a subgraph of the current graph

           Arguments:
            | ``subvertices`` -- The vertices that should be retained.
            | ``normalize`` -- Whether or not the vertices should renumbered and
                 reduced to the given set of subvertices. When True, also the
                 edges are sorted. It the end, this means that new order of the
                 edges does not depend on the original order, but only on the
                 order of the argument subvertices.
                 This option is False by default. When False, only edges will be
                 discarded, but the retained data remain unchanged. Also the
                 parameter num_vertices is not affected.

           The returned graph will have an attribute ``old_edge_indexes`` that
           relates the positions of the new and the old edges as follows::

             >>> self.edges[result._old_edge_indexes[i]] = result.edges[i]

           In derived classes, the following should be supported::

             >>> self.edge_property[result._old_edge_indexes[i]] = result.edge_property[i]

           When ``normalize==True``, also the vertices are affected and the
           derived classes should make sure that the following works::

             >>> self.vertex_property[result._old_vertex_indexes[i]] = result.vertex_property[i]

           The attribute ``old_vertex_indexes`` is only constructed when
           ``normalize==True``.
        """
        if normalize:
            revorder = dict((j, i) for i, j in enumerate(subvertices))
            new_edges = []
            old_edge_indexes = []
            for counter, (i, j) in enumerate(self.edges):
                new_i = revorder.get(i)
                if new_i is None:
                    continue
                new_j = revorder.get(j)
                if new_j is None:
                    continue
                new_edges.append((new_i, new_j))
                old_edge_indexes.append(counter)
            # sort the edges
            order = list(range(len(new_edges)))
            # argsort in pure python
            order.sort( key=(lambda i: tuple(sorted(new_edges[i]))) )
            new_edges = [new_edges[i] for i in order]
            old_edge_indexes = [old_edge_indexes[i] for i in order]

            result = Graph(new_edges, num_vertices=len(subvertices))
            result._old_vertex_indexes = np.array(subvertices, dtype=int)
            #result.new_vertex_indexes = revorder
            result._old_edge_indexes = np.array(old_edge_indexes, dtype=int)
        else:
            subvertices = set(subvertices)
            old_edge_indexes = np.array([
                i for i, edge in enumerate(self.edges)
                if edge.issubset(subvertices)
            ], dtype=int)
            new_edges = tuple(self.edges[i] for i in old_edge_indexes)
            result = Graph(new_edges, self.num_vertices)
            result._old_edge_indexes = old_edge_indexes
            # no need for old and new vertex_indexes because they remain the
            # same.
        return result

    def get_vertex_fingerprints(self, vertex_strings, edge_strings, num_iter=None):
        """Return an array with fingerprints for each vertex"""
        import hashlib
        def str2array(x):
            """convert a hash string to a numpy array of bytes"""
            if len(x) == 0:
                return np.zeros(0, np.ubyte)
            elif sys.version_info[0] == 2:
                return np.frombuffer(x, np.ubyte)
            else:
                return np.frombuffer(x.encode(), np.ubyte)
        hashrow = lambda x: np.frombuffer(hashlib.sha1(x.data).digest(), np.ubyte)
        # initialization
        result = np.zeros((self.num_vertices, 20), np.ubyte)
        for i in range(self.num_vertices):
            result[i] = hashrow(str2array(vertex_strings[i]))
        for i in range(self.num_edges):
            a, b = self.edges[i]
            tmp = hashrow(str2array(edge_strings[i]))
            result[a] += tmp
            result[b] += tmp
        work = result.copy()
        # iterations
        if num_iter is None:
            num_iter = self.max_distance
        for i in range(num_iter):
            for a, b in self.edges:
                work[a] += result[b]
                work[b] += result[a]
            #for a in xrange(self.num_vertices):
            #    for b in xrange(self.num_vertices):
            #        work[a] += hashrow(result[b]*self.distances[a, b])
            for a in range(self.num_vertices):
                result[a] = hashrow(work[a])
        return result

    def get_halfs(self, vertex1, vertex2):
        """Split the graph in two halfs by cutting the edge: vertex1-vertex2

           If this is not possible (due to loops connecting both ends), a
           GraphError is raised.

           Returns the vertices in both halfs.
        """
        def grow(origin, other):
            frontier = set(self.neighbors[origin])
            frontier.discard(other)
            result = set([origin])
            while len(frontier) > 0:
                pivot = frontier.pop()
                if pivot == other:
                    raise GraphError("The graph can not be separated in two halfs "
                                     "by disconnecting vertex1 and vertex2.")
                pivot_neighbors = set(self.neighbors[pivot])
                pivot_neighbors -= result
                frontier |= pivot_neighbors
                result.add(pivot)
            return result

        vertex1_part = grow(vertex1, vertex2)
        vertex2_part = grow(vertex2, vertex1)
        return vertex1_part, vertex2_part

    def get_part(self, vertex_in, vertices_border):
        """List all vertices that are connected to vertex_in, but are not
           included in or 'behind' vertices_border.
        """
        vertices_new = set(self.neighbors[vertex_in])
        vertices_part = set([vertex_in])

        while len(vertices_new) > 0:
            pivot = vertices_new.pop()
            if pivot in vertices_border:
                continue
            vertices_part.add(pivot)
            pivot_neighbors = set(self.neighbors[pivot])
            pivot_neighbors -= vertices_part
            vertices_new |= pivot_neighbors

        return vertices_part

    def get_halfs_double(self, vertex_a1, vertex_b1, vertex_a2, vertex_b2):
        """Compute the two parts separated by ``(vertex_a1, vertex_b1)`` and ``(vertex_a2, vertex_b2)``

           Raise a GraphError when ``(vertex_a1, vertex_b1)`` and
           ``(vertex_a2, vertex_b2)`` do not separate the graph in two
           disconnected parts. The edges must be neighbors. If not a GraphError
           is raised. The for vertices must not coincide or a GraphError is
           raised.

           Returns the vertices of the two halfs and the four 'hinge' vertices
           in the correct order, i.e. both ``vertex_a1`` and ``vertex_a2`` are
           in the first half and both ``vertex_b1`` and ``vertex_b2`` are in the
           second half.
        """
        if vertex_a1 not in self.neighbors[vertex_b1]:
            raise GraphError("vertex_a1 must be a neighbor of vertex_b1.")
        if vertex_a2 not in self.neighbors[vertex_b2]:
            raise GraphError("vertex_a2 must be a neighbor of vertex_b2.")

        # find vertex_a_part (and possibly switch vertex_a2, vertex_b2)
        vertex_a_new = set(self.neighbors[vertex_a1])
        vertex_a_new.discard(vertex_b1)
        if vertex_a1 == vertex_b2:
            # we now that we have to swap vertex_a2 and vertex_b2. The algo
            # below will fail otherwise in this 'exotic' case.
            vertex_a2, vertex_b2 = vertex_b2, vertex_a2
            #vertex_a_new.discard(vertex_a2) # in case there is overlap
        if vertex_a1 == vertex_a2:
            vertex_a_new.discard(vertex_b2) # in case there is overlap
        vertex_a_part = set([vertex_a1])

        touched = False # True if (the switched) vertex_a2 has been reached.
        while len(vertex_a_new) > 0:
            pivot = vertex_a_new.pop()
            if pivot == vertex_b1:
                raise GraphError("The graph can not be separated in two halfs. "
                                 "vertex_b1 reached by vertex_a1.")
            vertex_a_part.add(pivot)
            # create a new set that we can modify
            pivot_neighbors = set(self.neighbors[pivot])
            pivot_neighbors -= vertex_a_part
            if pivot == vertex_a2 or pivot == vertex_b2:
                if pivot == vertex_b2:
                    if touched:
                        raise GraphError("The graph can not be separated in "
                                         "two halfs. vertex_b2 reached by "
                                         "vertex_a1.")
                    else:
                        # put them in the correct order
                        vertex_a2, vertex_b2 = vertex_b2, vertex_a2
                pivot_neighbors.discard(vertex_b2)
                touched = True
            vertex_a_new |= pivot_neighbors

        if vertex_a2 not in vertex_a_part:
            raise GraphError("The graph can not be separated in two halfs. "
                             "vertex_a1 can not reach vertex_a2 trough "
                             "vertex_a_part")

        # find vertex_b_part: easy, is just the rest ...
        #vertex_b_part = set(xrange(self.num_vertices)) - vertex_a_part

        # ... but we also want that there is a path in vertex_b_part from
        # vertex_b1 to vertex_b2
        if vertex_b1 == vertex_b2:
            closed = True
        else:
            vertex_b_new = set(self.neighbors[vertex_b1])
            vertex_b_new.discard(vertex_a1)
            vertex_b_part = set([vertex_b1])

            closed = False
            while len(vertex_b_new) > 0:
                pivot = vertex_b_new.pop()
                if pivot == vertex_b2:
                    closed = True
                    break
                pivot_neighbors = set(self.neighbors[pivot])
                pivot_neighbors -= vertex_b_part
                vertex_b_new |= pivot_neighbors
                vertex_b_part.add(pivot)

        if not closed:
            raise GraphError("The graph can not be separated in two halfs. "
                             "vertex_b1 can not reach vertex_b2 trough "
                             "vertex_b_part")

        # finaly compute the real vertex_b_part, the former loop might break
        # early for efficiency.
        vertex_b_part = set(range(self.num_vertices)) - vertex_a_part

        # done!
        return vertex_a_part, vertex_b_part, \
               (vertex_a1, vertex_b1, vertex_a2, vertex_b2)

    def full_match(self, other):
        """Find the mapping between vertex indexes in self and other.

           This also works on disconnected graphs. Derived classes should just
           implement get_vertex_string and get_edge_string to make this method
           aware of the different nature of certain vertices. In case molecules,
           this would make the algorithm sensitive to atom numbers etc.
        """
        # we need normalize subgraphs because these graphs are used as patterns.
        graphs0 = [
            self.get_subgraph(group, normalize=True)
            for group in self.independent_vertices
        ]
        graphs1 = [
            other.get_subgraph(group)
            for group in other.independent_vertices
        ]

        if len(graphs0) != len(graphs1):
            return

        matches = []

        for graph0 in graphs0:
            pattern = EqualPattern(graph0)
            found_match = False
            for i, graph1 in enumerate(graphs1):
                local_matches = list(GraphSearch(pattern)(graph1, one_match=True))
                if len(local_matches) == 1:
                    match = local_matches[0]
                    # we need to restore the relation between the normalized
                    # graph0 and its original indexes
                    old_to_new = OneToOne((
                        (j, i) for i, j
                        in enumerate(graph0._old_vertex_indexes)
                    ))
                    matches.append(match * old_to_new)
                    del graphs1[i]
                    found_match = True
                    break
            if not found_match:
                return

        result = OneToOne()
        for match in matches:
            result.add_relations(match.forward.items())
        return result class FileFormatError(Exception):
    """Is raised when unexpected data is encountered while reading a file"""
    pass
import numpy as np
from molmod.periodic import periodic
from molmod.units import unified
from molmod.graphs import CriteriaSet, GraphSearch
from molmod.molecular_graphs import MolecularGraph, BondPattern, \
    BendingAnglePattern, DihedralAnglePattern, OutOfPlanePattern, \
    HasNumNeighbors
from molmod.graphs import Graph
from molmod.io.common import FileFormatError
# -*- coding: utf-8 -*-
# MolMod is a collection of molecular modelling tools for python.
# Copyright (C) 2007 - 2019 Toon Verstraelen <Toon.Verstraelen@UGent.be>, Center
# for Molecular Modeling (CMM), Ghent University, Ghent, Belgium; all rights
# reserved unless otherwise stated.
#
# This file is part of MolMod.
#
# MolMod is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 3
# of the License, or (at your option) any later version.
#
# MolMod is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>
#
# --
"""Tools for reading and writing PSF (Protein Structure File) files

   This format is orignally developed in conjunction with the CHARMM program,
   but is also used by CP2K as a generic format to define the molecular bond
   graph and other topological aspects of a molecular model. This module just
   creates files that can be read with CP2K.

   Due to the lack of public definition of a PSF file format, several variations
   exist. Therefore we do not make any attempt to follow the original format
   strictly, nor one of these variations. It would be more interesting to
   define a new public standard for molecular topologies, which is also suitable
   for non-protein systems.
"""


from __future__ import print_function, division




__all__ = ["PSFFile"]


class PSFFile(object):
    "A very simplistic and limited implementation of the PSF file format"

    def __init__(self, filename=None):
        """
           Argument:
            | ``filename``  --  When not given, an empty data structure is
                                created, otherwise the file is loaded from disk
        """
        if filename is None:
            self.clear()
        else:
            self.read_from_file(filename)

    def clear(self):
        """Clear the contents of the data structure"""
        self.title = None
        self.numbers = np.zeros(0, int)
        self.atom_types = [] # the atom_types in the second column, used to associate ff parameters
        self.charges = [] # ff charges
        self.names = [] # a name that is unique for the molecule composition and connectivity
        self.molecules = np.zeros(0, int) # a counter for each molecule
        self.bonds = np.zeros((0, 2), int)
        self.bends = np.zeros((0, 3), int)
        self.dihedrals = np.zeros((0, 4), int)
        self.impropers = np.zeros((0, 4), int)

        self.name_cache = {}

    def read_from_file(self, filename):
        """Load a PSF file"""
        self.clear()
        with open(filename) as f:
            # A) check the first line
            line = next(f)
            if not line.startswith("PSF"):
                raise FileFormatError("Error while reading: A PSF file must start with a line 'PSF'.")
            # B) read in all the sections, without interpreting them
            current_section = None
            sections = {}
            for line in f:
                line = line.strip()
                if line == "":
                    continue
                elif "!N" in line:
                    words = line.split()
                    current_section = []
                    section_name = words[1][2:]
                    if section_name.endswith(":"):
                        section_name = section_name[:-1]
                    sections[section_name] = current_section
                else:
                    current_section.append(line)
        # C) interpret the supported sections
        # C.1) The title
        self.title = sections['TITLE'][0]
        molecules = []
        numbers = []
        # C.2) The atoms and molecules
        for line in sections['ATOM']:
            words = line.split()
            self.atom_types.append(words[5])
            self.charges.append(float(words[6]))
            self.names.append(words[3])
            molecules.append(int(words[2]))

            atom = periodic[words[4]]