import torch
import numpy as np
from datasets import Dataset
# IntervalStrategy를 명시적으로 임포트하여 버전 충돌을 방지합니다.
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer, IntervalStrategy
from sklearn.metrics import accuracy_score, f1_score

# 1. GPU/CPU 장치 설정
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"사용 장치: {device}")

# 2. 모델 및 토크나이저 로드 (경량 모델 사용)
MODEL_NAME = "microsoft/xtremedistil-l12-h384-uncased"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
# AutoModelForSequenceClassification을 로드하여 분류층을 추가합니다.
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2) 
# num_labels=2: 이진 분류 (긍정: 1, 부정: 0)

# 3. 가상 데이터셋 준비 및 전처리
# 실제 사용 시에는 이 부분을 당신의 한국어 리뷰 데이터로 대체해야 합니다.
data = {
    'text': [
        "매너가 정말 좋으세요! 기분 좋은 거래였습니다.",
        "물건 상태가 별로고 답변도 너무 느렸어요.",
        "빠른 응답과 깔끔한 거래 감사합니다.",
        "가격이 너무 비싸네요. 비추입니다.",
        "오늘도 만족스러운 중고 거래였습니다.",
        "시간 약속 안 지키고 연락도 잘 안 되네요.",
        "친절함 덕분에 거래 과정이 순조로웠습니다.",
        "판매글과 실제 제품이 달라서 실망했습니다.",
    ],
    'label': [1, 0, 1, 0, 1, 0, 1, 0] # 1: 긍정, 0: 부정
}
raw_dataset = Dataset.from_dict(data)

# 데이터셋을 학습(train)과 평가(test) 세트로 분할 (8개 중 4개씩 분할)
train_test_split = raw_dataset.train_test_split(test_size=0.5, seed=42)
train_dataset = train_test_split['train']
eval_dataset = train_test_split['test']

def tokenize_function(examples):
    # 입력 텍스트를 토큰화하고, 경량 모델에 맞게 max_length를 지정합니다.
    return tokenizer(examples['text'], truncation=True, padding='max_length', max_length=128)

# 데이터셋에 토크나이저 적용 및 PyTorch 텐서 형식으로 지정
tokenized_train_dataset = train_dataset.map(tokenize_function, batched=True).with_format("torch")
tokenized_eval_dataset = eval_dataset.map(tokenize_function, batched=True).with_format("torch")

# 4. 평가 지표 함수 정의
def compute_metrics(p):
    predictions = np.argmax(p.predictions, axis=1)
    acc = accuracy_score(p.label_ids, predictions)
    f1 = f1_score(p.label_ids, predictions, average='binary') # 긍정(1)에 대한 F1-Score
    return {"accuracy": acc, "f1": f1}

# 5. 학습 설정 (TrainingArguments)
OUTPUT_DIR = "./xtreme-distil-review-classifier" # 모델 저장 경로
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    num_train_epochs=5,                     # 경량 모델이므로 에포크 수를 약간 늘렸습니다.
    per_device_train_batch_size=8,          # 배치 크기
    per_device_eval_batch_size=8,           
    warmup_steps=500,                       
    weight_decay=0.01,                      
    logging_dir='./logs',                   
    logging_steps=10,
    
    # 평가 및 저장 전략을 'EPOCH'으로 통일하여 load_best_model_at_end를 활성화합니다.
    eval_strategy=IntervalStrategy.EPOCH, 
    save_strategy=IntervalStrategy.EPOCH,       
    
    load_best_model_at_end=True,            
    fp16=torch.cuda.is_available(),         
)

# 6. Trainer 객체 생성 및 학습 시작
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train_dataset,
    eval_dataset=tokenized_eval_dataset,
    compute_metrics=compute_metrics,
)

print("\n--- 파인 튜닝 시작 (XTREME-Distil 모델) ---")
trainer.train()

# 7. 최종 모델 저장
print(f"\n--- 파인 튜닝 완료, 모델을 {OUTPUT_DIR}에 저장 중 ---")
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR) 

print("모델 저장 완료. 이제 저장된 모델을 로드하여 바로 사용할 수 있습니다.")