"""Matplotlib views of a plate, its modes, and a sound. Every function draws
into ``ax`` (a new figure by default) and returns the axes. Requires
``matplotlib`` (``pip install cymsim[plot]``)."""

from __future__ import annotations

import numpy as np

from .plate import Plate
from .sound import Sound
from .spectrogram import Spectrogram

__all__ = ["plate", "mode", "frame", "signal", "spectrogram", "decay_times", "force", "animate"]


def _ax(ax, figsize=(6, 5)):
    import matplotlib.pyplot as plt

    if ax is None:
        _, ax = plt.subplots(figsize=figsize)
    return ax


def _extent(pl: Plate):
    g = pl.grid
    return [g.x[0] * 1e3 - g.hx * 500, g.x[-1] * 1e3 + g.hx * 500, g.y[0] * 1e3 - g.hy * 500, g.y[-1] * 1e3 + g.hy * 500]


def field(pl: Plate, values: np.ndarray, ax=None, *, cmap="viridis", label: str = "", symmetric: bool = False, strike=None, listener=None):
    """A color map of a grid array over the plate (outside masked), in mm coordinates."""
    ax = _ax(ax)
    v = np.ma.masked_where(~pl.grid.mask, values)
    kw = {}
    if symmetric:
        m = float(np.abs(v).max()) or 1.0
        kw = {"vmin": -m, "vmax": m}
    im = ax.imshow(v, origin="lower", extent=_extent(pl), cmap=cmap, interpolation="nearest", **kw)
    ax.figure.colorbar(im, ax=ax, label=label)
    if strike is not None:
        x, y = pl.strike_point(strike)
        ax.plot(x * 1e3, y * 1e3, "o", mfc="none", mec="yellow", mew=2, ms=10)
    if listener is not None:
        lx, ly, _ = listener.position
        ax.plot(lx * 1e3, ly * 1e3, "v", color="white", ms=8, clip_on=False)
    ax.set_xlabel("x [mm]")
    ax.set_ylabel("y [mm]")
    ax.set_aspect("equal")
    return ax


def plate(pl: Plate, ax=None, what: str = "auto", strike=None, listener=None):
    """The plate: its profile (when it has one) or its thickness, with the strike point and the listener when given."""
    if what == "auto":
        what = "thickness" if pl.is_flat else "profile"
    if what == "profile":
        ax = field(pl, pl.w0 * 1e3, ax, cmap="cividis", label="stress-free height [mm]", strike=strike, listener=listener)
    elif what == "thickness":
        ax = field(pl, pl.h * 1e3, ax, cmap="magma", label="thickness [mm]", strike=strike, listener=listener)
    else:
        raise ValueError("what must be 'auto', 'profile' or 'thickness'")
    ax.set_title(pl.name)
    return ax


def mode(modes, m: int, ax=None, amplitude: float = 1e-3):
    """The transverse shape of mode ``m``, scaled to the given peak displacement [m]."""
    w = modes.w[m]
    w = w / (np.abs(w).max() or 1.0) * amplitude
    ax = field(modes.plate, w * 1e3, ax, cmap="RdBu_r", label="w [mm]", symmetric=True)
    ax.set_title(f"mode {m}: {modes.freqs[m]:.1f} Hz, T60 {modes.t60[m]:.2f} s")
    return ax


def frame(sound: Sound, index: int, ax=None, scale: float | None = None):
    """Movie frame ``index`` of the transverse displacement; ``scale`` fixes the color range [m]."""
    if sound.frames is None:
        raise ValueError("this sound has no movie frames; run with movie=True")
    w = sound.frames[index]
    ax = _ax(ax)
    v = np.ma.masked_where(~sound.plate.grid.mask, w * 1e3)
    m = (scale or float(np.abs(sound.frames).max()) or 1e-9) * 1e3
    im = ax.imshow(v, origin="lower", extent=_extent(sound.plate), cmap="RdBu_r", vmin=-m, vmax=m, interpolation="nearest")
    ax.figure.colorbar(im, ax=ax, label="w [mm]")
    ax.set_title(f"t = {sound.frame_time(index) * 1e3:.1f} ms")
    ax.set_xlabel("x [mm]")
    ax.set_ylabel("y [mm]")
    ax.set_aspect("equal")
    return ax


