import argparse
import gzip
import logging
import multiprocessing as mp
import pathlib
import pickle

import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.gridspec import GridSpec
from matplotlib.ticker import MultipleLocator
import numpy as np
from scipy import interpolate

from .olaflow import OFModel

parser = argparse.ArgumentParser(description="Post-process olaflow results")
parser.add_argument("-v", "--verbose", action="count", default=0)
parser.add_argument(
    "-o",
    "--output",
    type=pathlib.Path,
    help="Output directory for pickled data",
    required=True,
)
parser.add_argument(
    "-m",
    "--max",
    help="Only compute maximum rather than animation",
    action="store_true",
)
parser.add_argument(
    "-i",
    "--initial",
    help="Only compute initial domain",
    action="store_true",
)
args = parser.parse_args()

logging.basicConfig(level=max((10, 20 - 10 * args.verbose)))
log = logging.getLogger("ola_post")

log.info("Animating olaFlow output")
out = args.output
out.mkdir(parents=True, exist_ok=True)

with (
    path.open("rb")
    if (path := out.joinpath("pickle")).exists()
    else gzip.open(path.with_suffix(".gz"), "rb")
) as f:
    model = pickle.load(f)

x0, idx0 = np.unique(model.x.astype(np.half), return_inverse=True)
z0, idz0 = np.unique(model.z.astype(np.half), return_inverse=True)

ix0 = np.argsort(x0)
iz0 = np.argsort(z0)[::-1]

X, Z = np.meshgrid(x0, z0)

P = np.full((model.t.size, *X.shape), np.nan)
P[:, iz0[idz0], ix0[idx0]] = model.fields["porosity"]

AW = np.full((model.t.size, *X.shape), np.nan)
AW[:, iz0[idz0], ix0[idx0]] = model.fields["alpha.water"]

U = np.full((model.t.size, *X.shape), np.nan)
U[:, iz0[idz0], ix0[idx0]] = np.linalg.norm(model.fields["U"], axis=1)

fig = plt.figure(
    figsize=(15 / 2.54, 4 / 2.54), dpi=200, constrained_layout=True
)
gs = GridSpec(
    3 if not args.initial else 1,
    1,
    figure=fig,
    height_ratios=[1, 0.1, 0.1] if not args.initial else [1],
)
ax = fig.add_subplot(gs[0])
if not args.initial:
    cax1 = fig.add_subplot(gs[1])
    cax2 = fig.add_subplot(gs[2])
aw_m = ax.imshow(
    AW[0],
    vmin=0,
    vmax=1,
    extent=(x0.min(), x0.max(), z0.min(), z0.max()),
    cmap="Blues",
    zorder=1,
)
p_m = ax.imshow(
    P[1],
    vmin=0,
    vmax=1,
    extent=(x0.min(), x0.max(), z0.min(), z0.max()),
    cmap="Greys_r",
    alpha=(np.nan_to_num(1 - P[1]) / 2).clip(0, 1),
    zorder=1.1,
)
ax.axhline(4.5, ls="-.", lw=1, c="k", alpha=0.2, zorder=1.2)


if not args.initial:
    fig.colorbar(
        aw_m, label=r"$\alpha_w$", cax=cax1, shrink=0.6, orientation="horizontal"
    )
    fig.colorbar(p_m, label=r"Porosity", cax=cax2, shrink=0.6, orientation="horizontal")
ax.set(xlabel="x (m)", ylabel="z (m)", aspect="equal", facecolor="#000000")
ax.grid(c="k", alpha=0.2)
ax.xaxis.set_minor_locator(MultipleLocator(5))
ax.yaxis.set_minor_locator(MultipleLocator(5))


figU = plt.figure()
gsU = GridSpec(
    2 if args.max else 3,
    1,
    figure=figU,
    height_ratios=[1, 0.05] if args.max else [1, 0.05, 0.05],
)
axU = figU.add_subplot(gsU[0])
caxu1 = figU.add_subplot(gsU[1])
if not args.max:
    caxu2 = figU.add_subplot(gsU[2])
u_m = axU.imshow(
    U[0],
    cmap="BuPu",
    vmin=0,
    vmax=20,
    extent=(x0.min(), x0.max(), z0.min(), z0.max()),
    zorder=1,
    alpha=np.nan_to_num(AW[0]).clip(0, 1),
)
ur_m = axU.imshow(
    U[0],
    cmap="YlOrBr",
    vmin=0,
    vmax=20,
    extent=(x0.min(), x0.max(), z0.min(), z0.max()),
    zorder=1,
    alpha=1 - np.nan_to_num(AW[0]).clip(0, 1),
)
# aw_u = axU.contour(X, Z, AW[0], levels=(.5,))
axU.set(xlabel="x (m)", ylabel="z (m)", aspect="equal", facecolor="#bebebe")
axU.grid(c="k", alpha=0.2)
figU.colorbar(u_m, label=r"$U_w$", cax=caxu1, shrink=0.6, orientation="horizontal")


if args.max:
    aw_m.set_array(AW.max(axis=0))

    u_m.set_array(np.nanmax(np.where(AW > 0.5, U, np.nan), axis=0, initial=0))
    u_m.set_alpha(1)
    u_m.set_cmap("hot_r")
    ur_m.remove()

    fig.savefig(out.joinpath("max_aw.pdf"))
    figU.savefig(out.joinpath("max_U.pdf"))
elif args.initial:
    aw_m.set_array(AW[0])
    ax.vlines(-20, -15, 15, color="k", lw=1, ls="--", label="Measurements")
    ax.text(-20, 15, "Measurements", ha="right", va="bottom")

    fig.savefig(out.joinpath("aw_t0.pdf"))
    fig.savefig(out.joinpath("aw_t0.jpg"), dpi=200)
else:
    fig.set(figwidth=19.2, figheight=10.8, dpi=100)
    figU.set(figwidth=19.2, figheight=10.8, dpi=100)

    figU.colorbar(ur_m, label=r"$U_a$", cax=caxu2, shrink=0.6, orientation="horizontal")

    tit = ax.text(
        0.5,
        0.95,
        f"t={model.t[0]}s",
        horizontalalignment="center",
        verticalalignment="top",
        transform=ax.transAxes,
    )
    titU = axU.text(
        0.5,
        0.95,
        f"t={model.t[0]}s",
        horizontalalignment="center",
        verticalalignment="top",
        transform=axU.transAxes,
    )

    def anim(i):
        tit.set_text(f"t={model.t[i]}s")
        aw_m.set_array(AW[i])

    def animU(i):
        titU.set_text(f"t={model.t[i]}s")
        u_m.set_array(U[i])
        u_m.set_alpha(np.nan_to_num(AW[i]).clip(0, 1))
        ur_m.set_array(U[i])
        ur_m.set_alpha(1 - np.nan_to_num(AW[i]).clip(0, 1))

    ani = animation.FuncAnimation(fig, anim, frames=model.t.size, interval=1 / 24)
    aniU = animation.FuncAnimation(figU, animU, frames=model.t.size, interval=1 / 24)

    ani.save(out.joinpath("anim.mp4"), fps=24)
    aniU.save(out.joinpath("animU.mp4"), fps=24)