import scgenerator as sc
from plotapp import PlotApp
import numpy as np
from PySide6 import QtWidgets


def get_disp(w_c, wl, ind, beta):
    disp = sc.fiber.dispersion_from_coefficients(w_c, list(beta))
    return sc.units.D_ps_nm_km.inv(sc.fiber.beta2_to_dispersion_parameter(disp[ind], wl * 1e-9))


def main():
    beta2_coefficients = [
        -1.183e-26,
        8.1038e-41,
        -9.5205e-56,
        2.0737e-70,
        -5.3943e-85,
        1.3486e-99,
        -2.5495e-114,
        3.0524e-129,
        -1.714e-144,
    ]
    ranges = {}
    for i, b in enumerate(beta2_coefficients):
        i += 2
        ranges[f"beta{i:02}"] = np.linspace(-np.abs(b) * 10, np.abs(b) * 10, 3000)

    wl0 = 835e-9
    dt = 1.6e-15
    t_num = 4096
    t = sc.math.tspace(t_num=t_num, dt=dt)
    w_c = sc.math.wspace(t)
    w = w_c + sc.units.m_rads(wl0)
    wl = sc.units.m_rads(w)

    wl, ind, _ = sc.PlotRange(wl0 * 0.5e9, wl0 * 3e9, "nm").sort_axis(wl)
    disp0 = get_disp(w_c, wl, ind, beta2_coefficients)

    with PlotApp(**ranges) as app:

        def reset():
            for i, r in enumerate(app.params.values()):
                r.value = beta2_coefficients[i]

        def print_beta():
            keys = sorted(app.params.keys())
            values = [app.params[k].value for k in keys]
            print("[\n" + ",\n".join(format(v, "g") for v in values) + "\n]")

        def set_zero():
            for i, r in enumerate(app.params.values()):
                r.value = 0.0

        print_beta()

        app[0].set_line_data("original", wl, disp0)
        app[0].horizontal_line("zero", 0, color="k")
        app[0].vertical_line("pump", wl0 * 1e9, color="k")

        reset_btn = QtWidgets.QPushButton("Reset")
        reset_btn.clicked.connect(reset)
        app.params_layout.addWidget(reset_btn, *divmod(len(app.params), 2))
        reset()

        print_btn = QtWidgets.QPushButton("Print")
        print_btn.clicked.connect(print_beta)
        app.params_layout.addWidget(print_btn, *divmod(len(app.params) + 1, 2))

        zero_btn = QtWidgets.QPushButton("Zero")
        zero_btn.clicked.connect(set_zero)
        app.params_layout.addWidget(zero_btn, *divmod(len(app.params) + 2, 2))

        @app.update
        def draw(
            beta02,
            beta03,
            beta04,
            beta05,
            beta06,
            beta07,
            beta08,
            beta09,
            beta10,
        ):
            disp = get_disp(
                w_c,
                wl,
                ind,
                [
                    beta02,
                    beta03,
                    beta04,
                    beta05,
                    beta06,
                    beta07,
                    beta08,
                    beta09,
                    beta10,
                ],
            )
            app[0].set_line_data("dispersion", wl, disp)


if __name__ == "__main__":
    main()