import os
import torch
from dlkt.akt import AKT
from dlkt.sakt import SAKT
from dlkt.dkt import DKT
from dlkt.sparsekt import SparseKT
from dlkt.dkt_backbone import DKT_Backbone
from dlkt.akt_backbone import AKT_Backbone
from dlkt.sakt_backbone import SAKT_Backbone

from model.llmkt import *
from model.llmkt_meo import *
from load_llm import *

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


def try_makedirs(path_):
    if not os.path.isdir(path_):
        try:
            os.makedirs(path_)
        except FileExistsError:
            pass


def get_file_name_identifier(params):
    words = params.model.split('_')
    model_type = words[0]
    if model_type == 'dkt':
        file_name = [['_b', params.batch_size], ['_gn', params.maxgradnorm], ['_lr', params.lr],
                     ['_s', params.seed], ['_sl', params.seqlen], ['_dm', params.d_model], ['_ts', params.train_set],  ['_h', params.hidden_dim], ['_do', params.dropout], ['_l2', params.l2]]
    elif model_type == 'dktplus':
        file_name = [['_b', params.batch_size], ['_gn', params.maxgradnorm], ['_lr', params.lr],
                     ['_s', params.seed], ['_sl', params.seqlen], ['_dm', params.d_model], ['_ts', params.train_set],  ['_h', params.hidden_dim], ['_do', params.dropout], ['_l2', params.l2], ['_r', params.lamda_r], ['_w1', params.lamda_w1], ['_w2', params.lamda_w2]]
    elif model_type == 'dkvmn':
        file_name = [['_b', params.batch_size], ['_gn', params.maxgradnorm], ['_lr', params.lr],
                     ['_s', params.seed], ['_sl', params.seqlen], ['_q', params.q_embed_dim], ['_qa', params.qa_embed_dim], ['_ts', params.train_set], ['_m', params.memory_size], ['_l2', params.l2]]
    elif model_type in {'akt', 'sakt',"SparseKT","Denoise","LLM4KT"}:    
        file_name = [['_b', params.batch_size], ['_nb', params.n_block], ['_gn', params.maxgradnorm], ['_lr', params.lr],
                     ['_s', params.seed], ['_sl', params.seqlen], ['_do', params.dropout], ['_dm', params.d_model], ['_ts', params.train_set], ['_kq', params.kq_same], ['_l2', params.l2]]
    return file_name


def model_isPid_type(model_name):
    words = model_name.split('_')
    is_pid = True if 'pid' in words else False
    return is_pid, words[0]


def load_model(params):
    words = params.model.split('_')
    model_type = words[0]
    is_cid = words[1] == 'cid'
    if is_cid:
        params.n_pid = -1

    if model_type in {'akt'}:
        model = AKT(n_question=params.n_question, n_pid=params.n_pid, n_blocks=params.n_block, d_model=params.d_model,
                    dropout=params.dropout, kq_same=params.kq_same, model_type=model_type, noise_ratio=params.noise_ratio, l2=params.l2).to(device)
    elif model_type in {'sakt'}:
        model = SAKT(params).to(device)
        
    elif model_type in {'dkt'}:
        model = DKT(params).to(device)
    
    elif model_type in {'SparseKT'}:
        model = SparseKT(params).to(device)
        
    elif model_type in {'LLM4KT'}:
        """
        LLM4KT Model Loading
        
        Configuration options based on paper:
        - params.plm: Tracer LLM selection
          * qwen_id: Qwen3-4B (default, LLM4KT-4.6B)
          * qwen_base_id: Qwen3-1.7B (medium variant)
          * qwen_small_id: Qwen3-0.6B (small variant)
        
        - params.embed_lm: Inter LLM selection (set in main.py)
          * qwen_embed: Qwen3-Embed-0.6B (default, best performance)
          * bert: BERT-base (110M, faster alternative)
          * bert_large: BERT-large (340M)
          * qwen_embed_4b: Qwen3-Embed-4B (tested but not optimal)
        """
        
        # ==================== LLM4KT-4.6B (Default Configuration) ====================
        if params.plm == "qwen_id":
            # Inter LLM: Qwen3-Embed-0.6B (600M) + Tracer LLM: Qwen3-4B (4B)
            # Total parameters: ~4.6B
            if params.embed_lm == "qwen_embed":
                tokenizer, item_model = load_qwen_embd_model()  # Qwen3-Embed-0.6B
            elif params.embed_lm == "bert":
                tokenizer, item_model = load_bert_model()  # BERT-base (110M)
            elif params.embed_lm == "bert_large":
                tokenizer, item_model = load_bert_large_model()  # BERT-large (340M)
            elif params.embed_lm == "qwen_embed_4b":
                tokenizer, item_model = load_qwen_embd_4b_model()  # Qwen3-Embed-4B
            else:
                raise ValueError(f"Unknown embed_lm: {params.embed_lm}")
            
            _, user_model = load_qwen_lora_model()  # Qwen3-4B as Tracer LLM
            model = LLM4KT_ID_IU(item_model=item_model, item_tokenizer=tokenizer, user_model=user_model).to(device)
        
        # ==================== LLM4KT with Qwen3-1.7B as Tracer LLM ====================
        elif params.plm == "qwen_base_id":
            # Inter LLM: Qwen3-Embed-0.6B (600M) + Tracer LLM: Qwen3-1.7B
            if params.embed_lm == "qwen_embed":
                tokenizer, item_model = load_qwen_embd_model()
            elif params.embed_lm == "bert":
                tokenizer, item_model = load_bert_model()
            elif params.embed_lm == "bert_large":
                tokenizer, item_model = load_bert_large_model()
            else:
                raise ValueError(f"Unknown embed_lm: {params.embed_lm}")
            
            _, user_model = load_qwen_base_model()  # Qwen3-1.7B as Tracer LLM
            model = LLM4KT_ID_IU(item_model=item_model, item_tokenizer=tokenizer, user_model=user_model).to(device)
        
        # ==================== LLM4KT with Qwen3-0.6B as Tracer LLM ====================
        elif params.plm == "qwen_small_id":
            # Inter LLM: Qwen3-Embed-0.6B (600M) + Tracer LLM: Qwen3-0.6B
            if params.embed_lm == "qwen_embed":
                tokenizer, item_model = load_qwen_embd_model()
            elif params.embed_lm == "bert":
                tokenizer, item_model = load_bert_model()
            elif params.embed_lm == "bert_large":
                tokenizer, item_model = load_bert_large_model()
            else:
                raise ValueError(f"Unknown embed_lm: {params.embed_lm}")
            
            _, user_model = load_qwen_model()  # Qwen3-0.6B as Tracer LLM
            model = LLM4KT_ID_IU(item_model=item_model, item_tokenizer=tokenizer, user_model=user_model).to(device)
        
        else:
            raise ValueError(f"Unknown plm configuration: {params.plm}. Supported options: qwen_id, qwen_base_id, qwen_small_id")
        

    return model
