"""Plot a field saved by run.py (python plot.py results/<case>-<backend>): slices of Re(u) and level in dB,
plus an animation of the real pressure field Re(u exp(-i omega t)) over one period. Works for the room
and the cymbal problems."""

import argparse
import json
import os

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation, PillowWriter

SLICES = [("xy", "x", "y", 2), ("xz", "x", "z", 1), ("yz", "y", "z", 0)]


def masked(a):
    return np.ma.masked_invalid(a)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("out", help="a results directory written by run.py")
    ap.add_argument("--frames", type=int, default=32)
    args = ap.parse_args()

    d = np.load(os.path.join(args.out, "field.npz"))
    with open(os.path.join(args.out, "summary.json")) as f:
        if json.load(f).get("problem", "room") == "cymbal":
            return plot_cymbal(d, args)
    g, x0, k = d["grid"], d["x0"], float(d["k"])
    ext = [g[0], g[-1], g[0], g[-1]]
    us = {n: d[f"u_{n}"] for n, *_ in SLICES}

    # Color limits from the field more than half a wavelength from the source
    A, B = np.meshgrid(g, g, indexing="xy")
    far = []
    for n, lx, ly, fixed in SLICES:
        r = np.hypot(A - x0["xyz".index(lx)], B - x0["xyz".index(ly)])
        u = us[n]
        far.append(np.abs(u[np.isfinite(u) & (r > np.pi / k)]))
    finite = np.concatenate(far)
    umax = np.percentile(finite, 99)
    ref = np.median(finite)

    fig, axes = plt.subplots(2, 3, figsize=(13, 8.4), constrained_layout=True)
    for col, (n, lx, ly, fixed) in enumerate(SLICES):
        u = us[n]
        ax = axes[0, col]
        im0 = ax.imshow(masked(u.real), origin="lower", extent=ext, cmap="RdBu_r",
                        vmin=-umax, vmax=umax, interpolation="bilinear")
        ax.set_title(f"Re u, {n}-plane at {'xyz'[fixed]} = {x0[fixed]:.2f}")
        ax = axes[1, col]
        with np.errstate(divide="ignore", invalid="ignore"):
            level = 20 * np.log10(np.abs(u) / ref)
        im1 = ax.imshow(masked(level), origin="lower", extent=ext, cmap="Blues_r",
                        vmin=-30, vmax=15, interpolation="bilinear")
        ax.set_title(f"level (dB re median), {n}-plane")
        for ax in axes[:, col]:
            ax.plot(x0["xyz".index(lx)], x0["xyz".index(ly)], marker="o", ms=5,
                    mfc="none", mec="0.15", mew=1)
            ax.set_xlabel(lx)
            ax.set_ylabel(ly)
            ax.set_aspect("equal")
    fig.colorbar(im0, ax=axes[0, :], shrink=0.85, label="Re u")
    fig.colorbar(im1, ax=axes[1, :], shrink=0.85, label="dB")
    wl = 2 * np.pi / k
    fig.suptitle(f"Point source in a rounded cubical room: k = {k:g} "
                 f"(room width = {2 * float(d['a']) / wl:.2f} wavelengths), "
                 f"wall admittance {-float(d['lam']):g}")
    fig.savefig(os.path.join(args.out, "slices.png"), dpi=130)
    plt.close(fig)

    # Animation: Re(u e^{-i w t}) over one period on the three planes
    fig, axes = plt.subplots(1, 3, figsize=(13, 4.6), constrained_layout=True)
    ims = []
    for ax, (n, lx, ly, fixed) in zip(axes, SLICES):
        im = ax.imshow(masked(us[n].real), origin="lower", extent=ext, cmap="RdBu_r",
                       vmin=-umax, vmax=umax, interpolation="bilinear")
        ax.set_title(f"{n}-plane")
        ax.set_xlabel(lx)
        ax.set_ylabel(ly)
        ims.append(im)
    fig.colorbar(ims[-1], ax=axes, shrink=0.85, label="Re(u e^{-iωt})")
    title = fig.suptitle("")

    def frame(i):
        ph = np.exp(-2j * np.pi * i / args.frames)
        for im, (n, *_) in zip(ims, SLICES):
            im.set_data(masked((us[n] * ph).real))
        title.set_text(f"pressure over one period, t/T = {i / args.frames:.2f}")
        return ims

    anim = FuncAnimation(fig, frame, frames=args.frames, blit=False)
    anim.save(os.path.join(args.out, "animation.gif"), writer=PillowWriter(fps=12), dpi=80)
    plt.close(fig)
    print(f"wrote {args.out}/slices.png and {args.out}/animation.gif")


CYMBAL_SLICES = ["xz", "yz", "xy"]


