"""Mimi Neural Audio Codec - Hugging Face Space

A Gradio demo for the Kyutai Mimi neural audio codec (1.1 kbps @ 12.5 Hz).
Supports encoding audio to compact token representations and decoding back.

Paper: http://kyutai.org/Moshi.pdf
Model: gcxrightsify/mimi on Hugging Face
"""

import gradio as gr
import numpy as np
import torch
import torchaudio
import tempfile
import json
import os

# HF token for private model access (set as Space secret)
HF_TOKEN = os.environ.get("HF_TOKEN")


# =============================================================================
# Mimi Engine (standalone, no sunalink dependency)
# =============================================================================

MIMI_DOWNSAMPLE_FACTOR = 1920  # 24kHz @ 12.5 codes/sec
CONTEXT_FRAMES = 8  # Left context for chunked processing
CROSSFADE_SAMPLES = 240  # 10ms @ 24kHz


class MimiCodec:
    """Minimal Mimi codec wrapper for the Space."""

    def __init__(self):
        self._model = None
        self._feature_extractor = None
        self._device = None
        self._sampling_rate = None

    def _ensure_loaded(self):
        if self._model is None:
            from transformers import MimiModel, AutoFeatureExtractor

            self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            self._model = MimiModel.from_pretrained("gcxrightsify/mimi", token=HF_TOKEN)
            self._feature_extractor = AutoFeatureExtractor.from_pretrained("gcxrightsify/mimi", token=HF_TOKEN)
            self._model.to(self._device)
            self._model.eval()
            self._sampling_rate = self._feature_extractor.sampling_rate
            print(f"Mimi loaded on {self._device}")

    def encode(self, audio: np.ndarray, sample_rate: int) -> dict:
        """Encode audio to Mimi tokens."""
        self._ensure_loaded()

        # Convert to mono if stereo
        if audio.ndim == 2:
            audio = audio.mean(axis=0)

        # Resample if needed (high-quality Kaiser-windowed sinc interpolation)
        if sample_rate != self._sampling_rate:
            audio_torch = torch.from_numpy(audio).float()
            resampler = torchaudio.transforms.Resample(
                orig_freq=sample_rate,
                new_freq=self._sampling_rate,
                resampling_method="sinc_interp_kaiser",
                lowpass_filter_width=64,  # More taps = better quality
            )
            audio = resampler(audio_torch).numpy()

        # Prepare input
        inputs = self._feature_extractor(
            raw_audio=audio,
            sampling_rate=self._sampling_rate,
            return_tensors="pt"
        )
        input_values = inputs["input_values"].to(self._device)

        # Encode with chunking for memory efficiency
        import math
        n_samples = input_values.shape[-1]
        total_frames = math.ceil(n_samples / MIMI_DOWNSAMPLE_FACTOR)
        pad = total_frames * MIMI_DOWNSAMPLE_FACTOR - n_samples
        if pad:
            input_values = torch.nn.functional.pad(input_values, (0, pad))

        chunk_length = 2.0  # seconds
        chunk_frames = max(1, int(chunk_length * self._sampling_rate / MIMI_DOWNSAMPLE_FACTOR))

        chunk_codes = []
        for start_f in range(0, total_frames, chunk_frames):
            end_f = min(start_f + chunk_frames, total_frames)
            ctx_f = min(CONTEXT_FRAMES, start_f)

            seg = input_values[
                :, :,
                (start_f - ctx_f) * MIMI_DOWNSAMPLE_FACTOR:
                end_f * MIMI_DOWNSAMPLE_FACTOR
            ]

            with torch.no_grad():
                codes = self._model.encode(seg).audio_codes[0]

            chunk_codes.append(codes[:, ctx_f:])

        audio_codes = torch.cat(chunk_codes, dim=-1).cpu().numpy()

        return {
            "codes": audio_codes.tolist(),
            "shape": list(audio_codes.shape),
            "sample_rate": self._sampling_rate,
            "original_samples": n_samples,
        }

    def decode(self, encoded: dict) -> tuple[np.ndarray, int]:
        """Decode Mimi tokens back to audio."""
        self._ensure_loaded()

        audio_codes = np.array(encoded["codes"], dtype=np.int64)
        audio_codes_tensor = torch.from_numpy(audio_codes).to(self._device)

        chunk_length = 2.0
        chunk_frames = max(1, int(chunk_length * self._sampling_rate / MIMI_DOWNSAMPLE_FACTOR))
        total_frames = audio_codes_tensor.shape[-1]

        decoded_chunks = []
        for start_f in range(0, total_frames, chunk_frames):
            end_f = min(start_f + chunk_frames, total_frames)
            ctx_f = min(CONTEXT_FRAMES, start_f)

            seg_codes = audio_codes_tensor[:, start_f - ctx_f:end_f]

            with torch.no_grad():
                decoded = self._model.decode(seg_codes.unsqueeze(0)).audio_values

            if ctx_f == 0:
                decoded_chunks.append(decoded)
                continue

            cut = ctx_f * MIMI_DOWNSAMPLE_FACTOR - CROSSFADE_SAMPLES
            decoded = decoded[..., cut:]

            fade_in = torch.linspace(0.0, 1.0, CROSSFADE_SAMPLES, device=decoded.device)
            prev = decoded_chunks[-1]
            prev[..., -CROSSFADE_SAMPLES:] = (
                prev[..., -CROSSFADE_SAMPLES:] * (1.0 - fade_in)
                + decoded[..., :CROSSFADE_SAMPLES] * fade_in
            )
            decoded_chunks.append(decoded[..., CROSSFADE_SAMPLES:])

        audio_values = torch.cat(decoded_chunks, dim=-1)

        # Trim to original length
        original_samples = encoded.get("original_samples", audio_values.shape[-1])
        audio_values = audio_values[..., :original_samples]

        audio_array = audio_values.squeeze().cpu().numpy()
        return audio_array, encoded["sample_rate"]


