"""
Exact solutions for our hybrid model as well as numerical comparison
"""

import numpy as np
from numba import njit
from QuantumClassicalDynamics.split_op_pauli_like1D import SplitOpPauliLike1D


#############################################################################################
#
# Code generated by Maple (exact_solution.mw) begins
#
#############################################################################################

@njit
def compiled_D11(t, p, q, omega, alpha, beta):
    """
    The 11 matrix element of our hybrid density matrix
    :param t: time
    :param p: momentum
    :param q: coordinate
    :param beta: inverse temperature
    :param omega: frequency
    :param alpha: classical-quantum constant
    :return: D11
    """
    t5 = np.sqrt(0.1e1 / np.pi * omega / beta)
    t6 = omega ** 2
    t8 = np.sqrt(t6 + alpha)
    t9 = t * t8
    t10 = np.cos(t9)
    t13 = np.sin(t9)
    t15 = t13 * q * t8 + t10 * p
    t16 = t15 ** 2
    t19 = 0.1e1 / t8
    t21 = -t19 * t13 * p + t10 * q
    t22 = t21 ** 2
    t24 = t22 * t6 + t16
    t25 = t24 / 2
    t26 = t25 * beta
    t27 = 1 + t26
    t28 = np.exp(-t26)
    t29 = t28 * t27
    t31 = np.sqrt(1 - t29)
    t33 = 0.1e1 / t24
    t36 = np.sqrt(t6 - alpha)
    t37 = t * t36
    t38 = np.cos(t37)
    t41 = np.sin(t37)
    t43 = t41 * q * t36 + t38 * p
    t44 = t43 ** 2
    t47 = 0.1e1 / t36
    t49 = -t47 * t41 * p + t38 * q
    t50 = t49 ** 2
    t52 = t50 * t6 + t44
    t53 = t52 / 2
    t54 = t53 * beta
    t55 = 1 + t54
    t56 = np.exp(-t54)
    t57 = t56 * t55
    t59 = np.sqrt(1 - t57)
    t61 = 0.1e1 / t52
    t64 = np.abs(t33 * t31 * t5 + t61 * t59 * t5)
    t65 = t64 ** 2
    t66 = np.sqrt(2)
    t67 = t5 * t66
    t73 = t33 * t31 * t67 / 2 + t61 * t59 * t67 / 2
    t75 = np.conjugate(t24)
    t76 = t75 ** 2
    t78 = np.conjugate(t31)
    t79 = t78 / t76
    t80 = np.conjugate(t15)
    t82 = np.conjugate(t13 * t8)
    t84 = np.conjugate(t21)
    t85 = t84 * t6
    t86 = np.conjugate(t9)
    t87 = np.cos(t86)
    t89 = t82 * t80 + t87 * t85
    t93 = np.conjugate(t29)
    t95 = np.sqrt(1 - t93)
    t97 = 0.1e1 / t95 / t75
    t99 = np.conjugate(t25)
    t100 = t99 * beta
    t101 = np.exp(-t100)
    t104 = beta * (1 + t100)
    t112 = np.conjugate(t57)
    t114 = np.sqrt(1 - t112)
    t115 = 0.1e1 / t114
    t116 = np.conjugate(t43)
    t118 = np.conjugate(t41 * t36)
    t120 = np.conjugate(t49)
    t121 = t120 * t6
    t122 = np.conjugate(t37)
    t123 = np.cos(t122)
    t125 = t118 * t116 + t123 * t121
    t127 = np.conjugate(t53)
    t128 = t127 * beta
    t129 = np.exp(-t128)
    t132 = beta * (1 + t128)
    t137 = np.conjugate(t52)
    t138 = 0.1e1 / t137
    t141 = np.conjugate(t59)
    t142 = t137 ** 2
    t144 = 0.1e1 / t142 * t141
    t155 = np.conjugate(t19 * t13)
    t157 = -t155 * t85 + t87 * t80
    t171 = np.conjugate(t47 * t41)
    t173 = t123 * t116 - t171 * t121
    t187 = (-2 * t157 * t79 + (-t101 * t157 * beta + t101 * t157 * t104) * t97 / 2) * t67 / 2 + (
                t138 * (-t129 * t173 * beta + t129 * t173 * t132) * t115 / 2 - 2 * t173 * t144) * t67 / 2
    t190 = np.real(((-2 * t89 * t79 + (-t101 * t89 * beta + t101 * t89 * t104) * t97 / 2) * t67 / 2 + (t138 * (
                -t129 * t125 * beta + t129 * t125 * t132) * t115 / 2 - 2 * t125 * t144) * t67 / 2) * q * t73 + t187 * p * t73)
    t197 = t10 * t21 * t6 + t13 * t8 * t15
    t207 = t24 ** 2
    t220 = t41 * t36 * t43 + t38 * t49 * t6
    t230 = t52 ** 2
    t239 = np.imag(t187 * ((t28 * t197 * beta * t27 - t28 * t197 * beta) * t33 / t31 * t67 / 4 - t197 / t207 * t31 * t67 + (
                t56 * t220 * beta * t55 - t56 * t220 * beta) * t61 / t59 * t67 / 4 - t220 / t230 * t59 * t67))
    _D11 = t65 + t190 - 2 * t239

    return _D11

