external_fields.jl

"""
Gaussian laser pulse electromagnetic field.

Represents a focused Gaussian beam with a temporal envelope.
The beam propagates along the z-direction with waist w₀ at z=0.

Parameters:
- λ: wavelength
- a₀: normalized vector potential (a0 kwarg)
- w₀: beam waist (defaults to 75λ)
- T0: pulse duration parameter
- τ0: temporal width parameter
"""
@component function GaussLaser(; name, wavelength=nothing, frequency=nothing, a0=10.0, beam_waist=nothing, polarization = :linear, ref_frame)
    if wavelength === nothing && frequency === nothing
        wavelength = 1.0
    end
    if wavelength !== nothing && frequency !== nothing
        error("Specify either wavelength or frequency, not both")
    end

    # New interface with spacetime
    @named field_dynamics = EMFieldDynamics(; ref_frame)

    # Get spacetime variables and constants from parent scope
    @unpack c, m_e, q_e = ref_frame
    iv = ModelingToolkit.get_iv(ref_frame)

    # Create local position and time variables
    @variables x(iv)[1:4] t(iv)

    @unpack E, B = field_dynamics

    @parameters begin
        λ, [guess = 1.0]
        a₀ = a0
        ω, [guess = 1.0]
        k
        E₀
        w₀
        z_R
        T0 = 100
        τ0
    end

    # Fixed parameters
    if polarization == :linear
        ξx = 1.0 + 0im
        ξy = 0 + 0im
    elseif polarization == :circular
        ξx, ξy = (1/√2, im/√2) .|> complex
    else
        error("polarization $polarization not supported.")
    end

    ϕ₀ = 0
    t₀ = 5T0
    z₀ = 0

    @variables wz(iv) z(iv) r(iv)

    eqs = [
        # Extract z-coordinate from 4-position
        z ~ x[4]
        # Beam width as function of z
        wz ~ w₀ * √(1 + (z / z_R)^2)
        # Radial distance from beam axis
        r ~ hypot(x[2], x[3])

        # Electric field components
        E[1] ~ real(
            ξx * E₀ * w₀ / wz *
            exp(
                -(r / wz)^2 + im * (-(r^2 * z) / (z_R * wz^2) + atan(z, z_R) - k * z + ϕ₀),
            ) *
            exp(im * ω * t) *
            exp(-(((t - t₀) - (z - z₀) / c) / τ0)^2),
        )
        E[2] ~ real(
            ξy * E₀ * w₀ / wz *
            exp(
                -(r / wz)^2 + im * (-(r^2 * z) / (z_R * wz^2) + atan(z, z_R) - k * z + ϕ₀),
            ) *
            exp(im * ω * t) *
            exp(-(((t - t₀) - (z - z₀) / c) / τ0)^2),
        )
        E[3] ~ real(
            2im / (k * wz^2) *
            (1 + im * (z / z_R)) *
            (x[2] * E[1] + x[3] * E[2]) *
            exp(im * ω * t) *
            exp(-(((t - t₀) - (z - z₀) / c) / τ0)^2),
        )

        # Magnetic field components
        B[1] ~ -E[2] / c
        B[2] ~ E[1] / c
        B[3] ~ real(
            2im / (k * c * wz^2) *
            (1 + im * (z / z_R)) *
            (x[3] * E[1] - x[2] * E[2]) *
            exp(im * ω * t) *
            exp(-(((t - t₀) - (z - z₀) / c) / τ0)^2),
        )
    ]

    initialization_eqs = [
        ω ~ 2π * c / λ
        E₀ ~ a₀ * m_e * c * ω / abs(q_e)
        z_R ~ w₀^2 * k / 2
        k ~ 2π / λ
    ]
    bindings = [
        ω => missing
        λ => missing
        k => missing
        E₀ => missing
        z_R => missing
        w₀ => missing
    ]

    initial_conditions = Pair{SymbolicT,Any}[τ0 => 10 / ω]
    push!(initial_conditions, w₀ => (beam_waist === nothing ? 75λ : beam_waist))
    if wavelength !== nothing
        push!(initial_conditions, λ => wavelength)
    end
    if frequency !== nothing
        push!(initial_conditions, ω => frequency)
    end

    sys = System(eqs, iv, [x, t, wz, z, r], [λ, a₀, ω, k, E₀, w₀, z_R, T0, τ0];
        name,
        systems=[ref_frame],
        initial_conditions,
        initialization_eqs,
        bindings
    )

    extend(sys, field_dynamics)
