"""A simple web interactive chat demo based on gradio."""

import os
import time
import gradio as gr
import numpy as np
import spaces
import torch

import os
import lightning as L
import torch
import time
import spaces
from snac import SNAC
from litgpt import Tokenizer
from litgpt.utils import (
    num_parameters,
)
from litgpt.generate.base import (
    generate_AA,
    generate_ASR,
    generate_TA,
    generate_TT,
    generate_AT,
    generate_TA_BATCH,
)
from typing import Any, Literal, Optional
import soundfile as sf
from litgpt.model import GPT, Config
from lightning.fabric.utilities.load import _lazy_load as lazy_load
from utils.snac_utils import layershift, reconscruct_snac, reconstruct_tensors, get_time_str
from utils.snac_utils import get_snac
import whisper
from tqdm import tqdm
from huggingface_hub import snapshot_download
from litgpt.generate.base import sample


device = "cuda" if torch.cuda.is_available() else "cpu"
ckpt_dir = "./checkpoint"


OUT_CHUNK = 4096
OUT_RATE = 24000
OUT_CHANNELS = 1

# TODO
text_vocabsize = 151936
text_specialtokens = 64
audio_vocabsize = 4096
audio_specialtokens = 64

padded_text_vocabsize = text_vocabsize + text_specialtokens
padded_audio_vocabsize = audio_vocabsize + audio_specialtokens

_eot = text_vocabsize
_pad_t = text_vocabsize + 1
_input_t = text_vocabsize + 2
_answer_t = text_vocabsize + 3
_asr = text_vocabsize + 4

_eoa = audio_vocabsize
_pad_a = audio_vocabsize + 1
_input_a = audio_vocabsize + 2
_answer_a = audio_vocabsize + 3
_split = audio_vocabsize + 4


def download_model(ckpt_dir):
    repo_id = "gpt-omni/mini-omni"
    snapshot_download(repo_id, local_dir=ckpt_dir, revision="main")


if not os.path.exists(ckpt_dir):
    print(f"checkpoint directory {ckpt_dir} not found, downloading from huggingface")
    download_model(ckpt_dir)


snacmodel = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval().to(device)
whispermodel = whisper.load_model("small").to(device)
whispermodel.eval()
text_tokenizer = Tokenizer(ckpt_dir)
# fabric = L.Fabric(devices=1, strategy="auto")
config = Config.from_file(ckpt_dir + "/model_config.yaml")
config.post_adapter = False

model = GPT(config, device=device)

state_dict = lazy_load(ckpt_dir + "/lit_model.pth")
model.load_state_dict(state_dict, strict=True)
model = model.to(device)
model.eval()


@spaces.GPU
def get_input_ids_whisper_ATBatch(mel, leng, whispermodel, device):
    # with torch.no_grad():
    mel = mel.unsqueeze(0).to(device)
    # audio_feature = whisper.decode(whispermodel,mel, options).audio_features
    audio_feature = whispermodel.embed_audio(mel)[0][:leng]
    T = audio_feature.size(0)
    input_ids_AA = []
    for i in range(7):
        input_ids_item = []
        input_ids_item.append(layershift(_input_a, i))
        input_ids_item += [layershift(_pad_a, i)] * T
        input_ids_item += [(layershift(_eoa, i)), layershift(_answer_a, i)]
        input_ids_AA.append(torch.tensor(input_ids_item))
    input_id_T = torch.tensor([_input_t] + [_pad_t] * T + [_eot, _answer_t])
    input_ids_AA.append(input_id_T)

    input_ids_AT = []
    for i in range(7):
        input_ids_item = []
        input_ids_item.append(layershift(_input_a, i))
        input_ids_item += [layershift(_pad_a, i)] * T
        input_ids_item += [(layershift(_eoa, i)), layershift(_pad_a, i)]
        input_ids_AT.append(torch.tensor(input_ids_item))
    input_id_T = torch.tensor([_input_t] + [_pad_t] * T + [_eot, _answer_t])
    input_ids_AT.append(input_id_T)

    input_ids = [input_ids_AA, input_ids_AT]
    stacked_inputids = [[] for _ in range(8)]
    for i in range(2):
        for j in range(8):
            stacked_inputids[j].append(input_ids[i][j])
    stacked_inputids = [torch.stack(tensors) for tensors in stacked_inputids]
    return torch.stack([audio_feature, audio_feature]), stacked_inputids

    
@spaces.GPU
def next_token_batch(
    model: GPT,
    audio_features: torch.tensor,
    input_ids: list,
    whisper_lens: int,
    task: list,
    input_pos: torch.Tensor,
    **kwargs: Any,
) -> torch.Tensor:
    input_pos = input_pos.to(model.device)
    input_ids = [input_id.to(model.device) for input_id in input_ids]
    logits_a, logit_t = model(
        audio_features, input_ids, input_pos, whisper_lens=whisper_lens, task=task
    )

    for i in range(7):
        logits_a[i] = logits_a[i][0].unsqueeze(0)
    logit_t = logit_t[1].unsqueeze(0)

    next_audio_tokens = []
    for logit_a in logits_a:
        next_a = sample(logit_a, **kwargs).to(dtype=input_ids[0].dtype)
        next_audio_tokens.append(next_a)
    next_t = sample(logit_t, **kwargs).to(dtype=input_ids[0].dtype)
    return next_audio_tokens, next_t

    