def signal(sound: Sound, ax=None, tmax: float | None = None):
    """The listener pressure against time."""
    ax = _ax(ax, (8, 3))
    t = sound.time
    sel = slice(None) if tmax is None else t <= tmax
    ax.plot(t[sel], sound.pressure[sel], lw=0.5)
    ax.set_xlabel("time [s]")
    ax.set_ylabel("pressure [Pa]")
    ax.set_title(f"{sound.plate.name}, {sound.info.get('model', '')} model, peak {sound.peak_spl:.0f} dB SPL")
    return ax


def spectrogram(s: Sound | Spectrogram, ax=None, *, db_range: float = 90.0, fmin: float = 20.0, fmax: float | None = None, log_freq: bool = True, modes=None, window: int = 2048, hop: int = 256):
    """The spectrogram of a sound (or a precomputed ``Spectrogram``), with the
    modal frequencies drawn as ticks on the left edge when ``modes`` is given."""
    sp = s.spectrogram(window, hop) if isinstance(s, Sound) else s
    ax = _ax(ax, (8, 4))
    fmax = fmax or sp.freqs[-1]
    sel = (sp.freqs >= (fmin if log_freq else 0)) & (sp.freqs <= fmax)
    ax.pcolormesh(sp.times, sp.freqs[sel], sp.db[sel], cmap="inferno", vmin=-db_range, vmax=0, shading="nearest", rasterized=True)
    if log_freq:
        ax.set_yscale("log")
        ax.set_ylim(fmin, fmax)
    if modes is not None:
        f = modes.freqs[(modes.freqs >= fmin) & (modes.freqs <= fmax)]
        ax.plot(np.full(len(f), sp.times[0]), f, "_", color="cyan", ms=6, mew=1)
    ax.set_xlabel("time [s]")
    ax.set_ylabel("frequency [Hz]")
    return ax


def decay_times(modes, ax=None):
    """Decay time of every mode against frequency, total and by mechanism."""
    ax = _ax(ax, (8, 4))
    by = modes.t60_by_mechanism()
    ax.plot(modes.freqs, modes.t60, "k.", ms=4, label="total")
    for name, marker in (("radiation", "b."), ("material", "g."), ("thermal", "r.")):
        ax.plot(modes.freqs, by[name], marker, ms=2, alpha=0.6, label=name)
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_ylim(0.01, 1e4)
    ax.set_xlabel("frequency [Hz]")
    ax.set_ylabel("T60 [s]")
    ax.legend()
    return ax


def force(sound: Sound, ax=None):
    """The stick force during the contact (linear model only, which records it)."""
    if "force" not in sound.info:
        raise ValueError("the force history is recorded by the linear model only")
    ax = _ax(ax, (6, 3))
    f = sound.info["force"]
    ax.plot(np.arange(len(f)) * sound.info["force_dt"] * 1e3, f)
    ax.set_xlabel("time [ms]")
    ax.set_ylabel("force [N]")
    ax.set_title(f"peak {sound.info['peak_force']:.0f} N, contact {sound.info['contact_time'] * 1e3:.2f} ms")
    return ax


def animate(sound: Sound, interval_ms: float | None = None, scale: float | None = None, stride: int = 1):
    """A ``matplotlib.animation.FuncAnimation`` of the movie frames (display with ``HTML(anim.to_jshtml())`` in a notebook)."""
    import matplotlib.pyplot as plt
    from matplotlib.animation import FuncAnimation

    if sound.frames is None:
        raise ValueError("this sound has no movie frames; run with movie=True")
    fig, ax = plt.subplots(figsize=(6, 5))
    m = (scale or float(np.abs(sound.frames).max()) or 1e-9) * 1e3
    mask = ~sound.plate.grid.mask
    im = ax.imshow(np.ma.masked_where(mask, sound.frames[0] * 1e3), origin="lower", extent=_extent(sound.plate), cmap="RdBu_r", vmin=-m, vmax=m, interpolation="nearest")
    fig.colorbar(im, ax=ax, label="w [mm]")
    ax.set_aspect("equal")
    title = ax.set_title("")

    def update(i):
        im.set_data(np.ma.masked_where(mask, sound.frames[i] * 1e3))
        title.set_text(f"t = {sound.frame_time(i) * 1e3:.1f} ms")
        return im, title

    frames = range(0, len(sound.frames), stride)
    return FuncAnimation(fig, update, frames=frames, interval=interval_ms or 1000 * stride / sound.fps, blit=False)