# Adapted for diffusers from multimodal-art-projection/YuE at commit ef1936f2ee39fe8de486a0f47a481c95f8d4da87.
# Licensed under Apache-2.0; see LICENSE.
from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from numbers import Integral

import torch

from .protocol import CODEC_OFFSET, CODEC_SIZE, CONTEXT, MUSIC_END, chunk_ranges


@dataclass
class Chunk:
    ar_tokens: list[int]
    noise: torch.Tensor


class YuE2PrefixKVCache:
    """Holds a chunk's token-prefix keys and values; later acoustic calls attend to the prefix without extending it."""

    def __init__(self):
        self.keys, self.values = [], []

    def get_seq_length(self):
        return self.keys[0].shape[1] if self.keys else 0

    def update(self, key, value, layer_idx):
        if layer_idx == len(self.keys):
            self.keys.append(key)
            self.values.append(value)
            return key, value
        return torch.cat((self.keys[layer_idx], key), dim=1), torch.cat((self.values[layer_idx], value), dim=1)


def _integers(values, name):
    result = list(values)
    if not result or any(isinstance(v, bool) or not isinstance(v, Integral) for v in result):
        raise ValueError(f"{name} must be a nonempty sequence of integer token IDs")
    return [int(v) for v in result]


def song_chunks(prefix, codec, seed, context=CONTEXT):
    """Draw the whole song's CPU FP32 noise once, then split it at the release's chunk boundaries."""
    prefix = _integers(prefix, "prefix")
    codec = _integers(codec, "codec")
    if min(prefix) < 0 or min(codec) < 0 or max(codec) >= CODEC_SIZE:
        raise ValueError("Token IDs are outside their allowed vocabulary")
    if isinstance(context, bool) or not isinstance(context, Integral) or not 1 <= context <= CONTEXT:
        raise ValueError(f"context must be an integer in 1..{CONTEXT}")
    ranges = chunk_ranges(len(codec), len(prefix), int(context))
    generator = torch.Generator(device="cpu").manual_seed(int(seed))
    noise = torch.randn((len(codec), 64), dtype=torch.float32, device="cpu", generator=generator)
    return [
        Chunk(prefix + [value + CODEC_OFFSET for value in codec[a:b]] + [MUSIC_END], noise[a:b]) for a, b in ranges
    ]


@torch.inference_mode()
def solve_midpoint(
    transformer,
    kv_cache,
    noise,
    device,
    steps=32,
    cancelled: Callable[[], bool] | None = None,
    on_progress: Callable[[int, int], None] | None = None,
):
    """Integrate the flow from t=1 (noise) to t=0 with the midpoint method; returns CPU FP32 [frames, 64] latents."""
    if isinstance(steps, bool) or not isinstance(steps, Integral) or steps < 1:
        raise ValueError("steps must be a positive integer")
    if not torch.isfinite(noise).all():
        raise ValueError("Acoustic noise contains non-finite values")

    def velocity(state, t):
        # The model takes time in logit space; logit(1) is clamped to 20.
        raw = torch.logit(torch.tensor(t, dtype=torch.float64, device="cpu")).clamp(-20, 20).item()
        return transformer(latents=state[None], timestep=raw, kv_cache=kv_cache).sample[0]

    state = noise.to(device=device, dtype=transformer.dtype)
    dt = 1.0 / steps
    for step in range(steps):
        if cancelled is not None and cancelled():
            raise InterruptedError("Cancelled during acoustic flow matching")
        t = 1.0 - step * dt
        mid = state - velocity(state, t) * (dt / 2)
        if cancelled is not None and cancelled():
            raise InterruptedError("Cancelled during acoustic flow matching")
        state = state - velocity(mid, t - dt / 2) * dt
        if on_progress is not None:
            on_progress(step + 1, int(steps))
    result = state.float().cpu()
    if not torch.isfinite(result).all():
        raise FloatingPointError("Acoustic flow matching produced non-finite latents")
    return result