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

import matplotlib.pyplot as plt
import matplotlib.animation as animation
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("-c", "--config", default="config.ini")
args = parser.parse_args()

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

log.info("Animating olaFlow output")
config = configparser.ConfigParser()
config.read(args.config)
out = pathlib.Path(config.get("post", "out"))
out.mkdir(parents=True, exist_ok=True)

with gzip.open(out.joinpath("pickle.gz"), "rb") as f:
    model = pickle.load(f)

x0 = config.getfloat("post", "x")
z0 = config.getfloat("post", "z")
i0 = np.argmin(np.abs((model.x - x0) + 1j * (model.z - z0)))

X, Z = np.meshgrid(np.unique(model.x), np.unique(model.z))

C = np.where(
    (model.x[:, None, None].astype(np.single) == X[None, :, :].astype(np.single))
    & (model.z[:, None, None].astype(np.single) == Z[None, :, :].astype(np.single))
)

P = np.full((model.t.size, *X.shape), np.nan)
P[:, C[1], C[2]] = model.fields["porosity"][:, C[0]]

AW = np.full((model.t.size, *X.shape), np.nan)
AW[:, C[1], C[2]] = model.fields["alpha.water"][:, C[0]]

U = np.full((model.t.size, *X.shape), np.nan)
U[:, C[1], C[2]] = np.linalg.norm(model.fields["U"], axis=1)[:, C[0]]

fig, ax = plt.subplots(figsize=(19.2, 10.8), dpi=100)
tit = ax.text(
    0.5,
    0.95,
    f"t={model.t[0]}s",
    horizontalalignment="center",
    verticalalignment="top",
    transform=ax.transAxes,
)
aw_m = ax.pcolormesh(X, Z, AW[0], vmin=0, vmax=1, cmap="Blues", zorder=1)
ax.pcolormesh(
    X,
    Z,
    P[1],
    vmin=0,
    vmax=1,
    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)


fig.colorbar(aw_m)
ax.set(xlabel="x (m)", ylabel="z (m)", aspect="equal", facecolor="#bebebe")
ax.grid(c="k", alpha=0.2)


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



figU, axU = plt.subplots(figsize=(19.2, 10.8), dpi=100)
u_m = axU.pcolormesh(
    X, Z, U[0], cmap="BuPu", vmin=0, vmax=np.nanquantile(U, .99), zorder=1, alpha=np.nan_to_num(AW[0]).clip(0, 1)
)
ur_m = axU.pcolormesh(
    X, Z, U[0], cmap="YlOrBr", vmin=0, vmax=np.nanquantile(U, .99), zorder=1, alpha=1-np.nan_to_num(AW[0]).clip(0, 1)
)
# aw_u = axU.contour(X, Z, AW[0], levels=(.5,))
figU.colorbar(u_m)
axU.set(xlabel="x (m)", ylabel="z (m)", aspect="equal", facecolor="#bebebe")
axU.grid(c="k", alpha=0.2)
titU = axU.text(
    0.5,
    0.95,
    f"t={model.t[0]}s",
    horizontalalignment="center",
    verticalalignment="top",
    transform=axU.transAxes,
)


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)