def uniaxial_reflectivity(q: NDArray, layers: NDArray, tensor: NDArray, energy: float):
    """
    EMpy implementation of the uniaxial 4x4 matrix formalism.

    for calculating reflectivity from a stratified
    medium.

    Uses the implementation developed by FSRStools -
    https://github.com/ddietze/FSRStools - written by Daniel Dietze

    Parameters
    ----------
    q: np.ndarray
        q values for which to calculate the reflectivity.
        Units: 1/Angstroms

    layers: np.ndarray
        coefficients required for the calculation, has shape (2 + N, 4),
        where N is the number of layers
        - layers[0, 1] - SLD of fronting (/1e-6 Angstrom**-2)
        - layers[0, 2] - iSLD of fronting (/1e-6 Angstrom**-2)
        - layers[N, 0] - thickness of layer N
        - layers[N, 1] - SLD of layer N (/1e-6 Angstrom**-2)
        - layers[N, 2] - iSLD of layer N (/1e-6 Angstrom**-2)
        - layers[N, 3] - roughness between layer N-1/N
        - layers[-1, 1] - SLD of backing (/1e-6 Angstrom**-2)
        - layers[-1, 2] - iSLD of backing (/1e-6 Angstrom**-2)
        - layers[-1, 3] - roughness between backing and last layer

    tensor: np.ndarray
        contains the 1x3x3 dimensions
        First dimension may change in teh fiture to account for multi-energy
        currently it will just cycle
    scale: float
        Multiply all reflectivities by this value.
    bkg: float
        Linear background to be added to all reflectivities

    Returns
    -------
    Reflectivity: np.ndarray
        Calculated reflectivity values for each q value.
    """
    # Plane of incidence - required to define polarization vectors
    OpticAxis = np.array([0.0, 0.0, 1.0])
    phi = 0  # This is a uniaxial calculation

    ##Organize qvals into proper order
    qvals = np.asarray(q)
    flatq = qvals.ravel()
    numpnts = flatq.size  # Number of q-points

    ##Grab the number of layers
    nlayers = layers.shape[0]

    ##hc has been converted into eV*Angstroms
    wl = (
        hc / energy
    )  ##calculate the wavelength array in Aangstroms for layer calculations
    k0 = 2 * np.pi / (wl)

    # Convert optical constants into dielectric tensor
    tensor = np.conj(np.eye(3) - 2 * tensor[:, :, :])  # From tensor[:,0,:,:]

    # freq = 2*np.pi * c/wls #Angular frequency
    theta_exp = np.zeros(numpnts, dtype=float)
    theta_exp = np.pi / 2 - np.arcsin((flatq[:] / (2 * k0)).clip(-1, 1))

    ##Generate arrays of data for calculating transfer matrix
    ##Scalar values ~~
    ## Special cases!
    ##kx is constant for each wavelength but changes with angle
    ## Dimensionality ##
    ## (angle)
    # kx = np.zeros(numpnts, dtype=complex)
    # ky = np.zeros(numpnts, dtype=complex) #Used to keep the k vectors three
    # components later on for cross / dot products
    kx = k0 * np.sin(theta_exp) * np.cos(phi)
    ky = k0 * np.sin(theta_exp) * np.sin(phi)

    ## Calculate the eigenvalues corresponding to kz ~~ Each one has 4 solutions
    ## Dimensionality ##
    ## (angle, #layer, solution)
    kz = np.zeros((numpnts, nlayers, 4), dtype=complex)

    ## Calculate the eignvectors corresponding to each kz ~~ polarization of D and H
    ## Dimensionality ##
    ## (angle, #layers, solution, vector)
    Dpol = np.zeros(
        (numpnts, nlayers, 4, 3), dtype=complex
    )  ##The polarization of the displacement field
    Hpol = np.zeros(
        (numpnts, nlayers, 4, 3), dtype=complex
    )  ##The polarization of the magnetic field

    # Cycle through the layers and calculate kz
    for j, epsilon in enumerate(
        tensor
    ):  # Each layer will have a different epsilon and subsequent kz
        kz[:, j, :] = calculate_kz_uni(epsilon, kx, ky, k0, opticaxis=OpticAxis)
        Dpol[:, j, :, :], Hpol[:, j, :, :] = calculate_Dpol_uni(
            epsilon,
            kx,
            ky,
            kz[:, j, :],
            k0,
            opticaxis=OpticAxis,  # type: ignore
        )

    ##Make matrices for the transfer matrix calculation
    ##Dimensionality ##
    ##(angles, #layers, Matrix (4,4)

    ## Propogation Matrix
    P = calculate_P(
        numpnts, nlayers, kz[:, :, :], layers[:, 0]
    )  ##layers[k,0] is the thicknes of layer k
    ##Nevot-Croche roughness matrix
    W = calculate_W(numpnts, nlayers, kz[:, :, :], kz[:, :, :], layers[:, 3])
    ##Dynamic Matrix and inverse
    D, Di = calculate_D(numpnts, nlayers, Dpol[:, :, :, :], Hpol[:, :, :, :])

    ##Calculate the full system transfer matrix
    ##Dimensionality ##
    ##(angles, Matrix (4,4))
    M = np.ones((numpnts, 4, 4), dtype=complex)
    # Make a (numpnts x 4x4) identity matrix for the TMM -
    M = np.einsum("...ij,ij->...ij", M, np.identity(4))  # type: ignore
    M = calculate_TMM(numpnts, nlayers, M, D, Di, P, W)
    ##Calculate the final outputs and organize into the appropriate waves for later
    refl, tran = calculate_output(numpnts, M)

    return refl, tran, kx, ky, kz, Dpol, Hpol, D, Di, P, W, M

