# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import division, unicode_literals

import warnings
import subprocess
import numpy as np
import os.path

from pymatgen.core import Structure, Lattice, PeriodicSite, Molecule
from pymatgen.util.coord import lattice_points_in_supercell
from pymatgen.vis.structure_vtk import EL_COLORS

from monty.json import MSONable
from monty.os.path import which
from operator import itemgetter
from collections import namedtuple
from scipy.spatial import KDTree

try:
    import networkx as nx
    from networkx.readwrite import json_graph
    from networkx.drawing.nx_agraph import write_dot
except ImportError:
    raise ImportError("This module requires the NetworkX "
                      "graph library to be installed.")

import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

__author__ = "Matthew Horton"
__version__ = "0.1"
__maintainer__ = "Matthew Horton"
__email__ = "mkhorton@lbl.gov"
__status__ = "Beta"
__date__ = "August 2017"

ConnectedSite = namedtuple('ConnectedSite', 'periodic_site, jimage, index, weight, dist')

class StructureGraph(MSONable):

    def __init__(self, structure, graph_data):
        """
        If constructing this class manually, use the `with_empty_graph`
        method or `with_local_env_strategy` method (using an algorithm
        provided by the `local_env` module, such as O'Keeffe).

        This class that contains connection information:
        relationships between sites represented by a Graph structure,
        and an associated structure object.

        This class uses the NetworkX package to store and operate
        on the graph itself, but contains a lot of helper methods
        to make associating a graph with a given crystallographic
        structure easier.

        Use cases for this include storing bonding information,
        NMR J-couplings, Heisenberg exchange parameters, etc.

        For periodic graphs, class stores information on the graph
        edges of what lattice image the edge belongs to.

        :param *args: same as in :class: `pymatgen.core.Structure`
        :param graph_data: dict containing graph information in
        dict format, not intended to be constructed manually
        """

        self.structure = structure
        self.graph = nx.readwrite.json_graph.adjacency_graph(graph_data)

        # tidy up edge attr dicts, reading to/from json duplicates
        # information
        for u, v, k, d in self.graph.edges(keys=True, data=True):
            if 'id' in d:
                del d['id']
            if 'key' in d:
                del d['key']

    @classmethod
    def with_empty_graph(cls, structure, name="bonds",
                        edge_weight_name=None,
                         edge_weight_units=None):
        """
        Constructor for StructureGraph, returns a StructureGraph
        object with an empty graph (no edges, only nodes defined
        that correspond to Sites in Structure).

        :param structure (Structure):
        :param name (str): name of graph, e.g. "bonds"
        :param edge_weight_name (str): name of edge weights,
        e.g. "bond_length" or "exchange_constant"
        :param edge_weight_units (str): name of edge weight units
        e.g. "Å" or "eV"
        :return (StructureGraph):
        """

        if edge_weight_name and (edge_weight_units is None):
            raise ValueError("Please specify units associated "
                             "with your edge weights. Can be "
                             "empty string if arbitrary or "
                             "dimensionless.")

        # construct graph with one node per site
        # graph attributes don't change behavior of graph,
        # they're just for book-keeping
        graph = nx.MultiDiGraph(edge_weight_name=edge_weight_name,
                                edge_weight_units=edge_weight_units,
                                name=name)
        graph.add_nodes_from(range(len(structure)))

        graph_data = json_graph.adjacency_data(graph)

        return cls(structure, graph_data=graph_data)

    @staticmethod
    def with_local_env_strategy(structure, strategy):
        """
        Constructor for StructureGraph, using a strategy
        from  :Class: `pymatgen.analysis.local_env`.

        :param structure: Structure object
        :param strategy: an instance of a
         :Class: `pymatgen.analysis.local_env.NearNeighbors`
         object
        :return:
        """

        sg = StructureGraph.with_empty_graph(structure, name="bonds",
                                             edge_weight_name="weight",
                                             edge_weight_units="")

        for n in range(len(structure)):
            neighbors = strategy.get_nn_info(structure, n)
            for neighbor in neighbors:

                # local_env will always try to add two edges
                # for any one bond, one from site u to site v
                # and another form site v to site u: this is
                # harmless, so warn_duplicates=False
                sg.add_edge(from_index=n,
                            from_jimage=(0, 0, 0),
                            to_index=neighbor['site_index'],
                            to_jimage=neighbor['image'],
                            weight=neighbor['weight'],
                            warn_duplicates=False)

        return sg

    @property
    def name(self):
        """
        :return: Name of graph
        """
        return self.graph.graph['name']

    @property
    def edge_weight_name(self):
        """
        :return: Name of the edge weight property of graph
        """
        return self.graph.graph['edge_weight_name']

    @property
    def edge_weight_unit(self):
        """
        :return: Units of the edge weight property of graph
        """
        return self.graph.graph['edge_weight_units']

    def add_edge(self, from_index, to_index,
                 from_jimage=(0, 0, 0), to_jimage=None,
                 weight=None, warn_duplicates=True):
        """
        Add edge to graph.

        Since physically a 'bond' (or other connection
        between sites) doesn't have a direction, from_index,
        from_jimage can be swapped with to_index, to_jimage.

        However, images will always always be shifted so that
        from_index < to_index and from_jimage becomes (0, 0, 0).

        :param from_index: index of site connecting from
        :param to_index: index of site connecting to
        :param from_jimage (tuple of ints): lattice vector of periodic
        image, e.g. (1, 0, 0) for periodic image in +x direction
        :param to_jimage (tuple of ints): lattice vector of image
        :param weight (float): e.g. bond length
        :param warn_duplicates (bool): if True, will warn if
        trying to add duplicate edges (duplicate edges will not
        be added in either case)
        :return:
        """

        # this is not necessary for the class to work, but
        # just makes it neater
        if to_index < from_index:
            to_index, from_index = from_index, to_index
            to_jimage, from_jimage = from_jimage, to_jimage

        # constrain all from_jimages to be (0, 0, 0),
        # initial version of this class worked even if
        # from_jimage != (0, 0, 0), but making this
        # assumption simplifies logic later
        if not np.array_equal(from_jimage, (0, 0, 0)):
            shift = from_jimage
            from_jimage = np.subtract(from_jimage, shift)
            to_jimage = np.subtract(to_jimage, shift)

        # automatic detection of to_jimage if user doesn't specify
        # will try and detect all equivalent images and add multiple
        # edges if appropriate
        if to_jimage is None:
            # assume we want the closest site
            warnings.warn("Please specify to_jimage to be unambiguous, "
                          "trying to automatically detect.")
            dist, to_jimage = self.structure[from_index].distance_and_image(self.structure[to_index])
            if dist == 0:
                # this will happen when from_index == to_index,
                # typically in primitive single-atom lattices
                images = [1, 0, 0], [0, 1, 0], [0, 0, 1]
                dists = []
                for image in images:
                    dists.append(self.structure[from_index].distance_and_image(self.structure[from_index],
                                                                               jimage=image)[0])
                dist = min(dists)
            equiv_sites = self.structure.get_neighbors_in_shell(self.structure[from_index].coords,
                                                                dist,
                                                                dist*0.01,
                                                                include_index=True)
            for site, dist, to_index in equiv_sites:
                to_jimage = np.subtract(site.frac_coords, self.structure[from_index].frac_coords)
                to_jimage = to_jimage.astype(int)
                self.add_edge(from_index=from_index, from_jimage=(0, 0, 0),
                              to_jimage=to_jimage, to_index=to_index)
            return

        from_jimage, to_jimage = tuple(from_jimage), tuple(to_jimage)

        # check we're not trying to add a duplicate edge
        # there should only ever be at most one edge
        # between a given (site, jimage) pair and another
        # (site, jimage) pair
        existing_edge_data = self.graph.get_edge_data(from_index, to_index)
        if existing_edge_data:
            for key, d in existing_edge_data.items():
                if d["to_jimage"] == to_jimage:
                    if warn_duplicates:
                        warnings.warn("Trying to add an edge that already exists from "
                                      "site {} to site {} in {}.".format(from_index,
                                                                         to_index,
                                                                         to_jimage))
                    return

        if weight:
            self.graph.add_edge(from_index, to_index,
                                from_jimage=from_jimage, to_jimage=to_jimage,
                                weight=weight)
        else:
            self.graph.add_edge(from_index, to_index,
                                from_jimage=from_jimage, to_jimage=to_jimage)

    def get_connected_sites(self, n, jimage=(0, 0, 0)):
        """
        Returns a named tuple of neighbors of site n:
        periodic_site, jimage, index, weight.
        Index is the index of the corresponding site
        in the original structure, weight can be
        None if not defined.
        :param n: index of Site in Structure
        :param jimage: lattice vector of site
        :return: list of ConnectedSite tuples,
        sorted by closest first
        """

        connected_sites = set()

        out_edges = [(u, v, d, 'out') for u, v, d in self.graph.out_edges(n, data=True)]
        in_edges = [(u, v, d, 'in') for u, v, d in self.graph.in_edges(n, data=True)]

        for u, v, d, dir in out_edges + in_edges:

            to_jimage = d['to_jimage']

            if dir == 'in':
                u, v = v, u
                to_jimage = np.multiply(-1, to_jimage)

            site_d = self.structure[v].as_dict()
            site_d['abc'] = np.add(site_d['abc'], to_jimage).tolist()
            to_jimage = tuple(map(int, np.add(to_jimage, jimage)))
            periodic_site = PeriodicSite.from_dict(site_d)

            weight = d.get('weight', None)

            # from_site if jimage arg != (0, 0, 0)
            relative_jimage = np.subtract(to_jimage, jimage)
            dist = self.structure[u].distance(self.structure[v], jimage=relative_jimage)

            connected_site = ConnectedSite(periodic_site=periodic_site,
                                           jimage=to_jimage,
                                           index=v,
                                           weight=weight,
                                           dist=dist)

            connected_sites.add(connected_site)

        # return list sorted by closest sites first
        connected_sites = list(connected_sites)
        connected_sites.sort(key=lambda x: x.dist)

        return connected_sites

    def get_coordination_of_site(self, n):
        """
        Returns the number of neighbors of site n.
        In graph terms, simply returns degree
        of node corresponding to site n.
        :param n: index of site
        :return (int):
        """
        return self.graph.degree(n)

    def draw_graph_to_file(self, filename="graph",
                           diff=None,
                           hide_unconnected_nodes=False,
                           hide_image_edges=True,
                           edge_colors=False,
                           node_labels=False,
                           weight_labels=False,
                           image_labels=False,
                           color_scheme="VESTA",
                           keep_dot=False,
                           algo="fdp"):
        """
        Draws graph using GraphViz.

        The networkx graph object itself can also be drawn
        with networkx's in-built graph drawing methods, but
        note that this might give misleading results for
        multigraphs (edges are super-imposed on each other).

        If visualization is difficult to interpret,
        `hide_image_edges` can help, especially in larger
        graphs.

        :param filename: filename to output, will detect filetype
        from extension (any graphviz filetype supported, such as
        pdf or png)
        :param diff (StructureGraph): an additional graph to
        compare with, will color edges red that do not exist in diff
        and edges green that are in diff graph but not in the
        reference graph
        :param hide_unconnected_nodes: if True, hide unconnected
        nodes
        :param hide_image_edges: if True, do not draw edges that
        go through periodic boundaries
        :param edge_colors (bool): if True, use node colors to
        color edges
        :param node_labels (bool): if True, label nodes with
        species and site index
        :param weight_labels (bool): if True, label edges with
        weights
        :param image_labels (bool): if True, label edges with
        their periodic images (usually only used for debugging,
        edges to periodic images always appear as dashed lines)
        :param color_scheme (str): "VESTA" or "JMOL"
        :param keep_dot (bool): keep GraphViz .dot file for later
        visualization
        :param algo: any graphviz algo, "neato" (for simple graphs)
        or "fdp" (for more crowded graphs) usually give good outputs
        :return:
        """

        if not which(algo):
            raise RuntimeError("StructureGraph graph drawing requires "
                               "GraphViz binaries to be in the path.")

        # Developer note: NetworkX also has methods for drawing
        # graphs using matplotlib, these also work here. However,
        # a dedicated tool like GraphViz allows for much easier
        # control over graph appearance and also correctly displays
        # mutli-graphs (matplotlib can superimpose multiple edges).

        g = self.graph.copy()

        g.graph = {'nodesep': 10.0, 'dpi': 300, 'overlap': "false"}

        # add display options for nodes
        for n in g.nodes():

            # get label by species name
            label = "{}({})".format(str(self.structure[n].specie), n) if node_labels else ""

            # use standard color scheme for nodes
            c = EL_COLORS[color_scheme].get(str(self.structure[n].specie.symbol), [0, 0, 0])

            # get contrasting font color
            # magic numbers account for perceived luminescence
            # https://stackoverflow.com/questions/1855884/determine-font-color-based-on-background-color
            fontcolor = '#000000' if 1 - (c[0] * 0.299 + c[1] * 0.587
                                          + c[2] * 0.114) / 255 < 0.5 else '#ffffff'

            # convert color to hex string
            color = "#{:02x}{:02x}{:02x}".format(c[0], c[1], c[2])

            g.add_node(n, fillcolor=color, fontcolor=fontcolor, label=label,
                       fontname="Helvetica-bold", style="filled", shape="circle")

        edges_to_delete = []

        # add display options for edges
        for u, v, k, d in g.edges(keys=True, data=True):

            # retrieve from/to images, set as origin if not defined
            to_image = d['to_jimage']

            # set edge style
            d['style'] = "solid"
            if to_image != (0, 0, 0):
                d['style'] = "dashed"
                if hide_image_edges:
                    edges_to_delete.append((u, v, k))

            # don't show edge directions
            d['arrowhead'] = "none"

            # only add labels for images that are not the origin
            if image_labels:
                d['headlabel'] = "" if to_image == (0, 0, 0) else "to {}".format((to_image))
                d['arrowhead'] = "normal" if d['headlabel'] else "none"

            # optionally color edges using node colors
            color_u = g.node[u]['fillcolor']
            color_v = g.node[v]['fillcolor']
            d['color_uv'] = "{};0.5:{};0.5".format(color_u, color_v) if edge_colors else "#000000"

            # optionally add weights to graph
            if weight_labels:
                units = g.graph.get('edge_weight_units', "")
                if d.get('weight'):
                    d['label'] = "{:.2f} {}".format(d['weight'], units)

            # update edge with our new style attributes
            g.edges[u, v, k].update(d)

        # optionally remove periodic image edges,
        # these can be confusing due to periodic boundaries
        if hide_image_edges:
            for edge_to_delete in edges_to_delete:
                g.remove_edge(*edge_to_delete)

        # optionally hide unconnected nodes,
        # these can appear when removing periodic edges
        if hide_unconnected_nodes:
            g = g.subgraph([n for n in g.degree() if g.degree()[n] != 0])

        # optionally highlight differences with another graph
        if diff:
            diff = self.diff(diff, strict=True)
            green_edges = []
            red_edges = []
            for u, v, k, d in g.edges(keys=True, data=True):
                if (u, v, d['to_jimage']) in diff['self']:
                    # edge has been deleted
                    red_edges.append((u, v, k))
                elif (u, v, d['to_jimage']) in diff['other']:
                    # edge has been added
                    green_edges.append((u, v, k))
            for u, v, k in green_edges:
                g.edges[u, v, k].update({'color_uv': '#00ff00'})
            for u, v, k in red_edges:
                g.edges[u, v, k].update({'color_uv': '#ff0000'})

        basename, extension = os.path.splitext(filename)
        extension = extension[1:]

        write_dot(g, basename+".dot")

        with open(filename, "w") as f:

            args = [algo, "-T", extension, basename+".dot"]
            rs = subprocess.Popen(args,
                                  stdout=f,
                                  stdin=subprocess.PIPE, close_fds=True)
            rs.communicate()
            if rs.returncode != 0:
                raise RuntimeError("{} exited with return code {}.".format(algo, rs.returncode))

        if not keep_dot:
            os.remove(basename+".dot")

    def as_dict(self):
        """
        As in :Class: `pymatgen.core.Structure` except
        with using `to_dict_of_dicts` from NetworkX
        to store graph information.
        """

        d = {"@module": self.__class__.__module__,
             "@class": self.__class__.__name__,
             "structure": self.structure.as_dict(),
             "graphs": json_graph.adjacency_data(self.graph)}

        return d

    @classmethod
    def from_dict(cls, d):
        """
        As in :Class: `pymatgen.core.Structure` except
        restoring graphs using `from_dict_of_dicts`
        from NetworkX to restore graph information.
        """
        s = Structure.from_dict(d['structure'])
        return cls(s, d['graphs'])

    def __mul__(self, scaling_matrix):
        """
        Replicates the graph, creating a supercell,
        intelligently joining together
        edges that lie on periodic boundaries.
        In principle, any operations on the expanded
        graph could also be done on the original
        graph, but a larger graph can be easier to
        visualize and reason about.
        :param scaling_matrix: same as Structure.__mul__
        :return:
        """

        # Developer note: a different approach was also trialed, using
        # a simple Graph (instead of MultiDiGraph), with node indices
        # representing both site index and periodic image. Here, the
        # number of nodes != number of sites in the Structure. This
        # approach has many benefits, but made it more difficult to
        # keep the graph in sync with its corresponding Structure.

        # Broadly, it would be easier to multiply the Structure
        # *before* generating the StructureGraph, but this isn't
        # possible when generating the graph using critic2 from
        # charge density.

        # Multiplication works by looking for the expected position
        # of an image node, and seeing if that node exists in the
        # supercell. If it does, the edge is updated. This is more
        # computationally expensive than just keeping track of the
        # which new lattice images present, but should hopefully be
        # easier to extend to a general 3x3 scaling matrix.

        # code adapted from Structure.__mul__
        scale_matrix = np.array(scaling_matrix, np.int16)
        if scale_matrix.shape != (3, 3):
            scale_matrix = np.array(scale_matrix * np.eye(3), np.int16)
        else:
            # TODO: test __mul__ with full 3x3 scaling matrices
            raise NotImplementedError('Not tested with 3x3 scaling matrices yet.')
        new_lattice = Lattice(np.dot(scale_matrix, self.structure.lattice.matrix))

        f_lat = lattice_points_in_supercell(scale_matrix)
        c_lat = new_lattice.get_cartesian_coords(f_lat)

        new_sites = []
        new_graphs = []

        for v in c_lat:

            # create a map of nodes from original graph to its image
            mapping = {n: n + len(new_sites) for n in range(len(self.structure))}

            for idx, site in enumerate(self.structure):

                s = PeriodicSite(site.species_and_occu, site.coords + v,
                                 new_lattice, properties=site.properties,
                                 coords_are_cartesian=True, to_unit_cell=False)

                new_sites.append(s)

            new_graphs.append(nx.relabel_nodes(self.graph, mapping, copy=True))

        new_structure = Structure.from_sites(new_sites)

        # merge all graphs into one big graph
        new_g = nx.MultiDiGraph()
        for new_graph in new_graphs:
            new_g = nx.union(new_g, new_graph)

        edges_to_remove = []  # tuple of (u, v, k)
        edges_to_add = []  # tuple of (u, v, attr_dict)

        # list of new edges inside supercell
        # for duplicate checking
        edges_inside_supercell = [{u, v} for u, v, d in new_g.edges(data=True)
                                  if d['to_jimage'] == (0, 0, 0)]
        new_periodic_images = []

        orig_lattice = self.structure.lattice

        # use k-d tree to match given position to an
        # existing Site in Structure
        kd_tree = KDTree(new_structure.cart_coords)

        # tolerance in Å for sites to be considered equal
        # this could probably be a lot smaller
        tol = 0.05

        for u, v, k, d in new_g.edges(keys=True, data=True):

            to_jimage = d['to_jimage']  # for node v

            # reduce unnecessary checking
            if to_jimage != (0, 0, 0):

                # get index in original site
                n_u = u % len(self.structure)
                n_v = v % len(self.structure)

                # get fractional co-ordinates of where atoms defined
                # by edge are expected to be, relative to original
                # lattice (keeping original lattice has
                # significant benefits)
                v_image_frac = np.add(self.structure[n_v].frac_coords, to_jimage)
                u_frac = self.structure[n_u].frac_coords

                # using the position of node u as a reference,
                # get relative Cartesian co-ordinates of where
                # atoms defined by edge are expected to be
                v_image_cart = orig_lattice.get_cartesian_coords(v_image_frac)
                u_cart = orig_lattice.get_cartesian_coords(u_frac)
                v_rel = np.subtract(v_image_cart, u_cart)

                # now retrieve position of node v in
                # new supercell, and get absolute Cartesian
                # co-ordinates of where atoms defined by edge
                # are expected to be
                v_expec = new_structure[u].coords + v_rel

                # now search in new structure for these atoms
                # query returns (distance, index)
                v_present = kd_tree.query(v_expec)
                v_present = v_present[1] if v_present[0] <= tol else None

                # check if image sites now present in supercell
                # and if so, delete old edge that went through
                # periodic boundary
                if v_present is not None:

                    new_u = u
                    new_v = v_present
                    new_d = d.copy()

                    # node now inside supercell
                    new_d['to_jimage'] = (0, 0, 0)

                    edges_to_remove.append((u, v, k))

                    # make sure we don't try to add duplicate edges
                    # will remove two edges for everyone one we add
                    if {new_u, new_v} not in edges_inside_supercell:

                        # normalize direction
                        if new_v < new_u:
                            new_u, new_v = new_v, new_u

                        edges_inside_supercell.append({new_u, new_v})
                        edges_to_add.append((new_u, new_v, new_d))

                else:

                    # want to find new_v such that we have
                    # full periodic boundary conditions
                    # so that nodes on one side of supercell
                    # are connected to nodes on opposite side

                    v_expec_frac = new_structure.lattice.get_fractional_coords(v_expec)

                    # find new to_jimage
                    # use np.around to fix issues with finite precision leading to incorrect image
                    v_expec_image = np.around(v_expec_frac, decimals=3)
                    v_expec_image = v_expec_image - v_expec_image%1

                    v_expec_frac = np.subtract(v_expec_frac, v_expec_image)
                    v_expec = new_structure.lattice.get_cartesian_coords(v_expec_frac)
                    v_present = kd_tree.query(v_expec)
                    v_present = v_present[1] if v_present[0] <= tol else None

                    if v_present is not None:

                        new_u = u
                        new_v = v_present
                        new_d = d.copy()
                        new_to_jimage = tuple(map(int, v_expec_image))

                        # normalize direction
                        if new_v < new_u:
                            new_u, new_v = new_v, new_u
                            new_to_jimage = tuple(np.multiply(-1, d['to_jimage']).astype(int))

                        new_d['to_jimage'] = new_to_jimage

                        edges_to_remove.append((u, v, k))

                        if (new_u, new_v, new_to_jimage) not in new_periodic_images:
                            edges_to_add.append((new_u, new_v, new_d))
                            new_periodic_images.append((new_u, new_v, new_to_jimage))

        logger.debug("Removing {} edges, adding {} new edges.".format(len(edges_to_remove),
                                                                      len(edges_to_add)))

        # add/delete marked edges
        for edges_to_remove in edges_to_remove:
            new_g.remove_edge(*edges_to_remove)
        for (u, v, d) in edges_to_add:
            new_g.add_edge(u, v, **d)

        # return new instance of StructureGraph with supercell
        d = {"@module": self.__class__.__module__,
             "@class": self.__class__.__name__,
             "structure": new_structure.as_dict(),
             "graphs": json_graph.adjacency_data(new_g)}

        sg = StructureGraph.from_dict(d)

        return sg

    def __rmul__(self, other):
        return self.__mul__(other)

    def _edges_to_string(self, g):

        header = "from    to  to_image    "
        header_line = "----  ----  ------------"
        edge_weight_name = g.graph["edge_weight_name"]
        if edge_weight_name:
            print_weights = ["weight"]
            edge_label = g.graph["edge_weight_name"]
            edge_weight_units = g.graph["edge_weight_units"]
            if edge_weight_units:
                edge_label += " ({})".format(edge_weight_units)
            header += "  {}".format(edge_label)
            header_line += "  {}".format("-"*max([18, len(edge_label)]))
        else:
            print_weights = False

        s = header + "\n" + header_line + "\n"

        edges = list(g.edges(data=True))

        # sort edges for consistent ordering
        edges.sort(key=itemgetter(0,1))

        if print_weights:
            for u, v, data in edges:
                s += "{:4}  {:4}  {:12}  {:.12E}\n".format(u, v, str(data.get("to_jimage", (0, 0, 0))),
                                                           data.get("weight", 0))
        else:
            for u, v, data in edges:
                s += "{:4}  {:4}  {:12}\n".format(u, v,
                                                  str(data.get("to_jimage", (0, 0, 0))))

        return s

    def __str__(self):
        s = "Structure Graph"
        s += "\nStructure: \n{}".format(self.structure.__str__())
        s += "\nGraph: {}\n".format(self.name)
        s += self._edges_to_string(self.graph)
        return s

    def __repr__(self):
        s = "Structure Graph"
        s += "\nStructure: \n{}".format(self.structure.__repr__())
        s += "\nGraph: {}\n".format(self.name)
        s += self._edges_to_string(self.graph)
        return s

    def __len__(self):
        """
        :return: length of Structure / number of nodes in graph
        """
        return len(self.structure)

    def sort(self, key=None, reverse=False):
        """
        Same as Structure.sort(), also remaps nodes in graph.
        :param key:
        :param reverse:
        :return:
        """

        old_structure = self.structure.copy()

        # sort Structure
        self.structure._sites = sorted(self.structure._sites, key=key, reverse=reverse)

        # apply Structure ordering to graph
        mapping = {idx:self.structure.index(site) for idx, site in enumerate(old_structure)}
        self.graph = nx.relabel_nodes(self.graph, mapping, copy=True)

        # normalize directions of edges
        edges_to_remove = []
        edges_to_add = []
        for u, v, k, d in self.graph.edges(keys=True, data=True):
            if v < u:
                new_v, new_u, new_d = u, v, d.copy()
                new_d['to_jimage'] = tuple(np.multiply(-1, d['to_jimage']).astype(int))
                edges_to_remove.append((u,v,k))
                edges_to_add.append((new_u, new_v, new_d))

        # add/delete marked edges
        for edges_to_remove in edges_to_remove:
            self.graph.remove_edge(*edges_to_remove)
        for (u, v, d) in edges_to_add:
            self.graph.add_edge(u, v, **d)

    def __copy__(self):
        return StructureGraph.from_dict(self.as_dict())

    def __eq__(self, other):
        """
        Two StructureGraphs are equal if they have equal Structures,
        and have the same edges between Sites. Edge weights can be
        different and StructureGraphs can still be considered equal.

        :param other: StructureGraph
        :return (bool):
        """

        # sort for consistent node indices
        # PeriodicSite should have a proper __hash__() value,
        # using its frac_coords as a convenient key
        mapping = {tuple(site.frac_coords):self.structure.index(site) for site in other.structure}
        other_sorted = other.__copy__()
        other_sorted.sort(key=lambda site: mapping[tuple(site.frac_coords)])

        edges = {(u, v, d['to_jimage'])
                 for u, v, d in self.graph.edges(keys=False, data=True)}

        edges_other = {(u, v, d['to_jimage'])
                       for u, v, d in other_sorted.graph.edges(keys=False, data=True)}

        return (edges == edges_other) and \
               (self.structure == other_sorted.structure)

    def diff(self, other, strict=True):
        """
        Compares two StructureGraphs. Returns dict with
        keys 'self', 'other', 'both' with edges that are
        present in only one StructureGraph ('self' and
        'other'), and edges that are present in both.

        The Jaccard distance is a simple measure of the
        dissimilarity between two StructureGraphs (ignoring
        edge weights), and is defined by 1 - (size of the
        intersection / size of the union) of the sets of
        edges. This is returned with key 'dist'.

        Important note: all node indices are in terms
        of the StructureGraph this method is called
        from, not the 'other' StructureGraph: there
        is no guarantee the node indices will be the
        same if the underlying Structures are ordered
        differently.

        :param other: StructureGraph
        :param strict: if False, will compare bonds
        from different Structures, with node indices
        replaced by Specie strings, will not count
        number of occurrences of bonds
        :return:
        """

        if self.structure != other.structure and strict:
            return ValueError("Meaningless to compare StructureGraphs if "
                              "corresponding Structures are different.")

        if strict:

            # sort for consistent node indices
            # PeriodicSite should have a proper __hash__() value,
            # using its frac_coords as a convenient key
            mapping = {tuple(site.frac_coords):self.structure.index(site) for site in other.structure}
            other_sorted = other.__copy__()
            other_sorted.sort(key=lambda site: mapping[tuple(site.frac_coords)])

            edges = {(u, v, d['to_jimage'])
                     for u, v, d in self.graph.edges(keys=False, data=True)}

            edges_other = {(u, v, d['to_jimage'])
                           for u, v, d in other_sorted.graph.edges(keys=False, data=True)}

        else:

            edges = {(str(self.structure[u].specie),
                      str(self.structure[v].specie))
                     for u, v, d in self.graph.edges(keys=False, data=True)}

            edges_other = {(str(other.structure[u].specie),
                            str(other.structure[v].specie))
                           for u, v, d in other.graph.edges(keys=False, data=True)}

        if len(edges) == 0 and len(edges_other) == 0:
            jaccard_dist = 0  # by definition
        else:
            jaccard_dist = 1 - len(edges.intersection(edges_other)) / len(edges.union(edges_other))

        return {
            'self': edges - edges_other,
            'other': edges_other - edges,
            'both': edges.intersection(edges_other),
            'dist': jaccard_dist
        }

    def get_subgraphs_as_molecules(self, use_weights=False):
        """
        Retrieve subgraphs as molecules, useful for extracting
        molecules from periodic crystals.

        Will only return unique molecules, not any duplicates
        present in the crystal (a duplicate defined as an
        isomorphic subgraph).

        :param use_weights (bool): If True, only treat subgraphs
        as isomorphic if edges have the same weights. Typically,
        this means molecules will need to have the same bond
        lengths to be defined as duplicates, otherwise bond
        lengths can differ. This is a fairly robust approach,
        but will treat e.g. enantiomers as being duplicates.

        :return: list of unique Molecules in Structure
        """

        # creating a supercell is an easy way to extract
        # molecules (and not, e.g., layers of a 2D crystal)
        # without adding extra logic
        if getattr(self, '_supercell_sg', None) is None:
            self._supercell_sg = supercell_sg = self*(3,3,3)

        # make undirected to find connected subgraphs
        supercell_sg.graph = nx.Graph(supercell_sg.graph)

        # find subgraphs
        all_subgraphs = list(nx.connected_component_subgraphs(supercell_sg.graph))

        # discount subgraphs that lie across *supercell* boundaries
        # these will subgraphs representing crystals
        molecule_subgraphs = []
        for subgraph in all_subgraphs:
            intersects_boundary = any([d['to_jimage'] != (0, 0, 0)
                                      for u, v, d in subgraph.edges(data=True)])
            if not intersects_boundary:
                molecule_subgraphs.append(subgraph)

        # add specie names to graph to be able to test for isomorphism
        for subgraph in molecule_subgraphs:
            for n in subgraph:
                subgraph.add_node(n, specie=str(supercell_sg.structure[n].specie))

        # now define how we test for isomorphism
        def node_match(n1, n2):
            return n1['specie'] == n2['specie']
        def edge_match(e1, e2):
            if use_weights:
                return e1['weight'] == e2['weight']
            else:
                return True

        # prune duplicate subgraphs
        unique_subgraphs = []
        for subgraph in molecule_subgraphs:

            already_present = [nx.is_isomorphic(subgraph, g,
                                                node_match=node_match,
                                                edge_match=edge_match)
                               for g in unique_subgraphs]

            if not any(already_present):
                unique_subgraphs.append(subgraph)

        # get Molecule objects for each subgraph
        molecules = []
        for subgraph in unique_subgraphs:

            coords = [supercell_sg.structure[n].coords for n
                      in subgraph.nodes()]
            species = [supercell_sg.structure[n].specie for n
                      in subgraph.nodes()]

            molecule = Molecule(species, coords)

            # shift so origin is at center of mass
            molecule = molecule.get_centered_molecule()

            molecules.append(molecule)

        return molecules