def load_audio(path):
    audio = whisper.load_audio(path)
    duration_ms = (len(audio) / 16000) * 1000
    audio = whisper.pad_or_trim(audio)
    mel = whisper.log_mel_spectrogram(audio)
    return mel, int(duration_ms / 20) + 1

    
@spaces.GPU
def generate_audio_data(snac_tokens, snacmodel, device=None):
    audio = reconstruct_tensors(snac_tokens, device)
    with torch.inference_mode():
        audio_hat = snacmodel.decode(audio)
    audio_data = audio_hat.cpu().numpy().astype(np.float64) * 32768.0
    audio_data = audio_data.astype(np.int16)
    audio_data = audio_data.tobytes()
    return audio_data


@torch.inference_mode()
def run_AT_batch_stream(
                        audio_path,
                        stream_stride=4,
                        max_returned_tokens=2048,
                        temperature=0.9,
                        top_k=1,
                        top_p=1.0,
                        eos_id_a=_eoa,
                        eos_id_t=_eot,
    ):

    assert os.path.exists(audio_path), f"audio file {audio_path} not found"

    model.set_kv_cache(batch_size=2)

    mel, leng = load_audio(audio_path)
    audio_feature, input_ids = get_input_ids_whisper_ATBatch(mel, leng, whispermodel, device)
    T = input_ids[0].size(1)
    # device = input_ids[0].device

    assert max_returned_tokens > T, f"max_returned_tokens {max_returned_tokens} should be greater than audio length {T}"

    if model.max_seq_length < max_returned_tokens - 1:
        raise NotImplementedError(
            f"max_seq_length {model.max_seq_length} needs to be >= {max_returned_tokens - 1}"
        )

    input_pos = torch.tensor([T], device=device)
    list_output = [[] for i in range(8)]
    tokens_A, token_T = next_token_batch(
        model,
        audio_feature.to(torch.float32).to(model.device),
        input_ids,
        [T - 3, T - 3],
        ["A1T2", "A1T2"],
        input_pos=torch.arange(0, T, device=device),
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
    )

    for i in range(7):
        list_output[i].append(tokens_A[i].tolist()[0])
    list_output[7].append(token_T.tolist()[0])

    model_input_ids = [[] for i in range(8)]
    for i in range(7):
        tokens_A[i] = tokens_A[i].clone() + padded_text_vocabsize + i * padded_audio_vocabsize
        model_input_ids[i].append(tokens_A[i].clone().to(device).to(torch.int32))
        model_input_ids[i].append(torch.tensor([layershift(4097, i)], device=device))
        model_input_ids[i] = torch.stack(model_input_ids[i])

    model_input_ids[-1].append(token_T.clone().to(torch.int32))
    model_input_ids[-1].append(token_T.clone().to(torch.int32))
    model_input_ids[-1] = torch.stack(model_input_ids[-1])

    text_end = False
    index = 1
    nums_generate = stream_stride
    begin_generate = False
    current_index = 0
    for _ in tqdm(range(2, max_returned_tokens - T + 1)):
        tokens_A, token_T = next_token_batch(
            model,
            None,
            model_input_ids,
            None,
            None,
            input_pos=input_pos,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
        )

        if text_end:
            token_T = torch.tensor([_pad_t], device=device)

        if tokens_A[-1] == eos_id_a:
            break

        if token_T == eos_id_t:
            text_end = True

        for i in range(7):
            list_output[i].append(tokens_A[i].tolist()[0])
        list_output[7].append(token_T.tolist()[0])

        model_input_ids = [[] for i in range(8)]
        for i in range(7):
            tokens_A[i] = tokens_A[i].clone() +padded_text_vocabsize + i * padded_audio_vocabsize
            model_input_ids[i].append(tokens_A[i].clone().to(device).to(torch.int32))
            model_input_ids[i].append(
                torch.tensor([layershift(4097, i)], device=device)
            )
            model_input_ids[i] = torch.stack(model_input_ids[i])

        model_input_ids[-1].append(token_T.clone().to(torch.int32))
        model_input_ids[-1].append(token_T.clone().to(torch.int32))
        model_input_ids[-1] = torch.stack(model_input_ids[-1])

        if index == 7:
            begin_generate = True

        if begin_generate:
            current_index += 1
            if current_index == nums_generate:
                current_index = 0
                snac = get_snac(list_output, index, nums_generate)
                audio_stream = generate_audio_data(snac, snacmodel, device)
                yield audio_stream

        input_pos = input_pos.add_(1)
        index += 1
    text = text_tokenizer.decode(torch.tensor(list_output[-1]))
    print(f"text output: {text}")
    model.clear_kv_cache()
    return list_output

    
for chunk in run_AT_batch_stream('./data/samples/output1.wav'):
    pass


def process_audio(audio):
    filepath = audio
    print(f"filepath: {filepath}")
    if filepath is None:
        return

    cnt = 0
    tik = time.time()
    for chunk in run_AT_batch_stream(filepath):
        # Convert chunk to numpy array
        if cnt == 0:
            print(f"first chunk time cost: {time.time() - tik:.3f}")
        cnt += 1
        audio_data = np.frombuffer(chunk, dtype=np.int16)
        audio_data = audio_data.reshape(-1, OUT_CHANNELS)
        yield OUT_RATE, audio_data.astype(np.int16)


demo = gr.Interface(
    process_audio,
    inputs=gr.Audio(type="filepath", label="Microphone"),
    outputs=[gr.Audio(label="Response", streaming=True, autoplay=True)],
    title="Chat Mini-Omni Demo",
    # live=True,
)
demo.queue()
demo.launch()