"""Plot the room 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."""

import argparse
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"))
    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")


if __name__ == "__main__":
    main()