# Master Python module

import torch
import math
from typing import Dict, Tuple

Tensor = torch.Tensor

def master_mono_piano(mono: Tensor,
                      sr: int = 48000,
                      hp_cut: float = 20.0,
                      low_shelf_db: float = -1.5,
                      low_shelf_fc: float = 250.0,
                      presence_boost_db: float = 1.6,
                      presence_fc: float = 3500.0,
                      compressor_thresh_db: float = -6.0,
                      compressor_ratio: float = 2.5,
                      compressor_attack_ms: float = 8.0,
                      compressor_release_ms: float = 80.0,
                      compressor_makeup_db: float = 0.0,
                      stereo_spread_ms: float = 6.0,
                      stereo_spread_level_db: float = -12.0,
                      reverb_size_sec: float = 0.6,
                      reverb_mix: float = 0.06,
                      limiter_ceiling_db: float = -0.3,
                      dithering_bits: int = 24,
                      device: torch.device = None,
                      dtype: torch.dtype = None,
                      iir_impulse_len: int = 2048,        # length used to approximate IIRs as FIRs (tradeoff speed/accuracy)
                    ) -> Tuple[Tensor, Dict]:
    
    """
    Optimized mastering pipeline using FIR approximations for small biquads (computed once)
    and single-shot FFT convolutions for all FIRs (presence IR, IIR-approx, reverb).
    """

    # --- Setup, sanitize ---
    if mono.ndim == 2 and mono.shape[0] == 1:
        x = mono.squeeze(0)
    elif mono.ndim == 1:
        x = mono
    else:
        raise ValueError("mono must be shape [1, N] or [N]")

    device = device or x.device
    # prefer float32 for fastest throughput unless user explicitly provided float64
    if dtype is None:
        dtype = x.dtype if x.dtype in (torch.float32, torch.float64) else torch.float32
    x = x.to(device=device, dtype=dtype, copy=False)
    N = x.shape[-1]
    eps = 1e-12

    diagnostics: Dict = {}
    def db2lin(d): return 10.0 ** (d / 20.0)
    def lin2db(v): return 20.0 * math.log10(max(v, 1e-12))

    # --- Helper: design biquad coefficients (as before) ---
    def design_butter_hp(fc, fs, Q=0.7071):
        omega = 2.0 * math.pi * fc / fs
        alpha = math.sin(omega) / (2.0 * Q)
        cosw = math.cos(omega)
        b0 =  (1 + cosw) / 2.0
        b1 = -(1 + cosw)
        b2 =  (1 + cosw) / 2.0
        a0 = 1 + alpha
        a1 = -2 * cosw
        a2 = 1 - alpha
        return b0, b1, b2, a0, a1, a2

    def design_low_shelf(fc, fs, gain_db, Q=0.7071):
        A = 10 ** (gain_db / 40.0)
        w0 = 2.0 * math.pi * fc / fs
        alpha = math.sin(w0) / (2.0 * Q)
        cosw = math.cos(w0)
        b0 =    A*( (A+1) - (A-1)*cosw + 2*math.sqrt(A)*alpha )
        b1 =  2*A*( (A-1) - (A+1)*cosw )
        b2 =    A*( (A+1) - (A-1)*cosw - 2*math.sqrt(A)*alpha )
        a0 =       (A+1) + (A-1)*cosw + 2*math.sqrt(A)*alpha
        a1 =  -2*( (A-1) + (A+1)*cosw )
        a2 =       (A+1) + (A-1)*cosw - 2*math.sqrt(A)*alpha
        return b0, b1, b2, a0, a1, a2

    # --- Utility: convert a 2nd-order IIR (b,a) to an FIR impulse response of length L
    # compute impulse response by stepping the IIR for L samples (this loop runs only for L ~ 1-4k)
    def iir_to_fir(b0, b1, b2, a0, a1, a2, L, device, dtype):
        # normalize
        b0n, b1n, b2n = b0 / a0, b1 / a0, b2 / a0
        a1n, a2n = a1 / a0, a2 / a0
        # compute impulse response on CPU or device (run on device if small L and device != cpu)
        run_device = device if (device.type != 'cuda' or L <= 8192) else device
        h = torch.zeros(L, device=run_device, dtype=dtype)
        x_prev1 = 0.0
        x_prev2 = 0.0
        y_prev1 = 0.0
        y_prev2 = 0.0
        # feed delta(0)=1, others 0
        for n in range(L):
            xv = 1.0 if n == 0 else 0.0
            yv = b0n * xv + b1n * x_prev1 + b2n * x_prev2 - a1n * y_prev1 - a2n * y_prev2
            h[n] = yv
            x_prev2 = x_prev1
            x_prev1 = xv
            y_prev2 = y_prev1
            y_prev1 = yv
        # ensure on main device
        if h.device != device:
            h = h.to(device=device, dtype=dtype)
        return h

    # --- 1) Build IIR -> FIR approximations for HP and shelf (single impulse responses) ---
    # keep impulse length small (configurable) to balance cost/accuracy
    b0,b1,b2,a0,a1,a2 = design_butter_hp(hp_cut, sr)
    hp_ir = iir_to_fir(b0,b1,b2,a0,a1,a2, iir_impulse_len, device, dtype)

    b0s,b1s,b2s,a0s,a1s,a2s = design_low_shelf(low_shelf_fc, sr, low_shelf_db)
    shelf_ir = iir_to_fir(b0s,b1s,b2s,a0s,a1s,a2s, iir_impulse_len, device, dtype)

    # cascade IIRs by convolving their IRs (use FFT convolution)
    def fft_convolve_full(sig, kernel, device, dtype):
        n = sig.shape[-1]
        k = kernel.shape[-1]
        out_len = n + k - 1
        size = 1 << ((out_len - 1).bit_length())
        # cast to complex-friendly dtype (float32/64)
        S = torch.fft.rfft(sig, n=size)
        K = torch.fft.rfft(kernel, n=size)
        Y = S * K
        y = torch.fft.irfft(Y, n=size)[:out_len]
        # return same length as input (valid-ish) by trimming convolution to center-left aligned (like original)
        return y[:n]

    # apply HP then shelf by convolving with cascaded IR (hp_ir * shelf_ir)
    casc_ir = torch.fft.irfft(torch.fft.rfft(hp_ir, n=1<<((hp_ir.numel()+shelf_ir.numel()-2).bit_length())) * 
                               torch.fft.rfft(shelf_ir, n=1<<((hp_ir.numel()+shelf_ir.numel()-2).bit_length())))
    # safer: just convolve via fft_convolve_full on device
    casc_ir = fft_convolve_full(hp_ir, shelf_ir, device, dtype)[:max(hp_ir.numel(), shelf_ir.numel())]
    # normalize tiny numerical offsets
    casc_ir = casc_ir / (casc_ir.abs().sum().clamp(min=eps))

    # apply cascade IR to input (single FFT conv)
    x_hp_shelf = fft_convolve_full(x, casc_ir, device, dtype)

    # --- 2) Presence boost (small FIR) same approach but small kernel conv1d is very fast ---
    pres_len = min(256, max(65, int(sr * 0.0045)))  # ~4.5 ms
    t_idx = torch.arange(pres_len, device=device, dtype=dtype) - (pres_len - 1) / 2.0
    h = (torch.sinc(2.0 * presence_fc / sr * t_idx) * torch.hann_window(pres_len, device=device, dtype=dtype))
    h = h / (h.abs().sum() + eps)
    gain_lin = db2lin(presence_boost_db)
    presence_ir = (gain_lin - 1.0) * h
    presence_ir[(pres_len - 1) // 2] += 1.0
    # conv small kernel with conv1d (fast)
    x_eq = torch.nn.functional.conv1d(x_hp_shelf.view(1,1,-1), presence_ir.view(1,1,-1), padding=(pres_len-1)//2).view(-1)

    diagnostics.update({
        "hp_cut": hp_cut,
        "low_shelf_db": low_shelf_db,
        "presence_db": presence_boost_db,
        "presence_len": pres_len,
        "iir_impulse_len": iir_impulse_len,
    })

    # --- 3) Fast RMS compressor (vectorized with downsampled detector) ---
    sig = x_eq
    attack_tc = math.exp(-1.0 / max(1.0, (compressor_attack_ms * sr / 1000.0)))
    release_tc = math.exp(-1.0 / max(1.0, (compressor_release_ms * sr / 1000.0)))
    sq = sig * sig

    ds = max(1, int(sr // 4000))  # detector rate ~4kHz
    if ds > 1:
        # pad to multiple of ds
        pad = (-sq.shape[-1]) % ds
        if pad:
            sq_pad = torch.nn.functional.pad(sq, (0, pad))
        else:
            sq_pad = sq
        sq_ds = sq_pad.view(-1).reshape(-1, ds).mean(dim=1)
    else:
        sq_ds = sq

    # recurrence on small downsampled vector executed on CPU (cheap)
    sq_ds_cpu = sq_ds.detach().cpu()
    env_ds = torch.empty_like(sq_ds_cpu)
    s_val = float(sq_ds_cpu[0].item())
    a = attack_tc
    r = release_tc
    for i in range(sq_ds_cpu.shape[0]):
        v = float(sq_ds_cpu[i].item())
        if v > s_val:
            s_val = a * s_val + (1.0 - a) * v
        else:
            s_val = r * s_val + (1.0 - r) * v
        env_ds[i] = s_val
    env_ds = env_ds.to(device=device, dtype=dtype)

    if ds > 1:
        env = env_ds.repeat_interleave(ds)[:N]
    else:
        env = env_ds

    rms_env = torch.sqrt(torch.clamp(env, min=eps))
    lvl_db = 20.0 * torch.log10(torch.clamp(rms_env, min=1e-12))
    knee = 3.0
    over = lvl_db - compressor_thresh_db
    # soft knee
    zero = torch.zeros_like(over)
    gain_reduction_db = torch.where(
        over <= -knee,
        zero,
        torch.where(
            over >= knee,
            compressor_thresh_db + (over / compressor_ratio) - lvl_db,
            - ((1.0 - 1.0/compressor_ratio) * (over + knee)**2) / (4.0 * knee)
        )
    ).clamp(max=0.0)
    gain_lin = 10.0 ** (gain_reduction_db / 20.0)
    if ds > 1:
        gain_full = gain_lin.repeat_interleave(ds)[:N]
    else:
        gain_full = gain_lin
    makeup = db2lin(compressor_makeup_db)
    comp_out = sig * gain_full * makeup

    diagnostics.update({
        "compressor_thresh_db": compressor_thresh_db,
        "compressor_ratio": compressor_ratio,
        "compressor_attack_ms": compressor_attack_ms,
        "compressor_release_ms": compressor_release_ms,
        "compressor_makeup_db": compressor_makeup_db,
        "detector_downsample": ds,
        "avg_reduction_db": float((20.0 * torch.log10((gain_full.mean().clamp(min=1e-12))).item())),
    })

    # --- 4) Stereo widening with fractional sub-sample delays (vectorized) ---
    spread_samples = max(1e-4, stereo_spread_ms * sr / 1000.0)
    left_delay = spread_samples * 0.5
    right_delay = -spread_samples * 0.3333

    def fractional_delay_vec(sig, delay):
        n = sig.shape[-1]
        idx = torch.arange(n, device=device, dtype=dtype)
        pos = idx - delay
        pos_floor = pos.floor().long()
        pos_ceil = pos_floor + 1
        frac = (pos - pos_floor.to(dtype))
        pos_floor = pos_floor.clamp(0, n-1)
        pos_ceil = pos_ceil.clamp(0, n-1)
        s_floor = sig[pos_floor]
        s_ceil = sig[pos_ceil]
        return s_floor * (1.0 - frac) + s_ceil * frac

    left = 0.985 * comp_out + 0.015 * fractional_delay_vec(comp_out, left_delay)
    right = 0.985 * comp_out + 0.015 * fractional_delay_vec(comp_out, right_delay)

    mid = 0.5 * (left + right)
    side = 0.5 * (left - right)
    side = side * db2lin(stereo_spread_level_db)
    left = mid + side
    right = mid - side

    # --- 5) Reverb: build IR and FFT-convolve (single-shot) ---
    reverb_len = int(min(int(sr * reverb_size_sec), 65536))
    reverb_len = max(reverb_len, int(0.02 * sr))
    t = torch.arange(reverb_len, device=device, dtype=dtype)
    tail = torch.exp(-t / (reverb_size_sec * sr + 1e-12))
    taps_ms = [12, 23, 37, 53, 79]
    ir = torch.zeros(reverb_len, device=device, dtype=dtype)
    for i, tm in enumerate(taps_ms):
        idx = int(round(sr * tm / 1000.0))
        if idx < reverb_len:
            ir[idx] += (0.5 ** (i + 1))
    ir += 0.15 * tail
    ir = ir / (ir.abs().sum() + eps)

    left_rev = fft_convolve_full(left, ir, device, dtype)
    right_rev = fft_convolve_full(right, ir, device, dtype)
    left = (1.0 - reverb_mix) * left + reverb_mix * left_rev
    right = (1.0 - reverb_mix) * right + reverb_mix * right_rev

    diagnostics.update({
        "reverb_size_sec": reverb_size_sec,
        "reverb_mix": reverb_mix,
        "reverb_len": reverb_len,
    })

    # --- 6) Soft limiter ---
    def soft_limiter(x_chan, ceiling_db):
        ceiling_lin = db2lin(ceiling_db)
        peak = x_chan.abs().max().clamp(min=eps)
        if peak <= ceiling_lin:
            return x_chan
        scaled = x_chan * (ceiling_lin / peak)
        out = torch.tanh(scaled * 1.25) / 1.25
        out = out / out.abs().max().clamp(min=eps) * ceiling_lin
        return out

    left = soft_limiter(left, limiter_ceiling_db)
    right = soft_limiter(right, limiter_ceiling_db)

    # final safety scaling
    peak_val = max(left.abs().max().item(), right.abs().max().item())
    if peak_val > 0.999:
        scale = 0.999 / peak_val
        left = left * scale
        right = right * scale
        diagnostics['final_scale'] = float(scale)
    else:
        diagnostics['final_scale'] = 1.0

    # --- 7) Deterministic TPDF dithering (vectorized LCG) ---
    def vectorized_lcg(sz, seed):
        a = 1103515245
        c = 12345
        mod = 2**31
        seeds = (torch.arange(sz, device=device, dtype=torch.int64) * 1664525 + int(seed)) & (mod - 1)
        vals = (a * seeds + c) & (mod - 1)
        floats = (vals.to(dtype) / float(mod)) - 0.5
        return floats

    if 1 <= dithering_bits <= 32:
        q = 1.0 / (2 ** (dithering_bits - 1))
        seed = (N ^ sr ^ 0x9e3779b1) & 0xffffffff
        na = vectorized_lcg(N, seed)
        nb = vectorized_lcg(N, seed ^ 0x6a09e667)
        tpdf = (na - nb) * q
        left = left + 0.5 * tpdf
        right = right + 0.5 * tpdf

    # --- Output and diagnostics ---
    stereo = torch.stack([left.to(torch.float32), right.to(torch.float32)], dim=0)
    stereo = stereo.clamp(-1.0 + 1e-9, 1.0 - 1e-9)

    left_peak = left.abs().max().item()
    right_peak = right.abs().max().item()
    left_rms = math.sqrt(float(torch.mean(left * left).item()))
    right_rms = math.sqrt(float(torch.mean(right * right).item()))
    diagnostics.update({
        "left_peak": left_peak, "right_peak": right_peak,
        "left_peak_db": lin2db(left_peak), "right_peak_db": lin2db(right_peak),
        "left_rms_db": lin2db(left_rms), "right_rms_db": lin2db(right_rms),
        "num_samples": N, "sample_rate": sr,
    })

    return stereo, diagnostics