# Global codec instance (load at startup to avoid first-request latency)
codec = MimiCodec()
codec._ensure_loaded()


# =============================================================================
# Gradio Interface Functions
# =============================================================================

def process_audio(audio_input):
    """Encode and decode audio, return comparison stats."""
    if audio_input is None:
        return None, None, None, "Please upload an audio file."

    sample_rate, audio_data = audio_input

    # Convert to float32 normalized
    if audio_data.dtype == np.int16:
        audio_data = audio_data.astype(np.float32) / 32768.0
    elif audio_data.dtype == np.int32:
        audio_data = audio_data.astype(np.float32) / 2147483648.0

    # Ensure correct shape (samples,) or (channels, samples)
    if audio_data.ndim == 2 and audio_data.shape[0] > audio_data.shape[1]:
        audio_data = audio_data.T  # (samples, channels) -> (channels, samples)

    try:
        # Encode
        encoded = codec.encode(audio_data, sample_rate)

        # Decode
        reconstructed, out_sr = codec.decode(encoded)

        # Calculate stats
        codes_array = np.array(encoded["codes"])
        n_codebooks, n_frames = codes_array.shape
        original_bytes = audio_data.size * 4  # float32
        compressed_bits = n_codebooks * n_frames * 11  # 2048-entry codebook = 11 bits
        compressed_bytes = compressed_bits / 8
        compression_ratio = original_bytes / compressed_bytes

        duration = len(audio_data.flatten()) / sample_rate if audio_data.ndim == 1 else audio_data.shape[-1] / sample_rate
        bitrate = compressed_bits / duration / 1000  # kbps

        stats = f"""## Compression Statistics

| Metric | Value |
|--------|-------|
| Original size | {original_bytes:,} bytes |
| Compressed size | {int(compressed_bytes):,} bytes |
| Compression ratio | {compression_ratio:.1f}x |
| Bitrate | {bitrate:.2f} kbps |
| Duration | {duration:.2f}s |
| Codebooks | {n_codebooks} |
| Frames | {n_frames} |
| Frame rate | {n_frames / duration:.1f} Hz |
"""

        # Save tokens for download
        tokens_json = json.dumps(encoded, indent=2)
        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
            f.write(tokens_json)
            tokens_path = f.name

        return (out_sr, reconstructed), tokens_path, stats, ""

    except Exception as e:
        return None, None, None, f"Error: {str(e)}"


