import argparse
import configparser
import logging
import pathlib
import sys

import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import numpy as np
import scipy.signal as sgl
from scipy import fft

parser = argparse.ArgumentParser(description="Pre-process time-series")
parser.add_argument("-v", "--verbose", action="count", default=0)
parser.add_argument("-c", "--config", default="config.ini")
args = parser.parse_args()

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

log.info("Starting time-series pre-processing")
config = configparser.ConfigParser()
config.read(args.config)

inp_root = pathlib.Path(config.get("inp", "root"), "cerema/raw")
out_root = pathlib.Path(config.get("out", "root"))
out_root.mkdir(exist_ok=True)

raw_ts = []
#for tsi in sorted(inp_root.glob("2017022817*.raw")):
for tsi in sorted(inp_root.glob("*.raw")):
    raw_ts.append(
        np.loadtxt(
            tsi,
            dtype=[("state", int), ("z", float), ("y", float), ("x", float)],
            delimiter=",",
            max_rows=2304,
        )
    )
    log.debug(f"Loading <{tsi}>")
n = len(raw_ts)
raw_ts = np.concatenate(raw_ts)
log.debug(f"{raw_ts=}")

t0 = np.linspace(0, 30 * 60 * n, 2304 * n, endpoint=False)
t = (t0 * 1e3).astype("timedelta64[ms]") + np.datetime64("2017-02-28T00:00")

if (errs := np.count_nonzero(raw_ts["state"])) != 0:
    log.warning(f"{errs} transmission errors!")
    log.debug(f"{dict(zip(*np.unique(raw_ts['state'], return_counts=True)))}")
    # log.debug(f"{t[raw_ts['state'] != 0]}")

sos = sgl.butter(1, 0.2, btype="lowpass", fs=2305 / (30 * 60), output="sos")
z = sgl.sosfiltfilt(sos, raw_ts["z"]*1e-2)
cr0 = np.where(np.diff(np.sign(z)))[0]

wave = np.fromiter(
    (
        np.max(np.abs(z[cr0[i - 1] : cr0[i]])) + np.max(np.abs(z[cr0[i] : cr0[i + 1]]))
        #(np.max(np.abs(raw_ts["z"][cr0[i - 1] : cr0[i]])) + np.max(np.abs(raw_ts["z"][cr0[i] : cr0[i + 1]]))) * 1e-2
        for i in range(1, len(cr0) - 1)
    ),
    dtype=np.single,
)

log.debug(f"{wave=}")
log.debug(f"{t=}")

# plt.plot(t[cr0[1:-1]], wave)


dt = 30 * 60 / 2304
# Mlims = (int(5 / dt), int(30 / dt))
N = t.size // 24
s0 = 2 * dt
dj = 0.5
J = 1 / dj * np.log2(N * dt / s0)
j = np.arange(0, J)
sj = s0 * 2 ** (j * dj)
Tj = 2 * sj * np.pi / 5
# sj = s0 * np.arange(1, 2 ** (J * dj))
Mw = sj / dt
Mlims = sj[[0, -1]]
M = (np.abs(sgl.cwt(raw_ts["z"]*1e-2, sgl.morlet2, Mw))/np.std(raw_ts["z"]*1e-2))**2
# M = np.abs(sgl.cwt(z, sgl.morlet, Mw))
v = np.max(np.abs(M))

fig, ax = plt.subplots()
# ax2 = ax.twinx()
# ax.plot(t0, raw_ts["z"], lw=.5, c="k", alpha=.2)
# ax.plot(t0, z, ls="-.", lw=.25, alpha=.2, c="k")
st = raw_ts["state"][raw_ts["state"] != 0]
c = np.asarray(["g", "b", "r"])
# ax.vlines(t0[raw_ts["state"] != 0], -20, 20, colors=c[np.where(st != 777, st, 0)])
# ax.set(xlabel="t (s)", ylabel="z (cm)")
# ax.set(xlim=(17 * 3600 + 20 * 60, 17 * 3600 + 30 * 60))
ax.grid(c="k", alpha=0.2)
ax.set(zorder=1, frame_on=False)
ax.semilogy()
a = [t0[0], t0[-1], *Mlims]
# c = ax.imshow(M, extent=a, aspect="auto", cmap="plasma", vmin=0)
c = ax.contourf(t, Tj, M, cmap="Greys", vmin=0, vmax=v)
fig.colorbar(c)