@njit
def compiled_D12(t, p, q, omega, alpha, beta):
    """
    The 12 matrix element of our hybrid density matrix
    :param t: time
    :param p: momentum
    :param q: coordinate
    :param beta: inverse temperature
    :param omega: frequency
    :param alpha: classical-quantum constant
    :return: D12
    """
    t1 = np.sqrt(2)
    t6 = np.sqrt(0.1e1 / np.pi * omega / beta)
    t7 = t6 * t1
    t8 = omega ** 2
    t10 = np.sqrt(t8 + alpha)
    t11 = t * t10
    t12 = np.cos(t11)
    t15 = np.sin(t11)
    t17 = t15 * q * t10 + t12 * p
    t18 = t17 ** 2
    t21 = 0.1e1 / t10
    t23 = -t21 * t15 * p + t12 * q
    t24 = t23 ** 2
    t26 = t24 * t8 + t18
    t27 = t26 / 2
    t28 = t27 * beta
    t29 = 1 + t28
    t30 = np.exp(-t28)
    t31 = t30 * t29
    t33 = np.sqrt(1 - t31)
    t34 = 0.1e1 / t26
    t36 = t34 * t33 * t7
    t38 = np.sqrt(t8 - alpha)
    t39 = t * t38
    t40 = np.cos(t39)
    t43 = np.sin(t39)
    t45 = t43 * q * t38 + t40 * p
    t46 = t45 ** 2
    t49 = 0.1e1 / t38
    t51 = -t49 * t43 * p + t40 * q
    t52 = t51 ** 2
    t54 = t52 * t8 + t46
    t55 = t54 / 2
    t56 = t55 * beta
    t57 = 1 + t56
    t58 = np.exp(-t56)
    t59 = t58 * t57
    t61 = np.sqrt(1 - t59)
    t62 = 0.1e1 / t54
    t64 = t62 * t61 * t7
    t66 = t36 / 2 + t64 / 2
    t69 = np.conjugate(-t64 / 2 + t36 / 2)
    t73 = t34 / t33
    t76 = t23 * t8
    t78 = t15 * t10 * t17 + t12 * t76
    t81 = beta * t29
    t88 = t26 ** 2
    t90 = 0.1e1 / t88 * t33
    t96 = t62 / t61
    t99 = t51 * t8
    t101 = t43 * t38 * t45 + t40 * t99
    t104 = beta * t57
    t111 = t54 ** 2
    t113 = 0.1e1 / t111 * t61
    t118 = (-t30 * t78 * beta + t30 * t78 * t81) * t73 * t7 / 4 - t78 * t90 * t7 + (
                -t58 * t101 * beta + t58 * t101 * t104) * t96 * t7 / 4 - t101 * t113 * t7
    t119 = np.conjugate(t59)
    t121 = np.sqrt(1 - t119)
    t122 = 0.1e1 / t121
    t123 = np.conjugate(t45)
    t124 = np.conjugate(t39)
    t125 = np.cos(t124)
    t127 = np.conjugate(t51)
    t128 = t127 * t8
    t129 = t49 * t43
    t130 = np.conjugate(t129)
    t132 = t125 * t123 - t130 * t128
    t134 = np.conjugate(t55)
    t135 = t134 * beta
    t136 = np.exp(-t135)
    t139 = beta * (1 + t135)
    t144 = np.conjugate(t54)
    t145 = 0.1e1 / t144
    t148 = np.conjugate(t61)
    t149 = t144 ** 2
    t151 = 0.1e1 / t149 * t148
    t156 = np.conjugate(t26)
    t157 = t156 ** 2
    t159 = np.conjugate(t33)
    t160 = t159 / t157
    t161 = np.conjugate(t17)
    t162 = np.conjugate(t11)
    t163 = np.cos(t162)
    t165 = np.conjugate(t23)
    t166 = t165 * t8
    t167 = t21 * t15
    t168 = np.conjugate(t167)
    t170 = t163 * t161 - t168 * t166
    t174 = np.conjugate(t31)
    t176 = np.sqrt(1 - t174)
    t178 = 0.1e1 / t176 / t156
    t180 = np.conjugate(t27)
    t181 = t180 * beta
    t182 = np.exp(-t181)
    t185 = beta * (1 + t181)
    t194 = -(t145 * (-t136 * t132 * beta + t136 * t132 * t139) * t122 / 2 - 2 * t132 * t151) * t7 / 2 + (
                -2 * t170 * t160 + (-t182 * t170 * beta + t182 * t170 * t185) * t178 / 2) * t7 / 2
    t197 = np.conjugate(t43 * t38)
    t200 = t197 * t123 + t125 * t128
    t214 = np.conjugate(t15 * t10)
    t217 = t214 * t161 + t163 * t166
    t230 = -(t145 * (-t136 * t200 * beta + t136 * t200 * t139) * t122 / 2 - 2 * t200 * t151) * t7 / 2 + (
                -2 * t217 * t160 + (-t182 * t217 * beta + t182 * t217 * t185) * t178 / 2) * t7 / 2
    t233 = t12 * t17 - t167 * t76
    t248 = -t129 * t99 + t40 * t45
    t261 = (-t30 * t233 * beta + t30 * t233 * t81) * t73 * t7 / 4 - t233 * t90 * t7 + (
                -t58 * t248 * beta + t58 * t248 * t104) * t96 * t7 / 4 - t248 * t113 * t7
    _D12 = 2 * t69 * t66 + complex(0, 1) * (t194 * t118 - t261 * t230) + (t194 * p + t230 * q) * t66 / 2 + (
                t261 * p + t118 * q) * t69 / 2


    return _D12