end


"""
Plane wave electromagnetic field.

For a plane wave with normalized vector potential a₀ = eA/(mc²),
electrons can exhibit figure-8 motion when a₀ ~ 1.

Reference: Sarachik & Schappert, Phys. Rev. D 1, 2738 (1970)
"""
@component function PlaneWave(; name, amplitude=1.0, wavelength=nothing, frequency=nothing, k_direction=[0,0,1], polarization=[1,0,0], ref_frame)
    if wavelength === nothing && frequency === nothing
        frequency = 1.0
    end
    if wavelength !== nothing && frequency !== nothing
        error("Specify either wavelength or frequency, not both")
    end

    # New interface with spacetime
    @named field_dynamics = EMFieldDynamics(; ref_frame)

    # Get spacetime variables and constants from parent scope
    @unpack c, m_e, q_e = ref_frame
    iv = ModelingToolkit.get_iv(ref_frame)

    # Create local position and time variables
    @variables x(iv)[1:4]
    if nameof(iv) == :τ
        @variables t(iv)
    else
        t = iv
    end

    @unpack E, B = field_dynamics

    @parameters begin
        A = amplitude
        ω, [guess = 1.0]
        k_dir[1:3] = k_direction
        pol[1:3] = polarization
        λ, [guess = 1.0]
    end

    # Normalize k direction to get unit vector k̂
    k_norm = sqrt(k_dir[1]^2 + k_dir[2]^2 + k_dir[3]^2)
    k̂ = [k_dir[1] / k_norm, k_dir[2] / k_norm, k_dir[3] / k_norm]

    # Normalize polarization vector (user must ensure it's perpendicular to k)
    pol_norm = sqrt(pol[1]^2 + pol[2]^2 + pol[3]^2)
    ê = [pol[1] / pol_norm, pol[2] / pol_norm, pol[3] / pol_norm]

    # B-field direction: k̂ × ê
    b̂ = [
        k̂[2] * ê[3] - k̂[3] * ê[2],
        k̂[3] * ê[1] - k̂[1] * ê[3],
        k̂[1] * ê[2] - k̂[2] * ê[1]
    ]

    # Spatial position from 4-position (x = [ct, x, y, z] or [t, x, y, z])
    x⃗ = [x[2], x[3], x[4]]

    # Phase: k·r - ωt where |k| = ω/c (dispersion relation)
    phase = (ω / c) * (k̂[1] * x⃗[1] + k̂[2] * x⃗[2] + k̂[3] * x⃗[3]) - ω * t

    eqs = [
        E[1] ~ A * ê[1] * cos(phase)
        E[2] ~ A * ê[2] * cos(phase)
        E[3] ~ A * ê[3] * cos(phase)
        B[1] ~ (A / c) * b̂[1] * cos(phase)
        B[2] ~ (A / c) * b̂[2] * cos(phase)
        B[3] ~ (A / c) * b̂[3] * cos(phase)
    ]

    initialization_eqs = [
        λ ~ (2π * c) / ω
    ]
    bindings = [
        λ => missing
        ω => missing
    ]

    initial_conditions = Pair{SymbolicT,Any}[]
    if wavelength !== nothing
        push!(initial_conditions, λ => wavelength)
    end
    if frequency !== nothing
        push!(initial_conditions, ω => frequency)
    end

    vars = nameof(iv) == :τ ? [x, t] : [x]

    sys = System(eqs, iv, vars, [A, ω, k_dir, pol, λ]; name, systems=[ref_frame], initialization_eqs, bindings, initial_conditions)

    extend(sys, field_dynamics)
end

"""
Uniform electromagnetic field component.

In crossed E and B fields with E⊥B and |E| < |B|c,
particles drift with velocity v_drift = E×B/B²

Reference: Jackson, "Classical Electrodynamics", Section 12.4
"""
@component function UniformField(; name, E_field=[0,0,1], B_field=[0,0,0], ref_frame)
    @named field_dynamics = EMFieldDynamics(; ref_frame)
    @unpack E, B = field_dynamics

    # Get spacetime variables and constants from parent scope
    @unpack c, m_e, q_e = ref_frame
    iv = ModelingToolkit.get_iv(ref_frame)

    # Create local position and time variables
    @variables x(iv)[1:4] t(iv)

    @parameters E₀[1:3]=E_field B₀[1:3]=B_field

    eqs = [
        E ~ E₀,
        B ~ B₀
    ]

    sys = System(eqs, iv, [x, t], [E₀, B₀]; name, systems=[ref_frame])

    extend(sys, field_dynamics)
end

"""
Laguerre-Gauss laser beam electromagnetic field.

Represents a focused beam with orbital angular momentum (OAM).
The beam is characterized by radial index p and azimuthal index m.

Parameters:
- λ: wavelength
- a₀: normalized vector potential (a0 kwarg)
- w₀: beam waist (defaults to 75λ)
- radial_index (p): radial mode number, p ≥ 0
- azimuthal_index (m): azimuthal mode number (orbital angular momentum)
- temporal_profile: :gaussian (pulsed) or :constant (CW)
- temporal_width: pulse width for gaussian profile (defaults to 100.0)
- focus_position: focal position along z-axis (defaults to 0.0)

Reference: Allen et al., Phys. Rev. A 45, 8185 (1992)
"""
@component function LaguerreGaussLaser(;
    name,
    wavelength=nothing,
    frequency=nothing,
    a0=10.0,
    beam_waist=nothing,
    radial_index=0,      # p
    azimuthal_index=1,   # m
    ref_frame,
    temporal_profile=:gaussian,  # :gaussian or :constant
    temporal_width=nothing,      # pulse width (for gaussian profile)
    focus_position=nothing,       # focal position along z-axis
    polarization = :linear
)
    if wavelength === nothing && frequency === nothing
        wavelength = 1.0
    end
    if wavelength !== nothing && frequency !== nothing
        error("Specify either wavelength or frequency, not both")
    end

    # New interface with spacetime
    @named field_dynamics = EMFieldDynamics(; ref_frame)

    # Get spacetime variables from parent scope
    @unpack c, m_e, q_e = ref_frame
    iv = ModelingToolkit.get_iv(ref_frame)

    # Create local position and time variables
    @variables x(iv)[1:4]
    if nameof(iv) == :τ
        @variables t(iv)
    else
        t = iv
    end

    @unpack E, B = field_dynamics

    # Helper: compute |m|
    mₐ = abs(azimuthal_index)
    sgn = sign(azimuthal_index)

    # Compute normalization factor using Pochhammer symbol (rising factorial)
    # Nₚₘ = √((p+1)_{|m|}) where (x)_n is the Pochhammer symbol
    Npm_val = sqrt(pochhammer(radial_index + 1, mₐ))

    params = @parameters begin
        λ, [guess = 1.0]
        a₀ = a0
        ω, [guess = 1.0]
        k
        E₀
        w₀
        z_R
        τ0 = temporal_width === nothing ? 100.0 : temporal_width

        # Laguerre-Gauss quantum numbers
        p = radial_index
        m = azimuthal_index

        # Normalization factor (computed from p and m)
        Nₚₘ = Npm_val
    end

    if polarization == :linear
        ξx = 1.0 + 0im
        ξy = 0 + 0im
    elseif polarization == :circular
        ξx, ξy = (1/√2, im/√2) .|> complex
    else
        error("polarization $polarization not supported.")
    end

    # Fixed parameters (computed values, not symbolic parameters)
    ϕ₀ = 0
    t₀ = 0  # Center pulse at t=0
    z₀ = focus_position === nothing ? 0.0 : focus_position

    # Derived variables
    @variables begin
        wz(iv)      # Beam width
        z(iv)       # Propagation coordinate
        r(iv)       # Radial distance
        θ(iv)       # Azimuthal angle
        σ(iv)       # Normalized radial coordinate squared
        rwz(iv)     # Scaled radial coordinate
        env(iv)     # Temporal envelope factor
    end

    # Compute temporal envelope expression based on profile type
    env_expr = if temporal_profile == :constant
        1.0  # No temporal envelope
    elseif temporal_profile == :gaussian
        exp(-(((t - t₀) - (z - z₀) / c) / τ0)^2)
    else
        error("Unknown temporal_profile: $temporal_profile. Use :constant or :gaussian")
    end

    # We need to factor out the complex expressions so that we don't have equations with complex lhs
    # since MTK doesn't handle that well

    # Base Gaussian field (before polarization and LG modification)
    E_g = E₀ * w₀ / wz *
        exp(-(r / wz)^2) *
        exp(im * (-(r^2 * z) / (z_R * wz^2) + atan(z, z_R) - k * z + ϕ₀)) *
        exp(im * ω * t) * env

    # Laguerre-Gauss phase factor
    lg_phase = exp(im * ((2p + mₐ) * atan(z, z_R) - m * θ + ϕ₀))

    # NEgexp term used in Ez and Bz
    NEgexp = Nₚₘ * E_g * lg_phase

    # Complex transverse field components (used in Ez, Bz)
    Ex_complex = ξx * E_g * Nₚₘ * rwz^mₐ * _₁F₁(-p, mₐ + 1, 2σ) * lg_phase
    Ey_complex = ξy * E_g * Nₚₘ * rwz^mₐ * _₁F₁(-p, mₐ + 1, 2σ) * lg_phase

    eqs = [
        # Extract coordinates from 4-position
        z ~ x[4]
        r ~ hypot(x[2], x[3])
        θ ~ atan(x[3], x[2])

        # Beam width as function of z
        wz ~ w₀ * √(1 + (z / z_R)^2)

        # Intermediate variables
        σ ~ (r / wz)^2
        rwz ~ r * √2 / wz

        # Temporal envelope
        env ~ env_expr

        # Electric field Ex component (Laguerre-Gauss)
        E[1] ~ real(Ex_complex)

        # Electric field Ey component
        E[2] ~ real(Ey_complex)

        # Electric field Ez component (longitudinal)
        E[3] ~ real(
            -im / k * (
                # Term 1: m-dependent term
                (azimuthal_index == 0 ? 0.0 :
                    mₐ * (ξx - im * sgn * ξy) *
                    (√2 / wz)^mₐ * r^(mₐ - 1) * _₁F₁(-p, mₐ + 1, 2σ) *
                    NEgexp * exp(im * sgn * θ)
                ) -
                # Term 2: transverse field coupling
                2 / (wz^2) * (1 + im * z / z_R) * (x[2] * Ex_complex + x[3] * Ey_complex) -
                # Term 3: p-dependent term
                (radial_index == 0 ? 0.0 :
                    4 * p / ((mₐ + 1) * wz^2) * (x[2] * ξx + x[3] * ξy) *
                    rwz^mₐ * _₁F₁(-p + 1, mₐ + 2, 2σ) * NEgexp
                )
            )
        )

        # Magnetic field components
        # For paraxial beam: By ≈ Ex/c, Bx ≈ -Ey/c
        B[1] ~ -E[2] / c
        B[2] ~ E[1] / c

        # Bz component (full Laguerre-Gauss formula)
        B[3] ~ real(
            -im / ω * (
                # Term 1: m-dependent term
                -(azimuthal_index == 0 ? 0.0 :
                    mₐ * (ξy + im * sgn * ξx) *
                    (√2 / wz)^mₐ * r^(mₐ - 1) * _₁F₁(-p, mₐ + 1, 2σ) *
                    NEgexp * exp(im * sgn * θ)
                ) +
                # Term 2: transverse field coupling
                2 / (wz^2) * (1 + im * z / z_R) * (x[2] * Ey_complex - x[3] * Ex_complex) +
                # Term 3: p-dependent term
                (radial_index == 0 ? 0.0 :
                    4 * p / ((mₐ + 1) * wz^2) * (x[2] * ξy - x[3] * ξx) *
                    rwz^mₐ * _₁F₁(-p + 1, mₐ + 2, 2σ) * NEgexp
                )
            )
        )
    ]
    initialization_eqs = [
        ω ~ 2π * c / λ
        k ~ 2π / λ
        z_R ~ π * w₀^2 / λ
        E₀ ~ a₀ * m_e * c * ω / abs(q_e)
    ]
    bindings = [
        ω => missing
        λ => missing
        k => missing
        E₀ => missing
        z_R => missing
        w₀ => missing
    ]

    initial_conditions = Pair{SymbolicT,Any}[]
    push!(initial_conditions, w₀ => (beam_waist === nothing ? 75λ : beam_waist))
    if wavelength !== nothing
        push!(initial_conditions, λ => wavelength)
    end
    if frequency !== nothing
        push!(initial_conditions, ω => frequency)
    end

    vars = if nameof(iv) == :τ
        [x, t, z, r, θ, wz, σ, rwz, env]
    else
        @variables τ(iv)
        [x, τ, z, r, θ, wz, σ, rwz, env]
    end

    sys = System(eqs, iv, vars, params;
        name, systems=[ref_frame], initialization_eqs, bindings, initial_conditions)
    extend(sys, field_dynamics)
end