# Bass Python module

import math
from typing import Tuple, Dict, Optional, Union

import numpy as np
import torch
import torch.nn.functional as F

TensorLike = Union[torch.Tensor, np.ndarray]

def enhance_bass(mono: TensorLike,
                 sr: int = 48000,
                 low_cutoff: float = 200.0,
                 ir_len: int = 4097,
                 low_gain_db: float = 6.0,
                 sub_mix: float = 0.7,
                 compressor_thresh_db: float = -24.0,
                 compressor_ratio: float = 3.0,
                 compressor_attack_ms: float = 8.0,
                 compressor_release_ms: float = 120.0,
                 makeup_db: float = 0.0,
                 drive: float = 1.2,
                 wet_mix: float = 0.9,
                 device: Optional[torch.device] = None,
                 dtype: torch.dtype = torch.float32
                 ) -> Tuple[torch.Tensor, Dict]:
    
    """
    Pronounced bass enhancement for solo piano.

    Parameters
    - mono: 1-D waveform (numpy or torch)
    - sr: sample rate
    - low_cutoff: cutoff for low band extraction (Hz)
    - ir_len: length of linear-phase lowpass IR (odd recommended)
    - low_gain_db: additional gain applied to sub component (dB)
    - sub_mix: mix between original low and generated sub (0..1)
    - compressor_*: low-band compressor settings
    - makeup_db: makeup gain for compressor (dB)
    - drive: saturation drive
    - wet_mix: overall wet mix of enhancement (0..1)
    - device, dtype: processing device and dtype

    Returns
    - enhanced: 1-D torch.Tensor (same length as input)
    - diagnostics: dict
    """
    
    # --- prepare tensor ---
    if isinstance(mono, np.ndarray):
        x = torch.from_numpy(mono.astype(np.float32))
    elif isinstance(mono, torch.Tensor):
        x = mono.clone()
    else:
        raise TypeError("mono must be a numpy array or torch.Tensor")

    if x.ndim != 1:
        raise ValueError("mono must be 1-D (mono)")

    device = device or x.device if isinstance(x, torch.Tensor) else torch.device('cpu')
    x = x.to(device=device, dtype=dtype, copy=False)
    N = x.shape[-1]
    eps = 1e-12

    # --- helpers: linear-phase lowpass IR builder ---
    def make_lowpass_ir(cutoff_hz, sr, length):
        if length % 2 == 0:
            length += 1
        t = torch.arange(length, device=device, dtype=dtype) - (length - 1) / 2.0
        sinc_arg = 2.0 * cutoff_hz / sr * t
        h = torch.sinc(sinc_arg)
        beta = 8.0
        win = torch.kaiser_window(length, beta=beta, device=device, dtype=dtype)
        h = h * win
        h = h / (h.sum() + eps)
        return h

    def next_pow2(x):
        return 1 << ((x - 1).bit_length())

    def fft_convolve_full(sig, kernel):
        n = sig.shape[-1]
        k = kernel.shape[-1]
        out_len = n + k - 1
        size = next_pow2(out_len)
        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 y[:n]

    # --- 1) extract low band ---
    lp = make_lowpass_ir(low_cutoff, sr, ir_len)
    low = fft_convolve_full(x, lp)

    # --- 2) low-band compressor (RMS-based recurrence on CPU for determinism) ---
    def lowband_compress(sig, thresh_db, ratio, attack_ms, release_ms, makeup_db):
        sig_cpu = sig.detach().cpu()
        Nloc = sig_cpu.shape[-1]
        # RMS window ~ 10 ms
        win_len = max(1, int(round(sr * 0.01)))
        sq = sig_cpu * sig_cpu
        pad = (win_len - 1) // 2
        sq_p = F.pad(sq.unsqueeze(0).unsqueeze(0), (pad, win_len - 1 - pad)).squeeze(0).squeeze(0)
        kernel = torch.ones(win_len) / win_len
        rms_sq = F.conv1d(sq_p.unsqueeze(0).unsqueeze(0), kernel.view(1,1,-1)).squeeze(0).squeeze(0)
        rms = torch.sqrt(rms_sq + 1e-12)
        a = math.exp(-1.0 / max(1.0, (attack_ms * sr / 1000.0)))
        r = math.exp(-1.0 / max(1.0, (release_ms * sr / 1000.0)))
        s_val = float(rms[0].item())
        gain = torch.empty_like(rms)
        thresh_lin = 10.0 ** (thresh_db / 20.0)
        for i in range(rms.shape[0]):
            v = float(rms[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
            if s_val <= thresh_lin:
                g = 1.0
            else:
                in_db = 20.0 * math.log10(max(s_val, 1e-12))
                over = in_db - thresh_db
                out_db = thresh_db + over / ratio
                g = 10.0 ** ((out_db - in_db) / 20.0)
            gain[i] = g
        gain = gain.to(device=device, dtype=dtype)
        makeup = 10.0 ** (makeup_db / 20.0)
        return sig * gain * makeup, {"avg_gain_db": 20.0 * math.log10(max(gain.mean().item(), 1e-12))}

    low_comp, comp_diag = lowband_compress(low, compressor_thresh_db, compressor_ratio,
                                           compressor_attack_ms, compressor_release_ms, makeup_db)

    # --- 3) subharmonic / harmonic enhancement ---
    # half-wave rectification -> lowpass -> scale -> saturate -> mix
    rect = torch.clamp(low_comp, min=0.0)
    lp_sub = make_lowpass_ir(200.0, sr, length=2049)
    rect_lp = fft_convolve_full(rect, lp_sub)
    sub_gain = 10.0 ** (low_gain_db / 20.0)
    sub = rect_lp * sub_gain

    # soft saturation
    def soft_saturate(sig, drive_val=1.0):
        return torch.tanh(sig * drive_val) / (torch.tanh(torch.tensor(1.0, device=device, dtype=dtype)))

    sat_low = soft_saturate(low_comp, drive)
    sat_sub = soft_saturate(sub, drive * 0.8)
    enhanced_low = (1.0 - sub_mix) * sat_low + sub_mix * sat_sub

    # band-limit the enhanced low to avoid adding out-of-band energy
    enhanced_low_band = fft_convolve_full(enhanced_low, lp)

    # scale wet component and mix back
    wet = enhanced_low_band  # already scaled by sub_gain earlier
    wet = wet * (10.0 ** (low_gain_db / 20.0))
    out = (1.0 - wet_mix) * x + wet_mix * (x + wet)

    # final gentle limiter
    peak = float(out.abs().max().item())
    if peak > 0.999:
        out = out * (0.999 / peak)
        final_scale = 0.999 / peak
    else:
        final_scale = 1.0

    diagnostics = {
        "sr": sr,
        "low_cutoff": low_cutoff,
        "ir_len": ir_len,
        "low_gain_db": low_gain_db,
        "sub_mix": sub_mix,
        "compressor": comp_diag,
        "input_peak": float(x.abs().max().item()),
        "output_peak": float(out.abs().max().item()),
        "final_scale": final_scale,
    }

    return out.to(dtype=dtype, device=device), diagnostics