import gradio as gr
import torch
import numpy as np
from encodec import EncodecModel
from encodec.utils import convert_audio

model_cache = {}

def get_model(model_choice, bandwidth):
    key = model_choice
    if key not in model_cache:
        if model_choice == "24kHz – Mono":
            m = EncodecModel.encodec_model_24khz()
        else:
            m = EncodecModel.encodec_model_48khz()
        model_cache[key] = m
    m = model_cache[key]
    m.set_target_bandwidth(float(bandwidth))
    return m

def process(audio_input, model_choice, bandwidth):
    if audio_input is None:
        raise gr.Error("Upload an audio file first.")

    sr, data = audio_input

    # normalize to float32
    if data.dtype == np.int16:
        data = data.astype(np.float32) / 32768.0
    elif data.dtype == np.int32:
        data = data.astype(np.float32) / 2147483648.0
    else:
        data = data.astype(np.float32)

    # shape to (channels, samples)
    if data.ndim == 1:
        wav = torch.from_numpy(data).unsqueeze(0)
    else:
        wav = torch.from_numpy(data.T)

    target_sr = 24000 if "24kHz" in model_choice else 48000
    target_ch = 1 if "24kHz" in model_choice else 2

    model = get_model(model_choice, bandwidth)
    wav = convert_audio(wav, sr, target_sr, target_ch)
    wav_batch = wav.unsqueeze(0)

    with torch.no_grad():
        encoded_frames = model.encode(wav_batch)
        decoded = model.decode(encoded_frames)

    decoded = decoded.squeeze(0)  # (channels, samples)

    # stats
    codes = encoded_frames[0][0]  # (batch, n_codebooks, time_frames)
    n_codebooks = codes.shape[1]
    time_frames = codes.shape[2]
    orig_samples = wav_batch.shape[-1]
    duration_s = orig_samples / target_sr
    bits = time_frames * n_codebooks * 10  # 1024 vocab = 10 bits per token
    actual_bps = bits / duration_s / 1000

    info = (
        f"Duration:        {duration_s:.2f}s\n"
        f"Orig samples:    {orig_samples:,}\n"
        f"Time frames:     {time_frames}\n"
        f"Codebooks:       {n_codebooks}\n"
        f"Total tokens:    {time_frames * n_codebooks:,}\n"
        f"Approx bitrate:  {actual_bps:.2f} kbps\n"
        f"Target BW:       {bandwidth} kbps"
    )

    # back to numpy int16
    out = decoded.numpy()
    if out.shape[0] == 1:
        out = out[0]
    else:
        out = out.T  # stereo: (samples, channels)

    out = np.clip(out, -1.0, 1.0)
    out = (out * 32767).astype(np.int16)

    return (target_sr, out), info


with gr.Blocks(title="EnCodec") as demo:
    gr.Markdown("## EnCodec – Neural Audio Codec")
    gr.Markdown(
        "Uses [Meta's EnCodec](https://github.com/facebookresearch/encodec) to compress "
        "then reconstruct audio. Lower bandwidth = more compression = more artifacts."
    )

    with gr.Row():
        with gr.Column(scale=1):
            audio_in = gr.Audio(label="Input Audio", type="numpy")

            model_choice = gr.Radio(
                choices=["24kHz – Mono", "48kHz – Stereo"],
                value="24kHz – Mono",
                label="Model",
            )
            bandwidth = gr.Dropdown(
                choices=["1.5", "3.0", "6.0", "12.0", "24.0"],
                value="6.0",
                label="Target Bandwidth (kbps)",
            )
            btn = gr.Button("Encode + Decode", variant="primary")

        with gr.Column(scale=1):
            audio_out = gr.Audio(label="Reconstructed Audio", type="numpy")
            info_box = gr.Textbox(label="Stats", lines=8, interactive=False)

    btn.click(
        fn=process,
        inputs=[audio_in, model_choice, bandwidth],
        outputs=[audio_out, info_box],
    )

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