import numpy,from scipy,import networkx,from networkx


# Copyright 2018 The Cirq Developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, cast, Dict, NamedTuple, Optional, Sequence, Tuple, TYPE_CHECKING, Union

import numpy as np

from cirq import protocols, value, linalg
from cirq._doc import document
from cirq.ops import common_gates, gate_features, named_qubit, pauli_gates
from cirq.ops.pauli_gates import Pauli
from cirq.type_workarounds import NotImplementedType

if TYPE_CHECKING:
    import cirq

PauliTransform = NamedTuple('PauliTransform', [('to', Pauli), ('flip', bool)])
document(PauliTransform, """+X, -X, +Y, -Y, +Z, or -Z.""")

def _to_pauli_transform(matrix: np.ndarray) -> Optional[PauliTransform]:
    """Converts matrix to PauliTransform.

    If matrix is not ±Pauli matrix, returns None.
    """
    for pauli in Pauli._XYZ:
        p = protocols.unitary(pauli)
        if np.allclose(matrix, p):
            return PauliTransform(pauli, False)
        if np.allclose(matrix, -p):
            return PauliTransform(pauli, True)
    return None

def _pretend_initialized() -> 'SingleQubitCliffordGate':
    # HACK: This is a workaround to fool mypy and pylint into correctly handling
    # class fields that can't be initialized until after the class is defined.
    pass