@njit
def compiled_D22(t, p, q, omega, alpha, beta):
    """
    The 22 matrix element of our hybrid density matrix
    :param t: time
    :param p: momentum
    :param q: coordinate
    :param beta: inverse temperature
    :param omega: frequency
    :param alpha: classical-quantum constant
    :return: D22
    """
    t5 = np.sqrt(0.1e1 / np.pi * omega / beta)
    t6 = omega ** 2
    t8 = np.sqrt(t6 - alpha)
    t9 = t * t8
    t10 = np.cos(t9)
    t13 = np.sin(t9)
    t15 = t13 * q * t8 + t10 * p
    t16 = t15 ** 2
    t19 = 0.1e1 / t8
    t21 = -t19 * t13 * p + t10 * q
    t22 = t21 ** 2
    t24 = t22 * t6 + t16
    t25 = t24 / 2
    t26 = t25 * beta
    t27 = 1 + t26
    t28 = np.exp(-t26)
    t29 = t28 * t27
    t31 = np.sqrt(1 - t29)
    t33 = 0.1e1 / t24
    t36 = np.sqrt(t6 + alpha)
    t37 = t * t36
    t38 = np.cos(t37)
    t41 = np.sin(t37)
    t43 = t41 * q * t36 + t38 * p
    t44 = t43 ** 2
    t47 = 0.1e1 / t36
    t49 = -t47 * t41 * p + t38 * q
    t50 = t49 ** 2
    t52 = t50 * t6 + t44
    t53 = t52 / 2
    t54 = t53 * beta
    t55 = 1 + t54
    t56 = np.exp(-t54)
    t57 = t56 * t55
    t59 = np.sqrt(1 - t57)
    t61 = 0.1e1 / t52
    t64 = np.abs(-t33 * t31 * t5 + t61 * t59 * t5)
    t65 = t64 ** 2
    t66 = np.sqrt(2)
    t67 = t5 * t66
    t73 = -t33 * t31 * t67 / 2 + t61 * t59 * t67 / 2
    t75 = np.conjugate(t29)
    t77 = np.sqrt(1 - t75)
    t78 = 0.1e1 / t77
    t79 = np.conjugate(t15)
    t81 = np.conjugate(t13 * t8)
    t83 = np.conjugate(t21)
    t84 = t83 * t6
    t85 = np.conjugate(t9)
    t86 = np.cos(t85)
    t88 = t81 * t79 + t86 * t84
    t90 = np.conjugate(t25)
    t91 = t90 * beta
    t92 = np.exp(-t91)
    t95 = beta * (1 + t91)
    t100 = np.conjugate(t24)
    t101 = 0.1e1 / t100
    t104 = np.conjugate(t31)
    t105 = t100 ** 2
    t107 = 0.1e1 / t105 * t104
    t112 = np.conjugate(t52)
    t113 = t112 ** 2
    t115 = np.conjugate(t59)
    t116 = t115 / t113
    t117 = np.conjugate(t43)
    t119 = np.conjugate(t41 * t36)
    t121 = np.conjugate(t49)
    t122 = t121 * t6
    t123 = np.conjugate(t37)
    t124 = np.cos(t123)
    t126 = t119 * t117 + t124 * t122
    t130 = np.conjugate(t57)
    t132 = np.sqrt(1 - t130)
    t134 = 0.1e1 / t132 / t112
    t136 = np.conjugate(t53)
    t137 = t136 * beta
    t138 = np.exp(-t137)
    t141 = beta * (1 + t137)
    t155 = np.conjugate(t19 * t13)
    t157 = -t155 * t84 + t86 * t79
    t172 = np.conjugate(t47 * t41)
    t174 = t124 * t117 - t172 * t122
    t187 = -(t101 * (-t92 * t157 * beta + t92 * t157 * t95) * t78 / 2 - 2 * t157 * t107) * t67 / 2 + (
                -2 * t174 * t116 + (-t138 * t174 * beta + t138 * t174 * t141) * t134 / 2) * t67 / 2
    t190 = np.real((-(t101 * (-t92 * t88 * beta + t92 * t88 * t95) * t78 / 2 - 2 * t88 * t107) * t67 / 2 + (
                -2 * t126 * t116 + (
                    -t138 * t126 * beta + t138 * t126 * t141) * t134 / 2) * t67 / 2) * q * t73 + t187 * p * t73)
    t197 = t10 * t21 * t6 + t13 * t8 * t15
    t207 = t24 ** 2
    t220 = t41 * t36 * t43 + t38 * t49 * t6
    t230 = t52 ** 2
    t239 = np.imag(t187 * (-(t28 * t197 * beta * t27 - t28 * t197 * beta) * t33 / t31 * t67 / 4 + t197 / t207 * t31 * t67 + (
                t56 * t220 * beta * t55 - t56 * t220 * beta) * t61 / t59 * t67 / 4 - t220 / t230 * t59 * t67))
    _D22 = t65 + t190 - 2 * t239

    return _D22