def calculate_kz_uni(ep, kx, ky, k0, opticaxis=(None)):
    """Calculate the z-component of the wavevector for uniaxial media.

    This function calculates the z-components of both ordinary and extraordinary waves
    in a uniaxial medium, given the permittivity tensor and the x,y components of the
    wavevector.

    Parameters
    ----------
    ep : numpy.ndarray
        3x3 permittivity tensor of the material
    kx : numpy.ndarray
        x-component of the wavevector
    ky : numpy.ndarray
        y-component of the wavevector
    k0 : float
        Free space wavevector magnitude
    opticaxis : tuple, optional
        Unit vector defining the optical axis direction. Defaults to [0.0, 1.0, 0.0]

    Returns
    -------
    numpy.ndarray
        Array of shape (len(kx), 4) containing the z-components of the wavevector:
        - [:, 0]: forward extraordinary wave
        - [:, 1]: backward extraordinary wave
        - [:, 2]: forward ordinary wave
        - [:, 3]: backward ordinary wave

    Notes
    -----
    The calculation assumes the material is uniaxial with the extraordinary axis
    aligned with the optic axis. The ordinary and extraordinary components are
    calculated using the standard dispersion relations for uniaxial media.
    """
    # Calculate ordinary and extraordinary components from the tensor

    if opticaxis is None:
        opticaxis = [0.0, 1.0, 0.0]
    e_o = ep[0, 0]
    e_e = ep[2, 2]
    nu = (e_e - e_o) / e_o  # intermediate birefringence from reference
    k_par = np.sqrt(kx**2 + ky**2)  # Magnitude of parallel component
    # l = [kx/k_par, ky/k_par, 0]

    kz_ord = np.zeros(len(kx), dtype=np.complex128)
    kz_extraord = np.zeros(len(kx), dtype=np.complex128)
    kz_out = np.zeros((len(kx), 4), dtype=np.complex128)

    # n = [0, 0, 1] #Normal vector
    # if not numpy.isclose(k_par, 0):
    #    l = [kx/k_par, ky/k_par, 0]
    #    assert numpy.isclose(numpy.dot(l, l), 1)
    # else:
    #    l = [0, 0, 0]

    # Dot product between optical axis and vector normal and perpindicular component
    na = 1  # numpy.dot(n, opticAxis)
    la = 0  # numpy.dot(l, opticAxis)

    kz_ord = np.sqrt(e_o * k0**2 - k_par[:] ** 2)  # , dtype=np.complex128)

    kz_extraord = (1 / (1 + nu * na**2)) * (
        -nu * k_par[:] * na * la
        + np.sqrt(
            e_o * k0**2 * (1 + nu) * (1 + nu * na**2)
            - k_par[:] ** 2 * (1 + nu * (la**2 + na**2))
        )
    )

    kz_out[:, 2] = kz_ord
    kz_out[:, 3] = -kz_ord
    kz_out[:, 0] = kz_extraord
    kz_out[:, 1] = -kz_extraord
    return kz_out

