import numpy as np
from numba import njit


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

@njit
def compiled_D11(t, p, q, omega, alpha, beta):
    """
    The 11 matrix element of the Aleksandrov-Gerasimenko-Kapral 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: numpy.array
    """
    t3 = omega ** 2
    t4 = q ** 2
    t5 = t4 * t3
    t7 = p ** 2
    t9 = complex(0, -1) * t5 * alpha + complex(0, 1) * t7 * alpha
    t11 = 2 * omega * t
    t12 = np.sin(t11)
    t14 = p * q
    t15 = np.cos(t11)
    t17 = t15 * alpha * t14
    t18 = t3 ** 2
    t20 = t4 * t18 * beta
    t21 = t4 * alpha
    t22 = t * t21
    t23 = t7 * beta
    t29 = alpha * (p * t - q) * p
    t35 = 1 / t3 / omega
    t38 = np.exp(t35 * (t12 * t9 - 2 * (
                complex(0, 1) * t17 + t20 + t3 * (complex(0, 1) * t22 + t23) + complex(0, 1) * t29) * omega) / 4)
    t49 = np.exp(t35 * (-t12 * t9 + 2 * (
                complex(0, 1) * t17 - t20 + t3 * (complex(0, 1) * t22 - t23) + complex(0, 1) * t29) * omega) / 4)
    t52 = t3 - alpha
    t53 = np.sqrt(t52)
    t54 = t * t53
    t55 = np.cos(t54)
    t56 = t55 ** 2
    t58 = np.sin(t54)
    t64 = t4 * t18
    t67 = 2 * t4 * t3 * alpha
    t68 = t7 * t3
    t69 = alpha ** 2
    t70 = t4 * t69
    t77 = np.exp(0.1e1 / t52 * beta * (
                t56 * (-t5 + t21 + t7) * alpha + 2 * t53 * t14 * alpha * t58 * t55 - t64 + t67 - t68 - t70) / 2)
    t80 = t3 + alpha
    t81 = np.sqrt(t80)
    t82 = t * t81
    t83 = np.cos(t82)
    t84 = t83 ** 2
    t86 = np.sin(t82)
    t97 = np.exp(-0.1e1 / t80 * (
                t84 * (-t5 - t21 + t7) * alpha + 2 * t81 * t14 * alpha * t86 * t83 + t64 + t67 + t68 + t70) * beta / 2)
    _D11 = 0.1e1 / np.pi * (t38 + t49 + t77 + t97) * omega * beta / 8

    return _D11


@njit
def compiled_D12(t, p, q, omega, alpha, beta):
    """
    The 12 matrix element of the Aleksandrov-Gerasimenko-Kapral 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: numpy.array
    """
    t3 = omega ** 2
    t4 = q ** 2
    t5 = t4 * t3
    t7 = p ** 2
    t9 = complex(0, -1) * t5 * alpha + complex(0, 1) * t7 * alpha
    t11 = 2 * omega * t
    t12 = np.sin(t11)
    t14 = p * q
    t15 = np.cos(t11)
    t17 = t15 * alpha * t14
    t18 = t3 ** 2
    t20 = t4 * t18 * beta
    t21 = t4 * alpha
    t22 = t * t21
    t23 = t7 * beta
    t29 = alpha * (p * t - q) * p
    t35 = 1 / t3 / omega
    t38 = np.exp(t35 * (t12 * t9 - 2 * (
                complex(0, 1) * t17 + t20 + t3 * (complex(0, 1) * t22 + t23) + complex(0, 1) * t29) * omega) / 4)
    t49 = np.exp(t35 * (-t12 * t9 + 2 * (
                complex(0, 1) * t17 - t20 + t3 * (complex(0, 1) * t22 - t23) + complex(0, 1) * t29) * omega) / 4)
    t52 = t3 - alpha
    t53 = np.sqrt(t52)
    t54 = t * t53
    t55 = np.cos(t54)
    t56 = t55 ** 2
    t58 = np.sin(t54)
    t64 = t4 * t18
    t67 = 2 * t4 * t3 * alpha
    t68 = t7 * t3
    t69 = alpha ** 2
    t70 = t4 * t69
    t77 = np.exp(0.1e1 / t52 * beta * (
                t56 * (-t5 + t21 + t7) * alpha + 2 * t53 * t14 * alpha * t58 * t55 - t64 + t67 - t68 - t70) / 2)
    t80 = t3 + alpha
    t81 = np.sqrt(t80)
    t82 = t * t81
    t83 = np.cos(t82)
    t84 = t83 ** 2
    t86 = np.sin(t82)
    t97 = np.exp(-0.1e1 / t80 * (
                t84 * (-t5 - t21 + t7) * alpha + 2 * t81 * t14 * alpha * t86 * t83 + t64 + t67 + t68 + t70) * beta / 2)
    _D12 = -0.1e1 / np.pi * (t38 - t49 + t77 - t97) * omega * beta / 8

    return _D12

