import argparse
import configparser
import logging
import pathlib

import matplotlib.pyplot as plt
import numpy as np
import scipy.signal as sgl

from .reflex3s import reflex3s
from .pa.reflex3S import reflex3S
from .pa.PUV import PUV_dam

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("xp.npy"))
t = np.load(inp.joinpath("tsec.npy"))

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

x0 = config.getint("post", "x0")
arg_x0 = np.abs(x - x0).argmin()
t0 = config.getfloat("post", "t0")
arg_t0 = np.abs(t - t0).argmin()
dt = np.diff(t).mean()
f = 1 / dt

idx = [arg_x0 - 5, arg_x0, arg_x0 + 7]
wl = watl[arg_t0:, idx]

# Reflex3S
res = reflex3S(*wl.T, *x[idx], botl[idx].mean(), f, 0.02, 0.2)
f_, ai_, ar_ = reflex3s(wl.T, x[idx], botl[idx].mean(), f)
ai, ar, _, _, _, _, fre = res

# Custom PUV
eta = watl[t > t0, arg_x0]
u = vel[t > t0, 0, arg_x0]

phi_eta = sgl.welch(eta, f)
phi_u = sgl.welch(u, f)
phi_eta_u = sgl.csd(eta, u, f)

R = np.sqrt(
    (np.abs(phi_eta[1]) + np.abs(phi_u[1]) - 2 * phi_eta_u[1].real)
    / (np.abs(phi_eta[1]) + np.abs(phi_u[1]) + 2 * phi_eta_u[1].real)
)

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

plt.plot(fre, np.asarray(ar) / np.asarray(ai))
plt.xlim((0, 0.3))
plt.ylim((0, 1))
plt.savefig(out.joinpath("reflex3s_pa.pdf"))
plt.clf()
plt.plot(f_, ar_ / ai_)
plt.plot(phi_eta[0], R)
plt.xlim((0, 0.3))
plt.ylim((0, 1))
plt.savefig(out.joinpath("reflex3s.pdf"))

# pressk = np.load(inp.joinpath("pressk.npy"))
# press = pressk[:, 1, :]
# res = PUV_dam(vel[arg_t0:, 0, 500], press[arg_t0:, 500], botl[500], f, botl[500]/2, botl[500]/2)
# print(res)