H13 = np.quantile(wave, 2 / 3)
Hs = 4*np.std(raw_ts["z"]*1e-2)
th = 2 * Hs
log.info(f"Threshold: {th}m")
bigw = np.where(wave > th)[0]
ym = 1.1 * np.max(np.abs(z))
nw = wave.size / 2
nlw = bigw.size
log.info(f"Number of waves: {nw}")
log.info(f"Number of waves >m: {nlw}")
log.info(f"Frequency: {nlw/nw:e}")
log.info(f"Frequency: 1/{nw/nlw:.0f}")
log.info(f"H1/3: {H13}m")
log.info(f"Hs: {Hs}")

if bigw.size > 32:
    log.warning(f"Number of large waves: {bigw.size}")
    sys.exit()

fig, ax_ = plt.subplots(2 * (bigw.size // 2), 2, figsize=(15/2.54, 4/3*10/2.54), constrained_layout=True)
for w, ax2, ax in zip(bigw, ax_[::2].flatten(), ax_[1::2].flatten()):
    i0 = cr0[w] - int(400 / dt)
    i1 = cr0[w + 2] + int(400 / dt)
    # a = [t0[i0], t0[i1], *Mlims]
    # c = ax2.imshow(M[:, i0:i1], extent=a, aspect="auto", cmap="Spectral", vmin=-v, vmax=+v)
    ws = np.ptp(raw_ts["z"][cr0[w]:cr0[w+2]]) * 1e-2
    log.info(f"Wave [{w}] size: {ws:.2f}m")
    c = ax2.contourf(
        t[i0:i1],
        Tj,
        M[:, i0:i1],
        cmap="Greys",
        vmin=0,
        levels=[1, 2.5, 5, 10, 20, 40],
        extend="both",
    )
    fig.colorbar(c, ax=ax2, label="NWPS")
    ax.plot(t[i0:i1], (raw_ts["z"]*1e-2)[i0:i1], c="k", lw=1)
    #ax.plot(t[i0:i1], z[i0:i1], c="k", lw=1, alpha=0.2, ls="-.")
    # ax.vlines(t[raw_ts["state"] != 0], -20, 20, colors=c[np.where(st != 777, st, 0)])
    ax.set(xlim=(t[i0], t[i1 - 1]), ylim=(-ym, ym))
    ax2.set(ylim=(2, 200))
    ax2.set(ylabel="T (s)")
    ax2.grid(c="k", alpha=0.2)
    ax2.semilogy()
    ax.grid(c="k", alpha=.2)
    #ax.axhline(0, c="k", alpha=0.2, lw=1, ls="-.")
    #ax.set(zorder=1, frame_on=False)
    ax.set_xlabel("t (s)", loc="left")
    ax.set_ylabel("z (m)")
    ax.axvspan(t[cr0[w]], t[cr0[w+2]], color="k", alpha=.1)

    locator = mdates.AutoDateLocator(minticks=3, maxticks=7)
    formatter = mdates.ConciseDateFormatter(locator)
    ax.xaxis.set_major_locator(locator)
    ax.xaxis.set_major_formatter(formatter)
    ax2.xaxis.set_major_locator(locator)
    ax2.xaxis.set_major_formatter(formatter)
    ax2.axes.set_xticklabels([])

    ax2.set_rasterization_zorder(1.5)

fig.savefig(out_root.joinpath(f"wavelet.pdf"), dpi=300)
fig.savefig(out_root.joinpath(f"wavelet.png"), dpi=200)

#fig, ax = plt.subplots(constrained_layout=True)
## ax.plot(fft.rfftfreq(raw_ts["z"].size, dt), np.abs(fft.rfft(raw_ts["z"])), c="k", alpha=.2, lw=1)
#ax.plot(*sgl.welch(raw_ts["z"], 1 / dt), c="k", alpha=0.2, label="PSD")
#ax.plot(1 / sj, N * np.mean(M, axis=1), c="k", label="CWT")
## ax.grid(color="k", alpha=.2)
#ax.set(xlabel="T (s)", ylabel="PSD")
## ax2.set(ylabel="Average Wavelet Transform")
#ax.set(xlim=1 / Mlims)
#ax.legend()

plt.show()