"""
Discrete Hankel transform tools for cylindrical / axial simulations.
"""

from __future__ import annotations

import numpy as np
from scipy.special import jn_zeros, jv


class HankelTransform:
    """
    Discrete Hankel Transform object.

    Precomputes transform matrix and grids for repeated use.

    Parameters
    ----------
    n_points : int
        Number of radial grid points.
    order : int, optional
        Order of Bessel function (default = 0).
    r_max : float, optional
        Maximum radius. If provided, physical grids are computed.
    """

    def __init__(self, n_points: int, order: int = 0, r_max: float | None = None):
        self.n = n_points
        self.order = order
        self.r_max = r_max

        # Bessel zeros
        self.zeros = jn_zeros(order, n_points)
        self.factor = jn_zeros(order, n_points + 1)[-1]

        # Construct kernel
        grid = np.outer(self.zeros, self.zeros) / self.factor
        kernel = jv(order, grid)

        weight = 1.0 / (jv(order + 1, self.zeros) ** 2)

        # Hankel matrix
        self.H = kernel * weight

        # Normalization factors
        self.forward_prefactor = 2.0 / self.factor**2
        self.backward_prefactor = 2.0

        # Physical grids (if r_max is provided)
        if r_max is not None:
            self.r = self.zeros * (r_max / self.factor)
            self.k = self.zeros / r_max

            Jp1 = jv(order + 1, self.zeros)

            self.measure = (self.r_max / self.factor)**2 * 2.0 / (Jp1**2)
            self.measure_k = self.r_max**2 * 2.0 / (Jp1**2)


        else:
            self.r = None
            self.k = None
            self.measure = None
            self.measure_k = None

    def forward(self, array: np.ndarray, axis: int = -1) -> np.ndarray:
        """
        Forward discrete Hankel transform.
        """
        arr = np.moveaxis(array, axis, 0)

        reshaped = arr.reshape(self.n, -1)
        transformed = self.H @ reshaped
        transformed = transformed.reshape(arr.shape)

        transformed = np.moveaxis(transformed, 0, axis)

        return self.forward_prefactor * transformed

    def backward(self, array: np.ndarray, axis: int = -1) -> np.ndarray:
        """
        Inverse discrete Hankel transform.
        """
        arr = np.moveaxis(array, axis, 0)

        reshaped = arr.reshape(self.n, -1)
        transformed = self.H @ reshaped
        transformed = transformed.reshape(arr.shape)

        transformed = np.moveaxis(transformed, 0, axis)

        return self.backward_prefactor * transformed