def plot_cymbal(d, args):
    """Slices through the evaluation box around the cymbal. u(x) is the pressure at the
    listener due to a unit source at x (reciprocity). The level is shown relative to
    the free field, 20 log10 |u / G|, which is 0 dB where the cymbal has no effect."""
    x0, k, freq = d["x0"], float(d["k"]), float(d["freq"])
    rz, R = d["outline"], float(d["radius"])
    names = [n for n in CYMBAL_SLICES if f"u_{n}" in d]
    panels = []
    for n in names:
        ga, gb = d[f"axes_{n}"]
        fixed, val = d[f"plane_{n}"]
        fixed = int(fixed)
        ia, ib = [i for i in range(3) if i != fixed]
        A, B = np.meshgrid(ga, gb, indexing="xy")
        pts = np.zeros((3, A.size))
        pts[ia], pts[ib], pts[fixed] = A.ravel(), B.ravel(), val
        r = np.linalg.norm(pts - x0[:, None], axis=0)
        G = (np.exp(1j * k * r) / (4 * np.pi * r)).reshape(A.shape)
        panels.append((n, ga, gb, ia, ib, fixed, val, d[f"u_{n}"], G, r.reshape(A.shape)))
    far = np.concatenate([np.abs(u[np.isfinite(u) & (r > np.pi / k)]) for *_, u, G, r in panels])
    umax = np.percentile(far, 99)

    def outline(ax, ia, ib, fixed, val):
        if fixed == 2:  # horizontal plane: the cymbal's rim, projected
            t = np.linspace(0, 2 * np.pi, 200)
            ax.plot(R * np.cos(t), R * np.sin(t), color="0.2", lw=0.8, ls="--")
        elif abs(val) < R:  # vertical plane: the cross section, at distance |val| from the axis
            r, z = rz
            m = r >= abs(val)
            s = np.sqrt(np.maximum(r[m] ** 2 - val ** 2, 0))
            for sign in (1, -1):
                ax.fill(sign * s, z[m], color="0.25", lw=0)

    fig, axes = plt.subplots(2, len(panels), figsize=(4.4 * len(panels), 8.4), constrained_layout=True,
                             squeeze=False)
    for col, (n, ga, gb, ia, ib, fixed, val, u, G, r) in enumerate(panels):
        ext = [ga[0], ga[-1], gb[0], gb[-1]]
        ax = axes[0, col]
        im0 = ax.imshow(masked(u.real), origin="lower", extent=ext, cmap="RdBu_r",
                        vmin=-umax, vmax=umax, interpolation="bilinear")
        ax.set_title(f"Re u, {n}-plane at {'xyz'[fixed]} = {val:.2f} m")
        ax = axes[1, col]
        with np.errstate(divide="ignore", invalid="ignore"):
            level = 20 * np.log10(np.abs(u) / np.abs(G))
        im1 = ax.imshow(masked(level), origin="lower", extent=ext, cmap="PuOr_r",
                        vmin=-15, vmax=15, interpolation="bilinear")
        ax.set_title(f"level re free field (dB), {n}-plane")
        for ax in axes[:, col]:
            outline(ax, ia, ib, fixed, val)
            if abs(x0[fixed] - val) < 1e-9:
                ax.plot(x0[ia], x0[ib], marker="o", ms=5, mfc="none", mec="0.15", mew=1)
            ax.set_xlabel(f"{'xyz'[ia]} (m)")
            ax.set_ylabel(f"{'xyz'[ib]} (m)")
            ax.set_aspect("equal")
    fig.colorbar(im0, ax=axes[0, :], shrink=0.85, label="Re u")
    fig.colorbar(im1, ax=axes[1, :], shrink=0.85, label="dB")
    wl = 2 * np.pi / k
    fig.suptitle(f"Listener at ({x0[0]:.2f}, {x0[1]:.2f}, {x0[2]:.2f}) m near a rigid cymbal, "
                 f"{freq:g} Hz (diameter = {2 * R / wl:.2f} wavelengths)\nBy reciprocity u(x) is the "
                 f"pressure at the listener due to a unit source at x")
    fig.savefig(os.path.join(args.out, "slices.png"), dpi=130)
    plt.close(fig)

    # Animation of the vertical plane through the listener
    n, ga, gb, ia, ib, fixed, val, u, G, r = panels[0]
    fig, ax = plt.subplots(figsize=(6, 5.4), constrained_layout=True)
    im = ax.imshow(masked(u.real), origin="lower", extent=[ga[0], ga[-1], gb[0], gb[-1]], cmap="RdBu_r",
                   vmin=-umax, vmax=umax, interpolation="bilinear")
    outline(ax, ia, ib, fixed, val)
    ax.set_xlabel(f"{'xyz'[ia]} (m)")
    ax.set_ylabel(f"{'xyz'[ib]} (m)")
    fig.colorbar(im, ax=ax, shrink=0.85, label="Re(u e^{-iωt})")
    title = ax.set_title("")

    def frame(i):
        im.set_data(masked((u * np.exp(-2j * np.pi * i / args.frames)).real))
        title.set_text(f"{n}-plane, t/T = {i / args.frames:.2f}")
        return [im]

    anim = FuncAnimation(fig, frame, frames=args.frames, blit=False)
    anim.save(os.path.join(args.out, "animation.gif"), writer=PillowWriter(fps=12), dpi=80)
    plt.close(fig)
    print(f"wrote {args.out}/slices.png and {args.out}/animation.gif")


if __name__ == "__main__":
    main()