from __future__ import annotations

import os
import sys
import warnings
from functools import cache
from typing import Any, NamedTuple

import click
import numpy as np
import scgenerator as sc
import tomli
from customfunc.app import PlotApp
from pydantic import BaseModel, ValidationError, confloat

DEFAULT_CONFIG_FILE = "config.toml"


class Config(BaseModel):
    wl_min: confloat(ge=100, le=1000)
    wl_max: confloat(ge=500, le=6000)
    wl_pump: confloat(ge=200, le=6000)
    rep_rate: confloat(gt=0)
    gas: str

    @classmethod
    def load(cls, config_file: str | None = None) -> Config:
        config_file = config_file or DEFAULT_CONFIG_FILE
        with open(config_file, "rb") as file:
            d = tomli.load(file)
        d = cls.default() | d
        try:
            return cls(**d)
        except ValidationError as e:
            s = f"invalid input in config file {config_file}:\n{e}"
            print(s)
            sys.exit(1)

    @classmethod
    def default(cls) -> dict[str, Any]:
        return dict(wl_min=160, wl_max=1600, wl_pump=800, rep_rate=8e3, gas="argon")


class LimitValues(NamedTuple):
    wl_zero_disp: float
    ion_lim: float
    sf_lim: float


def b2(w, n_eff):
    dw = w[1] - w[0]
    beta = sc.fiber.beta(w, n_eff)
    return sc.math.differentiate_arr(beta, 2, 4, dw)


def N_sf_max(
    wl: np.ndarray, t0: float, wl_zero_disp: float, gas: sc.materials.Gas, safety: float = 10.0
) -> np.ndarray:
    """
    maximum soliton number according to self focusing

    eq. S15 in Travers2019
    """
    delta_gas = gas.sellmeier.delta(wl, wl_zero_disp)
    return t0 * np.sqrt(wl / (safety * np.abs(delta_gas)))


def N_ion_max(
    wl: np.ndarray, t0: float, wl_zero_disp: float, gas: sc.materials.Gas, safety: float = 10.0
) -> np.ndarray:
    """
    eq. S16 in Travers2019
    """
    ind = sc.math.argclosest(wl, wl_zero_disp)
    f = np.gradient(np.gradient(gas.sellmeier.chi(wl), wl), wl)
    factor = (sc.math.u_nm(1, 1) / sc.units.c) ** 2 * (0.5 * wl / np.pi) ** 3
    delta = factor * (f / f[ind] - 1)
    denom = safety * np.pi * wl * np.abs(delta) * f[ind]
    return t0 * sc.math.u_nm(1, 1) * np.sqrt(gas.n2() * gas.barrier_suppression / denom)


def solition_num(
    t0: float, w0: float, beta2: float, n2: float, core_radius: float, peak_power: float
) -> float:
    gamma = sc.fiber.gamma_parameter(n2, w0, sc.fiber.A_eff_marcatili(core_radius))
    ld = sc.pulse.L_D(t0, beta2)
    return np.sqrt(gamma * ld * peak_power)


def energy(
    t0: float, w0: float, beta2: float, n2: float, core_radius: float, solition_num: float
) -> float:
    gamma = sc.fiber.gamma_parameter(n2, w0, sc.fiber.A_eff_marcatili(core_radius))
    peak_power = solition_num**2 * abs(beta2) / (t0**2 * gamma)
    return sc.pulse.P0_to_E0(peak_power, t0, "sech")


