import sys
import os
import os.path
import glob
import logging
import argparse
import numpy as np
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
import math
import json
from load_data import DATA, PID_DATA
from run_mix import train, test, valid
# from run_id import train, test, valid
# from run_deepspeed import train, test, valid
from utils import try_makedirs, load_model, get_file_name_identifier

import logging
import os
# from accelerate import Accelerator
# os.environ["DS_SKIP_CUDA_CHECK"] = "1"

def setup_experiment_logger(experiment_name, save_path):
    """
    设置实验日志记录器。

    参数:
    experiment_name (str): 实验的名称，用于日志文件名。
    save_path (str): 日志文件的保存路径。

    返回:
    logger (logging.Logger): 配置好的日志记录器。
    """
    # 确保保存路径存在
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    
    # 日志文件的完整路径
    log_file_path = os.path.join(save_path, f"{experiment_name}.log")
    
    # 创建一个日志记录器
    logger = logging.getLogger(experiment_name)
    logger.setLevel(logging.DEBUG)  # 设置日志级别为DEBUG，可以根据需要调整

    # 创建一个文件处理器，并设置级别为DEBUG
    file_handler = logging.FileHandler(log_file_path)
    file_handler.setLevel(logging.DEBUG)

    # 创建一个流处理器（控制台输出），并设置级别为ERROR
    stream_handler = logging.StreamHandler()
    stream_handler.setLevel(logging.ERROR)

    # 创建一个日志格式器
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(formatter)
    stream_handler.setFormatter(formatter)

    # 将处理器添加到日志记录器
    logger.addHandler(file_handler)
    logger.addHandler(stream_handler)

    return logger

# # 使用示例
# logger = setup_experiment_logger('my_experiment', '/path/to/log/directory')
# logger.info('This is an info message.')
# logger.error('This is an error message.')
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# assert torch.cuda.is_available(), "No Cuda available, AssertionError"
def freeze_user_llm(model):
    """冻结模型中的user_llm参数"""
    for name, param in model.named_parameters():
        if 'user_llm' in name:
            param.requires_grad = False
    return model

def freeze_item_llm(model):
    """冻结模型中的user_llm参数"""
    for name, param in model.named_parameters():
        if 'item_llm' in name:
            param.requires_grad = False
    return model

