# -*- coding: utf-8 -*-

# This code is part of Qiskit.
#
# (C) Copyright IBM 2017, 2018.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

# pylint: disable=invalid-name,ungrouped-imports,import-error

"""
Visualization functions for quantum states.
"""

from functools import reduce
import colorsys
import numpy as np
from scipy import linalg
from qiskit.quantum_info.operators.pauli import pauli_group, Pauli
from .matplotlib import HAS_MATPLOTLIB

if HAS_MATPLOTLIB:
    from matplotlib import get_backend
    from matplotlib.ticker import MaxNLocator
    from matplotlib import pyplot as plt
    from matplotlib.patches import FancyArrowPatch
    from matplotlib.patches import Circle
    import matplotlib.colors as mcolors
    from matplotlib.colors import Normalize, LightSource
    import matplotlib.gridspec as gridspec
    from mpl_toolkits.mplot3d import proj3d
    from mpl_toolkits.mplot3d.art3d import Poly3DCollection
    from qiskit.visualization.exceptions import VisualizationError
    from qiskit.visualization.bloch import Bloch
    from qiskit.visualization.utils import _validate_input_state


if HAS_MATPLOTLIB:
    class Arrow3D(FancyArrowPatch):
        """Standard 3D arrow."""

        def __init__(self, xs, ys, zs, *args, **kwargs):
            """Create arrow."""
            FancyArrowPatch.__init__(self, (0, 0), (0, 0), *args, **kwargs)
            self._verts3d = xs, ys, zs

        def draw(self, renderer):
            """Draw the arrow."""
            xs3d, ys3d, zs3d = self._verts3d
            xs, ys, _ = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
            self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
            FancyArrowPatch.draw(self, renderer)


def plot_state_hinton(rho, title='', figsize=None):
    """Plot a hinton diagram for the quanum state.

    Args:
        rho (ndarray): Numpy array for state vector or density matrix.
        title (str): a string that represents the plot title
        figsize (tuple): Figure size in inches.
    Returns:
         matplotlib.Figure: The matplotlib.Figure of the visualization

    Raises:
        ImportError: Requires matplotlib.
    """
    if not HAS_MATPLOTLIB:
        raise ImportError('Must have Matplotlib installed.')
    rho = _validate_input_state(rho)
    if figsize is None:
        figsize = (8, 5)
    num = int(np.log2(len(rho)))
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)
    max_weight = 2 ** np.ceil(np.log(np.abs(rho).max()) / np.log(2))
    datareal = np.real(rho)
    dataimag = np.imag(rho)
    column_names = [bin(i)[2:].zfill(num) for i in range(2**num)]
    row_names = [bin(i)[2:].zfill(num) for i in range(2**num)]
    lx = len(datareal[0])            # Work out matrix dimensions
    ly = len(datareal[:, 0])
    # Real
    ax1.patch.set_facecolor('gray')
    ax1.set_aspect('equal', 'box')
    ax1.xaxis.set_major_locator(plt.NullLocator())
    ax1.yaxis.set_major_locator(plt.NullLocator())

    for (x, y), w in np.ndenumerate(datareal):
        color = 'white' if w > 0 else 'black'
        size = np.sqrt(np.abs(w) / max_weight)
        rect = plt.Rectangle([x - size / 2, y - size / 2], size, size,
                             facecolor=color, edgecolor=color)
        ax1.add_patch(rect)

    ax1.set_xticks(np.arange(0, lx+0.5, 1))
    ax1.set_yticks(np.arange(0, ly+0.5, 1))
    ax1.set_yticklabels(row_names, fontsize=14)
    ax1.set_xticklabels(column_names, fontsize=14, rotation=90)
    ax1.autoscale_view()
    ax1.invert_yaxis()
    ax1.set_title('Re[$\\rho$]', fontsize=14)
    # Imaginary
    ax2.patch.set_facecolor('gray')
    ax2.set_aspect('equal', 'box')
    ax2.xaxis.set_major_locator(plt.NullLocator())
    ax2.yaxis.set_major_locator(plt.NullLocator())

    for (x, y), w in np.ndenumerate(dataimag):
        color = 'white' if w > 0 else 'black'
        size = np.sqrt(np.abs(w) / max_weight)
        rect = plt.Rectangle([x - size / 2, y - size / 2], size, size,
                             facecolor=color, edgecolor=color)
        ax2.add_patch(rect)

    ax2.set_xticks(np.arange(0, lx+0.5, 1))
    ax2.set_yticks(np.arange(0, ly+0.5, 1))
    ax2.set_yticklabels(row_names, fontsize=14)
    ax2.set_xticklabels(column_names, fontsize=14, rotation=90)

    ax2.autoscale_view()
    ax2.invert_yaxis()
    ax2.set_title('Im[$\\rho$]', fontsize=14)
    if title:
        fig.suptitle(title, fontsize=16)
    if get_backend() in ['module://ipykernel.pylab.backend_inline',
                         'nbAgg']:
        plt.close(fig)
    return fig


def plot_bloch_vector(bloch, title="", ax=None, figsize=None):
    """Plot the Bloch sphere.

    Plot a sphere, axes, the Bloch vector, and its projections onto each axis.

    Args:
        bloch (list[double]): array of three elements where [<x>, <y>, <z>]
        title (str): a string that represents the plot title
        ax (matplotlib.Axes): An Axes to use for rendering the bloch sphere
        figsize (tuple): Figure size in inches. Has no effect is passing `ax`.

    Returns:
        Figure: A matplotlib figure instance if `ax = None`.

    Raises:
        ImportError: Requires matplotlib.
    """
    if not HAS_MATPLOTLIB:
        raise ImportError('Must have Matplotlib installed.')
    if figsize is None:
        figsize = (5, 5)
    B = Bloch(axes=ax)
    B.add_vectors(bloch)
    B.render(title=title)
    if ax is None:
        fig = B.fig
        fig.set_size_inches(figsize[0], figsize[1])
        if get_backend() in ['module://ipykernel.pylab.backend_inline',
                             'nbAgg']: