from numbers import Number

import sympy as sp

x, y, z = sp.symbols("x, y, z", real=True, finite=True)
#: Beam waist in x-direction (collimated).
waist_x0 = sp.Symbol("W_{x0}", real=True, positive=True, finite=True, nonzero=True)
#: Beam waist in z-direction at the focus.
waist_z0 = sp.Symbol("W_{z0}", real=True, positive=True, finite=True, nonzero=True)
#: Offset from waist at origin

#: Power per beam. Total power is 2x.
power = sp.Symbol("P", real=True, positive=True, finite=True, nonzero=True)
#: Magnetic field gradients.
grad_r, grad_z = sp.symbols(
    "\\nabla{B_x}, \\nabla{B_z}", real=True, finite=True, nonzero=True
)
#: Real part of polarizability :math:`\alpha`
a = sp.Symbol("a", real=True, finite=True, nonzero=True)
d = sp.Symbol("d", real=True, positive=True, finite=True, nonzero=True)
mu_b = sp.Symbol("mu_b", real=True, positive=True, finite=True, nonzero=True)
#: Wavelength
wvl = sp.Symbol("lambda", real=True, positive=True, finite=True, nonzero=True)
#: Angle of the incident beams
theta = sp.Symbol("theta", real=True, positive=True, finite=True, nonzero=True)


def BeamWaistZ(length: sp.Symbol | Number) -> sp.Expr:  # noqa: N802
    """The beam is focused in z-direction, leading to a waist that
    depends on the distance :math:`l` from the lens.

    Depending on which beam (upper or lower), the distance from the
    lens :math:`l` is given by
    :math:`y \\cos(\\theta) \\mp z \\sin(\\theta)`

    :param length: Distance :math:`l` from the lens.
    """
    return waist_z0 * sp.sqrt(1 + (wvl * length / (sp.pi * waist_z0**2)) ** 2)


def UpperBeamIntensity() -> sp.Expr:  # noqa: N802
    """Intensity distribution of the upper beam"""
    length = y * sp.cos(theta) - z * sp.sin(theta)
    waist_z = BeamWaistZ(length)
    # Intensity of the beam at the origin
    i_0 = 2 * power / (sp.pi * waist_x0 * waist_z0)
    # Turn off automatic formatting to preserve readability
    # fmt: off
    return i_0 * waist_z0 / waist_z * sp.exp(
        - 2 * x**2 / waist_x0**2 
        - 2 * (z * sp.cos(theta) + y * sp.sin(theta)) ** 2 / waist_z**2
    )  # fmt: on


def LowerBeamIntensity() -> sp.Expr:  # noqa: N802
    """Intensity distribution of the lower beam"""
    length = y * sp.cos(theta) + z * sp.sin(theta)
    waist_z = BeamWaistZ(length)
    # Intensity of the beam at the origin
    i_0 = 2 * power / (sp.pi * waist_x0 * waist_z0)
    # Turn off automatic formatting to preserve readability
    # fmt: off
    return i_0 * waist_z0 / waist_z * sp.exp(
        - 2 * x**2 / waist_x0**2 
        - 2 * (z * sp.cos(theta) - y * sp.sin(theta)) ** 2 / waist_z**2
    )  # fmt: on


def TotalIntensity() -> sp.Expr:  # noqa: N802
    i_1 = UpperBeamIntensity()
    i_2 = LowerBeamIntensity()
    return i_1 + i_2 + 2 * sp.sqrt(i_1 * i_2) * sp.cos(2 * sp.pi * z / d)


def TotalPotential() -> sp.Expr:  # noqa: N802
    return -a * TotalIntensity()