"""USBM RI8507 / OSMRE blasting compliance chart.

Renders the velocity-vs-frequency compliance scatter Blastware draws on its Event
Report: each channel's significant waveform cycles as ``(frequency, peak
velocity)`` points on log-log axes against the regulatory limit curve(s). A point
below the curve passes; above fails.

Two pieces, kept separate so both can be reused/extended:
  * ``limit_at`` / ``limit_curve`` — the regulatory limit curve(s), as data.
  * ``channel_compliance_points`` — the per-cycle (freq, velocity) scatter, by
    the zero-crossing method (matches Blastware: each channel's cloud tops out
    at that channel's PPV).

Limit curves (USBM RI8507 Figure B-1 / OSM 30 CFR 816.67), drawn CONTINUOUS — a
constant-displacement bound (sloped, ``v = 2πf·d``) meets a constant-velocity
plateau at the frequency where they're equal, so there are no vertical steps
(matching how Blastware draws it). Two lines:
  * **Drywall** (modern gypsum board) — 0.75 in/s plateau (solid).
  * **Plaster** on wood lath (older homes) — 0.50 in/s plateau (dashed).
Both use a 0.030 in low-frequency displacement bound and rise through a 0.010 in
displacement bound to a 2.0 in/s high-frequency plateau. Values from USBM RI8507
(Appendix B) / 30 CFR 816.67; ⚠ confirm the exact shape against a Blastware
report before trusting for compliance.
"""
from __future__ import annotations

import math
from typing import Dict, Sequence, Tuple

import numpy as np
from matplotlib.ticker import FixedLocator, NullLocator

# curve name → (low-freq "ultimate" displacement in, mid velocity plateau in/s,
#               high-freq displacement in, high-freq velocity plateau in/s).
# RI8507 Fig B-1 (p.74): ultimate max displacement 0.030 in (< ~4 Hz), plateau
# 0.75 (Drywall) / 0.50 (plaster), rising diagonal at 0.008 in displacement up to
# a 2.0 in/s plateau reached at ~40 Hz.
_CURVES: Dict[str, Tuple[float, float, float, float]] = {
    "Drywall": (0.030, 0.75, 0.008, 2.00),
    "Plaster": (0.030, 0.50, 0.008, 2.00),
}
# how each curve is stroked on the chart
_CURVE_STYLE = {"Drywall": {"ls": "-", "lw": 1.0}, "Plaster": {"ls": "--", "lw": 0.9}}

STANDARDS = tuple(_CURVES)

# Blastware's channel markers/colours on the compliance chart.
_CHANNEL_STYLE = {
    "Tran": ("+", "#d62728"),   # red +
    "Vert": ("x", "#2ca02c"),   # green x
    "Long": ("o", "#1f77b4"),   # blue o
}


def limit_at(freq_hz: float, curve: str = "Drywall") -> float:
    """Max allowed PPV (in/s) at ``freq_hz`` for ``curve`` (continuous)."""
    d_low, v_mid, d_high, v_high = _CURVES[curve]
    f = max(freq_hz, 1.0)
    f_a = v_mid / (2.0 * math.pi * d_low)    # disp_low → vel_mid
    f_b = v_mid / (2.0 * math.pi * d_high)   # vel_mid → disp_high
    f_c = v_high / (2.0 * math.pi * d_high)  # disp_high → vel_high
    if f <= f_a:
        return 2.0 * math.pi * f * d_low
    if f <= f_b:
        return v_mid
    if f <= f_c:
        return 2.0 * math.pi * f * d_high
    return v_high


def limit_curve(curve: str = "Drywall", fmin: float = 1.0, fmax: float = 100.0, n: int = 400):
    """(freqs, limits) sampled across the band for plotting one curve."""
    freqs = np.logspace(np.log10(fmin), np.log10(fmax), n)
    return freqs, np.array([limit_at(f, curve) for f in freqs])


def channel_compliance_points(
    samples: Sequence[float], sps: float, fmin: float = 1.0, fmax: float = 100.0,
    vmin: float = 0.0,
) -> Tuple[np.ndarray, np.ndarray]:
    """Per-cycle (frequency, peak velocity) scatter for one channel.

    Zero-crossing method: split the trace at sign changes; each half-cycle
    contributes one point at ``(1/(2·half_period), max|amplitude|)``. Matches
    Blastware — the cloud's ceiling is the channel PPV. ``samples`` must be in the
    velocity unit you want plotted (in/s). Points outside ``[fmin, fmax]`` or at
    or below ``vmin`` are dropped.
    """
    x = np.asarray(samples, dtype=float)
    if x.size < 3:
        return np.empty(0), np.empty(0)
    zc = np.where(np.diff(np.signbit(x)))[0]
    freqs, vels = [], []
    for a, b in zip(zc[:-1], zc[1:]):
        half_period = (b - a) / sps
        if half_period <= 0:
            continue
        freqs.append(1.0 / (2.0 * half_period))
        vels.append(float(np.abs(x[a:b + 1]).max()))
    f = np.array(freqs)
    v = np.array(vels)
    keep = (f >= fmin) & (f <= fmax) & (v > vmin)
    return f[keep], v[keep]


def draw_compliance_chart(ax, channels: Dict[str, Sequence[float]], sps: float) -> None:
    """Draw the compliance chart (both limit curves + per-channel scatter)."""
    for name, style in _CURVE_STYLE.items():
        cf, cv = limit_curve(name)
        ax.plot(cf, cv, color="#333", zorder=3, **style)

    for ch, (marker, color) in _CHANNEL_STYLE.items():
        samples = channels.get(ch)
        if samples is None or len(samples) == 0:
            continue
        f, v = channel_compliance_points(samples, sps)
        ax.scatter(f, v, marker=marker, s=12, c=color, linewidths=0.7, zorder=4, label=ch)

    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlim(1, 100)
    ax.set_ylim(0.0394, 10)
    ax.set_box_aspect(1)   # square plot box (log-log compliance charts are square)
    xt = [1, 2, 5, 10, 20, 50, 100]
    yt = [0.0394, 0.05, 0.1, 0.2, 0.5, 1, 2, 5, 10]
    ax.xaxis.set_major_locator(FixedLocator(xt)); ax.xaxis.set_minor_locator(NullLocator())
    ax.yaxis.set_major_locator(FixedLocator(yt)); ax.yaxis.set_minor_locator(NullLocator())
    ax.set_xticklabels([str(v) for v in xt])
    ax.set_yticklabels([("%g" % v) for v in yt])
    ax.set_xlabel("Frequency (Hz)", fontsize=7)
    ax.set_ylabel("Velocity (in/s)", fontsize=7)
    ax.tick_params(labelsize=6)
    ax.grid(True, which="both", ls=":", lw=0.4, color="#ccc")