"""Spectral diagnostics of the radiated pressure: per-octave band envelopes, onset times and the
spectral centroid, which together show whether energy climbs from low to high frequency (the
nonlinear upward cascade that gives a crash its build-up).
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.signal import stft

DEFAULT_BANDS = (60, 125, 250, 500, 1000, 2000, 4000, 8000, 16000)


@dataclass
class BandEnvelope:
    lo: float
    hi: float
    onset: float            # first time the band reaches -20 dB of its own peak [s]
    peak_time: float        # time of the band's peak [s]
    level_db: float         # band peak relative to the loudest band [dB]
    envelope: np.ndarray    # band energy vs time (linear)


@dataclass
class CascadeAnalysis:
    times: np.ndarray
    freqs: np.ndarray
    power: np.ndarray       # STFT power (n_freq, n_time)
    bands: list[BandEnvelope]
    centroid: np.ndarray    # spectral centroid vs time [Hz]

    def centroid_at(self, t: float) -> float:
        return float(self.centroid[min(np.searchsorted(self.times, t), len(self.times) - 1)])


def analyze_cascade(p: np.ndarray, fs: int, bands=DEFAULT_BANDS, nperseg: int = 256,
                    noverlap: int = 224) -> CascadeAnalysis:
    """Band envelopes at roughly 6 ms resolution (for fs = 44.1 kHz and the default hop)."""
    f_, t_, Z = stft(p, fs, nperseg=nperseg, noverlap=noverlap)
    Pw = np.abs(Z) ** 2
    tot_max = Pw.sum(0).max() or 1.0
    out = []
    for lo, hi in zip(bands[:-1], bands[1:]):
        m = (f_ >= lo) & (f_ < hi)
        env = Pw[m].sum(0)
        if env.max() <= 0:
            continue
        env_db = 10 * np.log10(env / env.max() + 1e-30)
        out.append(BandEnvelope(lo, hi, float(t_[np.argmax(env_db > -20)]), float(t_[np.argmax(env)]),
                                float(10 * np.log10(env.max() / tot_max)), env))
    cen = (Pw * f_[:, None]).sum(0) / (Pw.sum(0) + 1e-30)
    return CascadeAnalysis(t_, f_, Pw, out, cen)


def format_cascade_report(a: CascadeAnalysis, label: str = "") -> str:
    lines = [f"[{label}] band          onset(-20 dB)   peak     level vs loudest band"]
    for b in a.bands:
        lines.append(f"  {b.lo:5.0f}-{b.hi:5.0f} Hz   {b.onset * 1e3:7.1f} ms   {b.peak_time * 1e3:7.1f} ms"
                     f"   {b.level_db:6.1f} dB")
    for tq in (0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5):
        if tq < a.times[-1]:
            lines.append(f"  spectral centroid at {tq * 1e3:5.0f} ms: {a.centroid_at(tq):7.0f} Hz")
    return "\n".join(lines)


def plot_cascade(a: CascadeAnalysis, path: str, title: str = "", fs: float = 44100.0, t_max: float = 0.5) -> str:
    """Spectrogram and band envelopes (log-log). Requires matplotlib."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    tot_max = a.power.sum(0).max() or 1.0
    sel = a.times <= t_max
    fig, ax = plt.subplots(1, 2, figsize=(13, 4.5))
    ax[0].pcolormesh(a.times[sel] * 1e3, a.freqs, 10 * np.log10(a.power[:, sel] / a.power.max() + 1e-12),
                     vmin=-90, vmax=0, shading="auto")
    ax[0].set_yscale("log")
    ax[0].set_ylim(50, fs / 2)
    ax[0].set_xlabel("t [ms]")
    ax[0].set_ylabel("f [Hz]")
    ax[0].set_title(title)
    for b in a.bands:
        ax[1].plot(a.times[sel] * 1e3, 10 * np.log10(b.envelope[sel] / tot_max + 1e-30), label=f"{b.lo:.0f}-{b.hi:.0f} Hz")
    ax[1].set_ylim(-90, 5)
    ax[1].set_xscale("log")
    ax[1].set_xlim(2, t_max * 1e3)
    ax[1].set_xlabel("t [ms]")
    ax[1].set_ylabel("band energy [dB re loudest]")
    ax[1].legend(fontsize=7, ncol=2)
    ax[1].grid(alpha=0.3)
    ax[1].set_title("band envelopes (a cascade shows staggered rises, high bands last)")
    plt.tight_layout()
    plt.savefig(path, dpi=120)
    plt.close(fig)
    return path