@value.value_equality
class SingleQubitCliffordGate(gate_features.SingleQubitGate):
    """Any single qubit Clifford rotation."""

    I = _pretend_initialized()
    H = _pretend_initialized()
    X = _pretend_initialized()
    Y = _pretend_initialized()
    Z = _pretend_initialized()
    X_sqrt = _pretend_initialized()
    Y_sqrt = _pretend_initialized()
    Z_sqrt = _pretend_initialized()
    X_nsqrt = _pretend_initialized()
    Y_nsqrt = _pretend_initialized()
    Z_nsqrt = _pretend_initialized()

    def __init__(
        self,
        *,
        _rotation_map: Dict[Pauli, PauliTransform],
        _inverse_map: Dict[Pauli, PauliTransform],
    ) -> None:
        self._rotation_map = _rotation_map
        self._inverse_map = _inverse_map

    @staticmethod
    def from_xz_map(
        x_to: Tuple[Pauli, bool], z_to: Tuple[Pauli, bool]
    ) -> 'SingleQubitCliffordGate':
        """Returns a SingleQubitCliffordGate for the specified transforms.
        The Y transform is derived from the X and Z.

        Args:
            x_to: Which Pauli to transform X to and if it should negate.
            z_to: Which Pauli to transform Z to and if it should negate.
        """
        return SingleQubitCliffordGate.from_double_map(x_to=x_to, z_to=z_to)

    @staticmethod
    def from_single_map(
        pauli_map_to: Optional[Dict[Pauli, Tuple[Pauli, bool]]] = None,
        *,
        x_to: Optional[Tuple[Pauli, bool]] = None,
        y_to: Optional[Tuple[Pauli, bool]] = None,
        z_to: Optional[Tuple[Pauli, bool]] = None,
    ) -> 'SingleQubitCliffordGate':
        """Returns a SingleQubitCliffordGate for the
        specified transform with a 90 or 180 degree rotation.

        The arguments are exclusive, only one may be specified.

        Args:
            pauli_map_to: A dictionary with a single key value pair describing
                the transform.
            x_to: The transform from cirq.X
            y_to: The transform from cirq.Y
            z_to: The transform from cirq.Z
        """
        rotation_map = SingleQubitCliffordGate._validate_map_input(
            1, pauli_map_to, x_to=x_to, y_to=y_to, z_to=z_to
        )
        ((trans_from, (trans_to, flip)),) = tuple(rotation_map.items())
        if trans_from == trans_to:
            trans_from2 = Pauli.by_relative_index(trans_to, 1)  # 1 or 2 work
            trans_to2 = Pauli.by_relative_index(trans_from, 1)
            flip2 = False
        else:
            trans_from2 = trans_to
            trans_to2 = trans_from
            flip2 = not flip
        rotation_map[trans_from2] = PauliTransform(trans_to2, flip2)
        return SingleQubitCliffordGate.from_double_map(
            cast(Dict[Pauli, Tuple[Pauli, bool]], rotation_map)
        )

    @staticmethod
    def from_double_map(
        pauli_map_to: Optional[Dict[Pauli, Tuple[Pauli, bool]]] = None,
        *,
        x_to: Optional[Tuple[Pauli, bool]] = None,
        y_to: Optional[Tuple[Pauli, bool]] = None,
        z_to: Optional[Tuple[Pauli, bool]] = None,
    ) -> 'SingleQubitCliffordGate':
        """Returns a SingleQubitCliffordGate for the
        specified transform with a 90 or 180 degree rotation.

        Either pauli_map_to or two of (x_to, y_to, z_to) may be specified.

        Args:
            pauli_map_to: A dictionary with two key value pairs describing
                two transforms.
            x_to: The transform from cirq.X
            y_to: The transform from cirq.Y
            z_to: The transform from cirq.Z
        """
        rotation_map = SingleQubitCliffordGate._validate_map_input(
            2, pauli_map_to, x_to=x_to, y_to=y_to, z_to=z_to
        )
        (from1, trans1), (from2, trans2) = tuple(rotation_map.items())
        from3 = from1.third(from2)
        to3 = trans1.to.third(trans2.to)
        flip3 = trans1.flip ^ trans2.flip ^ ((from1 < from2) != (trans1.to < trans2.to))
        rotation_map[from3] = PauliTransform(to3, flip3)
        inverse_map = {to: PauliTransform(frm, flip) for frm, (to, flip) in rotation_map.items()}
        return SingleQubitCliffordGate(_rotation_map=rotation_map, _inverse_map=inverse_map)

    @staticmethod
    def from_pauli(pauli: Pauli, sqrt: bool = False) -> 'SingleQubitCliffordGate':
        prev_pauli = Pauli.by_relative_index(pauli, -1)
        next_pauli = Pauli.by_relative_index(pauli, 1)
        if sqrt:
            rotation_map = {
                prev_pauli: PauliTransform(next_pauli, True),
                pauli: PauliTransform(pauli, False),
                next_pauli: PauliTransform(prev_pauli, False),
            }
        else:
            rotation_map = {
                prev_pauli: PauliTransform(prev_pauli, True),
                pauli: PauliTransform(pauli, False),
                next_pauli: PauliTransform(next_pauli, True),
            }
        inverse_map = {to: PauliTransform(frm, flip) for frm, (to, flip) in rotation_map.items()}
        return SingleQubitCliffordGate(_rotation_map=rotation_map, _inverse_map=inverse_map)

    @staticmethod
    def from_quarter_turns(pauli: Pauli, quarter_turns: int) -> 'SingleQubitCliffordGate':
        quarter_turns = quarter_turns % 4
        if quarter_turns == 0:
            return SingleQubitCliffordGate.I
        if quarter_turns == 1:
            return SingleQubitCliffordGate.from_pauli(pauli, True)
        if quarter_turns == 2:
            return SingleQubitCliffordGate.from_pauli(pauli)

        return SingleQubitCliffordGate.from_pauli(pauli, True) ** -1

    @staticmethod
    def _validate_map_input(
        required_transform_count: int,
        pauli_map_to: Optional[Dict[Pauli, Tuple[Pauli, bool]]],
        x_to: Optional[Tuple[Pauli, bool]],
        y_to: Optional[Tuple[Pauli, bool]],
        z_to: Optional[Tuple[Pauli, bool]],
    ) -> Dict[Pauli, PauliTransform]:
        if pauli_map_to is None:
            xyz_to = {pauli_gates.X: x_to, pauli_gates.Y: y_to, pauli_gates.Z: z_to}
            pauli_map_to = {
                cast(Pauli, p): trans for p, trans in xyz_to.items() if trans is not None
            }
        elif x_to is not None or y_to is not None or z_to is not None:
            raise ValueError(
                '{} can take either pauli_map_to or a combination'
                ' of x_to, y_to, and z_to but both were given'
            )
        if len(pauli_map_to) != required_transform_count:
            raise ValueError(
                'Method takes {} transform{} but {} {} given'.format(
                    required_transform_count,
                    '' if required_transform_count == 1 else 's',
                    len(pauli_map_to),
                    'was' if len(pauli_map_to) == 1 else 'were',
                )
            )
        if len(set((to for to, _ in pauli_map_to.values()))) != len(pauli_map_to):
            raise ValueError('A rotation cannot map two Paulis to the same')
        return {frm: PauliTransform(to, flip) for frm, (to, flip) in pauli_map_to.items()}

    @staticmethod
    def from_unitary(u: np.ndarray) -> Optional['SingleQubitCliffordGate']:
        """Creates Clifford gate with given unitary (up to global phase).

        Args:
            u: 2x2 unitary matrix of a Clifford gate.

        Returns:
            SingleQubitCliffordGate, whose matrix is equal to given matrix (up
            to global phase), or `None` if `u` is not a matrix of a single-qubit
            Clifford gate.
        """
        if u.shape != (2, 2) or not linalg.is_unitary(u):
            return None
        x = protocols.unitary(pauli_gates.X)
        z = protocols.unitary(pauli_gates.Z)
        x_to = _to_pauli_transform(u @ x @ u.conj().T)
        z_to = _to_pauli_transform(u @ z @ u.conj().T)
        if x_to is None or z_to is None:
            return None
        return SingleQubitCliffordGate.from_double_map({pauli_gates.X: x_to, pauli_gates.Z: z_to})

    def transform(self, pauli: Pauli) -> PauliTransform:
        return self._rotation_map[pauli]

    def _value_equality_values_(self):
        return (
            self.transform(pauli_gates.X),
            self.transform(pauli_gates.Y),
            self.transform(pauli_gates.Z),
        )

    def __pow__(self, exponent) -> 'SingleQubitCliffordGate':
        if exponent == 0.5 or exponent == -0.5:
            return SQRT_EXP_MAP[exponent][self]
        if exponent != -1:
            return NotImplemented

        return SingleQubitCliffordGate(
            _rotation_map=self._inverse_map, _inverse_map=self._rotation_map
        )

    def _commutes_(self, other: Any, atol: float) -> Union[bool, NotImplementedType]:
        if isinstance(other, SingleQubitCliffordGate):
            return self.commutes_with_single_qubit_gate(other)
        if isinstance(other, Pauli):
            return self.commutes_with_pauli(other)
        return NotImplemented

    def commutes_with_single_qubit_gate(self, gate: 'SingleQubitCliffordGate') -> bool:
        """Tests if the two circuits would be equivalent up to global phase:
        --self--gate-- and --gate--self--"""
        for pauli0 in (pauli_gates.X, pauli_gates.Z):
            pauli1, flip1 = self.transform(cast(Pauli, pauli0))
            pauli2, flip2 = gate.transform(cast(Pauli, pauli1))
            pauli3, flip3 = self._inverse_map[pauli2]
            pauli4, flip4 = gate._inverse_map[pauli3]
            if pauli4 != pauli0 or (flip1 ^ flip2 ^ flip3 ^ flip4):
                return False
        return True

    def commutes_with_pauli(self, pauli: Pauli) -> bool:
        to, flip = self.transform(pauli)
        return to == pauli and not flip

    def merged_with(self, second: 'SingleQubitCliffordGate') -> 'SingleQubitCliffordGate':
        """Returns a SingleQubitCliffordGate such that the circuits
            --output-- and --self--second--
        are equivalent up to global phase."""
        x_intermediate_pauli, x_flip1 = self.transform(pauli_gates.X)
        x_final_pauli, x_flip2 = second.transform(x_intermediate_pauli)
        z_intermediate_pauli, z_flip1 = self.transform(pauli_gates.Z)
        z_final_pauli, z_flip2 = second.transform(z_intermediate_pauli)
        return SingleQubitCliffordGate.from_xz_map(
            (x_final_pauli, x_flip1 ^ x_flip2), (z_final_pauli, z_flip1 ^ z_flip2)
        )

    def _has_unitary_(self) -> bool:
        return True

    def _unitary_(self) -> np.ndarray:
        mat = np.eye(2)
        qubit = named_qubit.NamedQubit('arbitrary')
        for op in protocols.decompose_once_with_qubits(self, (qubit,)):
            mat = protocols.unitary(op).dot(mat)
        return mat

    def _decompose_(self, qubits: Sequence['cirq.Qid']) -> 'cirq.OP_TREE':
        (qubit,) = qubits
        if self == SingleQubitCliffordGate.H:
            return (common_gates.H(qubit),)
        rotations = self.decompose_rotation()
        return tuple(r.on(qubit) ** (qt / 2) for r, qt in rotations)

    def decompose_rotation(self) -> Sequence[Tuple[Pauli, int]]:
        """Returns ((first_rotation_axis, first_rotation_quarter_turns), ...)

        This is a sequence of zero, one, or two rotations."""
        x_rot = self.transform(pauli_gates.X)
        y_rot = self.transform(pauli_gates.Y)
        z_rot = self.transform(pauli_gates.Z)
        whole_arr = (
            x_rot.to == pauli_gates.X,
            y_rot.to == pauli_gates.Y,
            z_rot.to == pauli_gates.Z,
        )
        num_whole = sum(whole_arr)
        flip_arr = (x_rot.flip, y_rot.flip, z_rot.flip)
        num_flip = sum(flip_arr)
        if num_whole == 3:
            if num_flip == 0:
                # Gate is identity
                return []

            # 180 rotation about some axis
            pauli = Pauli.by_index(flip_arr.index(False))
            return [(pauli, 2)]
        if num_whole == 1:
            index = whole_arr.index(True)
            pauli = Pauli.by_index(index)
            next_pauli = Pauli.by_index(index + 1)
            flip = flip_arr[index]
            output = []
            if flip:
                # 180 degree rotation
                output.append((next_pauli, 2))
            # 90 degree rotation about some axis
            if self.transform(next_pauli).flip:
                # Negative 90 degree rotation
                output.append((pauli, -1))
            else:
                # Positive 90 degree rotation
                output.append((pauli, 1))
            return output
        elif num_whole == 0:
            # Gate is a 120 degree rotation
            if x_rot.to == pauli_gates.Y:
                return [
                    (pauli_gates.X, -1 if y_rot.flip else 1),
                    (pauli_gates.Z, -1 if x_rot.flip else 1),
                ]

            return [
                (pauli_gates.Z, 1 if y_rot.flip else -1),
                (pauli_gates.X, 1 if z_rot.flip else -1),
            ]
        # coverage: ignore
        assert (
            False
        ), 'Impossible condition where this gate only rotates one Pauli to a different Pauli.'

    def equivalent_gate_before(self, after: 'SingleQubitCliffordGate') -> 'SingleQubitCliffordGate':
        """Returns a SingleQubitCliffordGate such that the circuits
            --output--self-- and --self--gate--
        are equivalent up to global phase."""
        return self.merged_with(after).merged_with(self ** -1)

    def __repr__(self) -> str:
        x = self.transform(pauli_gates.X)
        y = self.transform(pauli_gates.Y)
        z = self.transform(pauli_gates.Z)
        x_sign = '-' if x.flip else '+'
        y_sign = '-' if y.flip else '+'
        z_sign = '-' if z.flip else '+'
        return (
            f'cirq.SingleQubitCliffordGate(X:{x_sign}{x.to!s}, '
            f'Y:{y_sign}{y.to!s}, Z:{z_sign}{z.to!s})'
        )

    def _circuit_diagram_info_(
        self, args: 'cirq.CircuitDiagramInfoArgs'
    ) -> 'cirq.CircuitDiagramInfo':
        well_known_map = {
            SingleQubitCliffordGate.I: 'I',
            SingleQubitCliffordGate.H: 'H',
            SingleQubitCliffordGate.X: 'X',
            SingleQubitCliffordGate.Y: 'Y',
            SingleQubitCliffordGate.Z: 'Z',
            SingleQubitCliffordGate.X_sqrt: 'X',
            SingleQubitCliffordGate.Y_sqrt: 'Y',
            SingleQubitCliffordGate.Z_sqrt: 'Z',
            SingleQubitCliffordGate.X_nsqrt: 'X',
            SingleQubitCliffordGate.Y_nsqrt: 'Y',
            SingleQubitCliffordGate.Z_nsqrt: 'Z',
        }
        if self in well_known_map:
            symbol = well_known_map[self]
        else:
            rotations = self.decompose_rotation()
            symbol = '-'.join(str(r) + ('^' + str(qt / 2)) * (qt % 4 != 2) for r, qt in rotations)
            symbol = '({})'.format(symbol)
        return protocols.CircuitDiagramInfo(
            wire_symbols=(symbol,),
            exponent={
                SingleQubitCliffordGate.X_sqrt: 0.5,
                SingleQubitCliffordGate.Y_sqrt: 0.5,
                SingleQubitCliffordGate.Z_sqrt: 0.5,
                SingleQubitCliffordGate.X_nsqrt: -0.5,
                SingleQubitCliffordGate.Y_nsqrt: -0.5,
                SingleQubitCliffordGate.Z_nsqrt: -0.5,
            }.get(self, 1),
        )