def calculate_Dpol_uni(ep, kx, ky, kz, k0, opticaxis: list[float] | None = None):
    """Calculate electric and magnetic dipole polarizations for uniaxial materials.

    This function computes the electric and magnetic dipole polarization vectors for
    a uniaxial anisotropic material, given the permittivity tensor and wavevector
    components.

    Parameters
    ----------
    ep : numpy.ndarray
        3x3 permittivity tensor for the uniaxial material.
    kx : numpy.ndarray
        x-component of the wavevector.
    ky : numpy.ndarray
        y-component of the wavevector.
    kz : numpy.ndarray
        z-component of the wavevector.
    k0 : float
        Free-space wavevector magnitude.
    opticaxis : list[float] | None, optional
        Unit vector defining the optical axis. Default is ``[0.0, 1.0, 0.0]`` when
        omitted.

    Returns
    -------
    dpol_temp : numpy.ndarray
        Normalized electric dipole polarization vectors.
    hpol_temp : numpy.ndarray
        Normalized magnetic dipole polarization vectors.

    Notes
    -----
    The optical axis should not be collinear with the k-vector and should be a unit
    vector. The function handles ordinary and extraordinary waves in the uniaxial
    medium.
    """
    if opticaxis is None:
        opticaxis = [0.0, 1.0, 0.0]
    e_o = ep[0, 0]
    e_e = ep[2, 2]
    nu = (e_e - e_o) / e_o  # intermediate birefringence from reference

    kvec = np.zeros((len(kx), 4, 3), dtype=np.complex128)
    kdiv = np.zeros((len(kx), 4), dtype=np.complex128)
    dpol_temp = np.zeros((len(kx), 4, 3), dtype=np.complex128)
    hpol_temp = np.zeros((len(kx), 4, 3), dtype=np.complex128)

    # create k-vector
    kvec[:, :, 0] = kx[:, None]
    kvec[:, :, 1] = ky[:, None]
    kvec[:, :, 2] = kz

    kdiv: np.ndarray = np.sqrt(
        np.einsum("ijk,ijk->ij", kvec, kvec)  # type: ignore
    )  # Performs the commented out dot product calculation

    knorm = kvec / kdiv[:, :, None]  # (np.linalg.norm(kvec,axis=-1)[:,:,None])

    # calc propogation of k along optical axis
    kpol = np.dot(knorm, opticaxis)

    dpol_temp[:, 2, :] = np.cross(opticaxis[None, :], knorm[:, 2, :])  # type: ignore
    dpol_temp[:, 3, :] = np.cross(opticaxis[None, :], knorm[:, 3, :])  # type: ignore
    dpol_temp[:, 0, :] = np.subtract(
        opticaxis[None, :],  # type: ignore
        ((1 + nu) / (1 + nu * kpol[:, 0, None] ** 2))
        * kpol[:, 0, None]
        * knorm[:, 0, :],
    )
    dpol_temp[:, 1, :] = np.subtract(
        opticaxis[None, :],  # type: ignore
        ((1 + nu) / (1 + nu * kpol[:, 1, None] ** 2))
        * kpol[:, 1, None]
        * knorm[:, 1, :],
    )

    dpol_norm = np.linalg.norm(dpol_temp, axis=-1)
    dpol_temp /= dpol_norm[:, :, None] + TINY
    hpol_temp = np.cross(kvec, dpol_temp) * (1 / k0)  # type: ignore

    return dpol_temp, hpol_temp

