from typing import Any, Dict, Tuple
import torch
import torch.nn.functional as F
import numpy as np
from collections import defaultdict, Counter
import heapq

SAMPLE_RATE = 16000
WINDOW_SIZE = 512


def chunked(iterable, size):
    """Helper to yield successive chunks of given size from iterable."""
    for i in range(0, len(iterable), size):
        yield iterable[i : i + size]


class AudioFeatureExtractor:
    def __init__(
        self,
        sample_rate=SAMPLE_RATE,
        n_fft=WINDOW_SIZE,
        hop_length=None,
        fan_value=5,
        min_match_ratio: float = 0.1,
        device="cuda",
        kernal_size=5,
        min_prune_size=500,
        max_prune_ratio=0.8,
    ):
        """
        GPU-based Shazam with optional query‐match threshold.
        Args:
          …
          min_query_matches: only accept a match if vote count ≥ this threshold
        """
        self.sample_rate = sample_rate
        self.n_fft = n_fft
        self.hop_length = hop_length or n_fft // 4
        self.fan_value = fan_value
        self.min_match_ratio = min_match_ratio

        self.window = torch.hann_window(self.n_fft, device=device)
        self.min_hash_time_delta = 0
        self.max_hash_time_delta = 200
        self.min_hash_freq_delta = -30
        self.max_hash_freq_delta = 30
        self.kernal_size = kernal_size
        self.device = device
        self.min_prune_size = min_prune_size
        self.max_prune_ratio = max_prune_ratio

    def fingerprint(self, waveform, delta_compress: bool = False):
        """
        Compute fingerprints for a single audio waveform.

        Args:
            waveform (1D array-like): audio samples

        Returns:
            List of (hash, time_offset) tuples.
        """
        # Move signal to GPU
        sig = torch.tensor(waveform, dtype=torch.float32, device=self.device)

        # Compute complex STFT on GPU
        spec = torch.stft(
            sig,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            win_length=self.n_fft,
            window=self.window,
            center=False,
            return_complex=True,
        )

        # Magnitude spectrogram
        mag = spec.abs()

        # Local max pooling for peak detection
        padding = self.kernal_size // 2
        max_pooled = F.max_pool2d(
            mag.unsqueeze(0).unsqueeze(0),
            kernel_size=(self.kernal_size, self.kernal_size),
            stride=1,
            padding=(padding, padding),
        )

        # amp_min = torch.quantile(mag.flatten(), 0.99).item()
        amp_min = 10
        peaks = (mag.unsqueeze(0).unsqueeze(0) == max_pooled) & (mag.unsqueeze(0).unsqueeze(0) > amp_min)
        peaks = peaks.squeeze().cpu().numpy()

        freqs, times = np.where(peaks)
        # after you compute freqs, times:
        peak_list = list(zip(times, freqs))
        if len(peak_list) == 0:
            return []

        peak_list.sort(key=lambda x: x[0])  # sort by time
        times_sorted, freqs_sorted = zip(*peak_list)

        # now pair in temporal order
        hashes = []
        for i in range(len(times_sorted)):
            for j in range(1, self.fan_value):
                if i + j < len(times_sorted):
                    t1, f1 = times_sorted[i], freqs_sorted[i]
                    t2, f2 = times_sorted[i + j], freqs_sorted[i + j]
                    dt = t2 - t1
                    df = f2 - f1
                    if (
                        self.min_hash_time_delta <= dt <= self.max_hash_time_delta
                        and self.min_hash_freq_delta <= df <= self.max_hash_freq_delta
                    ):
                        h = (f1 << 24) | (f2 << 16) | dt
                        hashes.append((h, t1))

        if not delta_compress:
            return hashes

        # 3) aggregate into a list of (h, [t1, t2, …])
        grouping = defaultdict(list)
        for h, t in hashes:
            grouping[h].append(t)
        return [(h, grouping[h]) for h in grouping]