def decode_from_tokens(tokens_file):
    """Decode audio from uploaded token JSON file."""
    if tokens_file is None:
        return None, "Please upload a tokens JSON file."

    try:
        with open(tokens_file.name, 'r') as f:
            encoded = json.load(f)

        audio, sr = codec.decode(encoded)
        return (sr, audio), ""

    except Exception as e:
        return None, f"Error decoding: {str(e)}"


# =============================================================================
# Gradio App
# =============================================================================

with gr.Blocks(title="Mimi Audio Codec", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
# Mimi Neural Audio Codec

A demo of [Kyutai's Mimi](https://huggingface.co/gcxrightsify/mimi) neural audio codec - achieving **1.1 kbps** compression at 12.5 Hz frame rate.

Mimi uses a causal architecture with 8 RVQ codebooks (2048 entries each), making it suitable for real-time streaming applications.

**Paper:** [Moshi: A Speech-Text Foundation Model for Real-Time Dialogue](http://kyutai.org/Moshi.pdf)
    """)

    with gr.Tab("Encode & Decode"):
        with gr.Row():
            with gr.Column():
                audio_input = gr.Audio(
                    label="Input Audio",
                    type="numpy",
                    sources=["upload", "microphone"],
                )
                encode_btn = gr.Button("Encode & Decode", variant="primary")

            with gr.Column():
                audio_output = gr.Audio(label="Reconstructed Audio", type="numpy")
                tokens_download = gr.File(label="Download Tokens (JSON)")

        stats_output = gr.Markdown(label="Statistics")
        error_output = gr.Markdown(label="Errors", visible=True)

        encode_btn.click(
            fn=process_audio,
            inputs=[audio_input],
            outputs=[audio_output, tokens_download, stats_output, error_output],
        )

    with gr.Tab("Decode from Tokens"):
        gr.Markdown("Upload a previously saved tokens JSON file to decode back to audio.")

        with gr.Row():
            tokens_upload = gr.File(label="Upload Tokens JSON", file_types=[".json"])
            decode_btn = gr.Button("Decode", variant="primary")

        decoded_audio = gr.Audio(label="Decoded Audio", type="numpy")
        decode_error = gr.Markdown()

        decode_btn.click(
            fn=decode_from_tokens,
            inputs=[tokens_upload],
            outputs=[decoded_audio, decode_error],
        )

    with gr.Tab("About"):
        gr.Markdown("""
## About Mimi

Mimi is a neural audio codec developed by Kyutai as part of the Moshi project. Key features:

- **Ultra-low bitrate**: 1.1 kbps (compared to 6-24 kbps for other neural codecs)
- **Causal architecture**: Suitable for real-time streaming
- **12.5 Hz frame rate**: 1920 samples per frame at 24 kHz
- **8 RVQ codebooks**: Each with 2048 entries (11 bits per code)

### Technical Details

| Parameter | Value |
|-----------|-------|
| Sample rate | 24,000 Hz |
| Frame rate | 12.5 Hz |
| Codebooks | 8 |
| Codebook size | 2048 |
| Bits per frame | 88 (8 x 11) |
| Bitrate | 1.1 kbps |

### How It Works

1. **Encoder**: Transforms audio into continuous embeddings using causal convolutions and transformers
2. **Quantizer**: Converts embeddings to discrete codes using Residual Vector Quantization (RVQ)
3. **Decoder**: Reconstructs audio from quantized embeddings

The causal architecture means Mimi can process audio in real-time without looking ahead, making it ideal for live streaming and interactive applications.

### Part of Sunalink

This Space demonstrates the Mimi codec as used in [Sunalink](https://github.com/rightsify/sunalink), a multi-modal compression engine supporting neural and classical codecs.
        """)

if __name__ == "__main__":
    demo.launch()