""" Visualizations for common meow-datatypes """

from collections.abc import Iterable
from typing import Any, Callable

import numpy as np

from meow.structures import _visualize_structures

try:
    import matplotlib.pyplot as plt  # fmt: skip


except ImportError:
    plt = None

try:
    import gdsfactory as gf  # fmt: skip
except ImportError:
    gf = None

try:
    from jaxlib.xla_extension import DeviceArray  # fmt: skip # type: ignore
except ImportError:
    DeviceArray = None


def _visualize_s_matrix(S, fmt=None, title=None, show=True, phase=False, ax=None):
    import matplotlib.pyplot as plt  # fmt: skip
    if phase:
        fmt = ".0f"
    else:
        fmt = ".3f"

    Z = np.abs(S)
    y, x = np.arange(Z.shape[0])[::-1], np.arange(Z.shape[1])
    Y, X = np.meshgrid(y, x)

    if ax:
        plt.sca(ax)
    else:
        plt.figure(figsize=(2 * x.shape[0] / 3, 2 * y.shape[0] / 3))

    plt.pcolormesh(X, Y, Z[::-1].T, cmap="Greys", vmin=0.0, vmax=2.0 * Z.max())

    coords_ = np.concatenate(
        [np.array([x[0] - 1]), x, np.array([x[-1] + 1])], axis=0, dtype=float
    )
    labels = ["" for _ in coords_]
    plt.xticks(coords_ + 0.5, labels)
    plt.xlim(coords_[0] + 0.5, coords_[-1] - 0.5)

    coords_ = np.concatenate(
        [np.array([y[0] + 1]), y, np.array([y[-1] - 1])], axis=0, dtype=float
    )
    coords_ = coords_[::-1]  # reverse
    labels = ["" for _ in coords_]
    plt.yticks(coords_ + 0.5, labels)
    plt.ylim(coords_[-1] - 0.5, coords_[0] + 0.5)
    plt.grid(True)

    for x, y, z in zip(X.ravel(), Y.ravel(), S[::-1].T.ravel()):
        if np.abs(z) > 0.0005:
            if phase:
                z = np.angle(z) * 180 / np.pi
            text = eval(f"f'{{z:{fmt}}}'")  # 😅
            text = text.replace("+", "\n+")
            text = text.replace("-", "\n-")
            if text[0] == "\n":
                text = text[1:]
            plt.text(x, y, text, ha="center", va="center", fontsize=8)

    if title is not None:
        plt.title(title)

    if show:
        plt.show()


def _visualize_s_pm_matrix(Spm, fmt=None, title=None, show=True, phase=False, ax=None):
    import matplotlib.pyplot as plt  # fmt: skip

    S, pm = Spm
    _visualize_s_matrix(S, fmt=fmt, title=title, show=False, phase=phase, ax=ax)
    num_left = len([p for p in pm if "left" in p])
    Z = np.abs(S)
    _, x = np.arange(Z.shape[0])[::-1], np.arange(Z.shape[1])

    plt.axvline(x[num_left] - 0.5, color="red")
    plt.axhline(x[num_left] - 0.5, color="red")

    if show:
        plt.show()


def _visualize_overlap_density(
    two_modes,
    conjugated=True,
    x_symmetry=False,
    y_symmetry=False,
    ax=None,
    n_cmap=None,
    mode_cmap=None,
    num_levels=8,
    show=True,
):
    import matplotlib.pyplot as plt  # fmt: skip

    from .mode import Mode  # fmt: skip

    mode1, mode2 = two_modes
    if conjugated:
        cross = mode1.Ex * mode2.Hy.conj() - mode1.Ey * mode2.Hx.conj()
    else:
        cross = mode1.Ex * mode2.Hy - mode1.Ey * mode2.Hx
    if x_symmetry:
        cross = 0.5 * (cross + cross[::-1])
    if y_symmetry:
        cross = 0.5 * (cross + cross[:, ::-1])
    zeros = np.zeros_like(cross)
    overlap = Mode(
        neff=mode1.neff,
        cs=mode1.cs,
        Ex=cross,
        Ey=zeros,
        Ez=zeros,
        Hx=zeros,
        Hy=zeros,
        Hz=zeros,
    )
    if ax is None:
        W, H = _figsize_visualize_mode(mode1.cs, 5)
        _, ax = plt.subplots(1, 3, figsize=(3 * W, H))

    field = "Ex" if mode1.te_fraction > 0.5 else "Hx"
    mode1._visualize(
        title=f"mode 1: {field}",
        fields=[field],
        ax=ax[0],
        n_cmap=n_cmap,
        mode_cmap=mode_cmap,
        num_levels=num_levels,
        show=False,
    )

    field = "Ex" if mode2.te_fraction > 0.5 else "Hx"
    mode2._visualize(
        title=f"mode 2: {field}",
        fields=[field],
        ax=ax[1],
        n_cmap=n_cmap,
        mode_cmap=mode_cmap,
        num_levels=num_levels,
        show=False,
    )

    title = "overlap density" + ("" if conjugated else " (no conjugations)")
    p = overlap._visualize(
        title=title,
        fields=["Ex"],
        ax=ax[2],
        n_cmap=n_cmap,
        mode_cmap=mode_cmap,
        num_levels=num_levels,
        show=False,
    )

    if show:
        plt.show()

    return p


