# Copyright (C) 2015-2023: The University of Edinburgh, United Kingdom
#                 Authors: Craig Warren, Antonis Giannopoulos, and John Hartley
#
# This file is part of gprMax.
#
# gprMax is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# gprMax is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with gprMax.  If not, see <http://www.gnu.org/licenses/>.

import numpy as np
cimport numpy as np

from cython.parallel import prange


# Use C-functions from 'complex.h' but doesn't work for Windows as it doesn't
#  support 'double complex' but instead defines its own type '_Dcomplex'.
#  https://docs.microsoft.com/en-us/cpp/c-runtime-library/complex-math-support?view=vs-2019
#  https://stackoverflow.com/questions/57837255/defining-dcomplex-externally-in-cython?rq=1
{% if not functions[0].iswin %}
cdef extern from "complex.h" nogil:
    double creal(double complex z)
    float crealf(float complex z)
{% endif %}


###############################################################
# Electric field updates - dispersive materials - multipole A #
###############################################################

{% for item in functions %}
cpdef void {{ item.name_a }}(
    int nx,
    int ny,
    int nz,
    int nthreads,
    int maxpoles,
    {{ item.field_type }}[:, ::1] updatecoeffsE,
    {{ item.dispersive_type }}[:, ::1] updatecoeffsdispersive,
    np.uint32_t[:, :, :, ::1] ID,
    {{ item.dispersive_type }}[:, :, :, ::1] Tx,
    {{ item.dispersive_type }}[:, :, :, ::1] Ty,
    {{ item.dispersive_type }}[:, :, :, ::1] Tz,
    {{ item.field_type }}[:, :, ::1] Ex,
    {{ item.field_type }}[:, :, ::1] Ey,
    {{ item.field_type }}[:, :, ::1] Ez,
    {{ item.field_type }}[:, :, ::1] Hx,
    {{ item.field_type }}[:, :, ::1] Hy,
    {{ item.field_type }}[:, :, ::1] Hz
):
    """Updates the electric field components when dispersive materials
        (with multiple poles) are present.

    Args:
        nx, ny, nz: int for grid size in cells.
        nthreads: int for number of threads to use.
        maxpoles: int for maximum number of poles.
        updatecoeffs, T, ID, E, H: memoryviews to access to update coeffients,
                                    temporary, ID and field component arrays.
    """

    cdef Py_ssize_t i, j, k, pole
    cdef int material
    cdef float phi = 0

    # Ex component
    if ny != 1 or nz != 1:
        for i in prange(0, nx, nogil=True, schedule='static', num_threads=nthreads):
            for j in range(1, ny):
                for k in range(1, nz):
                    material = ID[0, i, j, k]
                    phi = 0
                    for pole in range(maxpoles):
                        {% if 'complex' in item.dispersive_type %}
                        {% if item.iswin %}
                        phi = (phi + updatecoeffsdispersive[material, pole * 3].real
                               * Tx[pole, i, j, k].real)
                        {% else %}
                        phi = (phi + {{ item.real_part }}(updatecoeffsdispersive[material, pole * 3])
                               * {{ item.real_part }}(Tx[pole, i, j, k]))
                        {% endif %}
                        {% else %}
                        phi = phi + updatecoeffsdispersive[material, pole * 3] * Tx[pole, i, j, k]
                        {% endif %}
                        Tx[pole, i, j, k] = (updatecoeffsdispersive[material, 1 + (pole * 3)]
                                             * Tx[pole, i, j, k] + updatecoeffsdispersive[material, 2 + (pole * 3)]
                                             * Ex[i, j, k])
                    Ex[i, j, k] = (updatecoeffsE[material, 0] * Ex[i, j, k] + updatecoeffsE[material, 2]
                                   * (Hz[i, j, k] - Hz[i, j - 1, k]) - updatecoeffsE[material, 3]
                                   * (Hy[i, j, k] - Hy[i, j, k - 1]) - updatecoeffsE[material, 4] * phi)

    # Ey component
    if nx != 1 or nz != 1:
        for i in prange(1, nx, nogil=True, schedule='static', num_threads=nthreads):
            for j in range(0, ny):
                for k in range(1, nz):
                    material = ID[1, i, j, k]
                    phi = 0
                    for pole in range(maxpoles):
                        {% if 'complex' in item.dispersive_type %}
                        {% if item.iswin %}
                        phi = (phi + updatecoeffsdispersive[material, pole * 3].real
                               * Ty[pole, i, j, k].real)
                        {% else %}
                        phi = (phi + {{ item.real_part }}(updatecoeffsdispersive[material, pole * 3])
                               * {{ item.real_part }}(Ty[pole, i, j, k]))
                        {% endif %}
                        {% else %}
                        phi = phi + updatecoeffsdispersive[material, pole * 3] * Ty[pole, i, j, k]
                        {% endif %}
                        Ty[pole, i, j, k] = (updatecoeffsdispersive[material, 1 + (pole * 3)]
                                             * Ty[pole, i, j, k] + updatecoeffsdispersive[material, 2 + (pole * 3)]
                                             * Ey[i, j, k])
                    Ey[i, j, k] = (updatecoeffsE[material, 0] * Ey[i, j, k] + updatecoeffsE[material, 3]
                                   * (Hx[i, j, k] - Hx[i, j, k - 1]) - updatecoeffsE[material, 1]
                                   * (Hz[i, j, k] - Hz[i - 1, j, k]) - updatecoeffsE[material, 4] * phi)

    # Ez component
    if nx != 1 or ny != 1:
        for i in prange(1, nx, nogil=True, schedule='static', num_threads=nthreads):
            for j in range(1, ny):
                for k in range(0, nz):
                    material = ID[2, i, j, k]
                    phi = 0
                    for pole in range(maxpoles):
                        {% if 'complex' in item.dispersive_type %}
                        {% if item.iswin %}
                        phi = (phi + updatecoeffsdispersive[material, pole * 3].real
                               * Tz[pole, i, j, k].real)
                        {% else %}
                        phi = (phi + {{ item.real_part }}(updatecoeffsdispersive[material, pole * 3])
                               * {{ item.real_part }}(Tz[pole, i, j, k]))
                        {% endif %}
                        {% else %}
                        phi = phi + updatecoeffsdispersive[material, pole * 3] * Tz[pole, i, j, k]
                        {% endif %}
                        Tz[pole, i, j, k] = (updatecoeffsdispersive[material, 1 + (pole * 3)]
                                             * Tz[pole, i, j, k] + updatecoeffsdispersive[material, 2 + (pole * 3)]
                                             * Ez[i, j, k])
                    Ez[i, j, k] = (updatecoeffsE[material, 0] * Ez[i, j, k] + updatecoeffsE[material, 1]
                                   * (Hy[i, j, k] - Hy[i - 1, j, k]) - updatecoeffsE[material, 2]
                                   * (Hx[i, j, k] - Hx[i, j - 1, k]) - updatecoeffsE[material, 4] * phi)
{% endfor %}


