import torch
from torch import nn

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model, PeftModel
import pytorch_lightning as pl
from model import HubertXCNNEnoder

from torch.quantization import quantize_dynamic
import torch.jit as jit


if torch.cuda.is_available():
    # Set the device to CUDA
    device = "cuda"
else:
    # Set the device to CPU
    device = "cpu"

class SpeechLLMLightning(pl.LightningModule):
    def __init__(self, audio_enc_dim=512, llm_dim=2048, llm_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0", quantize=True):
        super().__init__()
        self.save_hyperparameters()

        self.audio_enc_dim = audio_enc_dim
        self.llm_dim = llm_dim
        self.llm_name = llm_name

        self.audio_encoder = HubertXCNNEnoder(self.audio_enc_dim, self.llm_dim)
        self.llm_tokenizer = AutoTokenizer.from_pretrained(self.llm_name)
        self.llm_tokenizer.pad_token = self.llm_tokenizer.eos_token
        self.llm_model = AutoModelForCausalLM.from_pretrained(
            self.llm_name,
            device_map=device,
        )

        peft_config = LoraConfig(
                r=4,
                lora_alpha=8,
                target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj', 'up_proj', 'down_proj', 'gate_proj'],
                lora_dropout=0.05,
            task_type="CAUSAL_LM",
            )

        self.llm_model = get_peft_model(self.llm_model, peft_config)
        self.llm_model.print_trainable_parameters()

        for param in self.llm_model.parameters():
            param.requires_grad = False
        
        self.audio_encoder.eval()
        self.llm_model.eval()

        if quantize:
            self.llm_model = jit.script(self.llm_model)
            self.llm_model = quantize_dynamic(
                                self.llm_model, {nn.Linear}, dtype=torch.qint8
            )


    def encode(self, mel, pre_tokenized_ids, post_tokenized_ids, output_tokenized_ids):
        batch_size = mel.shape[0]
        
        with torch.no_grad():
            speech_embeds = self.audio_encoder(mel)
        embedder = self.llm_model.model.model.embed_tokens
        pre_prompt_embeds = embedder(pre_tokenized_ids)
        post_prompt_embeds = embedder(post_tokenized_ids)
        output_prompt_embeds = embedder(output_tokenized_ids)

        combined_embeds = torch.cat([pre_prompt_embeds, speech_embeds, post_prompt_embeds, output_prompt_embeds], dim=1)
        atts = torch.ones(combined_embeds.size()[:-1], dtype=torch.long).to(combined_embeds.device)

        input_token_length = pre_tokenized_ids.shape[1] + speech_embeds.shape[1] + post_tokenized_ids.shape[1]
        label_ids = torch.cat([
            torch.ones([batch_size, input_token_length], device=combined_embeds.device)*-100,
            output_tokenized_ids
        ], 1).to(combined_embeds.device).to(torch.int64)
        return combined_embeds, atts, label_ids

    def forward(self, embeds, atts, label_ids):
        with torch.no_grad():
            return self.llm_model(
                inputs_embeds=embeds,
                attention_mask=atts,
                labels=label_ids,
            )