"""2D Mie scattering analytical solution for a dielectric cylinder (TM polarization).

Computes the exact scattered and total fields for a circular dielectric cylinder
under plane-wave illumination  u_inc = exp(i·k0·x).

Line-by-line translation of the validated MATLAB reference (result/mie.py).
"""

import numpy as np
from scipy.special import jv, hankel1
from typing import Optional, Tuple


def mie_scattered_field(
    points: np.ndarray,
    k0: float,
    eps_r: float,
    radius: float,
    cx: float = 0.5,
    cy: float = 0.5,
) -> np.ndarray:
    """Compute the scattered E_z field at arbitrary query points.

    The scattered field is  u_scat = u_total − u_inc, valid both inside and
    outside the cylinder.  This matches the FEM scattered-field formulation.

    Parameters
    ----------
    points : (N, 2) np.ndarray  — (x, y) coordinates
    k0     : float               — vacuum wavenumber
    eps_r  : float               — relative permittivity
    radius : float               — cylinder radius
    cx, cy : float               — cylinder centre

    Returns
    -------
    E_scat : (N,) np.complex128
    """
    m = np.sqrt(eps_r)
    k1 = k0 * m                         # wavenumber inside cylinder
    x_size = k0 * radius                # size parameter

    # ── polar coordinates relative to cylinder centre ──
    dx = points[:, 0] - cx
    dy = points[:, 1] - cy
    R = np.sqrt(dx * dx + dy * dy)
    Phi = np.arctan2(dy, dx)             # [-π, π], matches MATLAB cart2pol

    # ── Wiscombe truncation (matches MATLAB round(…)) ──
    N_trunc = int(np.round(x_size + 4.05 * x_size ** (1.0 / 3.0) + 2))
    N_trunc = max(N_trunc, 3)

    E_scat = np.zeros(len(points), dtype=np.complex128)
    E_int  = np.zeros(len(points), dtype=np.complex128)

    for n in range(-N_trunc, N_trunc + 1):
        # boundary values — matches MATLAB besselj / besselh(…, 1, …)
        J_nx   = jv(n, x_size)
        J_nmx  = jv(n, k1 * radius)
        H_nx   = hankel1(n, x_size)

        # derivatives via recurrence  Z'_n = ½ (Z_{n-1} − Z_{n+1})
        J_nx_p  = 0.5 * (jv(n - 1, x_size) - jv(n + 1, x_size))
        J_nmx_p = 0.5 * (jv(n - 1, k1 * radius) - jv(n + 1, k1 * radius))
        H_nx_p  = 0.5 * (hankel1(n - 1, x_size) - hankel1(n + 1, x_size))

        # TM scattering coefficient  a_n
        num_a = m * J_nx * J_nmx_p - J_nx_p * J_nmx
        den_a = J_nmx * H_nx_p - m * J_nmx_p * H_nx
        a_n = num_a / den_a

        # internal coefficient  c_n
        num_c = J_nx * H_nx_p - J_nx_p * H_nx   # Wronskian (2i/(π x) from theory)
        c_n = num_c / den_a

        # phase factor  iⁿ · exp(i·n·φ)
        phase = (1j) ** n * np.exp(1j * n * Phi)

        # scattered field (valid outside the cylinder)
        out = R >= radius
        if out.any():
            E_scat[out] += a_n * hankel1(n, k0 * R[out]) * phase[out]

        # internal total field (valid inside the cylinder)
        inside = R < radius
        if inside.any():
            E_int[inside] += c_n * jv(n, k1 * R[inside]) * phase[inside]

    # phase reference at cylinder centre  (matches MATLAB phase_shift)
    phase_shift = np.exp(1j * k0 * cx)
    E_scat *= phase_shift
    E_int  *= phase_shift

    # ── scattered field inside cylinder = internal total − incident ──
    E_inc = np.exp(1j * k0 * points[:, 0])
    inside = R < radius
    if inside.any():
        E_scat[inside] = E_int[inside] - E_inc[inside]

    return E_scat


