# FMMAX
# Copyright (C) 2025 Martin F. Schubert

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program 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 Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""Functions related to eigensolve with translated permittivity distributions.

Copyright (c) Martin F. Schubert
"""

import functools

import jax.numpy as jnp

from fmmax import basis, fft, fmm, fmm_matrices


def translate_layer_solve_result(
    solve_result: fmm.LayerSolveResult,
    dx: jnp.ndarray,
    dy: jnp.ndarray,
) -> fmm.LayerSolveResult:
    """Obtain solve result for layer with translated permittivity distribution.

    Rather than performing eigensolve on translated permittivity distribution, this
    function allows one to translate the solve result. This can result in significant
    compute savings for structures such as slanted gratings.

    Args:
        solve_result: The ``LayerSolveResult`` for the permittivity distribution prior
            to translation.
        dx: The translation of the permittivity distribution along the ``x`` direction.
        dy: The translation of the permittivity distribution along the ``y`` direction.

    Returns:
        The ``LayerSolveResult`` for the layer with translated permittivity.
    """
    assert dx.shape == dy.shape == ()

    _shift = functools.partial(
        _apply_shift_eigenvectors,
        expansion=solve_result.expansion,
        primitive_lattice_vectors=solve_result.primitive_lattice_vectors,
        dx=dx,
        dy=dy,
    )

    ev_te, ev_tm = jnp.split(solve_result.eigenvectors, 2, axis=-2)
    shifted_eigenvectors = jnp.concatenate([_shift(ev_te), _shift(ev_tm)], axis=-2)

    _shift_toeplitz = functools.partial(
        _apply_shift_toeplitz,
        expansion=solve_result.expansion,
        primitive_lattice_vectors=solve_result.primitive_lattice_vectors,
        dx=dx,
        dy=dy,
    )

    def _shift_block_toeplitz(matrix: jnp.ndarray) -> jnp.ndarray:
        m0, m1 = jnp.split(matrix, 2, axis=-2)
        m00, m01 = jnp.split(m0, 2, axis=-1)
        m10, m11 = jnp.split(m1, 2, axis=-1)
        return jnp.block(
            [
                [_shift_toeplitz(m00), _shift_toeplitz(m01)],
                [_shift_toeplitz(m10), _shift_toeplitz(m11)],
            ],
        )

    shifted_transverse_permittivity_matrix = _shift_block_toeplitz(
        solve_result.transverse_permittivity_matrix
    )
    shifted_transverse_permeability_matrix = _shift_block_toeplitz(
        solve_result.transverse_permeability_matrix
    )

    shifted_z_permittivity_matrix = _shift_toeplitz(solve_result.z_permittivity_matrix)

    omega_script_k_matrix = fmm_matrices.omega_script_k_matrix_patterned(
        wavelength=solve_result.wavelength,
        z_permittivity_matrix=shifted_z_permittivity_matrix,
        transverse_permeability_matrix=shifted_transverse_permeability_matrix,
        transverse_wavevectors=basis.transverse_wavevectors(
            in_plane_wavevector=solve_result.in_plane_wavevector,
            primitive_lattice_vectors=solve_result.primitive_lattice_vectors,
            expansion=solve_result.expansion,
        ),
    ).astype(solve_result.eigenvectors.dtype)

    return fmm.LayerSolveResult(
        wavelength=solve_result.wavelength,
        in_plane_wavevector=solve_result.in_plane_wavevector,
        primitive_lattice_vectors=solve_result.primitive_lattice_vectors,
        expansion=solve_result.expansion,
        eigenvalues=solve_result.eigenvalues,
        eigenvectors=shifted_eigenvectors,
        omega_script_k_matrix=omega_script_k_matrix,
        z_permittivity_matrix=shifted_z_permittivity_matrix,
        inverse_z_permittivity_matrix=_shift_toeplitz(
            solve_result.inverse_z_permittivity_matrix
        ),
        transverse_permittivity_matrix=shifted_transverse_permittivity_matrix,
        z_permeability_matrix=_shift_toeplitz(solve_result.z_permeability_matrix),
        inverse_z_permeability_matrix=_shift_toeplitz(
            solve_result.inverse_z_permeability_matrix
        ),
        transverse_permeability_matrix=shifted_transverse_permeability_matrix,
        tangent_vector_field=solve_result.tangent_vector_field,
    )


def _apply_shift_eigenvectors(
    eigenvectors: jnp.ndarray,
    expansion: basis.Expansion,
    primitive_lattice_vectors: basis.LatticeVectors,
    dx: jnp.ndarray,
    dy: jnp.ndarray,
) -> jnp.ndarray:
    """Returns eigenvectors that result from translating a structure."""
    transverse_wavevectors = basis.transverse_wavevectors(
        in_plane_wavevector=jnp.zeros((2,)),
        primitive_lattice_vectors=primitive_lattice_vectors,
        expansion=expansion,
    )
    kx = transverse_wavevectors[:, 0]
    ky = transverse_wavevectors[:, 1]
    return eigenvectors * jnp.exp(-1j * (kx * dx + ky * dy))[..., jnp.newaxis]


def _apply_shift_toeplitz(
    mat: jnp.ndarray,
    expansion: basis.Expansion,
    primitive_lattice_vectors: basis.LatticeVectors,
    dx: jnp.ndarray,
    dy: jnp.ndarray,
) -> jnp.ndarray:
    """Return toeplitz matrix that results from translating a structure."""
    idx = fft._standard_toeplitz_indices(expansion)

    reciprocal_vectors = primitive_lattice_vectors.reciprocal
    u = reciprocal_vectors.u
    v = reciprocal_vectors.v
    kx = 2 * jnp.pi * (idx[:, :, 0] * u[0] + idx[:, :, 1] * v[0])
    ky = 2 * jnp.pi * (idx[:, :, 0] * u[1] + idx[:, :, 1] * v[1])
    return mat * jnp.exp(-1j * (kx * dx + ky * dy))