###############################################################
# Electric field updates - dispersive materials - multipole B #
###############################################################

{% for item in functions %}
cpdef void {{ item.name_b }}(
    int nx,
    int ny,
    int nz,
    int nthreads,
    int maxpoles,
    {{ item.dispersive_type }}[:, ::1] updatecoeffsdispersive,
    np.uint32_t[:, :, :, ::1] ID,
    {{ item.dispersive_type }}[:, :, :, ::1] Tx,
    {{ item.dispersive_type }}[:, :, :, ::1] Ty,
    {{ item.dispersive_type }}[:, :, :, ::1] Tz,
    {{ item.field_type }}[:, :, ::1] Ex,
    {{ item.field_type }}[:, :, ::1] Ey,
    {{ item.field_type }}[:, :, ::1] Ez
):
    """Updates a temporary dispersive material array when disperisive materials
        (with multiple poles) are present.

    Args:
        nx, ny, nz: int for grid size in cells.
        nthreads: int for number of threads to use.
        maxpoles: int for maximum number of poles.
        updatecoeffs, T, ID, E, H: memoryviews to access to update coeffients,
                                    temporary, ID and field component arrays.
    """

    cdef Py_ssize_t i, j, k, pole
    cdef int material

    # Ex component
    if ny != 1 or nz != 1:
        for i in prange(0, nx, nogil=True, schedule='static', num_threads=nthreads):
            for j in range(1, ny):
                for k in range(1, nz):
                    material = ID[0, i, j, k]
                    for pole in range(maxpoles):
                        Tx[pole, i, j, k] = (Tx[pole, i, j, k]
                                             - updatecoeffsdispersive[material, 2 + (pole * 3)]
                                             * Ex[i, j, k])

    # Ey component
    if nx != 1 or nz != 1:
        for i in prange(1, nx, nogil=True, schedule='static', num_threads=nthreads):
            for j in range(0, ny):
                for k in range(1, nz):
                    material = ID[1, i, j, k]
                    for pole in range(maxpoles):
                        Ty[pole, i, j, k] = (Ty[pole, i, j, k]
                                             - updatecoeffsdispersive[material, 2 + (pole * 3)]
                                             * Ey[i, j, k])

    # Ez component
    if nx != 1 or ny != 1:
        for i in prange(1, nx, nogil=True, schedule='static', num_threads=nthreads):
            for j in range(1, ny):
                for k in range(0, nz):
                    material = ID[2, i, j, k]
                    for pole in range(maxpoles):
                        Tz[pole, i, j, k] = (Tz[pole, i, j, k]
                                             - updatecoeffsdispersive[material, 2 + (pole * 3)]
                                             * Ez[i, j, k])

