import argparse
import gzip
from itertools import starmap
import logging
from multiprocessing import pool
import pathlib
import pickle
import sys

from cycler import cycler
import matplotlib.pyplot as plt
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(
    "-f",
    "--func",
    type=str,
    help="Post-process function to compare",
    default="graphUniform",
    choices=("graphUniform", "graphUniform2"),
)

args = parser.parse_args()

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

log.info("Plotting comparison of model output")


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


model = get_pickle(args.output)

figsize = 15 / 2.54, 6 / 2.54

fig, ax_ = plt.subplots(
    2,
    figsize=figsize,
    dpi=200,
    constrained_layout=True,
)

_ax = ax_[0]
v = np.nanmax(np.abs(np.where(
    model.post_fields[args.func]["alpha.water"] > 0.5,
    #np.linalg.norm(model.post_fields[args.func]["U"], axis=2),
    model.post_fields[args.func]["U"][..., 0],
    np.nan,
)))
v150 = np.nanmax(np.abs(np.where(
    (model.post_fields[args.func]["alpha.water"] > 0.5) & (model.t[:, None] > 170) & (model.t[:, None] < 200),
    #np.linalg.norm(model.post_fields[args.func]["U"], axis=2),
    model.post_fields[args.func]["U"][..., 0],
    np.nan,
)))
_data = model.post_fields[args.func]["U"][..., 0].T
#_c = _ax.contourf(
#    model.t,
#    model.post_fields[args.func]["x_U"],
#    _data,
#    cmap="PiYG",
#    #levels=[-15, -10, -5, -2, -1, 0, 1, 2, 5, 10, 15],
#    vmin=-np.nanmax(np.abs(_data)),
#    vmax=np.nanmax(np.abs(_data)),
#    extend="both",
#)
_c = _ax.imshow(
    _data[::-1],
    cmap="PiYG",
    alpha=np.clip(model.post_fields[args.func]["alpha.water"], 0, 1).T[::-1],
    extent=(
        model.t.min(),
        model.t.max(),
        model.post_fields[args.func]["x_U"].min(),
        model.post_fields[args.func]["x_U"].max(),
    ),
    vmin=-v150,
    vmax=v150,
    aspect="auto",
)
_ax.set(xlim=(0, 400))
_ax.set(facecolor="k")
_ax.xaxis.set_minor_locator(MultipleLocator(10))
_ax.yaxis.set_minor_locator(MultipleLocator(1))
_ax.set(ylabel="z (m)")
_ax.axes.set_xticklabels([])
fig.colorbar(_c, label=f"U (m/s)", ax=_ax)
log.info(f"Vitesse max: {v}m/s")
log.info(f"Vitesse max [170,200]: {v150}m/s")
log.info(f"Écart: {abs(np.nanmax(_data)-17.7)/17.7:%}")


x = model.post_fields[args.func]["x_U"]
i0 = np.argmin(np.abs(x - 5))
_data = model.post_fields[args.func]["U"][..., i0, 0]
_alpha = model.post_fields[args.func]["alpha.water"][..., i0]

ax = ax_[1]
ax.fill_between(model.t, np.where(_alpha > 0.5, _data, 0), lw=1, color="#898989", edgecolor="k")
#ax.autoscale(True, "x", True)
ax.set(xlim=(0, 400))
ax.set(xlabel="t (s)", ylabel="U (m/s)")
ax.grid(c="k", alpha=.2)
ax.xaxis.set_minor_locator(MultipleLocator(10))
ax.yaxis.set_minor_locator(MultipleLocator(2))

fig.savefig(
    args.output.joinpath(
        f"U_{args.func}.pdf"
    )
)
fig.savefig(
    args.output.joinpath(
        f"U_{args.func}.jpg"
    )
)