"""LR4 (Linkwitz-Riley 4th order) crossovers and shared helpers.

LR4 = two cascaded 2nd-order Butterworth sections. LP+HP of one split sums
to an allpass (flat magnitude). Multi-band tree splits apply each later
crossover's allpass to already-extracted bands so the full sum stays flat.
"""
import numpy as np
from scipy.signal import butter, sosfilt


def _lr4_pair(freq, sr):
    lo = butter(2, freq, btype="low", fs=sr, output="sos")
    hi = butter(2, freq, btype="high", fs=sr, output="sos")
    return lo, hi


def _apply2(sos, x):
    return sosfilt(sos, sosfilt(sos, x, axis=0), axis=0)


def lr4_split(audio, freq, sr):
    """Return (low, high); low+high has flat magnitude."""
    lo, hi = _lr4_pair(freq, sr)
    low = _apply2(lo, audio)
    high = _apply2(hi, audio)
    return low, high


def lr4_allpass(audio, freq, sr):
    low, high = lr4_split(audio, freq, sr)
    return low + high


def lr4_bands(audio, freqs, sr):
    """Split into len(freqs)+1 bands, phase-aligned so sum(bands) is flat."""
    bands = []
    rest = audio
    for f in freqs:
        low, rest = lr4_split(rest, f, sr)
        bands = [lr4_allpass(b, f, sr) for b in bands]
        bands.append(low)
    bands.append(rest)
    return bands


def bandpass_lr4(audio, f_lo, f_hi, sr):
    """Isolate [f_lo, f_hi] band plus the complementary residual (band, rest).
    band + rest has flat magnitude."""
    low, mid_high = lr4_split(audio, f_lo, sr)
    band, high = lr4_split(mid_high, f_hi, sr)
    rest = lr4_allpass(low, f_hi, sr) + high
    return band, rest


def run_pb(board, audio, sr):
    """Run a pedalboard chain on (n, ch) float64 audio, preserving shape/dtype."""
    mono = audio.ndim == 1
    x = audio[:, None] if mono else audio
    y = board(x.T.astype(np.float32), sr).T.astype(np.float64)
    if y.shape[0] != x.shape[0]:           # pedalboard never changes length, guard anyway
        y = y[: x.shape[0]]
    return y[:, 0] if mono else y
