#!/usr/bin/env python3
"""Mean velocity and turbulence intensity profiles (AeroSim) vs EN 1991-1-4, terrain category II."""

import argparse
import glob

import h5py
import matplotlib.pyplot as plt
import numpy as np

Z0 = 0.05
Z_MIN = 2.0
Z0_II = 0.05
K_I = 1.0
C_0 = 1.0
Z_REF = 10.0


def read_line(path, t_start):
    with h5py.File(path, "r") as f:
        z = f["Geometry"][:, 2]
        x = float(f["Geometry"][0, 0])
        keys = sorted(f["ux"].keys(), key=lambda k: float(k[1:]))
        keys = [k for k in keys if float(k[1:]) >= t_start]
        ux = np.stack([f["ux"][k][:] for k in keys])
        uy = np.stack([f["uy"][k][:] for k in keys])
        uz = np.stack([f["uz"][k][:] for k in keys])
    return x, z, ux, uy, uz


def profiles(ux, uy, uz):
    mag = np.hypot(ux.mean(axis=0), uy.mean(axis=0))
    with np.errstate(divide="ignore", invalid="ignore"):
        iu = ux.std(axis=0) / mag
    return mag, iu


def eurocode(z, v_b):
    zc = np.maximum(z, Z_MIN)
    k_r = 0.19 * (Z0 / Z0_II) ** 0.07
    v_m = k_r * np.log(zc / Z0) * C_0 * v_b
    i_v = K_I / (C_0 * np.log(zc / Z0))
    return v_m, i_v


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--dir", default=".")
    p.add_argument("--t-start", type=float, default=20.0)
    p.add_argument("--z-max", type=float, default=100.0)
    p.add_argument("--v-b", type=float, default=None)
    p.add_argument("-o", "--out", default="abl_cat2_vs_eurocode.png")
    args = p.parse_args()

    files = sorted(glob.glob(f"{args.dir}/*line*.inst.h5"))
    if not files:
        raise SystemExit("no line export found")

    data = []
    for f in files:
        x, z, ux, uy, uz = read_line(f, args.t_start)
        mag, iu = profiles(ux, uy, uz)
        data.append((x, z, mag, iu))
    data.sort(key=lambda d: d[0])

    v_b = args.v_b
    if v_b is None:
        ref = data[len(data) // 2]
        k_r = 0.19 * (Z0 / Z0_II) ** 0.07
        v_b = np.interp(Z_REF, ref[1], ref[2]) / (k_r * np.log(max(Z_REF, Z_MIN) / Z0) * C_0)

    z_ec = np.linspace(0.1, args.z_max, 400)
    vm_ec, iv_ec = eurocode(z_ec, v_b)

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 6), sharey=True)
    colors = plt.cm.viridis(np.linspace(0.15, 0.8, len(data)))

    for (x, z, mag, iu), c in zip(data, colors):
        m = z <= args.z_max
        label = f"AeroSim, x = {x:.0f} m"
        ax1.plot(mag[m], z[m], color=c, lw=1.6, label=label)
        ax2.plot(iu[m], z[m], color=c, lw=1.6)

    ec_label = f"EN 1991-1-4, cat. II ($v_b$ = {v_b:.1f} m/s)"
    ax1.plot(vm_ec, z_ec, "k--", lw=1.8, label=ec_label)
    ax2.plot(iv_ec, z_ec, "k--", lw=1.8)

    ax1.set_xlabel("$U$ [m/s]")
    ax1.set_ylabel("$z$ [m]")
    ax1.set_title("Mean velocity")
    ax2.set_xlabel("$I_u$ [-]")
    ax2.set_xlim(0, 0.5)
    ax2.set_title("Turbulence intensity")

    for ax in (ax1, ax2):
        ax.set_ylim(0, args.z_max)
        ax.grid(alpha=0.3)

    handles, labels = ax1.get_legend_handles_labels()
    fig.legend(handles, labels, loc="lower center", ncol=len(labels), frameon=False, bbox_to_anchor=(0.5, 0.0))

    fig.suptitle("Atmospheric boundary layer, terrain category II")
    fig.tight_layout(rect=(0, 0.06, 1, 1))
    fig.savefig(args.out, dpi=200)
    print(f"saved to {args.out}")


if __name__ == "__main__":
    main()