from collections import namedtuple
import itertools
from scipy.optimize import root

Args = namedtuple("Args", ["alpha", "w", "Te"])

ARG1 = Args(0.04, 0.00409, 5_000)
ARG2 = Args(0.032, 0.00537, 10_000)
ARG3 = Args(0.023, 0.00873, 40_000)


def L_stark(Ne: float, alpha: float, w: float, Te: float) -> float:
    """computes stark broadening

    Parameters
    ----------
    Ne : float
        free electron density in 1e15 / cm^3
    alpha : float
        static ion-broadening parameter
    w : float
        electron impact parameter in nm/1e16/cm^3
    Te : float
        electron temperature in K

    Returns
    -------
    float
        Stark broadening in nm
    """
    Ne = 1e15 * Ne
    return (
        2e-16
        * (1 + 1.75e-4 * (Ne ** 0.25) * alpha * (1 - 0.068 * (Te ** -0.5) * (Ne ** (1 / 6))))
        * w
        * Ne
    )


def main():
    for ls, args in itertools.product([0.001, 0.003, 0.009], [ARG1, ARG2, ARG3]):

        def to_root(Ne):
            return L_stark(Ne, *args) - ls

        print(ls, args.Te, root(to_root, 1.0).x)


if __name__ == "__main__":
    main()