def train_one_dataset(params, file_name, train_q_data, train_qa_data, train_pid, valid_q_data, valid_qa_data, valid_pid,logger):
    # ================================== model initialization ==================================
    # 加载DeepSpeed配置

    # config_path = "configs/deepspeed_config.json"
    # with open(config_path, "r") as f:
    #     ds_config = json.load(f)
    
    
    model = load_model(params)# akt
    # model = freeze_item_llm(model)
    
    
    
    # 初始化DeepSpeed
    # model, optimizer, _, _ = deepspeed.initialize(   
    #     model=model,
    #     model_parameters=model.parameters(),
    #     config=ds_config
    # )
    
    # model=load_model_sakt(params)
    if params.model in {"LLM4KT_pid","LLM4KT_cid"}:
        optimizer = torch.optim.AdamW(
        model.parameters(), 
        lr=params.lr, 
        betas=(0.9, 0.999), 
        eps=1e-8, 
        weight_decay=params.weight_decay)
        
        # optimizer_grouped_parameters = []
        # # 排除 item_llm 的参数
        # for name, param in model.named_parameters():
        #     if 'user_llm' not in name:
        #         optimizer_grouped_parameters.append(param)
        
        # optimizer = torch.optim.AdamW(
        #     optimizer_grouped_parameters,
        #     lr=params.lr,
        #     betas=(0.9, 0.999),
        #     eps=1e-8,
        #     weight_decay=params.weight_decay
        # )
    else:
        optimizer = torch.optim.Adam(
            model.parameters(), lr=params.lr, betas=(0.9, 0.999), eps=1e-8)

    print("\n")
    """
    开启CPU加速
    """
    # accelerator = Accelerator(cpu=True) 
    # model, optimizer = accelerator.prepare(model, optimizer)
    # ================================== start training ==================================
    all_train_loss = {}
    all_train_accuracy = {}
    all_train_auc = {}
    all_valid_loss = {}
    all_valid_accuracy = {}
    all_valid_auc = {}
    best_valid_auc = 0

    for idx in range(params.max_iter):
        # Train Model
        train_loss, train_accuracy, train_auc ,train_rmse= train(
            model, params, optimizer, train_q_data, train_qa_data, train_pid,  label='Train')
        # Validation step
        valid_loss, valid_accuracy, valid_auc,valid_rmse = valid(
            model,  params, optimizer, valid_q_data, valid_qa_data, valid_pid, label='Valid')
        logger.info('epoch: {}, train_auc:  {}, train_accuracy:  {}, train_rmse: {}, train_loss:  {} '.format(idx+1,train_auc, train_accuracy,train_rmse, train_loss))
        logger.info('valid_auc:  {}, valid_accuracy:  {},valid_rmse: {}, valid_loss:  {} '.format(valid_auc, valid_accuracy,valid_rmse, valid_loss))
        print('epoch', idx + 1)
        print("valid_auc\t", valid_auc, "\ttrain_auc\t", train_auc)
        print("valid_accuracy\t", valid_accuracy,
              "\ttrain_accuracy\t", train_accuracy)
        print("valid_rmse\t", valid_rmse,
              "\ttrain_rmse\t", train_rmse)
        print("valid_loss\t", valid_loss, "\ttrain_loss\t", train_loss) 

        try_makedirs('model')
        try_makedirs(os.path.join('model', params.model))
        try_makedirs(os.path.join('model', params.model, params.plm, params.save))

        all_valid_auc[idx + 1] = valid_auc
        all_train_auc[idx + 1] = train_auc
        all_valid_loss[idx + 1] = valid_loss
        all_train_loss[idx + 1] = train_loss
        all_valid_accuracy[idx + 1] = valid_accuracy
        all_train_accuracy[idx + 1] = train_accuracy

        # output the epoch with the best validation auc
        if valid_auc > best_valid_auc:
            path = os.path.join('model', params.model, params.plm,
                                params.save,  file_name) + '_*'
            for i in glob.glob(path):
                os.remove(i)
            best_valid_auc = valid_auc
            best_epoch = idx+1
            torch.save({'epoch': idx,
                        'model_state_dict': model.state_dict(),
                        'optimizer_state_dict': optimizer.state_dict(),
                        'loss': train_loss,
                        },
                       os.path.join('model', params.model, params.plm, params.save,
                                    file_name)+'_' + str(idx+1)
                       )
        if idx-best_epoch >5:
            break   

    try_makedirs('result')
    try_makedirs(os.path.join('result', params.model))
    try_makedirs(os.path.join('result', params.model, params.plm, params.save))
    f_save_log = open(os.path.join(
        'result', params.model, params.plm, params.save, file_name), 'w')
    f_save_log.write("valid_auc:\n" + str(all_valid_auc) + "\n\n")
    f_save_log.write("train_auc:\n" + str(all_train_auc) + "\n\n")
    f_save_log.write("valid_loss:\n" + str(all_valid_loss) + "\n\n")
    f_save_log.write("train_loss:\n" + str(all_train_loss) + "\n\n")
    f_save_log.write("valid_accuracy:\n" + str(all_valid_accuracy) + "\n\n")
    f_save_log.write("train_accuracy:\n" + str(all_train_accuracy) + "\n\n")
    f_save_log.close()
    return best_epoch


def test_one_dataset(params, file_name, test_q_data, test_qa_data, test_pid,  best_epoch,logger):
    print("\n\nStart testing ......................\n Best epoch:", best_epoch)
    model = load_model(params)

    checkpoint = torch.load(os.path.join(
        'model', params.model, params.plm, params.save, file_name) + '_'+str(best_epoch))
    model.load_state_dict(checkpoint['model_state_dict'])

    test_loss, test_accuracy, test_auc = test(
        model, params, None, test_q_data, test_qa_data, test_pid, label='Test')
    logger.info('test_auc:  {}, test_accuracy:  {}, test_loss:  {} '.format(test_auc, test_accuracy, test_loss))
    print("\ntest_auc\t", test_auc)
    print("test_accuracy\t", test_accuracy)
    print("test_loss\t", test_loss)

    # Now Delete all the models
    path = os.path.join('model', params.model, params.plm, params.save,  file_name) + '_*'
    for i in glob.glob(path):
        os.remove(i)