def calculate_D(numpnts, nlayers, Dpol, Hpol) -> tuple[np.ndarray, np.ndarray]:
    """
    Transfer matrices D and its inverse Di for a multilayer optical system.

    This function constructs the transfer matrix D and calculates its inverse Di for
    each point
    and layer in the system using polarization components from D-polarized and
    H-polarized fields.

    Parameters
    ----------
    numpnts : int
        Number of points (typically wavelengths or angles) to calculate for
    nlayers : int
        Number of layers in the optical system
    Dpol : ndarray
        D-polarization transfer matrix components with shape (numpnts, nlayers, 4, 4)
    Hpol : ndarray
        H-polarization transfer matrix components with shape (numpnts, nlayers, 4, 4)

    Returns
    -------
    list
        A list containing two elements:
        - D_Temp : ndarray of shape (numpnts, nlayers, 4, 4)
            The constructed transfer matrix D
        - Di_Temp : ndarray of shape (numpnts, nlayers, 4, 4)
            The inverse of D_Temp, calculated using np.linalg.inv() or np.linalg.pinv()
            as fallback

    Notes
    -----
    The function attempts to use numpy.linalg.inv() for matrix inversion first, and
    falls back
    to numpy.linalg.pinv() (pseudo-inverse) if the regular inversion fails.
    """
    D_Temp: np.ndarray = np.zeros((numpnts, nlayers, 4, 4), dtype=np.complex128)
    Di_Temp: np.ndarray = np.zeros((numpnts, nlayers, 4, 4), dtype=np.complex128)

    D_Temp[:, :, 0, 0] = Dpol[:, :, 0, 0]
    D_Temp[:, :, 0, 1] = Dpol[:, :, 1, 0]
    D_Temp[:, :, 0, 2] = Dpol[:, :, 2, 0]
    D_Temp[:, :, 0, 3] = Dpol[:, :, 3, 0]
    D_Temp[:, :, 1, 0] = Hpol[:, :, 0, 1]
    D_Temp[:, :, 1, 1] = Hpol[:, :, 1, 1]
    D_Temp[:, :, 1, 2] = Hpol[:, :, 2, 1]
    D_Temp[:, :, 1, 3] = Hpol[:, :, 3, 1]
    D_Temp[:, :, 2, 0] = Dpol[:, :, 0, 1]
    D_Temp[:, :, 2, 1] = Dpol[:, :, 1, 1]
    D_Temp[:, :, 2, 2] = Dpol[:, :, 2, 1]
    D_Temp[:, :, 2, 3] = Dpol[:, :, 3, 1]
    D_Temp[:, :, 3, 0] = Hpol[:, :, 0, 0]
    D_Temp[:, :, 3, 1] = Hpol[:, :, 1, 0]
    D_Temp[:, :, 3, 2] = Hpol[:, :, 2, 0]
    D_Temp[:, :, 3, 3] = Hpol[:, :, 3, 0]

    """
    for i in range(numpnts):
        for j in range(nlayers):
            Di_Temp[i,j,:,:] = np.linalg.pinv(D_Temp[i,j,:,:])
    """
    # Try running an a matrix inversion for the tranfer matrix.
    # If it fails, run a pseudo-inverse
    # Update 07/07/2021: I don't think the uniaxial calculation will error...changing
    # pinv to inv
    #                   for default calculation
    try:
        Di_Temp = np.linalg.inv(D_Temp)  # type: ignore
    except LinAlgError:
        Di_Temp = np.linalg.pinv(D_Temp)  # type: ignore

    return [D_Temp, Di_Temp]  # type: ignore

def calculate_P(_numpnts, _nlayers, kz, d):
    """
    Calculate the propagation matrix using the previously calculated values for kz.

    Parameters
    ----------
    _numpnts : int
        Number of points in the calculation grid (not used directly in this function).
    _nlayers : int
        Number of layers in the model (not used directly in this function).
    kz : ndarray
        Array of shape (numpnts, nlayers, 4) containing the z-component of the wavevector for each point, layer, and solution.
    d : ndarray
        Array of shape (nlayers,) containing the thickness of each layer.
    """
    # Create the diagonal components in the propogation matrix
    # Cast into a 4x4 version through redundent broadcasting
    diagonal_components: np.ndarray = np.exp(
        -1j * kz[:, :, :, None] * d[None, :, None, None]
    )
    # Element by element multiplication with the identity over each q-point
    P_temp = np.einsum("...jk,jk->...jk", diagonal_components, np.identity(4))  # type: ignore

    return P_temp