SingleQubitCliffordGate.I = SingleQubitCliffordGate.from_xz_map(
    (pauli_gates.X, False), (pauli_gates.Z, False)
)
SingleQubitCliffordGate.H = SingleQubitCliffordGate.from_xz_map(
    (pauli_gates.Z, False), (pauli_gates.X, False)
)
SingleQubitCliffordGate.X = SingleQubitCliffordGate.from_xz_map(
    (pauli_gates.X, False), (pauli_gates.Z, True)
)
SingleQubitCliffordGate.Y = SingleQubitCliffordGate.from_xz_map(
    (pauli_gates.X, True), (pauli_gates.Z, True)
)
SingleQubitCliffordGate.Z = SingleQubitCliffordGate.from_xz_map(
    (pauli_gates.X, True), (pauli_gates.Z, False)
)
SingleQubitCliffordGate.X_sqrt = SingleQubitCliffordGate.from_xz_map(
    (pauli_gates.X, False), (pauli_gates.Y, True)
)
SingleQubitCliffordGate.X_nsqrt = SingleQubitCliffordGate.from_xz_map(
    (pauli_gates.X, False), (pauli_gates.Y, False)
)
SingleQubitCliffordGate.Y_sqrt = SingleQubitCliffordGate.from_xz_map(
    (pauli_gates.Z, True), (pauli_gates.X, False)
)
SingleQubitCliffordGate.Y_nsqrt = SingleQubitCliffordGate.from_xz_map(
    (pauli_gates.Z, False), (pauli_gates.X, True)
)
SingleQubitCliffordGate.Z_sqrt = SingleQubitCliffordGate.from_xz_map(
    (pauli_gates.Y, False), (pauli_gates.Z, False)
)
SingleQubitCliffordGate.Z_nsqrt = SingleQubitCliffordGate.from_xz_map(
    (pauli_gates.Y, True), (pauli_gates.Z, False)
)

