import argparse
import configparser
import logging
import pathlib

import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator, LogLocator, NullFormatter
import numpy as np
import scipy.signal as sgl

parser = argparse.ArgumentParser(description="Post-process swash output")
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("post")

log.info("Starting post-processing")
config = configparser.ConfigParser()
config.read(args.config)

inp = pathlib.Path(config.get("post", "inp"))
root = pathlib.Path(config.get("swash", "out"))

log.info(f"Reading data from '{inp}'")
x = np.load(inp.joinpath("x.npy"))
t = np.load(inp.joinpath("t.npy")) * 1e-3

botl = np.load(inp.joinpath("botl.npy"))
watl = np.load(inp.joinpath("watl.npy"))
vel = np.load(inp.joinpath("vel.npy"))[0]

# Plotting
log.info("Plotting results")

vlim = np.nanmin(np.maximum(watl, -botl)), np.nanmax(np.maximum(watl, -botl))

x0 = np.linspace(-600, -200, 5)
i0 = np.argmin(np.abs(x[:, None] - x0), axis=0)

fig_x, ax = plt.subplots(
    5, 1, figsize=(15 / 2.54, 15/ 2.54), constrained_layout=True
)
dt = np.mean(np.diff(t))
N = t.size
s0 = 2 * dt
dj = 0.5
J = 1 / dj * np.log2(N * dt / s0)
j = np.arange(0, J)
sj = s0 * 2 ** (j * dj)
Mw = sj / dt
sig = np.std(watl[:, i0])
M = np.stack([(np.abs(sgl.cwt(watl[:, i], sgl.morlet2, Mw))/sig)**2 for i in i0])
v = np.max(M)
T = 2 * sj * np.pi / 5

for ax_x, M_, x_ in zip(ax.reshape(-1), M, x[i0]):
    c = ax_x.contourf(t, T, M_, cmap="Greys", vmin=0, levels=[1, 2.5, 5, 10, 20, 40], extend="both")
    fig_x.colorbar(c, ax=ax_x, label="NWPS")
    ax_x.grid(color="k", alpha=0.2)
    ax_x.text(
        0.95,
        0.95,
        f"x={x_:.0f}m",
        horizontalalignment="right",
        verticalalignment="top",
        transform=ax_x.transAxes,
        #c="w",
    )
    ax_x.semilogy()
    ax_x.autoscale(True, "both", True)
    ax_x.set_rasterization_zorder(1.5)
    ax_x.set(ylabel="T (s)", ylim=(sj[0], sj[-1]))

    ax_x.xaxis.set_minor_locator(MultipleLocator(100))
    ax_x.yaxis.set_major_locator(LogLocator(10, numticks=2**10))
    ax_x.yaxis.set_minor_locator(LogLocator(10, subs=np.arange(10), numticks=2**10))
    ax_x.yaxis.set_minor_formatter(NullFormatter())
    if ax_x != ax.reshape(-1)[-1]:
        ax_x.axes.set_xticklabels([])
    else:
        ax_x.set(xlabel="t (s)")

out = pathlib.Path(config.get("post", "out")).joinpath(f"trans")
log.info(f"Saving plots in '{out}'")
out.mkdir(parents=True, exist_ok=True)

fig_x.savefig(out.joinpath("wavelet.pdf"), dpi=300)
fig_x.savefig(out.joinpath("wavelet.jpg"), dpi=200)
fig_x.savefig(out.joinpath("wavelet.png"), dpi=200)

log.info("Finished post-processing")