# tools/automaster_app/pipeline/width.py
"""Frequency-dependent stereo width correction toward reference profile.
3 bands (split at profile width_bands_hz), S gain scaled toward target
S/(M+S) ratio, correction capped at +-50%."""
import numpy as np
from .filters import lr4_bands

MAX_CORRECTION = 0.50
BAND_NAMES = ["low", "mid", "high"]


def adjust_width(audio, sr, profile):
    if audio.ndim == 1 or audio.shape[1] == 1:
        return audio
    target = profile["band_width_ratio"]
    splits = [float(b) for b in profile.get("width_bands_hz", [250, 4000])]
    mid = (audio[:, 0] + audio[:, 1]) * 0.5
    side = (audio[:, 0] - audio[:, 1]) * 0.5
    ms = np.stack([mid, side], axis=1)
    bands = lr4_bands(ms, splits, sr)
    scales = []
    for name, b in zip(BAND_NAMES, bands):
        em, es = np.mean(b[:, 0] ** 2), np.mean(b[:, 1] ** 2)
        cur = es / (em + es + 1e-12)
        tgt = float(target[name])
        if cur < 1e-6:
            scale = 1.0                       # nothing to widen (pure mono band)
        else:
            # solve gain g for side energy so ratio moves to tgt, then cap
            want = tgt / max(1.0 - tgt, 1e-6) * em
            scale = np.sqrt(want / (es + 1e-12))
            scale = float(np.clip(scale, 1.0 - MAX_CORRECTION, 1.0 + MAX_CORRECTION))
        scales.append(scale)
    if all(abs(s - 1.0) < 1e-6 for s in scales):
        return audio
    out_m = np.zeros_like(mid)
    out_s = np.zeros_like(side)
    for b, scale in zip(bands, scales):
        out_m += b[:, 0]
        out_s += b[:, 1] * scale
    return np.stack([out_m + out_s, out_m - out_s], axis=1)