SQRT_EXP_MAP = {
    0.5: {
        SingleQubitCliffordGate.X: SingleQubitCliffordGate.X_sqrt,
        SingleQubitCliffordGate.Y: SingleQubitCliffordGate.Y_sqrt,
        SingleQubitCliffordGate.Z: SingleQubitCliffordGate.Z_sqrt,
    },
    -0.5: {
        SingleQubitCliffordGate.X: SingleQubitCliffordGate.X_nsqrt,
        SingleQubitCliffordGate.Y: SingleQubitCliffordGate.Y_nsqrt,
        SingleQubitCliffordGate.Z: SingleQubitCliffordGate.Z_nsqrt,
    },
}


# This file is part of QuTiP: Quantum Toolbox in Python.
#
#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
#       of its contributors may be used to endorse or promote products derived
#       from this software without specific prior written permission.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################
from collections.abc import Iterable
import numbers

import numpy as np

from qutip.qobj import Qobj
from qutip.qobjevo import QobjEvo
from qutip.qip.operations.gates import globalphase
from qutip.tensor import tensor
from qutip.mesolve import mesolve
from qutip.qip.circuit import QubitCircuit
from qutip.qip.device.processor import Processor

__all__ = ['ModelProcessor']

class ModelProcessor(Processor):
    """
    The base class for a circuit processor simulating a physical device,
    e.g cavityQED, spinchain.
    The available Hamiltonian of the system is predefined.
    The processor can simulate the evolution under the given
    control pulses either numerically or analytically.
    It cannot be used alone, please refer to the sub-classes.
    (Only additional attributes are documented here, for others please
    refer to the parent class :class:`.Processor`)

    Parameters
    ----------
    N: int
        The number of component systems.

    correct_global_phase: boolean, optional
        If true, the analytical solution will track the global phase. It
        has no effect on the numerical solution.

    t1: list or float
        Characterize the decoherence of amplitude damping for
        each qubit. A list of size `N` or a float for all qubits.

    t2: list of float
        Characterize the decoherence of dephasing for
        each qubit. A list of size `N` or a float for all qubits.

    Attributes
    ----------
    params: dict
        A Python dictionary contains the name and the value of the parameters
        in the physical realization, such as laser frequency, detuning etc.

    correct_global_phase: float
        Save the global phase, the analytical solution
        will track the global phase.
        It has no effect on the numerical solution.
    """
    def __init__(self, N, correct_global_phase=True, t1=None, t2=None):
        super(ModelProcessor, self).__init__(N, t1=t1, t2=t2)
        self.correct_global_phase = correct_global_phase
        self.global_phase = 0.
        self._params = {}

    def to_array(self, params, N):
        """
        Transfer a parameter to an array.
        """
        if isinstance(params, numbers.Real):
            return np.asarray([params] * N)
        elif isinstance(params, Iterable):
            return np.asarray(params)

    def set_up_params(self):
        """
        Save the parameters in the attribute `params` and check the validity.
        (Defined in subclasses)

        Notes
        -----
        All parameters will be multiplied by 2*pi for simplicity
        """
        raise NotImplementedError("Parameters should be defined in subclass.")

    @property
    def params(self):
        return self._params

    @params.setter
    def params(self, par):
        self._params = par

    def run_state(self, init_state=None, analytical=False, qc=None,
                  states=None, **kwargs):
        """
        If `analytical` is False, use :func:`qutip.mesolve` to
        calculate the time of the state evolution
        and return the result. Other arguments of mesolve can be
        given as keyword arguments.
        If `analytical` is True, calculate the propagator
        with matrix exponentiation and return a list of matrices.

        Parameters
        ----------
        init_state: Qobj
            Initial density matrix or state vector (ket).

        analytical: boolean
            If True, calculate the evolution with matrices exponentiation.

        qc: :class:`.QubitCircuit`, optional
            A quantum circuit. If given, it first calls the ``load_circuit``
            and then calculate the evolution.

        states: :class:`qutip.Qobj`, optional
         Old API, same as init_state.

        **kwargs
           Keyword arguments for the qutip solver.

        Returns
        -------
        evo_result: :class:`qutip.Result`
            If ``analytical`` is False,  an instance of the class
            :class:`qutip.Result` will be returned.

            If ``analytical`` is True, a list of matrices representation
            is returned.
        """
        if qc is not None:
            self.load_circuit(qc)
        return super(ModelProcessor, self).run_state(
            init_state=init_state, analytical=analytical,
            states=states, **kwargs)

    def get_ops_and_u(self):
        """
        Get the labels for each Hamiltonian.

        Returns
        -------
        ctrls: list
            The list of Hamiltonians
        coeffs: array_like
            The transposed pulse matrix
        """
        return (self.ctrls, self.get_full_coeffs().T)

    def pulse_matrix(self, dt=0.01):
        """
        Generates the pulse matrix for the desired physical system.

        Returns
        -------
        t, u, labels:
            Returns the total time and label for every operation.
        """
        ctrls = self.ctrls
        coeffs = self.get_full_coeffs().T

        # FIXME This might becomes a problem if new tlist other than
        # int the default pulses are added.
        tlist = self.get_full_tlist()
        dt_list = tlist[1:] - tlist[:-1]
        t_tot = tlist[-1]
        num_step = int(np.ceil(t_tot / dt))

        t = np.linspace(0, t_tot, num_step)
        u = np.zeros((len(ctrls), num_step))

        t_start = 0
        for n in range(len(dt_list)):
            t_idx_len = int(np.floor(dt_list[n] / dt))
            mm = 0
            for m in range(len(ctrls)):
                u[mm, t_start:(t_start + t_idx_len)] = (np.ones(t_idx_len) *
                                                        coeffs[n, m])
                mm += 1
            t_start += t_idx_len

        return t, u, self.get_operators_labels()


