"""Spectral normalized cross-correlation; offset = dance time - reference time."""
from dataclasses import dataclass

import librosa
import numpy as np
from scipy import signal

SR = 16000
HOP = 160


class SyncError(ValueError):
    """An actionable media or alignment error suitable for the UI."""


@dataclass(frozen=True)
class Alignment:
    offset: float
    confidence: float
    overlap: float
    refined: bool


def _features(audio):
    if len(audio) < 5 * SR:
        raise SyncError("Each recording must contain at least 5 seconds of audio.")
    if not np.isfinite(audio).all() or np.sqrt(np.mean(audio ** 2)) < 1e-5:
        raise SyncError("A recording is silent or contains invalid audio.")
    mel = librosa.feature.melspectrogram(
        y=audio, sr=SR, n_fft=1024, hop_length=HOP,
        n_mels=40, fmin=100, fmax=7500, power=2, center=False,
    )
    # A relative floor avoids letting tiny background fluctuations dominate silence.
    mel = np.log1p(mel / max(float(mel.max()) * 1e-4, 1e-12))
    deviation = mel.std(axis=1, keepdims=True)
    if np.median(deviation) < 0.05:
        raise SyncError("Audio has too little musical variation to align reliably.")
    return np.clip((mel - mel.mean(axis=1, keepdims=True)) /
                   np.maximum(deviation, 0.1), -3, 3).astype(np.float64)


def _scores(a, b, minimum):
    """Exact overlap-normalized, mean-centered multiband correlation for every lag."""
    n, m = a.shape[1], b.shape[1]
    lags = signal.correlation_lags(n, m)
    a0, b0 = np.maximum(lags, 0), np.maximum(-lags, 0)
    count = np.minimum(n - a0, m - b0)
    cross = signal.fftconvolve(a, b[:, ::-1], mode="full", axes=1).sum(axis=0)
    sums, energy = [], []
    for x, start in ((a, a0), (b, b0)):
        integral = np.pad(np.cumsum(x, axis=1), ((0, 0), (1, 0)))
        squared = np.pad(np.cumsum(x * x, axis=1), ((0, 0), (1, 0)))
        total = integral[:, start + count] - integral[:, start]
        sums.append(total)
        energy.append(np.maximum(0, (squared[:, start + count] -
                                     squared[:, start] - total ** 2 / count).sum(axis=0)))
    numerator = cross - (sums[0] * sums[1]).sum(axis=0) / count
    denominator = np.sqrt(energy[0] * energy[1])
    scores = np.divide(numerator, denominator, out=np.zeros_like(numerator),
                       where=denominator > 1e-8)
    scores[(count < minimum) | (denominator < 1e-8)] = -1
    return lags, scores, count


def _refine(dance, reference, coarse):
    """Use several independent waveform windows only when they agree."""
    sos = signal.butter(4, [150, 4500], btype="bandpass", fs=SR, output="sos")
    dance, reference = (signal.sosfilt(sos, x) for x in (dance, reference))
    lag = round(coarse * SR)
    start = max(0, -lag)
    end = min(len(reference), len(dance) - lag)
    radius, width = int(0.08 * SR), 2 * SR
    candidates = []
    for r in np.linspace(start + radius, end - width - radius, 5).astype(int):
        d = r + lag
        if r < 0 or d < radius or d + width + radius > len(dance):
            continue
        x = reference[r:r + width]
        y = dance[d - radius:d + width + radius]
        corr = signal.correlate(y, x, mode="valid", method="fft")
        energies = signal.convolve(y * y, np.ones(width), mode="valid", method="fft")
        norm = np.sqrt(np.maximum(energies, 1e-15) * max(np.dot(x, x), 1e-15))
        scores = np.abs(corr / norm)  # microphone polarity may be reversed
        peak = int(np.argmax(scores))
        if scores[peak] >= 0.18:
            candidates.append((lag + peak - radius) / SR)
    if len(candidates) >= 3 and np.ptp(candidates) < 0.015:
        return float(np.median(candidates)), True
    return coarse, False


def estimate_offset(dance: np.ndarray, reference: np.ndarray) -> Alignment:
    """Positive offset inserts reference silence; negative offset trims reference."""
    a, b = _features(dance), _features(reference)
    minimum = max((5 * SR - 1024) // HOP + 1, round(0.25 * min(a.shape[1], b.shape[1])))
    lags, scores, overlap = _scores(a, b, minimum)
    best = int(np.argmax(scores))
    peak = float(scores[best])
    outside = np.abs(lags - lags[best]) > round(0.5 * SR / HOP)
    rival = float(scores[outside].max(initial=-1))
    if peak < 0.40 or peak - rival < 0.055:
        raise SyncError(
            "No reliable unique match was found. Use recordings of the same song/version "
            "with at least 5 seconds of shared music and less background noise."
        )
    offset, refined = _refine(dance, reference, float(lags[best] * HOP / SR))
    return Alignment(offset, peak, float(overlap[best] * HOP / SR), refined)