"""Constant-Q frame and directional total variation that follows the singer's pitch.

CQTFrame: a painless non-stationary Gabor frame (Velasco et al. 2011) with log-spaced, Hann-shaped
frequency windows and one common hop, so every channel shares the same time grid. Windows are
normalised so that sum_k |G_k|^2 is constant on the covered band, which makes the frame tight there.
Coefficients use the absolute-time phase convention: a steady sinusoid at f rotates by 2*pi*f*hop/sr
per frame in every channel it touches, so the phase of the next frame is predictable.

On a log-frequency axis vibrato and glides shift every harmonic by the same number of channels, so a
sung note is smooth along one direction per frame, given by the pitch slope:

    (D_theta c)[k, t] = c[k + dk(t), t+1] * exp(-i * 2*pi * f_k' * hop / sr) - c[k, t]
    dk(t) = bpo * (log2 f0(t+1) - log2 f0(t))       (0 where unvoiced; |dk| capped)

c[k + dk] uses linear interpolation across channels; f_k' is the frequency halfway along the step.
With dk = 0 everywhere this is ordinary phase-predicted time differencing (fixed direction).
"""
from __future__ import annotations

import numpy as np


class CQTFrame:
    def __init__(self, T, sr=16000, fmin=80.0, fmax=4000.0, bpo=36, width=2.0, hop=None):
        self.T, self.sr, self.bpo = T, sr, bpo
        nb = int(np.floor(bpo * np.log2(fmax / fmin))) + 1
        self.fk = fmin * 2.0 ** (np.arange(nb) / bpo)
        self.K = nb
        fj = np.arange(T // 2 + 1) * sr / T                                 # rfft bin frequencies
        lf = np.log2(np.maximum(fj, 1e-9))
        win, lo, hi = [], [], []
        for k in range(nb):
            u = (lf - np.log2(self.fk[k])) * bpo / width                    # in units of half-width
            idx = np.nonzero(np.abs(u) < 1)[0]
            win.append(np.cos(np.pi / 2 * u[idx]) ** 2)                     # Hann bump in log-frequency
            lo.append(idx[0]); hi.append(idx[-1] + 1)
        # tight on the covered band: divide by sqrt(sum_k |G_k|^2)
        tot = np.zeros(T // 2 + 1)
        for k in range(nb):
            tot[lo[k]:hi[k]] += win[k] ** 2
        self.cover = tot > 1e-3 * tot.max()
        norm = np.sqrt(np.where(self.cover, tot, 1.0))
        self.win = [win[k] / norm[lo[k]:hi[k]] for k in range(nb)]
        self.lo, self.hi = np.array(lo), np.array(hi)
        support = int((self.hi - self.lo).max())
        if hop is None:                                                     # largest hop with M >= support dividing T
            hop = max(h for h in range(1, 513) if T % h == 0 and T // h >= support)
        assert T % hop == 0 and T // hop >= support, "hop must divide T and keep the frame painless"
        self.hop, self.M = hop, T // hop
        self.t = np.arange(self.M) * hop / sr
        self.pos = [np.arange(lo[k], hi[k]) % self.M for k in range(nb)]
        # one gather/scatter for all channels
        self._flat_bins = np.concatenate([np.arange(lo[k], hi[k]) for k in range(nb)])
        self._flat_win = np.concatenate(self.win)
        self._flat_row = np.concatenate([np.full(hi[k] - lo[k], k) for k in range(nb)])
        self._flat_pos = np.concatenate(self.pos)

    def fwd(self, x):
        """Real signal (T,) -> coefficients (K, M) complex."""
        X = np.fft.rfft(x)
        B = np.zeros((self.K, self.M), dtype=np.complex128)
        np.add.at(B, (self._flat_row, self._flat_pos), X[self._flat_bins] * self._flat_win)
        return np.fft.ifft(B, axis=1) * np.sqrt(self.M) / np.sqrt(self.T / 2)   # Parseval on the covered band

    def adj(self, C):
        """Adjoint (real inner products): coefficients (K, M) -> real signal (T,)."""
        B = np.fft.fft(C, axis=1) / np.sqrt(self.M)
        Y = np.zeros(self.T // 2 + 1, dtype=np.complex128)
        np.add.at(Y, self._flat_bins, np.conj(self._flat_win) * B[self._flat_row, self._flat_pos])
        full = np.zeros(self.T, dtype=np.complex128)
        full[: self.T // 2 + 1] = Y
        return np.real(np.fft.ifft(full)) * self.T / np.sqrt(self.T / 2)   # adjoint of rfft under Re<.,.>


class DirectionalTV:
    """K = D_theta: coefficients (K, M) -> differences (K, M-1). Linear; .T is its exact adjoint.

    follow=True : step along the pitch slope, and predict each coefficient's phase from the nearest
                  harmonic h*f0 of the tracked pitch (a sung partial rotates at h*f0, not at the
                  channel's centre frequency).
    follow=False: fixed direction (dk = 0).
    phase_ref:    if given, predict the phase from this reference's own phase advance along the same
                  step instead (the mixture's coefficients, or the current voice estimate's).
    """

    def __init__(self, frame: CQTFrame, f0_times, f0_hz, follow=True, max_shift=3.0, phase_ref=None):
        self.fr = frame
        K, M = frame.K, frame.M
        f0 = np.interp(frame.t, f0_times, np.where(f0_hz > 0, f0_hz, 0.0))
        voiced = np.interp(frame.t, f0_times, (f0_hz > 0).astype(float)) > 0.99
        if follow:
            g = np.where(voiced, np.log2(np.maximum(f0, 1e-9)), np.nan)
            dk = frame.bpo * (g[1:] - g[:-1])
            dk = np.where(np.isfinite(dk) & (np.abs(dk) <= max_shift), dk, 0.0)
        else:
            dk = np.zeros(M - 1)
        self.dk = dk
        kk = np.arange(K)[:, None] + dk[None, :]                           # target channel at t+1
        self.k0 = np.clip(np.floor(kk).astype(int), 0, K - 1)
        self.k1 = np.clip(self.k0 + 1, 0, K - 1)
        self.w1 = np.clip(kk - np.floor(kk), 0, 1)
        self.w0 = 1 - self.w1
        valid = (kk >= 0) & (kk <= K - 1)
        self.w0 *= valid; self.w1 *= valid
        fmid = frame.fk[:, None] * 2.0 ** (dk[None, :] / (2 * frame.bpo))   # centre frequency along the step
        if phase_ref is not None:
            # phase advance of a reference signal along the same step (k,t) -> (k+dk, t+1)
            nxt = phase_ref[:, 1:]
            ref_next = self.w0 * np.take_along_axis(nxt, self.k0, 0) + self.w1 * np.take_along_axis(nxt, self.k1, 0)
            adv = ref_next * np.conj(phase_ref[:, :-1])
            mag = np.abs(adv)
            fallback = np.exp(-2j * np.pi * fmid * frame.hop / frame.sr)
            self.rot = np.where(mag > 1e-12, np.conj(adv) / np.maximum(mag, 1e-12), fallback)
        elif follow:
            fm0 = 0.5 * (f0[1:] + f0[:-1])[None, :]
            h = np.round(fmid / np.maximum(fm0, 1e-9))
            near = (h >= 1) & (np.abs(fmid - h * fm0) < fmid * (2 ** (1.5 / frame.bpo) - 1)) & (voiced[1:] & voiced[:-1])[None, :]
            fpred = np.where(near, h * fm0, fmid)
            self.rot = np.exp(-2j * np.pi * fpred * frame.hop / frame.sr)
        else:
            self.rot = np.exp(-2j * np.pi * fmid * frame.hop / frame.sr)
        self.cols = np.arange(M - 1)[None, :]

    def __call__(self, C):
        nxt = C[:, 1:]
        interp = self.w0 * np.take_along_axis(nxt, self.k0, 0) + self.w1 * np.take_along_axis(nxt, self.k1, 0)
        return interp * self.rot - C[:, :-1]

    def T(self, Y):
        K, M = self.fr.K, self.fr.M
        C = np.zeros((K, M), dtype=np.complex128)
        C[:, :-1] -= Y
        Z = Y * np.conj(self.rot)
        nxt = np.zeros((K, M - 1), dtype=np.complex128)
        np.add.at(nxt, (self.k0, np.broadcast_to(self.cols, self.k0.shape)), self.w0 * Z)
        np.add.at(nxt, (self.k1, np.broadcast_to(self.cols, self.k1.shape)), self.w1 * Z)
        C[:, 1:] += nxt
        return C