"""Base class for shared forward model components."""

import jax.numpy as jnp
from jax import vmap


class BeamForwardModel:
    """Base class containing shared forward model components for beam simulation."""

    @staticmethod
    def triangle_wave(f: jnp.ndarray | float, t: jnp.ndarray | float, phi: float = 0.0):
        """
        Triangle wave function for rastering pattern.

        Parameters
        ----------
        f : float or array
            Frequency in Hz
        t : float or array
            Time
        phi : float
            Phase offset in radians

        Returns
        -------
        float or array
            Triangle wave value in [-1, 1]
        """
        phase = (f * t - 0.25 + phi / (2 * jnp.pi)) % 1.0
        return 4.0 * jnp.abs(phase - 0.5) - 1.0

    def raster_position(
        self,
        t: jnp.ndarray | float,
        ax: float,
        ay: float,
        fx: float,
        fy: float,
        cx: float,
        cy: float,
        phix: float,
        phiy: float,
    ):
        """
        Compute beam raster position at time t.

        Parameters
        ----------
        t : float or array
            Time
        ax, ay : float
            Rastering amplitudes (horizontal, vertical)
        fx, fy : float
            Rastering frequencies in kHz
        cx, cy : float
            Center offsets
        phix, phiy : float
            Phase offsets in radians

        Returns
        -------
        gx, gy : float or array
            Beam center position (x, y)
        """
        fx_hz = fx * 1000.0
        fy_hz = fy * 1000.0
        x_pos = ax * self.triangle_wave(fx_hz, t, phix) + cx
        y_pos = ay * self.triangle_wave(fy_hz, t, phiy) + cy
        return x_pos, y_pos

    def beam_intensity_single_t(
        self, x: jnp.ndarray, y: jnp.ndarray, t: float, k: jnp.ndarray
    ):
        """
        Compute beam intensity at a single time point.

        Parameters
        ----------
        x, y : jnp.ndarray
            2D coordinate grids
        t : float
            Time point
        k : jnp.ndarray
            Parameter vector [Ax, Ay, sigx, sigy, cx, cy, fx, fy, phix, phiy]

        Returns
        -------
        jnp.ndarray
            2D intensity image at time t
        """
        ax, ay, sigx, sigy, cx, cy, fx, fy, phix, phiy = k
        gx, gy = self.raster_position(t, ax, ay, fx, fy, cx, cy, phix, phiy)

        normalizing_const = 1.0 / (2.0 * jnp.pi * sigx * sigy)
        exponent = -0.5 * (((x - gx) / sigx) ** 2 + ((y - gy) / sigy) ** 2)
        return normalizing_const * jnp.exp(exponent)

    def beam_intensity_batch(
        self, x: jnp.ndarray, y: jnp.ndarray, t_vals: jnp.ndarray, k: jnp.ndarray
    ):
        """
        Compute beam intensity for all time points (vectorized).

        Parameters
        ----------
        x, y : jnp.ndarray
            2D coordinate grids
        t_vals : jnp.ndarray
            1D array of time points
        k : jnp.ndarray
            Parameter vector

        Returns
        -------
        jnp.ndarray
            3D array of intensity images (time, y, x)
        """
        batched_fn = vmap(self.beam_intensity_single_t, in_axes=(None, None, 0, None))
        return batched_fn(x, y, t_vals, k)

    def simulate_image(
        self, x: jnp.ndarray, y: jnp.ndarray, t_vals: jnp.ndarray, k: jnp.ndarray
    ):
        """
        Simulate accumulated beam image over time.

        Parameters
        ----------
        x, y : jnp.ndarray
            2D coordinate grids
        t_vals : jnp.ndarray
            1D array of time points
        k : jnp.ndarray
            Parameter vector [Ax, Ay, sigx, sigy, cx, cy, fx, fy, phix, phiy]

        Returns
        -------
        jnp.ndarray
            Normalized 2D intensity image
        """
        dt = t_vals[1] - t_vals[0]
        i_stack = self.beam_intensity_batch(x, y, t_vals, k)
        i_accumulated = jnp.sum(i_stack, axis=0) * dt
        return i_accumulated / jnp.max(i_accumulated)
