"""PERCH 2 embeddings for audio of any length, using perch_v2_embed.pt2 (needs only torch and huggingface_hub).

    from perch2_embed import load, embed
    core = load("cpu")                    # or "cuda", "mps", "auto"
    out = embed(core, wave)               # mono float waveform at 32 kHz
    out["avg"], out["max"], out["temporal"], out["times"]
"""
import math

import torch

REPO, FILE = "DAS-DeepAudioSegmenter/perch2-torch-embed", "perch_v2_embed.pt2"
WINDOW, FRAME = 160_000, 10_240  # 5 s window and 0.32 s frame step, in samples at the 32 kHz model rate


def load(device="cpu", path=None, revision=None):
    """Load the embedding core onto a device ("cpu", "cuda", "mps" or "auto"). Downloads and caches the file unless
    `path` points to a local copy."""
    from torch.export.passes import move_to_device_pass
    if path is None:
        from huggingface_hub import hf_hub_download
        path = hf_hub_download(REPO, FILE, revision=revision)
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
    return move_to_device_pass(torch.export.load(path), device).module()


def embed(core, wave, hop_frames=1, max_seconds=None, sample_rate=32_000, batch_size=32):
    """Embed a mono waveform of any length with the perch_v2_embed.pt2 core.

    wave         1-D float array or tensor. It is fed to the model unchanged, i.e. treated as 32 kHz audio.
    hop_frames   window hop in 0.32 s frames, 1-15 (1 = the model's own 0.32 s hop; larger = fewer windows).
    max_seconds  keep only the first max_seconds of audio (None = everything).
    sample_rate  rate of `wave`; used only to convert max_seconds and the returned times to seconds.
    batch_size   5 s windows per forward pass (memory use).

    Returns a dict with avg (1536,), max (1536,), temporal (frames, 1536) (mean over the 4 frequency positions),
    spatial (frames, 4, 1536), times (frames,) frame centres in seconds, and n_windows.
    """
    if not 1 <= hop_frames <= 15:
        raise ValueError("hop_frames must be 1-15: windows must overlap or touch")
    wave = torch.as_tensor(wave, dtype=torch.float32).flatten()
    if max_seconds is not None:
        wave = wave[: int(max_seconds * sample_rate)]
    length, hop = len(wave), hop_frames * FRAME
    if length == 0:
        raise ValueError("empty waveform")
    n = math.ceil(max(length - WINDOW, 0) / hop) + 1
    windows = torch.nn.functional.pad(wave, (0, (n - 1) * hop + WINDOW - length)).unfold(0, WINDOW, hop)
    device = next(core.parameters()).device
    with torch.inference_mode():
        chunks = torch.cat([core(windows[i:i + batch_size].to(device)).float().cpu() for i in range(0, n, batch_size)])
    # Overlap pooling. At hop_frames=1 this uses the full model's rule: windows after the third drop their first
    # 2 frames and all but the last two drop their last 2. At larger hops every overlapping frame is averaged.
    n_time = chunks.shape[1]
    total = (n - 1) * hop_frames + n_time
    sums, counts = chunks.new_zeros(total, *chunks.shape[2:]), chunks.new_zeros(total)
    maximum = torch.full(chunks.shape[3:], -torch.inf)
    trim = 2 if hop_frames == 1 else 0
    for i in range(n):
        lo = trim if i > 2 else 0
        hi = n_time - trim if i < n - 2 else n_time
        part, start = chunks[i, lo:hi], i * hop_frames
        sums[start + lo:start + hi] += part
        counts[start + lo:start + hi] += 1
        maximum = torch.maximum(maximum, part.amax(dim=(0, 1)))
    avg = sums.sum(dim=(0, 1)) / (counts.sum() * chunks.shape[2])
    spatial = sums / counts.clamp_min(1)[:, None, None]
    centers = (torch.arange(total) + 0.5) * FRAME
    keep = centers < length
    if not keep.any():  # audio shorter than half a frame: one frame, centred on the audio
        keep[0] = True
        centers[0] = length / 2
    return {"avg": avg, "max": maximum, "temporal": spatial[keep].mean(dim=1), "spatial": spatial[keep],
            "times": centers[keep] / sample_rate, "n_windows": n}