#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Authors: Tomas Cassanelli and Andrea Pinna

import logging
import os
import shutil
import sys
import yaml
import numpy as np
from astropy.io import ascii, fits
from scipy.constants import c as light_speed
from .aperture import illum_gauss, illum_pedestal

# ---------------------------------------------------------------------------- #

__all__ = [
    'get_run_config', 'init_output_dir', 'init_logger',
    'extract_data_pyoof', 'extract_data_effelsberg',
    'extract_real_data_srt', 'extract_synthetic_data_srt', 'extract_data_srt',
    'precompute_srt_opd', 'str2LaTeX',
    'store_data_csv', 'uv_ratio', 'illum_strings', 'store_data_ascii'
    ]

# ---------------------------------------------------------------------------- #

def get_run_config(yaml_config_file):
    """
    Read configuration parameters for the current run.
    """

    if not os.path.exists(yaml_config_file):
        sys.exit('Error! Configuration file "{}" does not exist!'.format(
            yaml_config_file))
    try:
        with open(yaml_config_file, 'r') as f:
            return yaml.load(f, Loader=yaml.SafeLoader)
    except:
        sys.exit('Cannot read configuration file "{}"!'.format(
            yaml_config_file))

# ---------------------------------------------------------------------------- #

def init_output_dir(config):
    """
    Create output directory.
    """

    # Create output directory
    if not os.path.exists(config['output']['output_dir']):
        os.makedirs(config['output']['output_dir'])
    elif config['output']['overwrite_dir']:
        shutil.rmtree(config['output']['output_dir'])
        os.mkdir(config['output']['output_dir'])
    else:
        try:
            os.mkdir(config['output']['output_dir'])
        except FileExistsError:
            sys.exit('Error! Directory "{}" already exists!'.format(
                config['output']['output_dir']))

# ---------------------------------------------------------------------------- #

def init_logger(config):
    """
    Initialize loggers to standard output and to file.
    """

    logger = logging.getLogger('pyoof')
    logger.setLevel(logging.DEBUG)

    log_file = os.path.join(config['output']['output_dir'],
                            '{}.log'.format(config['info']['label']))

    fh = logging.FileHandler(log_file)
    fh.setLevel(logging.DEBUG)

    ch = logging.StreamHandler()
    ch.setLevel(logging.INFO)

    formatter = logging.Formatter('%(asctime)s : %(levelname)s : %(message)s')
    fh.setFormatter(formatter)
    ch.setFormatter(formatter)

    logger.addHandler(fh)
    logger.addHandler(ch)

    return logger