{% endfor %}


#################################################################
# Electric field updates - dispersive materials - single pole A #
#################################################################

{% for item in functions %}
cpdef void {{ item.name_a_1 }}(
    int nx,
    int ny,
    int nz,
    int nthreads,
    int maxpoles,
    {{ item.field_type }}[:, ::1] updatecoeffsE,
    {{ item.dispersive_type }}[:, ::1] updatecoeffsdispersive,
    np.uint32_t[:, :, :, ::1] ID,
    {{ item.dispersive_type }}[:, :, :, ::1] Tx,
    {{ item.dispersive_type }}[:, :, :, ::1] Ty,
    {{ item.dispersive_type }}[:, :, :, ::1] Tz,
    {{ item.field_type }}[:, :, ::1] Ex,
    {{ item.field_type }}[:, :, ::1] Ey,
    {{ item.field_type }}[:, :, ::1] Ez,
    {{ item.field_type }}[:, :, ::1] Hx,
    {{ item.field_type }}[:, :, ::1] Hy,
    {{ item.field_type }}[:, :, ::1] Hz
):
    """Updates the electric field components when dispersive materials
        (with 1 pole) are present.

    Args:
        nx, ny, nz: int for grid size in cells.
        nthreads: int for number of threads to use.
        maxpoles: int for maximum number of poles.
        updatecoeffs, T, ID, E, H: memoryviews to access to update coeffients,
                                    temporary, ID and field component arrays.
    """

    cdef Py_ssize_t i, j, k
    cdef int material
    cdef float phi = 0

    # Ex component
    if ny != 1 or nz != 1:
        for i in prange(0, nx, nogil=True, schedule='static', num_threads=nthreads):
            for j in range(1, ny):
                for k in range(1, nz):
                    material = ID[0, i, j, k]
                    {% if 'complex' in item.dispersive_type %}
                    {% if item.iswin %}
                    phi = (updatecoeffsdispersive[material, 0].real
                           * Tx[0, i, j, k].real)
                    {% else %}
                    phi = ({{ item.real_part }}(updatecoeffsdispersive[material, 0])
                           * {{ item.real_part }}(Tx[0, i, j, k]))
                    {% endif %}
                    {% else %}
                    phi = updatecoeffsdispersive[material, 0] * Tx[0, i, j, k]
                    {% endif %}
                    Tx[0, i, j, k] = (updatecoeffsdispersive[material, 1] * Tx[0, i, j, k]
                                      + updatecoeffsdispersive[material, 2] * Ex[i, j, k])
                    Ex[i, j, k] = (updatecoeffsE[material, 0] * Ex[i, j, k] + updatecoeffsE[material, 2]
                                   * (Hz[i, j, k] - Hz[i, j - 1, k]) - updatecoeffsE[material, 3]
                                   * (Hy[i, j, k] - Hy[i, j, k - 1]) - updatecoeffsE[material, 4] * phi)

    # Ey component
    if nx != 1 or nz != 1:
        for i in prange(1, nx, nogil=True, schedule='static', num_threads=nthreads):
            for j in range(0, ny):
                for k in range(1, nz):
                    material = ID[1, i, j, k]
                    {% if 'complex' in item.dispersive_type %}
                    {% if item.iswin %}
                    phi = (updatecoeffsdispersive[material, 0].real
                           * Ty[0, i, j, k].real)
                    {% else %}
                    phi = ({{ item.real_part }}(updatecoeffsdispersive[material, 0])
                           * {{ item.real_part }}(Ty[0, i, j, k]))
                    {% endif %}
                    {% else %}
                    phi = updatecoeffsdispersive[material, 0] * Ty[0, i, j, k]
                    {% endif %}
                    Ty[0, i, j, k] = (updatecoeffsdispersive[material, 1] * Ty[0, i, j, k]
                                      + updatecoeffsdispersive[material, 2] * Ey[i, j, k])
                    Ey[i, j, k] = (updatecoeffsE[material, 0] * Ey[i, j, k] + updatecoeffsE[material, 3]
                                   * (Hx[i, j, k] - Hx[i, j, k - 1]) - updatecoeffsE[material, 1]
                                   * (Hz[i, j, k] - Hz[i - 1, j, k]) - updatecoeffsE[material, 4] * phi)

    # Ez component
    if nx != 1 or ny != 1:
        for i in prange(1, nx, nogil=True, schedule='static', num_threads=nthreads):
            for j in range(1, ny):
                for k in range(0, nz):
                    material = ID[2, i, j, k]
                    {% if 'complex' in item.dispersive_type %}
                    {% if item.iswin %}
                    phi = (updatecoeffsdispersive[material, 0].real
                           * Tz[0, i, j, k].real)
                    {% else %}
                    phi = ({{ item.real_part }}(updatecoeffsdispersive[material, 0])
                           * {{ item.real_part }}(Tz[0, i, j, k]))
                    {% endif %}
                    {% else %}
                    phi = updatecoeffsdispersive[material, 0] * Tz[0, i, j, k]
                    {% endif %}
                    Tz[0, i, j, k] = (updatecoeffsdispersive[material, 1] * Tz[0, i, j, k]
                                      + updatecoeffsdispersive[material, 2] * Ez[i, j, k])
                    Ez[i, j, k] = (updatecoeffsE[material, 0] * Ez[i, j, k] + updatecoeffsE[material, 1]
                                   * (Hy[i, j, k] - Hy[i - 1, j, k]) - updatecoeffsE[material, 2]
                                   * (Hx[i, j, k] - Hx[i, j - 1, k]) - updatecoeffsE[material, 4] * phi)

