import sys
sys.stdout.reconfigure(line_buffering=True)

import threading
import torch
import torchaudio
import soundfile as sf
import numpy as np
import gradio as gr
from transformers import RobertaTokenizer
from huggingface_hub import hf_hub_download
from pyharp import ModelCard, build_endpoint

from caco_torch.caco import create_caco_model

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

model = None
tokenizer = None
model_loading = True
model_error = None


def load_model():
    global model, tokenizer, model_loading, model_error
    try:
        m = create_caco_model()
        ckpt_path = hf_hub_download(repo_id="teamup-tech/cacophony-weights", filename="Cacophony_torch.pt")
        ckpt = torch.load(ckpt_path, map_location=DEVICE, weights_only=False)
        if "model_state_dict" in ckpt:
            m.load_state_dict(ckpt["model_state_dict"])
        elif "state_dict" in ckpt:
            m.load_state_dict(ckpt["state_dict"])
        else:
            m.load_state_dict(ckpt)
        model = m.to(DEVICE).eval()
        tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
        print("Model loaded.", flush=True)
    except Exception as e:
        model_error = str(e)
        print(f"Load error: {e}", flush=True)
    finally:
        model_loading = False


threading.Thread(target=load_model, daemon=True).start()


def audio_to_patches(path):
    """Convert an audio file to the patch format expected by the model.
    Ported from eval_caco_torch.py: compute_mel_spectrogram, spectrogram_to_patches,
    and prepare_audio_batch. Audio loading uses soundfile instead of torchaudio.load.
    """
    # Load, downmix to mono, resample to 16 kHz
    data, sr = sf.read(path, dtype='float32', always_2d=True)
    wav = torch.from_numpy(data.T)  # (channels, samples)
    if wav.shape[0] > 1:
        wav = wav.mean(0, keepdim=True)
    if sr != 16000:
        wav = torchaudio.functional.resample(wav, sr, 16000)
    audio = wav.squeeze()

    hop, win, n_fft, n_mels = 160, 400, 512, 128  # from compute_mel_spectrogram

    # TF's STFT uses ceil(len/hop) frames; pad to match that before computing STFT.
    n_frames = (len(audio) + hop - 1) // hop
    required_len = (n_frames - 1) * hop + n_fft
    audio = torch.nn.functional.pad(audio, (0, max(0, required_len - len(audio))))

    stft = torch.stft(audio, n_fft=n_fft, hop_length=hop, win_length=win,
                      window=torch.hann_window(win), return_complex=True, center=False)
    spec = torch.abs(stft).T  # (time, freq)
    mel_fb = torchaudio.functional.melscale_fbanks(
        n_freqs=n_fft // 2 + 1, f_min=0, f_max=8000, n_mels=n_mels, sample_rate=16000, norm=None)
    mel = torch.log(spec @ mel_fb + 1e-5).mul(0.2).add(0.9).numpy()

    pt, pf, max_patches = 16, 16, 512  # from spectrogram_to_patches
    mel = mel[:mel.shape[0] // pt * pt]
    nt, nf = mel.shape[0] // pt, n_mels // pf
    # reshape axes: (nt, pt, nf, pf) to (nt, nf, pt, pf) to (n_patches, patch_size)
    patches = mel.reshape(nt, pt, nf, pf).transpose(0, 2, 1, 3).reshape(-1, pt * pf)
    total = nt * nf

    if total > max_patches:  # clip > ~83s: keep first 512 patches
        patches, mask = patches[:max_patches], np.ones(max_patches, dtype=np.float32)
        idx = np.arange(max_patches)
    else:
        mask = (np.arange(max_patches) < total).astype(np.float32)
        idx = (mask * np.arange(max_patches)).astype(np.int64)
        patches = np.pad(patches, [[0, max_patches - total], [0, 0]])

    ti, fi = idx // nf, idx % nf

    def to_tensor(a):
        return torch.from_numpy(a.astype(np.float32)).unsqueeze(0).to(DEVICE)

    return {
        "audio_patches":   to_tensor(patches),
        "audio_time_inds": to_tensor(ti.astype(np.float32)),
        "audio_freq_inds": to_tensor(fi.astype(np.float32)),
        "audio_mask":      to_tensor(mask),
    }


@torch.inference_mode()
def process_fn(audio_path, temperature):
    if model_loading:
        raise gr.Error("Model is still loading, please try again in a moment.")
    if model is None:
        raise gr.Error(f"Model failed to load: {model_error}")

    batch = audio_to_patches(audio_path)
    _, audio_hidden = model.get_audio_embedding(
        audio_patches=batch["audio_patches"],
        audio_time_inds=batch["audio_time_inds"],
        audio_freq_inds=batch["audio_freq_inds"],
        audio_mask=batch["audio_mask"],
        deterministic=True, return_hidden_state=True, normalize=False,
    )

    # Decode token by token: seed with BOS, stop at EOS or 100 tokens.
    generated = [tokenizer.bos_token_id]
    for _ in range(100):
        ids = torch.tensor([generated], dtype=torch.long, device=DEVICE)
        mask = torch.ones(1, len(generated), device=DEVICE)
        logits = model.get_decoder_logits(
            audio_hidden_state=audio_hidden, audio_mask=batch["audio_mask"],
            text_input_ids=ids, text_mask=mask,
        )
        next_token = int(torch.multinomial(
            torch.softmax(logits[0, -1] / temperature, dim=-1), 1))
        if next_token == tokenizer.eos_token_id:
            break
        generated.append(next_token)

    return tokenizer.decode(generated[1:], skip_special_tokens=True)  # skip BOS


model_card = ModelCard(
    name="Cacophony",
    description="Generates a text description of the input audio.",
    author="Ge Zhu, Jordan Darefsky, Zhiyao Duan",
    tags=["audio", "captioning"],
)

with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input Audio").harp_required(True),
        gr.Slider(minimum=0.1, maximum=1.0, step=0.05, value=0.1,
                  label="Caption Creativity",
                  info="Lower = more focused, higher = more varied"),
    ]
    output_components = [
        gr.Textbox(label="Generated Caption"),
    ]
    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

demo.queue().launch()