"""
Exact solution of the Pauli equation with parameters:
"""
pauli_params = {
    "omega": 1.,
    "beta": 2.,
    "alpha": 0.95,
}

import numpy as np
from numba import njit

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

@njit
def compiled_W(t, p, q):
    """
    Pauli Wigner function
    :param t: time
    :param p:
    :param q:
    :return: numpy.array
    """

    t1 = np.sqrt(5)
    t2 = np.sqrt(39)
    t3 = t1 * t2
    t5 = t * t3 / 10
    t6 = np.cos(t5)
    t8 = np.sin(t5)
    t13 = t6 ** 2
    t14 = p ** 2
    t17 = q ** 2
    t23 = np.exp(
        -0.19e2 / 0.195e3 * q * p * t8 * t6 * t3 - 0.19e2 / 0.39e2 * t14 * t13 + 0.19e2 / 0.20e2 * t13 * t17 - 0.20e2 / 0.39e2 * t14 - 0.39e2 / 0.20e2 * t17)
    t26 = t * t1 / 10
    t27 = np.cos(t26)
    t29 = np.sin(t26)
    t30 = 100 * t29
    t32 = np.sqrt(complex(0, 10) * t27 * t1 + t30)
    t37 = np.sqrt(complex(0, -10) * t27 * t1 + t30)
    t40 = np.sqrt(19 * t13 + 20)
    t44 = t6 * t1
    t46 = 100 * t8
    t48 = np.sqrt(complex(0, 10) * t44 * t2 + t46)
    t52 = np.sqrt(complex(0, -10) * t44 * t2 + t46)
    t59 = t27 ** 2
    t67 = np.exp(
        0.19e2 / 0.5e1 * q * p * t29 * t27 * t1 + 19 * t14 * t59 - 0.19e2 / 0.20e2 * t17 * t59 - 20 * t14 - t17 / 20)
    t70 = np.sqrt(-19 * t59 + 20)
    W = 5 / t37 / t32 / np.pi / t52 / t48 * (t40 * t37 * t32 * t23 + t70 * t67 * t52 * t48) * t1

    return W


def compiled_rotated_rho_12(t):
    t1 = 39 ** (0.1e1 / 0.4e1)
    t2 = np.sqrt(5)
    t5 = t * t2 / 10
    t6 = np.cos(t5)
    t8 = np.sin(t5)
    t11 = np.sqrt(complex(0, -10) * t6 * t2 + 100 * t8)
    t13 = np.conjugate(0.1e1 / t11)
    t15 = np.sqrt(39)
    t19 = t * t2 * t15 / 10
    t20 = np.cos(t19)
    t23 = np.sin(t19)
    t26 = np.sqrt(complex(0, -10) * t20 * t2 * t15 + 100 * t23)
    t33 = t8 * t2
    t53 = np.sqrt(0.1e1 / (t23 * (complex(0, -2) * t6 * t2 - 20 * t8) + 2 * t20 * (complex(0, 1) * t33 - t6 / 2) * t15) * (
                t23 * (complex(0, -59) * t6 * t2 - 400 * t8) + 21 * t20 * (
                    complex(0, 1) * t33 - 0.20e2 / 0.21e2 * t6) * t15))
    rotated_rho_12 = -50 / t53 / t26 * t13 * t1

    return rotated_rho_12


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

class  SolPauli(object):
    """
    Exact solution for the Pauli equation with parameters

        hbar = 1, m = 1, omega = 1, alpha = 0.95, beta = 2 / (hbar * omega)
    """

    def __init__(self, *, p, q, **kwargs):
        """
        Constructor
        :param p: momenta
        :param q: coordinate
        :param kwargs: ignored
        """
        self.p = p
        self.q = q

        # 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.dpdq = self.dq * self.dp

        # rotation matrix
        self.U = np.array(
            [[1., 1.], [-1., 1.]]
        ) / np.sqrt(2)

    def classical_density(self, t):
        """
        Return the Pauli Wigner function
        :return: numpy.array
        """
        W = compiled_W(t, self.p, self.q)

        assert np.allclose(W.imag, 0), "Classical density must be real"
        assert np.allclose(W.sum() * self.dpdq, 1), "Classical density normalization is not correct"

        return W.real

    def quantum_density(self, t):
        """
        Return the marginal quantum density matrix Pauli
        :param t: time
        :return: numpy.array
        """
        rho_12 = compiled_rotated_rho_12(t)

        rotated_rho = np.array(
            [[0.5, rho_12], [np.conjugate(rho_12), 0.5]]
        )

        rho = self.U.conjugate().T.dot(rotated_rho.dot(self.U))

        # 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


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


if __name__ == '__main__':

    import matplotlib.pyplot as plt
    from wigner_normalize import WignerNormalize, WignerSymLogNorm

    p = np.linspace(-20, 20, 500)[:, np.newaxis]
    q = np.linspace(-20, 20, 500)[np.newaxis, :]

    W = compiled_W(5.7, p, q)

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

    plt.subplot(121)

    plt.title("real part")
    plt.imshow(W.real, **img_params)
    plt.colorbar()

    plt.subplot(122)

    plt.title("imag part")
    plt.imshow(W.imag, **img_params)
    plt.colorbar()

    plt.show()