# ---------------------------------------------------------------------------- #

def illum_strings(illum_func):
    """
    It assigns string labels to the illumination function. The `~pyoof` package
    has two standard illumination functions, `~pyoof.aperture.illum_pedestal`
    and `~pyoof.aperture.illum_gauss`.

    Parameters
    ----------
    illum_func : `function`
        Illumination function, :math:`E_\\mathrm{a}(x, y)`, to be evaluated
        with ``I_coeff``. The illumination functions available are
        `~pyoof.aperture.illum_pedestal` and `~pyoof.aperture.illum_gauss`.

    Returns
    -------
    illum_name : `str`
        String with the illumination function name.
    taper_name : `str`
        String with the illumination function taper.
    """

    # adding illumination function information
    if illum_func == illum_pedestal:
        illum_name = 'pedestal'
        taper_name = 'c_dB'
    elif illum_func == illum_gauss:
        illum_name = 'gauss'
        taper_name = 'sigma_dB'
    else:
        illum_name = 'manual'
        taper_name = 'taper_dB'

    return illum_name, taper_name

# ---------------------------------------------------------------------------- #

def extract_data_pyoof(pathfits):
    """
    Extracts data from the `~pyoof` default fits file OOF holography
    observations, ready to use for the least squares minimization (see
    `~pyoof.fit_beam`). The fits file has to have the following keys on its
    PrimaryHDU header: ``'FREQ'``, ``'WAVEL'``, ``'MEANEL'``, ``'OBJECT'`` and
    ``'DATE_OBS'``. Besides this three BinTableHDU are required for the data
    itself; ``MINUS OOF``, ``ZERO OOF`` and ``PLUS OOF``. The BinTableHDU
    header has to have the ``'DZ'`` key which includes the radial offset,
    :math:`d_z`. Finally the BinTableHDU has the data files ``'U'``, ``'V'``
    and ``'BEAM'``, which is the :math:`x`- and :math:`y`-axis position in
    radians and the ``'BEAM'`` in a flat array, in mJy.

    Parameters
    ----------
    pathfits : `str`
        Path to the fits file that contains the three beam maps pre-calibrated,
        using the correct PrimaryHDU and the three BinTableHDU (``MINUS OOF``,
        ``ZERO OOF`` and ``PLUS OOF``).

    Returns
    -------
    data_info : `list`
        It contains all extra data besides the beam map. The output
        corresponds to a list,
        ``[name, pthto, obs_object, obs_date, freq, wavel, d_z, meanel]``.
        These are, name of the fits file, paht of the fits file, observed
        object, observation date, frequency, wavelength, radial offset and
        mean elevation, respectively.
    data_obs : `list`
        It contains beam maps and :math:`x`-, and :math:`y`-axis
        (:math:`uv`-plane in Fourier space) data for the least squares
        minimization (see `~pyoof.fit_beam`). The list has the following order
        ``[beam_data, u_data, v_data]``. ``beam_data`` is the three beam
        observations, minus, zero and plus out-of-focus, in a flat array.
        ``u_data`` and ``v_data`` are the beam axes in a flat array.
    """

    hdulist = fits.open(pathfits)  # open fits file, pyoof format
    # path or directory where the fits file is located
    pthto = os.path.split(pathfits)[0]
    # name of the fit file to fit
    name = os.path.split(pathfits)[1][:-5]

    if not all(k in hdulist[0].header
               for k in ['FREQ', 'WAVEL', 'MEANEL', 'OBJECT', 'DATE_OBS']):
        raise TypeError('Not all necessary keys found in FITS header.')

    freq = hdulist[0].header['FREQ']
    wavel = hdulist[0].header['WAVEL']
    meanel = hdulist[0].header['MEANEL']
    obs_object = hdulist[0].header['OBJECT']
    obs_date = hdulist[0].header['DATE_OBS']

    beam_data = [hdulist[i].data['BEAM'] for i in range(1, 4)]
    u_data = [hdulist[i].data['U'] for i in range(1, 4)]
    v_data = [hdulist[i].data['V'] for i in range(1, 4)]
    d_z = [hdulist[i].header['DZ'] for i in range(1, 4)]

    data_file = [name, pthto]
    data_info = data_file + [obs_object, obs_date, freq, wavel, d_z, meanel]
    data_obs = [beam_data, u_data, v_data]

    return data_info, data_obs

# ---------------------------------------------------------------------------- #

def extract_data_srt(config, logger):
    """
    Read data from synthetic or real measurement files and generate the output
    FITS file in the requested format by pyoof.
    """

    data_files = [os.path.join(config['input']['input_dir'],
                               config['input']['oof_minus']),
                  os.path.join(config['input']['input_dir'],
                               config['input']['in_focus']),
                  os.path.join(config['input']['input_dir'],
                               config['input']['oof_plus'])]

    d_z = [-float(config['params']['delta_z']),
           0.0,
           float(config['params']['delta_z'])]

    wavelength = light_speed / float(config['params']['frequency'])  # Hz frequency

    if config['input']['real_data']:
        u, v, P = extract_real_data_srt(data_files, logger)
    else:
        u, v, P = extract_synthetic_data_srt(data_files, logger)

    u_to_save = [u[i].flatten() for i in range(3)]
    v_to_save = [v[i].flatten() for i in range(3)]
    p_to_save = [P[i].flatten() for i in range(3)]

    # Writing default fits file for OOF observations
    table_hdu0 = fits.BinTableHDU.from_columns([
        fits.Column(name='U', format='E', array=u_to_save[0]),
        fits.Column(name='V', format='E', array=v_to_save[0]),
        fits.Column(name='BEAM', format='E', array=p_to_save[0])
        ])

    table_hdu1 = fits.BinTableHDU.from_columns([
        fits.Column(name='U', format='E', array=u_to_save[1]),
        fits.Column(name='V', format='E', array=v_to_save[1]),
        fits.Column(name='BEAM', format='E', array=p_to_save[1])
        ])

    table_hdu2 = fits.BinTableHDU.from_columns([
        fits.Column(name='U', format='E', array=u_to_save[2]),
        fits.Column(name='V', format='E', array=v_to_save[2]),
        fits.Column(name='BEAM', format='E', array=p_to_save[2])
        ])

    fits_file = os.path.join(config['output']['output_dir'],
                             '{}.fits'.format(config['info']['label']))

    logger.info('Writing data to {}...'.format(fits_file))

    prihdr = fits.Header()
    prihdr['FREQ'] = float(config['params']['frequency'])
    prihdr['WAVEL'] = wavelength
    prihdr['MEANEL'] = 0
    prihdr['OBJECT'] = config['info']['label']
    prihdr['DATE_OBS'] = config['info']['observation_date']
    prihdr['COMMENT'] = config['info']['comment']
    prihdr['AUTHOR'] = config['info']['author']
    prihdu = fits.PrimaryHDU(header=prihdr)
    pyoof_fits = fits.HDUList([prihdu, table_hdu0, table_hdu1, table_hdu2])

    for i in range(3):
        pyoof_fits[i + 1].header['DZ'] = d_z[i]

    pyoof_fits[1].name = 'MINUS OOF'
    pyoof_fits[2].name = 'ZERO OOF'
    pyoof_fits[3].name = 'PLUS OOF'

    pyoof_fits.writeto(fits_file)
    logger.info('Done!')

    data_info, data_obs = extract_data_pyoof(fits_file)
    return data_info, data_obs

# ---------------------------------------------------------------------------- #

def extract_real_data_srt(data_files, logger):
    """
    Read data from input files containing real measurements and generate
    organized output variables.
    """

    # Load SRT data
    u_all, v_all, P_all = [], [], []

    for i in range(3):
        logger.info('Reading data from file "{}"...'.format(data_files[i]))
        u_data, v_data, beam_data = np.loadtxt(data_files[i], unpack=True)
        u = sorted(set(u_data))
        v = sorted(set(v_data))
        n_u = len(u)
        n_v = len(v)
        n_values = n_u * n_v
        uu, vv = np.meshgrid(u, v)
        P_tot = beam_data.reshape(n_u, n_v).transpose()

        u_all.append(uu)
        v_all.append(vv)
        P_all.append(P_tot)

    return u_all, v_all, P_all

# ---------------------------------------------------------------------------- #

def extract_synthetic_data_srt(data_files, logger):
    """
    Read data from input files containing synthetic measurements and generate
    organized output variables.
    """

    u_all, v_all, P_all = [], [], []

    for i in range(3):

        logger.info('Reading data from file "{}"...'.format(data_files[i]))
        with open(data_files[i], 'r') as f:
            data = f.readlines()[9:]

        grid_start, grid_end = [float(k) for k in data[0].split()[0::2]]
        n_u, n_v = [int(k) for k in data[1].split()[0:2]]
        n_values = n_u * n_v
        u = np.linspace(grid_start, grid_end, n_u)
        v = np.linspace(grid_start, grid_end, n_v)
        uu, vv = np.meshgrid(u, v)
        pr = np.zeros(n_values)
        pi = np.zeros(n_values)
        lines = data[2:]
        for j, line in enumerate(lines):
            pr[j], pi[j] = [float(k) for k in line.split()[0:2]]
        P_tot = np.power(pr, 2) + np.power(pi, 2)

        u_all.append(uu)
        v_all.append(vv)
        P_all.append(P_tot)

    return u_all, v_all, P_all

# ---------------------------------------------------------------------------- #

