import argparse
import configparser
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("Starting sws -> olaFlow converter")
config = configparser.ConfigParser()
config.read(args.config)
out = pathlib.Path(config.get("post", "pickle"))
out.parent.mkdir(parents=True, exist_ok=True)

with out.open("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)

a = mp.Process(target=ani.save, args=(out.parent.joinpath("anim.mp4"),), kwargs={"fps": 24})
aU = mp.Process(target=aniU.save, args=(out.parent.joinpath("animU.mp4"),), kwargs={"fps": 24})

a.start()
aU.start()

a.join()
aU.join()
#plt.show()