def calculate_W(numpnts, nlayers, kz1, kz2, r):
    """Calculate matrix W for uniaxial model fitting.

    This function calculates the W matrix used in uniaxial model fitting, which involves
    exponential terms of wavevectors for different layers.

    Parameters
    ----------
    numpnts : int
        Number of points in the calculation grid
    nlayers : int
        Number of layers in the model
    kz1 : ndarray
        Array of shape (numpnts, nlayers, 4) containing z-component of wavevector 1
    kz2 : ndarray
        Array of shape (numpnts, nlayers, 4) containing z-component of wavevector 2
    r : ndarray
        Array of shape (nlayers,) containing roughness parameters for each layer

    Returns
    -------
    ndarray
        4D array of shape (numpnts, nlayers, 4, 4) containing the W matrix elements
        with exponential terms for interface roughness

    Notes
    -----
    The function calculates exponential terms for interface roughness using the
    sum and difference of wavevectors, and arranges them in a 4x4 matrix format
    for each point and layer.
    """
    W_temp = np.zeros((numpnts, nlayers, 4, 4), dtype=np.complex128)
    eplus = np.zeros((numpnts, nlayers, 4), dtype=np.complex128)
    eminus = np.zeros((numpnts, nlayers, 4), dtype=np.complex128)

    kz2 = np.roll(
        kz2, 1, axis=1
    )  # Reindex to allow broadcasting in the next step....see commented loop
    # for j in range(nlayers):
    eplus[:, :, :] = np.exp(
        -((kz1[:, :, :] + kz2[:, :, :]) ** 2) * r[None, :, None] ** 2 / 2
    )
    eminus[:, :, :] = np.exp(
        -((kz1[:, :, :] - kz2[:, :, :]) ** 2) * r[None, :, None] ** 2 / 2
    )

    W_temp[:, :, 0, 0] = eminus[:, :, 0]
    W_temp[:, :, 0, 1] = eplus[:, :, 1]
    W_temp[:, :, 0, 2] = eminus[:, :, 2]
    W_temp[:, :, 0, 3] = eplus[:, :, 3]
    W_temp[:, :, 1, 0] = eplus[:, :, 0]
    W_temp[:, :, 1, 1] = eminus[:, :, 1]
    W_temp[:, :, 1, 2] = eplus[:, :, 2]
    W_temp[:, :, 1, 3] = eminus[:, :, 3]
    W_temp[:, :, 2, 0] = eminus[:, :, 0]
    W_temp[:, :, 2, 1] = eplus[:, :, 1]
    W_temp[:, :, 2, 2] = eminus[:, :, 2]
    W_temp[:, :, 2, 3] = eplus[:, :, 3]
    W_temp[:, :, 3, 0] = eplus[:, :, 0]
    W_temp[:, :, 3, 1] = eminus[:, :, 1]
    W_temp[:, :, 3, 2] = eplus[:, :, 2]
    W_temp[:, :, 3, 3] = eminus[:, :, 3]

    return W_temp

def calculate_TMM(_numpnts, nlayers, M, D, Di, P, W):
    """Calculate Transfer Matrix Method (TMM) for multilayered optical system.

    This function computes the total transfer matrix for a multilayer optical system
    using
    the Transfer Matrix Method (TMM). It iteratively multiplies matrices representing
    different
    layers to obtain the overall system response.

    Parameters
    ----------
    _numpnts : int
        Unused; retained for signature compatibility. Number of points is implied by array axes.
    nlayers : int
        Number of layers in the optical system.
    M : ndarray
        Initial transfer matrix, shape (numpnts, 2, 2).
    D : ndarray
        Dynamic matrices for each layer, shape (numpnts, nlayers, 2, 2).
    Di : ndarray
        Inverse dynamic matrices for each layer, shape (numpnts, nlayers-1, 2, 2).
    P : ndarray
        Propagation matrices for each layer, shape (numpnts, nlayers-1, 2, 2).
    W : ndarray
        Interface matrices between layers, shape (numpnts, nlayers, 2, 2).

    Returns
    -------
    ndarray
        Final transfer matrix M after considering all layers, shape (numpnts, 2, 2).

    Notes
    -----
    The function implements the TMM calculation following the sequence:
    1. Iterates through internal layers
    2. Multiplies appropriate matrices using Einstein summation
    3. Includes final layer calculation separately
    """
    for j in range(1, nlayers - 1):
        A = np.einsum("...ij,...jk ->...ik", Di[:, j - 1, :, :], D[:, j, :, :])
        B = A * W[:, j, :, :]
        C = np.einsum("...ij,...jk ->...ik", B, P[:, j, :, :])
        M[:, :, :] = np.einsum("...ij,...jk ->...ik", M[:, :, :], C)
    AA = np.einsum("...ij,...jk ->...ik", Di[:, -2, :, :], D[:, -1, :, :])
    BB = AA * W[:, -1, :, :]
    M[:, :, :] = np.einsum("...ij,...jk ->...ik", M[:, :, :], BB)
    return M

def calculate_output(
    numpnts: int, M_full: np.ndarray
) -> tuple[NDArray[np.float64], NDArray[np.complex128]]:
    """Extract reflectivity and transmittance amplitudes from the system matrix.

    Parameters
    ----------
    numpnts : int
        Number of angular or momentum samples.
    M_full : ndarray
        System transfer matrix with shape ``(numpnts, 4, 4)``.

    Returns
    -------
    refl : ndarray
        Real reflectivity values, shape ``(numpnts,)``.
    tran : ndarray
        Complex transmittance-related amplitudes, shape ``(numpnts,)``.

    Notes
    -----
    The full Yeh 4x4 extraction from ``M_full`` is not yet implemented in this
    port; the returned arrays are zero-filled so the module remains importable
    and documentable while the optical reduction is completed.
    """
    _ = M_full
    return (
        np.zeros(numpnts, dtype=np.float64),
        np.zeros(numpnts, dtype=np.complex128),
    )