def precompute_srt_opd(data_info, telgeo, resolution, box_factor, config, logger):
    """
    Precompute the optical path difference for the Sardinia Radio Telescope.
    """

    pr = telgeo[2]
    box_size = pr * box_factor

    x = np.linspace(-box_size, box_size, resolution)
    y = x
    x_grid, y_grid = np.meshgrid(x, y)

    # Cassegrain/Gregorian (at focus) telescope
    Fp = config['params']['focus_primary_reflector']  # Focus primary reflector m
    F = config['params']['total_focus']  # Total focus Gregorian telescope m
    r = np.sqrt(np.power(x_grid, 2) + np.power(y_grid, 2))  # polar coordinates radius
    a = r / (2 * Fp)
    b = r / (2 * F)

    # Polynomial fitting
    R = np.asarray([
        [-1.438998E-09, -2.706440E-10, -7.098885E-14],
        [ 1.467218E-07,  2.888850E-08,  7.767377E-12],
        [-5.904197E-06, -1.234467E-06, -3.330234E-10],
        [ 1.188731E-04,  2.717683E-05,  7.038221E-09],
        [-1.239164E-03, -3.357801E-04, -7.544711E-08],
        [ 6.017033E-03,  2.461091E-03,  3.776198E-07],
        [-1.323805E-02, -9.312688E-03, -7.093549E-07],
        [ 5.676379E-03,  8.636208E-04,  2.160941E-07]])

    # Degree of polynomial
    N1 = R.shape[0]
    N2 = R.shape[1]

    opd = [[], [], []]
    delta_opd = [[], [], []]
    for i_dz in range(3):
        d_z = data_info[6][i_dz]
        opd[i_dz] = d_z * ((1 - a ** 2) / (1 + a ** 2) + \
                           (1 - b ** 2) / (1 + b ** 2))

        coeff = np.zeros(N1)
        for k in range(N1-1, -1, -1):
            for i in range(N2-1, -1, -1):
                coeff[k] = coeff[k] + R[k, i] * np.power(d_z, N2-i)

        delta_opd[i_dz] = np.zeros(r.shape)
        for i in range(r.shape[0]):
            for j in range(r.shape[1]):
                for k in range(N1-1, -1, -1):
                    delta_opd[i_dz][i, j] = delta_opd[i_dz][i, j] + \
                                            coeff[k] * np.power(r[i, j], N1-k)

    if config['params']['residual_opd']:
        logger.info('Returning "opd + delta_opd"...')
        return [opd[0] + delta_opd[0], opd[1] + delta_opd[1], opd[2] + delta_opd[2]]
    else:
        logger.info('Returning "opd" without "delta_opd"...')
        return opd

# ---------------------------------------------------------------------------- #

def extract_data_effelsberg(pathfits):
    """
    Extracts data from the Effelsberg OOF holography observations, ready to
    use for the least squares minimization. This function will only work for
    the Effelsberg telescope beam maps.

    Parameters
    ----------
    pathfits : `str`
        Path to the fits file that contains the three beam maps pre-calibrated,
        from the Effelsberg telescope.

    Returns
    -------
    data_info : `list`
        It contains all extra data besides the beam map. The output
        corresponds to a list,
        ``[name, pthto, obs_object, obs_date, freq, wavel, d_z, meanel]``.
        These are, name of the fits file, paht of the fits file, observed
        object, observation date, frequency, wavelength, radial offset and
        mean elevation, respectively.
    data_obs : `list`
        It contains beam maps and :math:`x`-, and :math:`y`-axis
        (:math:`uv`-plane in Fourier space) data for the least squares
        minimization (see `~pyoof.fit_beam`). The list has the following order
        ``[beam_data, u_data, v_data]``. ``beam_data`` is the three beam
        observations, minus, zero and plus out-of-focus, in a flat array.
        ``u_data`` and ``v_data`` are the beam axes in a flat array.
    """

    # Opening fits file with astropy
    try:
        # main fits file with the OOF holography format
        hdulist = fits.open(pathfits)

        # Observation frequency
        freq = hdulist[0].header['FREQ']  # Hz
        wavel = light_speed / freq

        # Mean elevation
        meanel = hdulist[0].header['MEANEL']  # Degrees
        obs_object = hdulist[0].header['OBJECT']  # observed object
        obs_date = hdulist[0].header['DATE_OBS']  # observation date
        d_z = [hdulist[i].header['DZ'] for i in range(1, 4)][::-1]

        beam_data = [hdulist[i].data['fnu'] for i in range(1, 4)][::-1]
        u_data = [hdulist[i].data['DX'] for i in range(1, 4)][::-1]
        v_data = [hdulist[i].data['DY'] for i in range(1, 4)][::-1]

    except FileNotFoundError:
        print('Fits file does not exists in directory: ' + pathfits)
    except NameError:
        print('Fits file does not have the OOF holography format')

    else:
        pass

    # Permuting the position to provide same as main_functions
    beam_data.insert(1, beam_data.pop(2))
    u_data.insert(1, u_data.pop(2))
    v_data.insert(1, v_data.pop(2))
    d_z.insert(1, d_z.pop(2))

    # path or directory where the fits file is located
    pthto = os.path.split(pathfits)[0]
    # name of the fit file to fit
    name = os.path.split(pathfits)[1][:-5]

    data_info = [name, pthto, obs_object, obs_date, freq, wavel, d_z, meanel]
    data_obs = [beam_data, u_data, v_data]

    return data_info, data_obs

# ---------------------------------------------------------------------------- #

def str2LaTeX(python_string):
    """
    Function that solves the underscore problem in a python string to
    :math:`\LaTeX` string.

    Parameters
    ----------
    python_string : `str`
        String that needs to be changed.

    Returns
    -------
    LaTeX_string : `str`
        String with the new underscore symbol.
    """

    string_list = list(python_string)
    for idx, string in enumerate(string_list):
        if string_list[idx] == '_':
            string_list[idx] = '\\_'

    LaTeX_string = ''.join(string_list)

    return LaTeX_string

# ---------------------------------------------------------------------------- #

def store_data_csv(name, name_dir, order, save_to_csv):
    """
    Stores all important information in a csv file after the least squares
    minimization has finished, `~pyoof.fit_beam`. All data will be stores in
    the ``pyoof_out/name`` directory, with ``name`` the name of the fits file.

    Parameters
    ----------
    name : `str`
        File name of the fits file to be optimized.
    name_dir : `str`
        Path to store all the csv files. The files will depend on the order of
        the Zernike circle polynomial.
    order : `int`
        Order used for the Zernike circle polynomial, :math:`n`.
    save_to_csv : `list`
        It contains all data that will be stored. The list must have the
        following order, ``[beam_data, u_data, v_data, res_optim, jac_optim,
        grad_optim, phase, cov_ptrue, corr_ptrue]``.
    """

    headers = [
        'Normalized beam', 'u vector radians', 'v vector radians', 'Residual',
        'Jacobian', 'Gradient', 'Phase primary reflector radians', 'Deformation Error',
        'Variance-Covariance matrix (first row fitted parameters idx)',
        'Correlation matrix (first row fitted parameters idx)'
        ]

    fnames = [
        '/beam_data.csv', '/u_data.csv', '/v_data.csv',
        '/res_n{}.csv'.format(order), '/jac_n{}.csv'.format(order),
        '/grad_n{}.csv'.format(order), '/phase_n{}.csv'.format(order),
        '/error_n{}.csv'.format(order),
        '/cov_n{}.csv'.format(order), '/corr_n{}.csv'.format(order)
        ]

    if order != 1:
        headers = headers[3:]
        fnames = fnames[3:]
        save_to_csv = save_to_csv[3:]

    for fname, header, file in zip(fnames, headers, save_to_csv):
        np.savetxt(
            fname=name_dir + fname,
            X=file,
            header=header + ' ' + name
            )

# ---------------------------------------------------------------------------- #

def store_data_ascii(name, name_dir, taper_name, order, params_solution,
                     params_init):
    """
    Stores in an ascii format the parameters found by the least squares
    minimization (see `~pyoof.fit_beam`).

    Parameters
    ----------
    name : `str`
        File name of the fits file to be optimized.
    name_dir : `str`
        Path to store all the csv files. The files will depend on the order of
        the Zernike circle polynomial.
    taper_name : `str`
        Name of the illumination function taper.
    order : `int`
        Order used for the Zernike circle polynomial, :math:`n`.
    params_solution : `~numpy.ndarray`
        Contains the best fitted parameters, the illumination function
        coefficients, ``I_coeff`` and the Zernike circle polynomial
        coefficients, ``K_coeff`` in one array.
    params_init : `~numpt.ndarray`
        Contains the initial parameters used in the least squares minimization
        to start finding the best fitted combination of them.
    """

    n = order
    N_K_coeff = (n + 1) * (n + 2) // 2

    # Making nice table :)
    ln = [(j, i) for i in range(0, n + 1) for j in range(-i, i + 1, 2)]
    L = np.array(ln)[:, 0]
    N = np.array(ln)[:, 1]

    params_names = ['i_amp', taper_name, 'x_0', 'y_0']
    for i in range(N_K_coeff):
        params_names.append('K({}, {})'.format(N[i], L[i]))

    # To store fit information and found parameters in ascii file
    ascii.write(table=[params_names, params_solution, params_init],
                output=os.path.join(name_dir, 'fitpar_n{}.csv'.format(n)),
                names=['parname', 'parfit', 'parinit'],
                comment='Fitted parameters {}'.format(name))

# ---------------------------------------------------------------------------- #

def uv_ratio(u, v):
    """
    Calculates the aspect ratio for the 3 power pattern plots, plus some
    corrections for the text on it. Used in the `function` `~pyoof.plot_beam`
    and `~pyoof.plot_data`

    Parameters
    ----------
    u : `~numpy.ndarray`
        Spatial frequencies from the power pattern, usually in degrees.
    v : `~numpy.ndarray`
        Spatial frequencies from the power pattern, usually in degrees.

    Returns
    -------
    plot_width : `float`
        Width for the power pattern figure.
    plot_height : `float`
        Height for the power pattern figure.
    """

    ratio = (v.max() - v.min()) / (3 * (u.max() - u.min()))
    width = 14
    height = width * (ratio) + 0.2

    return width, height

# ---------------------------------------------------------------------------- #