def mie_total_field(
    points: np.ndarray,
    k0: float,
    eps_r: float,
    radius: float,
    cx: float = 0.5,
    cy: float = 0.5,
) -> np.ndarray:
    """Compute the total E_z field.

    Outside:  u_inc + u_scat
    Inside:   internal field (refracted wave)
    """
    m = np.sqrt(eps_r)
    k1 = k0 * m
    x_size = k0 * radius

    dx = points[:, 0] - cx
    dy = points[:, 1] - cy
    R = np.sqrt(dx * dx + dy * dy)
    Phi = np.arctan2(dy, dx)

    N_trunc = int(np.round(x_size + 4.05 * x_size ** (1.0 / 3.0) + 2))
    N_trunc = max(N_trunc, 3)

    E_scat = np.zeros(len(points), dtype=np.complex128)
    E_int  = np.zeros(len(points), dtype=np.complex128)

    for n in range(-N_trunc, N_trunc + 1):
        J_nx   = jv(n, x_size)
        J_nmx  = jv(n, k1 * radius)
        H_nx   = hankel1(n, x_size)

        J_nx_p  = 0.5 * (jv(n - 1, x_size) - jv(n + 1, x_size))
        J_nmx_p = 0.5 * (jv(n - 1, k1 * radius) - jv(n + 1, k1 * radius))
        H_nx_p  = 0.5 * (hankel1(n - 1, x_size) - hankel1(n + 1, x_size))

        num_a = m * J_nx * J_nmx_p - J_nx_p * J_nmx
        den_a = J_nmx * H_nx_p - m * J_nmx_p * H_nx
        a_n = num_a / den_a

        num_c = J_nx * H_nx_p - J_nx_p * H_nx
        c_n = num_c / den_a

        phase = (1j) ** n * np.exp(1j * n * Phi)

        out = R >= radius
        if out.any():
            E_scat[out] += a_n * hankel1(n, k0 * R[out]) * phase[out]

        inside = R < radius
        if inside.any():
            E_int[inside] += c_n * jv(n, k1 * R[inside]) * phase[inside]

    phase_shift = np.exp(1j * k0 * cx)
    E_scat *= phase_shift
    E_int  *= phase_shift

    E_inc = np.exp(1j * k0 * points[:, 0])

    E_total = np.zeros(len(points), dtype=np.complex128)
    E_total[R >= radius] = E_inc[R >= radius] + E_scat[R >= radius]
    E_total[R < radius]  = E_int[R < radius]

    return E_total


def mie_grid_solution(
    k0: float,
    eps_r: float,
    radius: float,
    cx: float = 0.5,
    cy: float = 0.5,
    x_range: Tuple[float, float] = (0.0, 1.0),
    y_range: Tuple[float, float] = (0.0, 1.0),
    Nx: int = 400,
    Ny: int = 400,
) -> dict:
    """Compute Mie solution on a regular grid (for plotting / visual checks).

    Returns a dict with keys: X, Y, R, Phi, E_inc, E_scat, E_total.
    """
    x_vec = np.linspace(x_range[0], x_range[1], Nx)
    y_vec = np.linspace(y_range[0], y_range[1], Ny)
    X, Y = np.meshgrid(x_vec, y_vec)

    points = np.column_stack([X.ravel(), Y.ravel()])
    dx = points[:, 0] - cx
    dy = points[:, 1] - cy
    R = np.sqrt(dx * dx + dy * dy).reshape(Ny, Nx)
    Phi = np.arctan2(dy, dx).reshape(Ny, Nx)

    E_inc   = np.exp(1j * k0 * X)
    E_scat  = mie_scattered_field(points, k0, eps_r, radius, cx, cy).reshape(Ny, Nx)
    E_total = mie_total_field(points, k0, eps_r, radius, cx, cy).reshape(Ny, Nx)

    return {
        "X": X, "Y": Y, "R": R, "Phi": Phi,
        "E_inc": E_inc, "E_scat": E_scat, "E_total": E_total,
    }