def app(config_file: os.PathLike | None = None):
    config = Config.load(config_file)
    # grid stuff
    w = sc.w_from_wl(config.wl_min, config.wl_max, 2048)
    wl = sc.units.m.inv(w)
    wl_ind = sc.math.argclosest(wl, config.wl_pump * 1e-9)
    w0 = w[wl_ind]

    gas = sc.materials.Gas(config.gas)

    with PlotApp(
        f"Dispersion design with {config.gas.title()}",
        core_diameter_um=np.linspace(50, 300),
        pressure_mbar=np.geomspace(1, 2000),
        wall_thickness_um=np.geomspace(0.01, 10),
        n_tubes=np.arange(6, 16),
        gap_um=np.linspace(1, 15),
        t_fwhm_fs=np.linspace(10, 200, 96),
    ) as app:
        # initial setup
        app[0].horizontal_line("reference", 0, color="gray")
        app[0].set_xlabel("wavelength (nm)")
        app[0].set_ylabel("beta2 (fs^2/cm)")
        app.params["wall_thickness_um"].value = 1
        app.params["core_diameter_um"].value = 100
        app.params["pressure_mbar"].value = 500
        app.params["n_tubes"].value = 7
        app.params["gap_um"].value = 5
        app.params["t_fwhm_fs"].value = 100
        app[0].set_lim(ylim=(-4, 2))

        @cache
        def compute_max_energy(
            core_diameter_um: float, pressure_mbar: float, t_fwhm_fs: float
        ) -> LimitValues:
            t_fwhm = 1e-15 * t_fwhm_fs
            pressure = 1e2 * pressure_mbar
            core_radius = 0.5e-6 * core_diameter_um
            t0 = sc.pulse.fwhm_to_T0_fac["sech"] * t_fwhm

            disp = compute_capillary(core_diameter_um, pressure_mbar)
            wl_zero_disp = sc.math.all_zeros(wl, disp)
            if len(wl_zero_disp) == 0:
                return
            wl_zero_disp = wl_zero_disp[0]

            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                ion_limit = N_ion_max(wl, t0, wl_zero_disp, gas)[wl_ind]
                sf_limit = N_sf_max(wl, t0, wl_zero_disp, gas)[wl_ind]

            beta2 = disp[wl_ind]
            n2 = gas.n2(pressure=pressure)

            return LimitValues(
                wl_zero_disp,
                energy(t0, w0, beta2, n2, core_radius, ion_limit),
                energy(t0, w0, beta2, n2, core_radius, sf_limit),
            )

        @cache
        def compute_vincetti(
            wall_thickness_um: float,
            core_diameter_um: float,
            pressure_mbar: float,
            n_tubes: int,
            gap_um: float,
        ) -> np.ndarray:
            core_diameter = core_diameter_um * 1e-6
            wall_thickness = wall_thickness_um * 1e-6
            gap = gap_um * 1e-6
            pressure = pressure_mbar * 1e2

            tr = sc.fiber.tube_radius_from_gap(core_diameter / 2, gap, n_tubes)
            n_gas_2 = gas.sellmeier.n_gas_2(wl, None, pressure)
            n_eff_vinc = sc.fiber.n_eff_vincetti(
                wl, 800e-9, n_gas_2, wall_thickness, tr, gap, n_tubes
            )
            return b2(w, n_eff_vinc)

        @cache
        def compute_capillary(core_diameter_um: float, pressure_mbar: float) -> np.ndarray:
            pressure = pressure_mbar * 1e2
            core_diameter = core_diameter_um * 1e-6
            n_gas_2 = gas.sellmeier.n_gas_2(wl, None, pressure)
            n_eff_cap = sc.fiber.n_eff_marcatili(wl, n_gas_2, core_diameter / 2)
            return b2(w, n_eff_cap)

        @app.update
        def draw_vincetty(
            wall_thickness_um: float,
            core_diameter_um: float,
            pressure_mbar: float,
            n_tubes: int,
            gap_um: float,
        ):
            b2 = compute_vincetti(
                wall_thickness_um, core_diameter_um, pressure_mbar, n_tubes, gap_um
            )
            app[0].set_line_data("Vincetti", wl * 1e9, sc.units.beta2_fs_cm_inv(b2))

        @app.update
        def draw_capillary(core_diameter_um: float, pressure_mbar: float):
            b2 = compute_capillary(core_diameter_um, pressure_mbar)
            app[0].set_line_data("Capillary", wl * 1e9, sc.units.beta2_fs_cm_inv(b2))

        @app.update
        def draw_energy_limit(core_diameter_um: float, pressure_mbar: float, t_fwhm_fs: float):
            lim = compute_max_energy(core_diameter_um, pressure_mbar, t_fwhm_fs)
            if lim.ion_lim > lim.sf_lim:
                power = lim.sf_lim * 1e3 * config.rep_rate
                app[0].set_line_name(
                    "Capillary", f"Capillary, max energy = {power:.0f}mW (self-focusing)"
                )
            else:
                power = lim.ion_lim * 1e3 * config.rep_rate
                app[0].set_line_name(
                    "Capillary", f"Capillary, max energy = {power:.0f}mW (ionization)"
                )

            zdw = lim.wl_zero_disp * 1e9
            app[0].set_line_data("zdw", [zdw, zdw], [-3, 3])
            app[0].set_line_name("zdw", f"ZDW = {zdw:.0f}nm")


@click.command()
@click.option(
    "-c",
    "--config",
    default=DEFAULT_CONFIG_FILE,
    help="configuration file in TOML format",
    type=click.Path(exists=True, file_okay=True, dir_okay=False, resolve_path=True),
)
def main(config: os.PathLike):
    app(config)


if __name__ == "__main__" and not hasattr(sys, "ps1"):
    main()