"""read Sparameters from CSV file and returns sax model."""
import pathlib
from typing import Union

import jax
import jax.numpy as jnp
import numpy as np
import pandas as pd
from sax.typing_ import Float, Model
from typing_extensions import Literal

import gdsfactory as gf
from gdsfactory.simulation.get_sparameters_path import (
    get_sparameters_path_lumerical,
    get_sparameters_path_meep,
    get_sparameters_path_tidy3d,
)

wl_cband = np.linspace(1.500, 1.600, 128)

PathType = Union[str, pathlib.Path]

Simulator = Literal["lumerical", "meep", "tidy3d"]


def model_from_npz(
    filepath: Union[PathType, np.ndarray],
    xkey: str = "wavelengths",
    xunits: float = 1,
    prefix: str = "s",
) -> Model:
    """Returns a SAX Sparameters Model from a npz file.

    The SAX Model is a function that returns a SAX SDict interpolated over wavelength.

    Args:
        filepath: CSV Sparameters path or pandas DataFrame.
        wl: wavelength to interpolate (um).
        xkey: key for wavelengths in file.
        xunits: x units in um from the loaded file (um). 1 means 1um.
        prefix: for the sparameters column names in file.

    """
    sp = filepath if isinstance(filepath, np.ndarray) else np.load(filepath)
    keys = list(sp.keys())

    if xkey not in keys:
        raise ValueError(f"{xkey!r} not in {keys}")

    x = jnp.asarray(sp[xkey] * xunits)
    wl = jnp.asarray(wl_cband)

    # make sure x is sorted from low to high
    idxs = jnp.argsort(x)
    x = x[idxs]
    sp = {k: v[idxs] for k, v in sp.items()}

    @jax.jit
    def model(wl: Float = wl):
        S = {}
        zero = jnp.zeros_like(x)

        for key in sp:
            if not key.startswith("wav"):
                port_mode0, port_mode1 = key.split(",")
                port0, mode0 = port_mode0.split("@")
                port1, mode1 = port_mode1.split("@")
                S[(port0, port1)] = jnp.interp(wl, x, sp.get(key, zero))

        return S

    return model


def model_from_csv(
    filepath: Union[PathType, pd.DataFrame],
    xkey: str = "wavelengths",
    xunits: float = 1,
    prefix: str = "s",
) -> Model:
    """Returns a SAX Sparameters Model from a CSV file.

    The SAX Model is a function that returns a SAX SDict interpolated over wavelength.

    Args:
        filepath: CSV Sparameters path or pandas DataFrame.
        wl: wavelength to interpolate (um).
        xkey: key for wavelengths in file.
        xunits: x units in um from the loaded file (um). 1 means 1um.
        prefix: for the sparameters column names in file.
    """
    df = filepath if isinstance(filepath, pd.DataFrame) else pd.read_csv(filepath)
    assert isinstance(df, pd.DataFrame)
    df = df.reset_index()  # maybe there is useful info in the index...
    dic = dict(zip(df.columns, jnp.asarray(df.values.T)))
    keys = list(dic.keys())

    if xkey not in keys:
        raise ValueError(f"{xkey!r} not in {keys}")

    nsparameters = (len(keys) - 1) // 2
    nports = int(nsparameters**0.5)

    x = jnp.asarray(dic[xkey] * xunits)
    wl = jnp.asarray(wl_cband)

    # make sure x is sorted from low to high
    idxs = jnp.argsort(x)
    x = x[idxs]
    dic = {k: v[idxs] for k, v in dic.items()}

    @jax.jit
    def model(wl: Float = wl):
        S = {}
        zero = jnp.zeros_like(x)
        for i in range(1, nports + 1):
            for j in range(1, nports + 1):
                m = jnp.interp(wl, x, dic.get(f"{prefix}{i}{j}m", zero))
                a = jnp.interp(wl, x, dic.get(f"{prefix}{i}{j}a", zero))
                S[f"o{i}", f"o{j}"] = m * jnp.exp(1j * a)

        return S

    return model


def _demo_mmi_lumerical_csv() -> None:
    import matplotlib.pyplot as plt
    from plot_model import plot_model

    filepath = get_sparameters_path_lumerical(gf.c.mmi1x2)
    mmi = model_from_csv(filepath=filepath, xkey="wavelengths")
    plot_model(mmi)
    plt.show()


def model_from_component(component, simulator: Simulator, **kwargs) -> Model:
    """Returns SAX model from lumerical FDTD simulations.

    Args:
        component: to simulate.
        simulator: meep, lumerical or tidy3d.
        kwargs: simulator settings.

    """
    simulators = ["lumerical", "meep", "tidy3d"]

    if simulator == "lumerical":
        filepath = get_sparameters_path_lumerical(component=component, **kwargs)
    elif simulator == "meep":
        filepath = get_sparameters_path_meep(component=component, **kwargs)
    elif simulator == "tidy3d":
        filepath = get_sparameters_path_tidy3d(component=component, **kwargs)
    else:
        raise ValueError(f"{simulator!r} no in {simulators}")
    return model_from_csv(filepath=filepath)


model_from_component_lumerical = gf.partial(model_from_component, simulator="fdtd")


if __name__ == "__main__":
    # mmi1x2 = gf.partial(model_from_component_lumerical, component=gf.components.mmi1x2)
    # mmi2x2 = gf.partial(model_from_component_lumerical, component=gf.components.mmi2x2)

    # grating_coupler_elliptical = model_from_csv(
    #     filepath=sparameters_path / "grating_coupler_ellipti_9d85a0c6_18c08cac.csv"
    # )

    # model_factory = dict(
    #     mmi1x2=mmi1x2, mmi2x2=mmi2x2, grating_coupler_elliptical=grating_coupler_elliptical
    # )
    # from gdsfactory.config import sparameters_path

    import matplotlib.pyplot as plt
    from plot_model import plot_model

    # s = from_csv(filepath=filepath, wavelength=wl_cband, xunits=1e-3)
    # s21 = s[("o1", "o3")]
    # plt.plot(wl_cband, np.abs(s21) ** 2)
    # plt.show()
    # filepath = get_sparameters_path_lumerical(gf.components.mmi1x2)
    # mmi = model_from_csv(filepath=filepath, xkey="wavelengths")
    # plot_model(mmi)
    # plt.show()
    # plot_model(grating_coupler_elliptical)
    # plt.show()
    # model = partial(
    #     from_csv, filepath=filepath, xunits=1, xkey="wavelengths", prefix="s"
    # )
    # sax.plot.plot_model(model)
    # plt.show()
    #  This looks correct
    # coupler_fdtd = gf.partial(
    #     sdict_from_csv,
    #     filepath=gf.config.sparameters_path / "coupler" / "coupler_G224n_L20_S220.csv",
    #     xkey="wavelength_nm",
    #     prefix="S",
    #     xunits=1e-3,
    # )
    # this looks wrong
    # coupler_fdtd = model_from_csv(
    #     filepath=sparameters_path / "coupler" / "coupler_G224n_L20_S220.csv",
    #     xkey="wavelength_nm",
    #     prefix="S",
    #     xunits=1e-3,
    # )
    # coupler_fdtd = model_from_npz(
    #     filepath=sparameters_path / "coupler" / "coupler_G224n_L20_S220.npz",
    #     xkey="wavelength_nm",
    #     prefix="S",
    #     xunits=1e-3,
    # )
    filepath = get_sparameters_path_tidy3d(gf.c.mmi1x2)
    coupler_fdtd = model_from_npz(filepath)
    plot_model(coupler_fdtd)
    plt.show()