def _visualize_gf_component(comp):
    import gdsfactory as gf  # fmt: skip

    gf.plot(comp)  # type: ignore


def _is_two_tuple(obj):
    if not isinstance(obj, tuple):
        return False
    try:
        x, y = obj  # type: ignore
        return True
    except Exception:
        return False


def _figsize_visualize_mode(cs, W0):
    x_min, x_max = cs.mesh.x.min(), cs.mesh.x.max()
    y_min, y_max = cs.mesh.y.min(), cs.mesh.y.max()
    delta_x = x_max - x_min
    delta_y = y_max - y_min
    aspect = delta_y / delta_x
    W, H = W0 + 1, W0 * aspect + 1
    return W, H


def _visualize_modes(
    modes,
    n_cmap=None,
    mode_cmap=None,
    num_levels=8,
    operation=lambda x: np.abs(x) ** 2,
    show=True,
    plot_width=6.4,
    fields=("Ex", "Hx"),
    ax=None,
):
    import matplotlib.pyplot as plt  # fmt: skip

    num_modes = len(modes)
    cs = modes[0].cs
    W, H = _figsize_visualize_mode(cs, plot_width)

    if ax is None:
        fig, ax = plt.subplots(
            num_modes,
            2,
            figsize=(2 * W, num_modes * H),
            sharex=True,
            sharey=True,
            squeeze=False,
        )
    else:
        fig = None
    for i, m in enumerate(modes):
        m._visualize(
            title=None,
            title_prefix=f"m{i}: ",
            fields=list(fields),
            ax=ax[i],
            n_cmap=n_cmap,
            mode_cmap=mode_cmap,
            num_levels=num_levels,
            operation=operation,
            show=False,
        )
    if fig is not None:
        fig.subplots_adjust(hspace=0, wspace=2 / (2 * W))
    if show:
        plt.show()


def _visualize_base_model(obj, **kwargs):
    return obj._visualize(**kwargs)


def _is_mode_list(obj: Any) -> bool:
    from .mode import Mode  # fmt: skip
    return isinstance(obj, Iterable) and all(isinstance(o, Mode) for o in obj)


def _is_structure_3d_list(obj: Any) -> bool:
    from .structures import Structure3D  # fmt: skip
    return isinstance(obj, Iterable) and all(isinstance(o, Structure3D) for o in obj)


def _is_base_model(obj: Any) -> bool:
    from .base_model import BaseModel  # fmt: skip
    return isinstance(obj, BaseModel)


def _is_mode_overlap(obj: Any) -> bool:
    from .mode import Mode  # fmt: skip
    return _is_two_tuple(obj) and all(isinstance(o, Mode) for o in obj)


def _is_s_matrix(obj: Any):
    return (
        (
            isinstance(obj, np.ndarray)
            or (DeviceArray is not None and isinstance(obj, DeviceArray))
        )
        and obj.ndim == 2
        and obj.shape[0] > 1
        and obj.shape[1] > 1
    )


def _is_s_pm_matrix(obj):
    return _is_two_tuple(obj) and _is_s_matrix(obj[0]) and isinstance(obj[1], dict)


def _is_gf_component(obj):
    return gf is not None and isinstance(obj, gf.Component)


VISUALIZATION_MAPPING: dict[Callable, Callable] = {
    _is_base_model: _visualize_base_model,
    _is_mode_list: _visualize_modes,
    _is_structure_3d_list: _visualize_structures,
    _is_mode_overlap: _visualize_overlap_density,
    _is_s_matrix: _visualize_s_matrix,
    _is_s_pm_matrix: _visualize_s_pm_matrix,
    _is_gf_component: _visualize_gf_component,
}


def visualize(obj: Any, **kwargs: Any):
    """visualize any meow object

    Args:
        obj: the meow object to visualize
        **kwargs: extra configuration to visualize the object

    Note:
        Most meow objects have a `._visualize` method.
        Check out its help to see which kwargs are accepted.
    """
    if plt is None:
        print(obj)
        return

    for check_func, vis_func in VISUALIZATION_MAPPING.items():
        if check_func(obj):
            return vis_func(obj, **kwargs)
    print(obj)


vis = visualize  # shorthand for visualize
