import os
import torch
from datasets import load_dataset
from transformers import (
    AutoTokenizer, 
    AutoModelForCausalLM, 
    TrainingArguments, 
    Trainer, 
    DataCollatorForLanguageModeling
)
from peft import LoraConfig, get_peft_model
import spaces
import time

# Đường dẫn lưu checkpoint
CHECKPOINT_DIR = "./checkpoints"
if not os.path.exists(CHECKPOINT_DIR):
    os.makedirs(CHECKPOINT_DIR)

# Tải Dataset
dataset = load_dataset('vntc/wiki-mini-corpus')

# Chia Dataset thành train và validation
split_dataset = dataset['train'].train_test_split(test_size=0.1, seed=42)
train_dataset = split_dataset['train']
validation_dataset = split_dataset['test']

# Tiền Xử Lý Văn Bản
@spaces.GPU(duration=180, queue=False)
def preprocess_function(examples):
    passages = [passage.lower().strip() for passage in examples['passage']]
    return {'passage': passages}

processed_train = train_dataset.map(preprocess_function, batched=True, remove_columns=['id', 'metadata'])
processed_validation = validation_dataset.map(preprocess_function, batched=True, remove_columns=['id', 'metadata'])

# Tokenization
model_name = "meta-llama/Llama-3.2-3B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Đảm bảo tokenizer có pad_token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

@spaces.GPU(duration=180, queue=False)
def tokenize_function(examples):
    return tokenizer(
        examples['passage'],
        padding='max_length',
        truncation=True,
        max_length=512,
    )

tokenized_train = processed_train.map(tokenize_function, batched=True)
tokenized_validation = processed_validation.map(tokenize_function, batched=True)

# Thêm trường 'labels'
@spaces.GPU(duration=180, queue=False)
def add_labels(examples):
    examples['labels'] = examples['input_ids'].copy()
    return examples

tokenized_train = tokenized_train.map(add_labels, batched=True)
tokenized_validation = tokenized_validation.map(add_labels, batched=True)

# Loại bỏ các cột không cần thiết
tokenized_train = tokenized_train.remove_columns(['passage'])
tokenized_validation = tokenized_validation.remove_columns(['passage'])

# Định dạng dữ liệu cho PyTorch
tokenized_train.set_format('torch')
tokenized_validation.set_format('torch')

# Tạo DatasetDict
final_dataset = {
    'train': tokenized_train,
    'validation': tokenized_validation
}

# Tải và Cấu Hình Mô Hình với LoRA
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16,
    load_in_8bit=False
)

lora_config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "out_proj"],
    lora_dropout=0.1,
    bias="none",
)

model = get_peft_model(model, lora_config)
print(model)

# Cấu Hình TrainingArguments
training_args = TrainingArguments(
    output_dir=CHECKPOINT_DIR,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=8,
    num_train_epochs=3,
    learning_rate=3e-4,
    weight_decay=0.01,
    logging_steps=50,  # Giảm số bước logging
    evaluation_strategy="steps",  # Đánh giá sau mỗi vài bước
    eval_steps=100,  # Đánh giá sau mỗi 100 bước
    save_strategy="steps",  # Lưu checkpoint sau mỗi vài bước
    save_steps=100,  # Lưu checkpoint sau mỗi 100 bước
    save_total_limit=5,  # Giới hạn số lượng checkpoint lưu trữ
    fp16=True,
    report_to="none",
    load_best_model_at_end=True,
)

# Data Collator
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, 
    mlm=False,  # Vì bạn đang thực hiện Causal LM
    pad_to_multiple_of=8
)

# Tạo Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=final_dataset['train'],
    eval_dataset=final_dataset['validation'],
    tokenizer=tokenizer,
    data_collator=data_collator,
)

# Định Nghĩa Hàm Huấn Luyện với Decorator @spaces.GPU
@spaces.GPU(duration=180, queue=False)
def run_training():
    # Kiểm tra nếu có checkpoint
    checkpoints = [os.path.join(CHECKPOINT_DIR, d) for d in os.listdir(CHECKPOINT_DIR) if d.startswith('checkpoint')]
    if checkpoints:
        latest_checkpoint = max(checkpoints, key=os.path.getctime)
        print(f"Đang tiếp tục huấn luyện từ checkpoint: {latest_checkpoint}")
        trainer.train(resume_from_checkpoint=latest_checkpoint)
    else:
        trainer.train()
    
    # Lưu checkpoint sau khi huấn luyện
    trainer.save_model(CHECKPOINT_DIR)
    return "Huấn luyện hoàn tất hoặc đã tiếp tục từ checkpoint."

# Hàm Tự Động Hóa Việc Gọi Lặp Lại `run_training`
@spaces.GPU(duration=180, queue=False)
def continuous_training(max_epochs=3):
    current_epoch = 0
    while current_epoch < max_epochs:
        result = run_training()
        print(result)
        
        # Giả định mỗi lần gọi huấn luyện tiến thêm 0.25 epoch (tùy thuộc vào tốc độ huấn luyện)
        current_epoch += 0.25
        print(f"Đã huấn luyện {current_epoch} / {max_epochs} epochs.")
        
        # Kiểm tra nếu đã đạt số epoch mong muốn
        if current_epoch >= max_epochs:
            print("Đã hoàn thành quá trình huấn luyện.")
            break
        
        # Chờ một khoảng thời gian trước khi gọi lại (tùy thuộc vào yêu cầu của hệ thống)
        time.sleep(1)  # Thời gian chờ có thể điều chỉnh

# Chạy quá trình huấn luyện liên tục
if __name__ == "__main__":
    continuous_training(max_epochs=3)