@njit
def compiled_D22(t, p, q, omega, alpha, beta):
    """
    The 12 matrix element of the Aleksandrov-Gerasimenko-Kapral 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: numpy.array
    """
    t3 = omega ** 2
    t4 = q ** 2
    t5 = t4 * t3
    t7 = p ** 2
    t9 = complex(0, -1) * t5 * alpha + complex(0, 1) * t7 * alpha
    t11 = 2 * omega * t
    t12 = np.sin(t11)
    t14 = p * q
    t15 = np.cos(t11)
    t17 = t15 * alpha * t14
    t18 = t3 ** 2
    t20 = t4 * t18 * beta
    t21 = t4 * alpha
    t22 = t * t21
    t23 = t7 * beta
    t29 = alpha * (p * t - q) * p
    t35 = 1 / t3 / omega
    t38 = np.exp(t35 * (t12 * t9 - 2 * (
                complex(0, 1) * t17 + t20 + t3 * (complex(0, 1) * t22 + t23) + complex(0, 1) * t29) * omega) / 4)
    t49 = np.exp(t35 * (-t12 * t9 + 2 * (
                complex(0, 1) * t17 - t20 + t3 * (complex(0, 1) * t22 - t23) + complex(0, 1) * t29) * omega) / 4)
    t52 = t3 - alpha
    t53 = np.sqrt(t52)
    t54 = t * t53
    t55 = np.cos(t54)
    t56 = t55 ** 2
    t58 = np.sin(t54)
    t64 = t4 * t18
    t67 = 2 * t4 * t3 * alpha
    t68 = t7 * t3
    t69 = alpha ** 2
    t70 = t4 * t69
    t77 = np.exp(0.1e1 / t52 * beta * (
                t56 * (-t5 + t21 + t7) * alpha + 2 * t53 * t14 * alpha * t58 * t55 - t64 + t67 - t68 - t70) / 2)
    t80 = t3 + alpha
    t81 = np.sqrt(t80)
    t82 = t * t81
    t83 = np.cos(t82)
    t84 = t83 ** 2
    t86 = np.sin(t82)
    t97 = np.exp(-0.1e1 / t80 * (
                t84 * (-t5 - t21 + t7) * alpha + 2 * t81 * t14 * alpha * t86 * t83 + t64 + t67 + t68 + t70) * beta / 2)
    _D22 = -0.1e1 / np.pi * (t38 + t49 - t77 - t97) * omega * beta / 8

    return _D22

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


class SolAGKEq(object):
    """
    Exact solution to the Aleksandrov-Gerasimenko-Kapral Eq for a quadratic hamiltonian as shown in the paper.

    The values of the parameters: 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

        self.omega = omega
        self.beta = beta
        self.alpha = alpha

        # save the Hamiltonian
        self.H0 = 0.5 * (self.p ** 2 + self.omega ** 2 * self.q ** 2)

    def D11(self, t):
        """
        The 11 element of the Aleksandrov-Gerasimenko-Kapral 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 Aleksandrov-Gerasimenko-Kapral 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 Aleksandrov-Gerasimenko-Kapral 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 Aleksandrov-Gerasimenko-Kapral 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."

        assert np.allclose(rho.imag, 0.), "Classical density must be real"

        return rho.real

    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

        assert np.allclose(rho, rho.conj().T), "The quantum density matrix must be Hermitian"

        # 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

    def energy(self):
        """
        Calculate the total energy, which is a conserved quantity.
        :return: float
        """
        return np.sum(
            self.H0 * (self._D11 + self._D22) + self.alpha * self.q ** 2 * self._D12.real
        ).real * self.dqdp

    def average_sigma_1(self):
        """
        To check weather the Bloch vector stays in the yz plane
        <sigma_1>
        :return: float
        """
        return 2 * self._D12.real.sum() * self.dqdp


########################################################################################################
#
#  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(-0.1, 0.1, 500)[:, np.newaxis]
            q = np.linspace(-0.1, 0.1, 500)[np.newaxis, :]

            self.hybrid = SolAGKEq(p=p, q=q, omega=1., beta=1e5, 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)
            )

            # List to save the total hybrid energy for testing (it must be a conserved quantity)
            self.total_energy = []

            # List to save the x component of the Bloch vector
            self.sigma_1 = []

            #################################################################
            #
            # 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., 1000000], [1, 0.5])
            ax.set_xlabel('time (a.u.)')
            ax.set_ylabel("quantum purity")

            self.time = []
            self.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"
            )
            ax.legend()

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

            ax = fig.add_subplot(224, projection='3d', azim=0,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 = 4000. * frame_num

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

            # Save total energy
            self.total_energy.append(self.hybrid.energy())

            # Save <sigma_1>
            self.sigma_1.append(self.hybrid.average_sigma_1())

            # 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)

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

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

            # plot Bloch vector
            self.bloch.clear()
            self.bloch.add_states(
                Qobj(self.hybrid.quantum_density())
            )
            self.bloch.make_sphere()

            #
            # self.pauli.propagate(100)

            return self.img_classical_density, self.quantum_purity_plot, self.bloch,

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

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

    plt.subplot(211)
    plt.plot(visualizer.time, visualizer.total_energy)
    plt.xlabel('$t$ (a.u.)')
    plt.ylabel('Total energy (must be conserved)')

    plt.subplot(212)
    plt.plot(visualizer.time, visualizer.sigma_1)
    plt.xlabel('$t$ (a.u.)')
    plt.ylabel('$\\langle \sigma_1 \\rangle$ (must be small)')

    plt.show()