from cmath import exp, sqrt
import numpy as np
from numba import cuda
from yasfpy.computers import Computer


class ComputerCUDA(Computer):
    # TODO: Implement data batching for GPUs with smaller memory
    @cuda.jit(fastmath=True)
    def particle_interaction(
        lmax: int,
        particle_number: int,
        idx: np.ndarray,
        x: np.ndarray,
        wx_real: np.ndarray,
        wx_imag: np.ndarray,
        translation_table: np.ndarray,
        plm: np.ndarray,
        sph_h: np.ndarray,
        e_j_dm_phi,
    ):
        """
        Perform particle interaction calculations on the GPU.

        Args:
            lmax (int): Maximum angular momentum quantum number.
            particle_number (int): Number of particles.
            idx (np.ndarray): Array containing particle indices.
            x (np.ndarray): Array of particle positions.
            wx_real (np.ndarray): Array to store the real part of the result.
            wx_imag (np.ndarray): Array to store the imaginary part of the result.
            translation_table (np.ndarray): Array containing translation table.
            plm (np.ndarray): Array containing associated Legendre polynomials.
            sph_h (np.ndarray): Array containing spherical harmonics.
            e_j_dm_phi (np.ndarray): Additional parameter for the calculation.

        Todo:
            Implement data batching for GPUs with smaller memory
        """
        j1, j2, w = cuda.grid(3)

        jmax = particle_number * 2 * lmax * (lmax + 2)
        channels = sph_h.shape[-1]

        if (j1 >= jmax) or (j2 >= jmax) or (w >= channels):
            return

        s1, n1, tau1, l1, m1 = idx[j1, :]
        s2, n2, tau2, l2, m2 = idx[j2, :]

        if s1 == s2:
            return

        delta_tau = abs(tau1 - tau2)
        delta_l = abs(l1 - l2)
        delta_m = abs(m1 - m2)

        p_dependent = complex(0)
        for p in range(max(delta_m, delta_l + delta_tau), l1 + l2 + 1):
            p_dependent += (
                translation_table[n2, n1, p]
                * plm[p * (p + 1) // 2 + delta_m, s1, s2]
                * sph_h[p, s1, s2, w]
            )
        p_dependent *= e_j_dm_phi[m2 - m1 + 2 * lmax, s1, s2] * x[j2]

        # atomic.add performs the += operation in sync
        cuda.atomic.add(wx_real, (j1, w), p_dependent.real)
        cuda.atomic.add(wx_imag, (j1, w), p_dependent.imag)

    @cuda.jit(fastmath=True)
    def compute_scattering_cross_section(
        lmax: int,
        particle_number: int,
        idx: np.ndarray,
        sfc: np.ndarray,
        translation_table: np.ndarray,
        plm: np.ndarray,
        sph_h: np.ndarray,
        e_j_dm_phi: np.ndarray,
        c_sca_real: np.ndarray,
        c_sca_imag: np.ndarray,
    ):
        """
        Compute the scattering cross section on the GPU using CUDA.

        Args:
            lmax (int): The maximum degree of the spherical harmonics expansion.
            particle_number (int): The number of particles.
            idx (np.ndarray): The index array.
            sfc (np.ndarray): The scattering form factor array.
            translation_table (np.ndarray): The translation table array.
            plm (np.ndarray): The associated Legendre polynomials array.
            sph_h (np.ndarray): The spherical harmonics array.
            e_j_dm_phi (np.ndarray): The phase factor array.
            c_sca_real (np.ndarray): The real part of the scattering cross section array.
            c_sca_imag (np.ndarray): The imaginary part of the scattering cross section array.
        """
        jmax = particle_number * 2 * lmax * (lmax + 2)
        channels = sph_h.shape[-1]

        j1, j2, w = cuda.grid(3)

        if (j1 >= jmax) or (j2 >= jmax) or (w >= channels):
            return

        s1, n1, _, _, m1 = idx[j1, :]
        s2, n2, _, _, m2 = idx[j2, :]

        delta_m = abs(m1 - m2)

        p_dependent = complex(0)
        for p in range(delta_m, 2 * lmax + 1):
            p_dependent += (
                translation_table[n2, n1, p]
                * plm[p * (p + 1) // 2 + delta_m, s1, s2]
                * sph_h[p, s1, s2, w]
            )
        p_dependent *= (
            sfc[s1, n1, w].conjugate()
            * e_j_dm_phi[m2 - m1 + 2 * lmax, s1, s2]
            * sfc[s2, n2, w]
        )

        # atomic.add performs the += operation in sync
        cuda.atomic.add(c_sca_real, w, p_dependent.real)
        cuda.atomic.add(c_sca_imag, w, p_dependent.imag)

    @cuda.jit(fastmath=True)
    def compute_radial_independent_scattered_field(
        lmax: int,
        particles_position: np.ndarray,
        idx: np.ndarray,
        sfc: np.ndarray,
        k_medium: np.ndarray,
        azimuthal_angles: np.ndarray,
        e_r: np.ndarray,
        e_phi: np.ndarray,
        e_theta: np.ndarray,
        pilm: np.ndarray,
        taulm: np.ndarray,
        e_1_sca_real: np.ndarray,
        e_1_sca_imag: np.ndarray,
    ):
        """
        Compute the radial independent scattered field using GPU acceleration.

        Args:
            lmax (int): The maximum degree of the spherical harmonics expansion.
            particles_position (np.ndarray): Array of particle positions.
            idx (np.ndarray): Array of indices for particle properties.
            sfc (np.ndarray): Array of scattering form factors.
            k_medium (np.ndarray): Array of wave numbers in the medium.
            azimuthal_angles (np.ndarray): Array of azimuthal angles.
            e_r (np.ndarray): Array of radial electric field components.
            e_phi (np.ndarray): Array of azimuthal electric field components.
            e_theta (np.ndarray): Array of polar electric field components.
            pilm (np.ndarray): Array of associated Legendre polynomials.
            taulm (np.ndarray): Array of tau coefficients.
            e_1_sca_real (np.ndarray): Array of real parts of the scattered electric field.
            e_1_sca_imag (np.ndarray): Array of imaginary parts of the scattered electric field.
        """
        j_idx, a_idx, w_idx = cuda.grid(3)

        jmax = particles_position.shape[0] * 2 * lmax * (lmax + 2)

        if (
            (j_idx >= jmax)
            or (a_idx >= azimuthal_angles.size)
            or (w_idx >= k_medium.size)
        ):
            return

        s, n, tau, l, m = idx[j_idx, :]

        # Temporary variable
        # If tau = 1 -> 1j**(tau-1) = 1, if tau = 2 -> 1j**(tau-1) = 1j
        # 1j**(-l-1) = (-1j)**(l+1) => both lead to the coefficient 1j**(tau-l-2)
        # k * <particle_position, e_r> is the phase shift due to the distance and relative position
        t = (
            1j ** (tau - l - 2)
            * sfc[s, n, w_idx]
            / sqrt(2 * l * (l + 1))
            * exp(
                1j
                * (
                    m * azimuthal_angles[a_idx]
                    - k_medium[w_idx]
                    * (
                        particles_position[s, 0] * e_r[a_idx, 0]
                        + particles_position[s, 1] * e_r[a_idx, 1]
                        + particles_position[s, 2] * e_r[a_idx, 2]
                    )
                )
            )
        )

        for c in range(3):
            if tau == 1:
                e_1_sca = t * (
                    e_theta[a_idx, c] * pilm[l, abs(m), a_idx] * 1j * m
                    - e_phi[a_idx, c] * taulm[l, abs(m), a_idx]
                )
            else:
                e_1_sca = t * (
                    e_phi[a_idx, c] * pilm[l, abs(m), a_idx] * 1j * m
                    + e_theta[a_idx, c] * taulm[l, abs(m), a_idx]
                )

            cuda.atomic.add(e_1_sca_real, (a_idx, c, w_idx), e_1_sca.real)
            cuda.atomic.add(e_1_sca_imag, (a_idx, c, w_idx), e_1_sca.imag)

    @cuda.jit(fastmath=True)
    def compute_electric_field_angle_components(
        lmax: int,
        particles_position: np.ndarray,
        idx: np.ndarray,
        sfc: np.ndarray,
        k_medium: np.ndarray,
        azimuthal_angles: np.ndarray,
        e_r: np.ndarray,
        pilm: np.ndarray,
        taulm: np.ndarray,
        e_field_theta_real: np.ndarray,
        e_field_theta_imag: np.ndarray,
        e_field_phi_real: np.ndarray,
        e_field_phi_imag: np.ndarray,
    ):
        """
        Compute the electric field angle components on the GPU.

        Args:
            lmax (int): The maximum angular momentum quantum number.
            particles_position (np.ndarray): Array of particle positions.
            idx (np.ndarray): Array of indices.
            sfc (np.ndarray): Array of scattering form factors.
            k_medium (np.ndarray): Array of medium wavevectors.
            azimuthal_angles (np.ndarray): Array of azimuthal angles.
            e_r (np.ndarray): Array of radial unit vectors.
            pilm (np.ndarray): Array of associated Legendre polynomials.
            taulm (np.ndarray): Array of tau coefficients.
            e_field_theta_real (np.ndarray): Array of real parts of electric field theta component.
            e_field_theta_imag (np.ndarray): Array of imaginary parts of electric field theta component.
            e_field_phi_real (np.ndarray): Array of real parts of electric field phi component.
            e_field_phi_imag (np.ndarray): Array of imaginary parts of electric field phi component.
        """
        j_idx, a_idx, w_idx = cuda.grid(3)
        jmax = particles_position.shape[0] * 2 * lmax * (lmax + 2)
        if (
            (j_idx >= jmax)
            or (a_idx >= azimuthal_angles.size)
            or (w_idx >= k_medium.size)
        ):
            return

        s, n, tau, l, m = idx[j_idx, :]

        t = (
            1j ** (tau - l - 2)
            * sfc[s, n, w_idx]
            / sqrt(2 * l * (l + 1))
            * exp(
                1j
                * (
                    m * azimuthal_angles[a_idx]
                    - k_medium[w_idx]
                    * (
                        particles_position[s, 0] * e_r[a_idx, 0]
                        + particles_position[s, 1] * e_r[a_idx, 1]
                        + particles_position[s, 2] * e_r[a_idx, 2]
                    )
                )
            )
        )

        if tau == 1:
            e_field_theta = t * pilm[l, abs(m), a_idx] * 1j * m
            e_field_phi = -t * taulm[l, abs(m), a_idx]
        else:
            e_field_theta = t * taulm[l, abs(m), a_idx]
            e_field_phi = t * pilm[l, abs(m), a_idx] * 1j * m

        cuda.atomic.add(e_field_theta_real, (a_idx, w_idx), e_field_theta.real)
        cuda.atomic.add(e_field_theta_imag, (a_idx, w_idx), e_field_theta.imag)
        cuda.atomic.add(e_field_phi_real, (a_idx, w_idx), e_field_phi.real)
        cuda.atomic.add(e_field_phi_imag, (a_idx, w_idx), e_field_phi.imag)

    @cuda.jit(fastmath=True)
    def compute_polarization_components(
        number_of_wavelengths: int,
        number_of_angles: int,
        e_field_theta_real: np.ndarray,
        e_field_theta_imag: np.ndarray,
        e_field_phi_real: np.ndarray,
        e_field_phi_imag: np.ndarray,
        intensity: np.ndarray,
        degree_of_polarization: np.ndarray,
        degree_of_linear_polarization: np.ndarray,
        degree_of_linear_polarization_q: np.ndarray,
        degree_of_linear_polarization_u: np.ndarray,
        degree_of_circular_polarization: np.ndarray,
    ):
        """
        Compute the polarization components using GPU acceleration.

        Args:
            number_of_wavelengths (int): Number of wavelengths.
            number_of_angles (int): Number of angles.
            e_field_theta_real (np.ndarray): Real part of the electric field in the theta direction.
            e_field_theta_imag (np.ndarray): Imaginary part of the electric field in the theta direction.
            e_field_phi_real (np.ndarray): Real part of the electric field in the phi direction.
            e_field_phi_imag (np.ndarray): Imaginary part of the electric field in the phi direction.
            intensity (np.ndarray): Array to store the intensity component.
            degree_of_polarization (np.ndarray): Array to store the degree of polarization component.
            degree_of_linear_polarization (np.ndarray): Array to store the degree of linear polarization component.
            degree_of_linear_polarization_q (np.ndarray): Array to store the degree of linear polarization (Q) component.
            degree_of_linear_polarization_u (np.ndarray): Array to store the degree of linear polarization (U) component.
            degree_of_circular_polarization (np.ndarray): Array to store the degree of circular polarization component.
        """
        a_idx, w_idx = cuda.grid(2)
        if (w_idx >= number_of_wavelengths) or (a_idx >= number_of_angles):
            return

        # Jones vector components (1,2,4)
        e_field_theta_abs = (
            e_field_theta_real[a_idx, w_idx] ** 2
            + e_field_theta_imag[a_idx, w_idx] ** 2
        )
        e_field_phi_abs = (
            e_field_phi_real[a_idx, w_idx] ** 2 + e_field_phi_imag[a_idx, w_idx] ** 2
        )
        e_field_angle_interaction_real = (
            e_field_theta_real[a_idx, w_idx] * e_field_phi_real[a_idx, w_idx]
            + e_field_theta_imag[a_idx, w_idx] * e_field_phi_imag[a_idx, w_idx]
        )
        e_field_angle_interaction_imag = (
            e_field_theta_imag[a_idx, w_idx] * e_field_phi_real[a_idx, w_idx]
            - e_field_theta_real[a_idx, w_idx] * e_field_phi_imag[a_idx, w_idx]
        )

        # Stokes components S = (I, Q, U, V)
        I = e_field_theta_abs + e_field_phi_abs
        Q = e_field_theta_abs - e_field_phi_abs
        U = -2 * e_field_angle_interaction_real
        V = 2 * e_field_angle_interaction_imag

        intensity[a_idx, w_idx] = I
        degree_of_polarization[a_idx, w_idx] = sqrt(Q**2 + U**2 + V**2).real / I
        degree_of_linear_polarization[a_idx, w_idx] = sqrt(Q**2 + U**2).real / I
        degree_of_linear_polarization_q[a_idx, w_idx] = -Q.real / I
        degree_of_linear_polarization_u[a_idx, w_idx] = U.real / I
        degree_of_circular_polarization[a_idx, w_idx] = V / I

    @cuda.jit(fastmath=True)
    def compute_field(
        lmax: int,
        idx: np.ndarray,
        size_parameter: np.ndarray,
        sph_h: np.ndarray,
        derivative: np.ndarray,
        e_j_dm_phi: np.ndarray,
        p_lm: np.ndarray,
        pi_lm: np.ndarray,
        tau_lm: np.ndarray,
        e_r: np.ndarray,
        e_theta: np.ndarray,
        e_phi: np.ndarray,
        scattered_field_coefficients: np.ndarray,
        field_real: np.ndarray,
        field_imag: np.ndarray,
    ):  # , initial_field_coefficients: np.ndarray, scatter_to_internal: np.ndarray):
        """
        Compute the field on the GPU using CUDA.

        Args:
            lmax (int): Maximum degree of the spherical harmonics.
            idx (np.ndarray): Array of indices.
            size_parameter (np.ndarray): Array of size parameters.
            sph_h (np.ndarray): Array of spherical harmonics.
            derivative (np.ndarray): Array of derivatives.
            e_j_dm_phi (np.ndarray): Array of phi-dependent terms.
            p_lm (np.ndarray): Array of Legendre polynomials.
            pi_lm (np.ndarray): Array of pi-dependent terms.
            tau_lm (np.ndarray): Array of tau-dependent terms.
            e_r (np.ndarray): Array of r-dependent terms.
            e_theta (np.ndarray): Array of theta-dependent terms.
            e_phi (np.ndarray): Array of phi-dependent terms.
            scattered_field_coefficients (np.ndarray): Array of scattered field coefficients.
            field_real (np.ndarray): Array to store the real part of the field.
            field_imag (np.ndarray): Array to store the imaginary part of the field.
        """
        jmax = sph_h.shape[1] * 2 * lmax * (lmax + 2)
        channels = sph_h.shape[-1]

        sampling_idx, j_idx, w = cuda.grid(3)

        if (sampling_idx >= sph_h.shape[2]) or (j_idx >= jmax) or (w >= channels):
            return

        particle_idx, n, tau, l, m = idx[j_idx, :]

        invariant = (
            1
            / sqrt(2 * (l + 1) * l)
            * e_j_dm_phi[m + 2 * lmax, particle_idx, sampling_idx]
        )

        for c in range(3):
            term = scattered_field_coefficients[particle_idx, n, w] * invariant

            # Calculate M
            if tau == 1:
                c_term_1 = (
                    pi_lm[l, abs(m), particle_idx, sampling_idx]
                    * e_theta[particle_idx, sampling_idx, c]
                    * 1j
                    * m
                )
                c_term_2 = (
                    tau_lm[l, abs(m), particle_idx, sampling_idx]
                    * e_phi[particle_idx, sampling_idx, c]
                )
                c_term = sph_h[l, particle_idx, sampling_idx, w] * (c_term_1 - c_term_2)

                term *= c_term

            # Calculate N
            else:
                p_term = (
                    l
                    * (l + 1)
                    / size_parameter[particle_idx, sampling_idx, w]
                    * sph_h[l, particle_idx, sampling_idx, w]
                )
                p_term *= (
                    p_lm[l, abs(m), particle_idx, sampling_idx]
                    * e_r[particle_idx, sampling_idx, c]
                )

                b_term_1 = (
                    derivative[l, particle_idx, sampling_idx, w]
                    / size_parameter[particle_idx, sampling_idx, w]
                )
                b_term_2 = (
                    tau_lm[l, abs(m), particle_idx, sampling_idx]
                    * e_theta[particle_idx, sampling_idx, c]
                )
                b_term_3 = (
                    pi_lm[l, abs(m), particle_idx, sampling_idx]
                    * e_phi[particle_idx, sampling_idx, c]
                    * 1j
                    * m
                )
                b_term = b_term_1 * (b_term_2 + b_term_3)

                term *= p_term + b_term

            cuda.atomic.add(field_real, (w, sampling_idx, c), term.real)
            cuda.atomic.add(field_imag, (w, sampling_idx, c), term.imag)
