import os, torch, torchaudio, gradio as gr

import spaces

from zonos.model import Zonos
from zonos.conditioning import make_cond_dict, supported_language_codes

device = "cuda"
MODEL_NAME = "Zyphra/Zonos-v0.1-transformer"
MODEL = Zonos.from_pretrained(MODEL_NAME, device=device)

# def _patch_cuda_props():
#     if torch.cuda.is_available():
#         for i in range(torch.cuda.device_count()):
#             p = torch.cuda.get_device_properties(i)
#             if not hasattr(p, "regs_per_multiprocessor"):
#                 setattr(p, "regs_per_multiprocessor", 65536)
#             if not hasattr(p, "max_threads_per_multi_processor"):
#                 setattr(p, "max_threads_per_multi_processor", 2048)


# _patch_cuda_props()

@spaces.GPU
def generate_audio(
    text,
    language,
    speaker_audio,
    e1,
    e2,
    e3,
    e4,
    e5,
    e6,
    e7,
    e8,
    clarity,
    fmax,
    pitch_std,
    speaking_rate,
    dnsmos_ovrl,
    cfg_scale,
    min_p,
    steps,
    seed,
    randomize_seed,
    progress=gr.Progress(),
):
    if randomize_seed:
        seed = torch.randint(0, 2**32 - 1, (1,)).item()
    torch.manual_seed(int(seed))

    speaker_embedding = None
    if speaker_audio is not None:
        wav, sr = torchaudio.load(speaker_audio)
        speaker_embedding = (
            MODEL.make_speaker_embedding(wav, sr).to(device, dtype=torch.bfloat16)
        )

    emotion_tensor = torch.tensor(
        [e1, e2, e3, e4, e5, e6, e7, e8], device=device, dtype=torch.float32
    )
    vq_tensor = torch.tensor([clarity] * 8, device=device, dtype=torch.float32).unsqueeze(
        0
    )

    cond_dict = make_cond_dict(
        text=text,
        language=language,
        speaker=speaker_embedding,
        emotion=emotion_tensor,
        vqscore_8=vq_tensor,
        fmax=float(fmax),
        pitch_std=float(pitch_std),
        speaking_rate=float(speaking_rate),
        dnsmos_ovrl=float(dnsmos_ovrl),
        device=device,
    )
    conditioning = MODEL.prepare_conditioning(cond_dict)

    estimated_total_steps = int(steps)

    def cb(_, step, __):
        progress((step, estimated_total_steps))
        return True

    codes = MODEL.generate(
        prefix_conditioning=conditioning,
        max_new_tokens=int(steps),
        cfg_scale=float(cfg_scale),
        batch_size=1,
        sampling_params=dict(min_p=float(min_p)),
        callback=cb,
    )

    wav_out = MODEL.autoencoder.decode(codes).cpu().detach()
    sr_out = MODEL.autoencoder.sampling_rate
    if wav_out.dim() == 2 and wav_out.size(0) > 1:
        wav_out = wav_out[0:1, :]
    return (sr_out, wav_out.squeeze().numpy()), seed


def build_interface():
    with gr.Blocks() as demo:
        text = gr.Textbox(label="text", value="hello, world!", lines=4, max_length=500)
        language = gr.Dropdown(choices=supported_language_codes, value="en-us", label="language")
        speaker_audio = gr.Audio(label="voice reference", type="filepath")

        clarity_slider = gr.Slider(0.5, 0.8, 0.8, 0.01, label="clarity")
        steps_slider = gr.Slider(1, 3000, 300, 1, label="steps")

        dnsmos_slider = gr.Slider(1.0, 5.0, 5.0, 0.1, label="quality")
        fmax_slider = gr.Slider(0, 24000, 24000, 1, label="fmax")
        pitch_std_slider = gr.Slider(0.0, 300.0, 30.0, 1, label="pitch std")
        speaking_rate_slider = gr.Slider(5.0, 30.0, 15.0, 0.1, label="rate")

        cfg_scale_slider = gr.Slider(1.0, 5.0, 2.5, 0.1, label="guidance")
        min_p_slider = gr.Slider(0.0, 1.0, 0.05, 0.01, label="min p")

        with gr.Row():
            e1 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="happy")
            e2 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="sad")
            e3 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="disgust")
            e4 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="fear")

        with gr.Row():
            e5 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="surprise")
            e6 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="anger")
            e7 = gr.Slider(0.0, 1.0, 0.0, 0.01, label="other")
            e8 = gr.Slider(0.0, 1.0, 1.0, 0.01, label="neutral")

        seed_number = gr.Number(label="seed", value=420, precision=0)
        randomize_seed_toggle = gr.Checkbox(label="randomize seed", value=True)

        generate_button = gr.Button("generate")
        output_audio = gr.Audio(label="output", type="numpy", autoplay=True)

        generate_button.click(
            fn=generate_audio,
            inputs=[
                text,
                language,
                speaker_audio,
                e1,
                e2,
                e3,
                e4,
                e5,
                e6,
                e7,
                e8,
                clarity_slider,
                fmax_slider,
                pitch_std_slider,
                speaking_rate_slider,
                dnsmos_slider,
                cfg_scale_slider,
                min_p_slider,
                steps_slider,
                seed_number,
                randomize_seed_toggle,
            ],
            outputs=[output_audio, seed_number],
        )
        
    return demo


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