#############################################################################################
#
# Code generated by Maple ends
#
#############################################################################################

class CAnalyticQCHybrid(object):
    """
    Equations for plotting the found analytical solution

    hbar = 1, m = 1,
    """
    def __init__(self, *, p, q, omega, alpha, beta):
        """
        Constructor
        :param p: momentum
        :param q: coordinate
        :param beta: inverse temperature
        :param omega: frequency
        :param alpha: classical-quantum constant
        """
        # save parameters
        self.p = p
        self.q = q
        self.beta = beta

        self._args = {
            "p":self.p, "q":self.q, "omega":omega, "alpha":alpha, "beta":beta,
        }

        # extract the volume element in the phase space
        p = p.reshape(-1)
        q = q.reshape(-1)

        self.dq = q[1] - q[0]
        self.dp = p[1] - p[0]

        self.dqdp = self.dq * self.dp

    def D11(self, t):
        """
        The 11 element of the density matrix
        :param t: time
        :return: numpy.array
        """
        self._D11 = compiled_D11(t=t, **self._args)
        return self._D11

    def D12(self, t):
        """
        The 12 element of the hybrid density matrix
        :param t: time
        :return: numpy.array
        """
        self._D12 = compiled_D12(t=t, **self._args)
        return self._D12

    def D22(self, t):
        """
        The 22 element of the hybrid density matrix
        :param t: time
        :return: numpy.array
        """
        self._D22 = compiled_D22(t=t, **self._args)
        return self._D22

    def calculate_D(self, t):
        """
        Save the components of the hybrid density matrix.
        The method must be called before all the other methods are called.
        :param t: time
        :return: self
        """
        self.D11(t)
        self.D12(t)
        self.D22(t)
        return self

    def classical_density(self):
        """
        Calculate the classical. This method must be called after self.calculate_D().
        :return: numpy.array
        """
        rho = self._D11 + self._D22

        # Check for the normalization condition
        assert np.allclose(rho.sum() * self.dqdp, 1), "Classical density must be normalized."

        # Check for the positivity of the density
        assert np.allclose(rho[rho < 0], 0), "Classical density mus be nonnegative"

        assert not np.isnan(rho).any(), "Classical density contains NaNs."

        # Make sure all values are real
        assert isinstance(rho[0, 0],  np.float), "Classical density must be real"

        return rho

    def quantum_density(self):
        """
        Get the quantum density matrix. This method must be called after self.calculate_D().
        :return: 2x2 numpy.array
        """
        d12 = self._D12.sum()

        rho = np.array(
            [[self._D11.sum(), d12], [d12.conjugate(), self._D22.sum()]]
        )

        rho *= self.dqdp

        # For the following consistency checks, get eigenvalues
        p = np.linalg.eigvalsh(rho)

        assert np.allclose(p[p < 0], 0), "Quantum density matrix must be non-negative"

        assert np.allclose(p.sum(), 1), "Trace of quantum density matrix must be one"

        return rho

    def quantum_purity(self):
        """
        Get the purity of the quantum density matrix. This method must be called after self.calculate_D().
        :return: float
        """
        rho = self.quantum_density()

        return rho.dot(rho).trace().real

