"""
LLM Model Loading Functions for LLM4KT

This module provides functions to load different LLM models used in the LLM4KT framework:
- Inter LLM (for semantic compression): BERT-base, BERT-large, Qwen3-Embed-0.6B, Qwen3-Embed-4B
- Tracer LLM (for temporal modeling): Qwen3-0.6B, Qwen3-1.7B, Qwen3-4B

All models support LoRA for parameter-efficient fine-tuning.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer, BertModel, BertTokenizer, AutoModelForCausalLM
from torch.cuda.amp import autocast, GradScaler
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model


# ==================== Inter LLM Models (for semantic compression) ====================

def load_qwen_embd_model():
    """
    Load Qwen3-Embed-0.6B model (default Inter LLM, best performance)
    Used as Inter LLM for compressing textual interaction descriptions
    Parameters: 600M
    """
    model_name_or_path = "PLMs/Qwen3-Embedding-0.6B"
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
    model = AutoModel.from_pretrained(model_name_or_path)
    lora_config = LoraConfig(
        r=32,  # LoRA rank
        lora_alpha=32,  # LoRA scaling factor
        target_modules=["q_proj", "v_proj"],  # Target attention modules
        lora_dropout=0.1
    )
    model = get_peft_model(model, lora_config)
    return tokenizer, model


def load_bert_model():
    """
    Load BERT-base model (110M parameters)
    Can be used as Inter LLM (smaller, faster alternative)
    """
    model_name_or_path = "PLMs/bert-base-uncased"
    tokenizer = BertTokenizer.from_pretrained(model_name_or_path)
    model = BertModel.from_pretrained(model_name_or_path)
    return tokenizer, model


def load_bert_large_model():
    """
    Load BERT-large model (340M parameters)
    Can be used as Inter LLM (larger capacity)
    """
    model_name_or_path = "PLMs/bert-large-uncased"
    tokenizer = BertTokenizer.from_pretrained(model_name_or_path)
    model = BertModel.from_pretrained(model_name_or_path)
    return tokenizer, model


# ==================== Tracer LLM Models (for temporal modeling) ====================

def load_qwen_model():
    """
    Load Qwen3-0.6B model (Tracer LLM, smallest variant)
    Used as Tracer LLM for temporal knowledge state modeling
    Parameters: 600M
    """
    model_name_or_path = "PLMs/Qwen3-0.6B"
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
    model = AutoModel.from_pretrained(model_name_or_path)
    lora_config = LoraConfig(
        r=128,  # LoRA rank
        lora_alpha=128,  # LoRA scaling factor
        target_modules=["q_proj", "v_proj"],  # Target attention modules
        lora_dropout=0.1
    )
    model = get_peft_model(model, lora_config)
    return tokenizer, model


def load_qwen_base_model():
    """
    Load Qwen3-1.7B model (Tracer LLM, medium variant)
    Used as Tracer LLM for temporal knowledge state modeling
    Parameters: 1.7B
    """
    model_name_or_path = "PLMs/Qwen3-1.7B"
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
    model = AutoModel.from_pretrained(model_name_or_path)
    lora_config = LoraConfig(
        r=32,  # LoRA rank
        lora_alpha=32,  # LoRA scaling factor
        target_modules=["q_proj", "v_proj"],  # Target attention modules
        lora_dropout=0.1
    )
    model = get_peft_model(model, lora_config)
    return tokenizer, model


def load_qwen_lora_model():
    """
    Load Qwen3-4B model (default Tracer LLM, best performance)
    Used as Tracer LLM for temporal knowledge state modeling
    Parameters: 4B
    """
    model_name_or_path = "PLMs/Qwen3-4B"
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
    model = AutoModel.from_pretrained(model_name_or_path)
    lora_config = LoraConfig(
        r=128,  # LoRA rank
        lora_alpha=128,  # LoRA scaling factor
        target_modules=["q_proj", "v_proj"],  # Target attention modules
        lora_dropout=0.1
    )
    model = get_peft_model(model, lora_config)
    return tokenizer, model


def load_qwen_embd_4b_model():
    """
    Load Qwen3-Embed-4B model (tested in paper but not optimal)
    Used as Inter LLM for semantic compression
    Parameters: 4B
    """
    model_name_or_path = "PLMs/Qwen3-Embedding-4B"
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
    model = AutoModel.from_pretrained(model_name_or_path)
    lora_config = LoraConfig(
        r=32,  # LoRA rank
        lora_alpha=32,  # LoRA scaling factor
        target_modules=["q_proj", "v_proj"],  # Target attention modules
        lora_dropout=0.1
    )
    model = get_peft_model(model, lora_config)
    return tokenizer, model