if __name__ == '__main__':

    # Parse Arguments
    parser = argparse.ArgumentParser(description='Script to test KT')
    # Basic Parameters
    parser.add_argument('--max_iter', type=int, default=20,
                        help='number of iterations')
    parser.add_argument('--train_set', type=int, default=1)
    parser.add_argument('--seed', type=int, default=224, help='default seed')

    # Common parameters
    parser.add_argument('--optim', type=str, default='adam',
                        help='Default Optimizer')
    parser.add_argument('--batch_size', type=int,
                        default=24, help='the batch size')
    parser.add_argument('--lr', type=float, default=1e-5,
                        help='learning rate')
    parser.add_argument('--maxgradnorm', type=float,
                        default=-1, help='maximum gradient norm')
    parser.add_argument('--final_fc_dim', type=int, default=512,
                        help='hidden state dim for final fc layer')
    parser.add_argument('--noise_ratio', type=float, default=0,
                        help='noise level')

    # AKT Specific Parameter
    parser.add_argument('--d_model', type=int, default=256,
                        help='Transformer d_model shape')
    parser.add_argument('--d_ff', type=int, default=1024,
                        help='Transformer d_ff shape')
    parser.add_argument('--dropout', type=float,
                        default=0.05, help='Dropout rate')
    parser.add_argument('--n_block', type=int, default=1,
                        help='number of blocks')
    parser.add_argument('--n_head', type=int, default=8,
                        help='number of heads in multihead attention')
    parser.add_argument('--kq_same', type=int, default=1)

    # AKT-R Specific Parameter
    parser.add_argument('--l2', type=float,
                        default=1e-5, help='l2 penalty for difficulty')

    # DKVMN Specific  Parameter
    parser.add_argument('--q_embed_dim', type=int, default=50,
                        help='question embedding dimensions')
    parser.add_argument('--qa_embed_dim', type=int, default=256,
                        help='answer and question embedding dimensions')
    parser.add_argument('--memory_size', type=int,
                        default=50, help='memory size')
    parser.add_argument('--init_std', type=float, default=0.1,
                        help='weight initialization std')
    # DKT Specific Parameter
    parser.add_argument('--hid_size', type=int, default=256)
    parser.add_argument('--hidden_dim', type=int, default=256)
    parser.add_argument('--num_hid_layers', type=int, default=4)
    # parser.add_argument('--drop_prob', type=float, default=0.05)
    
    #  SAKT
    parser.add_argument('--embed_size', type=int, default=256)
    parser.add_argument('--num_attn_layers', type=int, default=1)
    parser.add_argument('--num_heads', type=int, default=8)
    parser.add_argument('--encode_pos', action='store_true')
    parser.add_argument('--max_pos', type=int, default=10)
    parser.add_argument('--grad_clip', type=float, default=10)
    parser.add_argument("--drop_prob",type=float,default=0.1)

    #  TCN
    parser.add_argument('--levels', type=int, default=4, help='use gpu')
    parser.add_argument('--nhid', type=int, default=128, help='gpu')

    # training
    parser.add_argument('--num_workers', type=int, default=5, help='data loader num workers')
    parser.add_argument('--patience', type=int, default=3, help='early stopping patience')
    parser.add_argument('--itr', type=int, default=1, help='experiment times')
    parser.add_argument('--train_epochs', type=int, default=1, help='train epochs')
    
    parser.add_argument('--learning_rate', type=float, default=0.0005, help='optimizer learning rate')
    parser.add_argument('--weight_decay', type=float, default=0.01, help='weight decay')
    parser.add_argument('--zeta', type=float, default=0.5, help='trade off parameter zeta')
    parser.add_argument('--eta', type=float, default=1.0, help='trade off parameter eta')

    # Device
    parser.add_argument('--use_gpu', type=bool, default=True, help='use gpu')
    parser.add_argument('--gpu', type=int, default=0, help='gpu')
    
    
    # LoRA Configuration (for parameter-efficient fine-tuning)
    parser.add_argument('--lora_dropout', type=float, default=0.1, help='LoRA dropout rate')
    parser.add_argument('--lora_alpha', type=int, default=1, help='LoRA alpha (scaling factor)')   
    parser.add_argument('--lora_r', type=int, default=32, help='LoRA rank r')
    
    # Inter LLM (for semantic compression)
    # Options: qwen_embed (Qwen3-Embed-0.6B, default), bert (BERT-base), bert_large (BERT-large), qwen_embed_4b (Qwen3-Embed-4B)
    parser.add_argument('--embed_lm', type=str, default="qwen_embed", 
                       help='Inter LLM for compressing interaction descriptions')
    
    # Tracer LLM (for temporal modeling)
    # Options: qwen_id (Qwen3-4B, default), qwen_base_id (Qwen3-1.7B), qwen_small_id (Qwen3-0.6B)
    parser.add_argument('--plm', type=str, default="qwen_id", 
                       help='Tracer LLM for temporal knowledge state modeling')



    # Datasets and Model
    parser.add_argument('--model', type=str, default='LLM4KT_pid',
                        help="combination of akt/sakt/dkvmn/dkt (mandatory), pid/cid (mandatory) separated by underscore '_'. For example tf_pid")
    parser.add_argument('--dataset', type=str, default="algebra05_pid")
    # parser.add_argument('--dataset', type=str, default="statics")
    parser.add_argument('--lora', type=str, default="no", help='pretrained_language_model') # use the function run_id.py to train the model
    parser.add_argument('--share', type=str, default="no", help='pretrained_language_model') # use the function run_id.py to train the model
    parser.add_argument('--meo', type=str, default="yes", help='pretrained_language_model') # use the function run_id.py to train the model
    
    parser.add_argument('--concept_name', type=str, default="data/assist2009_pid/skill.json", help='parameter description')
    parser.add_argument('--ct', type=str, default="no", help='parameter description')

    parser.add_argument('--offset', type=float, default="1.0", help='dataset scales')

   
    params = parser.parse_args()
    dataset = params.dataset

    if dataset in {"assist2009_pid"}:
        # params.n_question = 110
        params.n_question = 200
        params.batch_size = 16
        params.seqlen = 100
        params.data_dir = 'DATASETS/'+dataset
        params.data_name = dataset
        # params.n_pid = 16891
        params.n_pid = 20000

    if dataset in {"assist2017_pid"}:
        params.batch_size = 16
        params.seqlen = 50
        params.data_dir = 'DATASETS/'+dataset
        params.data_name = dataset
        params.n_question = 200
        params.n_pid = 3162

    if dataset in {"assist2015"}:
        params.n_question = 100        
        params.batch_size = 24
        params.seqlen = 200
        params.data_dir = 'data/'+dataset
        params.data_name = dataset

    if dataset in {"statics"}:
        params.n_question = 1500
        params.batch_size = 16
        params.seqlen = 100
        params.data_dir = 'DATASETS/'+dataset
        params.data_name = dataset
        
    if dataset in {"algebra06_pid"}:
            # params.max_iter=2000
        params.n_question = 550
        params.n_pid = 1084
        params.batch_size = 16
        params.seqlen = 100
        params.data_dir = 'DATASETS/'+dataset
        params.data_name = dataset
                
        
    if dataset in {"algebra05_pid"}:
        # params.max_iter=2000
        params.n_question = 300
        params.n_pid = 1084
        params.batch_size = 16
        params.seqlen = 100
        params.data_dir = 'DATASETS/'+dataset
        params.data_name = dataset
        
    if dataset in {"spanish_pid"}:
            # params.max_iter=2000
        params.n_question = 300
        params.n_pid = 500
        params.batch_size = 16
        params.seqlen = 100
        params.data_dir = 'DATASETS/'+dataset
        params.data_name = dataset

    
    if dataset in {"ednet_pid"}:
        # params.max_iter=2000
        params.n_question = 1500
        params.n_pid = 500
        params.batch_size = 16
        params.seqlen = 100
        params.data_dir = 'DATASETS/'+dataset
        params.data_name = dataset
        
    params.save = params.data_name
    params.load = params.data_name

    # Setup
    if "pid" not in params.data_name:
        dat = DATA(n_question=params.n_question,
                   seqlen=params.seqlen, separate_char=',')
    else:
        dat = PID_DATA(n_question=params.n_question,
                       seqlen=params.seqlen, separate_char=',')
    seedNum = params.seed
    np.random.seed(seedNum)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    torch.manual_seed(seedNum)
    np.random.seed(seedNum)
    file_name_identifier = get_file_name_identifier(params)

    ###Train- Test
    d = vars(params)
    for key in d:
        print('\t', key, '\t', d[key])
    file_name = ''
    for item_ in file_name_identifier:
        file_name = file_name+item_[0] + str(item_[1])

    # ================================== 5-Fold Cross Validation ==================================
    # Perform 5-fold cross validation as described in the paper

    for fold_idx in range(1, 6):  # fold_1 to fold_5
        print(f"\n{'='*80}")
        print(f"Starting Fold {fold_idx}/5")
        print(f"{'='*80}\n")
        
        # Construct data paths for current fold
        train_data_path = params.data_dir + "/" + f"fold_{fold_idx}" + "/" + "train_processed.csv"
        valid_data_path = params.data_dir + "/" + f"fold_{fold_idx}" + "/" + "test_processed.csv"

        # Load data for current fold
        train_q_data, train_qa_data, train_pid = dat.load_data(train_data_path)
        valid_q_data, valid_qa_data, valid_pid = dat.load_data(valid_data_path)

        # Optional: Apply data sampling ratio for few-shot experiments
        # Uncomment if you want to use --offset parameter for data sampling
        # scale = params.offset
        # offset = math.ceil(train_q_data.shape[0] * scale)
        # idx = np.arange(offset)
        # train_q_data = train_q_data[idx]
        # train_qa_data = train_qa_data[idx]
        # train_pid = train_pid[idx]

        # Print data shapes
        print(f"Fold {fold_idx} Data Statistics:")
        print(f"  train_q_data.shape: {train_q_data.shape}")
        print(f"  train_qa_data.shape: {train_qa_data.shape}")
        print(f"  valid_q_data.shape: {valid_q_data.shape}")
        print(f"  valid_qa_data.shape: {valid_qa_data.shape}")
        print()
        
        # Setup logger for current fold
        logger = setup_experiment_logger(
            str(dataset) + "_fold_" + str(fold_idx) + "_plm_" + str(params.plm) + 
            "_embed_" + str(params.embed_lm) + "_bs_" + str(params.batch_size) + 
            "_sqlen_" + str(params.seqlen) + "_lora_128_" + str(params.lora) + 
            "_share_" + str(params.share) + "_meo_0.1_" + str(params.meo) + 
            "_ct_" + str(params.ct) + "_scale_" + str(params.offset), 
            'new_exp/' + params.model + "/" + str(params.plm) + "/"
        )
        
        logger.info(f'Training on Fold {fold_idx}/5 has begun.')
        
        # Train on current fold and get the best epoch
        best_epoch = train_one_dataset(
            params, file_name + f"_fold_{fold_idx}", 
            train_q_data, train_qa_data, train_pid, 
            valid_q_data, valid_qa_data, valid_pid, 
            logger
        )
        
        logger.info(f'Training on Fold {fold_idx}/5 is finished.')
        logger.info(f'Best epoch for Fold {fold_idx}: {best_epoch}')
        
        # Close logger for current fold
        logging.shutdown()
        
        print(f"\nCompleted Fold {fold_idx}/5")
        print(f"{'='*80}\n")
    
    print(f"\n{'='*80}")
    print("5-Fold Cross Validation Completed!")
    print(f"{'='*80}\n")