########################################################################################################
#
#  Plots
#
########################################################################################################

if __name__ == '__main__':

    import matplotlib.pyplot as plt
    from wigner_normalize import WignerNormalize, WignerSymLogNorm
    from matplotlib.animation import FuncAnimation
    from mpl_toolkits.mplot3d import Axes3D

    # Bloch sphere drawing tool
    from qutip import Bloch, Qobj

    ########################################################################################################
    #
    # Make animation
    #
    ########################################################################################################

    class CVisualizeAnim(object):
        """
        Class for drawing animation
        """
        def __init__(self, fig):

            #################################################################
            #
            # Initialize plotting exact plots tools
            #
            #################################################################

            # p = np.linspace(-20., 20., 1000)[:, np.newaxis]
            # q = np.linspace(-20., 20., 1000)[np.newaxis, :]
            #
            # self.hybrid = CAnalyticQCHybrid(p, q, kT=0.5)
            #
            # img_params = dict(
            #     extent=[q.min(), q.max(), p.min(), p.max()],
            #     origin='lower',
            #     cmap='seismic',
            #     # norm=WignerNormalize(vmin=-0.1, vmax=0.1)
            #     norm=WignerSymLogNorm(linthresh=1e-10, vmin=-0.01, vmax=0.1)
            # )

            p = np.linspace(-25, 25, 600)[:, np.newaxis]
            q = np.linspace(-25, 25, 600)[np.newaxis, :]

            self.hybrid = CAnalyticQCHybrid(p=p, q=q, omega=1., beta=2., alpha=0.95)

            img_params = dict(
                extent=[q.min(), q.max(), p.min(), p.max()],
                origin='lower',
                cmap='seismic',
                # norm=WignerNormalize(vmin=-0.1, vmax=0.1)
                norm=WignerSymLogNorm(linthresh=1e-10, vmin=-0.01, vmax=0.1)
            )

            #################################################################
            #
            # Initialize Pauli propagator
            #
            #################################################################

            self.pauli = SplitOpPauliLike1D(
                X_amplitude=2. * q.max(),
                X_gridDIM=512,
                dt=0.0001,
                K0="0.5 * P ** 2",
                V0="0.5 * X ** 2",
                V1="0.5 * 0.95 * X ** 2",

                # kT=self.hybrid.kT, # parameters controlling the width of the initial wavepacket
            ).set_wavefunction("exp(-0.5 * X ** 2)")

            #################################################################
            #
            # Initialize plotting facility
            #
            #################################################################

            self.fig = fig

            self.fig.suptitle(
                "Quantum-classical hybrid $m=1$, $\omega=1$, $\\alpha=0.95$, $\\beta={:.1e}$ (a.u.) (a.u.)".format(self.hybrid.beta)
            )

            self.ax = fig.add_subplot(221)
            self.ax.set_title('Classical density')

            # generate empty plots
            self.img_classical_density = self.ax.imshow([[0]], **img_params)

            self.ax.set_xlabel('$q$ (a.u.)')
            self.ax.set_ylabel('$p$ (a.u.)')

            ax = fig.add_subplot(222)

            ax.set_title('Quantum purity')
            self.quantum_purity_plot, = ax.plot([0., 40], [1, 0.5], label='hybrid')
            self.pauli_quantum_purity_plot, = ax.plot([0., 40], [1, 0.5], label='Pauli')

            ax.set_xlabel('time (a.u.)')
            ax.set_ylabel("quantum purity")
            ax.legend()

            self.time = []
            self.qpurity = []
            self.pauli_qpurity = []

            ax = fig.add_subplot(223)

            ax.set_title("Coordinate distribution")
            self.c_coordinate_distribution, = ax.semilogy(
                [self.hybrid.q.min(), self.hybrid.q.max()], [1e-11, 1e0], label="hybrid"
            )
            self.pauli_coordinate_distribution, = ax.semilogy(
                [self.hybrid.q.min(), self.hybrid.q.max()], [1e-11, 1e0], label="Pauli"
            )

            ax.legend()

            ax.set_xlabel('$q$ (a.u.)')
            ax.set_ylabel('Probability density')

            ax = fig.add_subplot(224, projection='3d', azim=90,elev=0)

            self.bloch = Bloch(axes=ax)
            self.bloch.make_sphere()

        def __call__(self, frame_num):
            """
            Draw a new frame
            :param frame_num: current frame number
            :return: image objects
            """
            # convert the frame number to time
            #t = 0.05 * frame_num
            t= self.pauli.t

            # calculate the hybrid density matrix
            self.hybrid.calculate_D(t)

            # plot the classical density
            c_rho = self.hybrid.classical_density()

            self.img_classical_density.set_array(c_rho)
            self.ax.set_title('Classical density \n $t = {:.1f}$ (a.u.)'.format(t))

            # plot the coordinate distribution for the classical density
            coordinate_marginal = c_rho.sum(axis=0)
            # coordinate_marginal *= self.hybrid.dp
            coordinate_marginal /= coordinate_marginal.max()

            self.c_coordinate_distribution.set_data(self.hybrid.q.reshape(-1), coordinate_marginal)
            coordinate_density = self.pauli.coordinate_density
            coordinate_density /= coordinate_density.max()
            self.pauli_coordinate_distribution.set_data(self.pauli.X, coordinate_density)

            # plot quantum purity
            self.time.append(t)
            self.qpurity.append(self.hybrid.quantum_purity())

            self.quantum_purity_plot.set_data(self.time, self.qpurity)

            # Get the Pauli density matrix
            rho12 =  np.sum(self.pauli.psi1 * self.pauli.psi2.conjugate())

            rho_pauli = np.array(
                [[np.sum(np.abs(self.pauli.psi1) ** 2), rho12],
                [rho12.conjugate(), np.sum(np.abs(self.pauli.psi2) ** 2)]]
            )
            rho_pauli *= self.pauli.dX

            # plot Pauli purity
            self.pauli_qpurity.append(
                rho_pauli.dot(rho_pauli).trace().real
            )

            self.pauli_quantum_purity_plot.set_data(self.time, self.pauli_qpurity)

            # plot Bloch vector
            self.bloch.clear()
            self.bloch.add_states(
                [Qobj(self.hybrid.quantum_density()),
                Qobj(rho_pauli)]
            )
            self.bloch.make_sphere()

            #
            self.pauli.propagate(1000)

            return self.img_classical_density, self.quantum_purity_plot, self.pauli_quantum_purity_plot,\
                   self.bloch, self.pauli_coordinate_distribution


    fig = plt.gcf()
    fig.set_size_inches(10, 10)
    visualizer = CVisualizeAnim(fig)
    animation = FuncAnimation(
        fig, visualizer, frames=np.arange(400), repeat=True, # blit=True
    )
    plt.show()

    # If you want to make a movie, comment "plt.show()" out and uncomment the lines bellow
    # Save animation into the file
    # animation.save('hybrid_animation.mp4')

    # check weather the energy is preserved during the
    energy = np.array(visualizer.pauli.hamiltonian_average)
    print("Hamiltonian preserved: {:.1e} %".format(np.real(1 - energy.min() / energy.max()) * 100.))

    #######################################################################################################

    #plt.imshow(hybrid.classical_density(0.5), **img_params)
    #plt.ylabel('$p$ (a.u.)')
    #plt.xlabel('$q$ (a.u.)')
    #plt.show()