{% endfor %}


#################################################################
# Electric field updates - dispersive materials - single pole B #
#################################################################

{% for item in functions %}
cpdef void {{ item.name_b_1 }}(
    int nx,
    int ny,
    int nz,
    int nthreads,
    int maxpoles,
    {{ item.dispersive_type }}[:, ::1] updatecoeffsdispersive,
    np.uint32_t[:, :, :, ::1] ID,
    {{ item.dispersive_type }}[:, :, :, ::1] Tx,
    {{ item.dispersive_type }}[:, :, :, ::1] Ty,
    {{ item.dispersive_type }}[:, :, :, ::1] Tz,
    {{ item.field_type }}[:, :, ::1] Ex,
    {{ item.field_type }}[:, :, ::1] Ey,
    {{ item.field_type }}[:, :, ::1] Ez
):
    """Updates a temporary dispersive material array when disperisive materials
        (with 1 pole) are present.

    Args:
        nx, ny, nz: int for grid size in cells.
        nthreads: int for number of threads to use.
        maxpoles: int for maximum number of poles.
        updatecoeffs, T, ID, E, H: memoryviews to access to update coeffients,
                                    temporary, ID and field component arrays.
    """

    cdef Py_ssize_t i, j, k
    cdef int material

    # Ex component
    if ny != 1 or nz != 1:
        for i in prange(0, nx, nogil=True, schedule='static', num_threads=nthreads):
            for j in range(1, ny):
                for k in range(1, nz):
                    material = ID[0, i, j, k]
                    Tx[0, i, j, k] = Tx[0, i, j, k] - updatecoeffsdispersive[material, 2] * Ex[i, j, k]

    # Ey component
    if nx != 1 or nz != 1:
        for i in prange(1, nx, nogil=True, schedule='static', num_threads=nthreads):
            for j in range(0, ny):
                for k in range(1, nz):
                    material = ID[1, i, j, k]
                    Ty[0, i, j, k] = Ty[0, i, j, k] - updatecoeffsdispersive[material, 2] * Ey[i, j, k]

    # Ez component
    if nx != 1 or ny != 1:
        for i in prange(1, nx, nogil=True, schedule='static', num_threads=nthreads):
            for j in range(1, ny):
                for k in range(0, nz):
                    material = ID[2, i, j, k]
                    Tz[0, i, j, k] = Tz[0, i, j, k] - updatecoeffsdispersive[material, 2] * Ez[i, j, k]
{% endfor %}