# 필요 라이브러리 임포트
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np # alpha가 list/ndarray일 경우를 위해 추가

class FocalLossSoftLabels(nn.Module):
    """
    다중 클래스 분류를 위한 Focal Loss 구현 (소프트 레이블 지원).
    Label smoothing이 적용된 경우에도 작동하도록 설계되었습니다.

    Args:
        alpha (float or list/tensor): 각 클래스에 대한 가중치 요소.
                                       리스트/텐서 형태 권장 (길이=num_classes).
                                       float일 경우 모든 클래스에 동일 적용. 기본값 0.25.
        gamma (float): 초점 파라미터. 기본값 2.0.
        reduction (str): 감소 방식 ('none' | 'mean' | 'sum'). 기본값 'mean'.
        epsilon (float): 수치 안정성을 위한 작은 값. 기본값 1e-7.
    """
    def __init__(self, alpha=0.25, gamma=2.0, reduction='mean', epsilon=1e-7):
        super(FocalLossSoftLabels, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction
        self.epsilon = epsilon # log(0) 방지

    def forward(self, inputs, targets):
        """
        Focal Loss 계산 (소프트 레이블 지원).

        Args:
            inputs (torch.Tensor): 모델 로짓 (N, C).
            targets (torch.Tensor): 실제 레이블 (N) 또는 소프트 레이블 (N, C).

        Returns:
            torch.Tensor: 계산된 Focal Loss.
        """
        num_classes = inputs.shape[1]

        # 1. 입력 로짓에 log_softmax 적용
        log_probs = F.log_softmax(inputs, dim=-1)
        # 예측 확률 계산
        probs = F.softmax(inputs, dim=-1)

        # 2. targets 형태 확인 및 변환
        if targets.ndim == 1: # 하드 레이블인 경우 (Label smoothing 미적용 시)
            # 원-핫 인코딩으로 변환 (Cross Entropy 계산 위해)
            targets_one_hot = F.one_hot(targets, num_classes=num_classes).float().to(inputs.device)
            # pt 계산을 위한 정수 인덱스
            target_indices = targets
        elif targets.ndim == 2 and targets.shape[1] == num_classes: # 소프트 레이블인 경우
            targets_one_hot = targets.float().to(inputs.device)
            # pt 계산 및 alpha 적용을 위해 가장 확률 높은 클래스 인덱스 사용
            target_indices = targets_one_hot.argmax(dim=1)
        else:
            raise ValueError("targets 텐서의 형태가 올바르지 않습니다. (N) 또는 (N, C)여야 합니다.")

        # 3. 기본 Cross Entropy 손실 계산 (소프트/하드 레이블 모두 처리)
        # ce_loss = - (targets_one_hot * log_probs).sum(dim=-1) # 수동 계산
        # 또는 PyTorch 함수 사용 (내부적으로 소프트 레이블 처리)
        # 참고: F.cross_entropy는 log_softmax + NLLLoss 이므로, log_probs를 직접 사용 시 NLLLoss 사용
        ce_loss_per_sample = F.cross_entropy(inputs, targets_one_hot, reduction='none')

        # 4. pt 계산 (예측 확률 중 '정답' 클래스에 해당하는 값)
        # target_indices를 사용하여 해당 클래스의 예측 확률(probs)을 가져옴
        pt = probs.gather(1, target_indices.unsqueeze(-1)).squeeze(-1)
        # 수치 안정성 보장 (pt가 0이나 1에 매우 가까워지는 것 방지)
        pt = pt.clamp(min=self.epsilon, max=1.0 - self.epsilon)

        # 5. alpha 가중치 계산 (alpha_t)
        alpha_t = None
        if isinstance(self.alpha, (float, int)):
            # 단일 float alpha -> 모든 클래스에 동일하게 적용하거나 특정 로직 추가 가능
            # 여기서는 우선 모든 샘플에 동일 alpha 적용 (리스트 형태 권장)
            # alpha_t = torch.full_like(ce_loss_per_sample, self.alpha) # 단순 적용 예시
            # 클래스별로 다른 alpha를 적용하고 싶다면 아래 리스트/텐서 방식 사용
             alpha_tensor = torch.full((num_classes,), self.alpha).to(inputs.device)
             alpha_t = alpha_tensor.gather(0, target_indices)

        elif isinstance(self.alpha, (list, np.ndarray, torch.Tensor)):
            # 리스트, ndarray, 텐서 형태 alpha 처리
            if not isinstance(self.alpha, torch.Tensor):
                alpha_tensor = torch.tensor(self.alpha).to(inputs.device)
            else:
                alpha_tensor = self.alpha.to(inputs.device)

            if alpha_tensor.shape[0] != num_classes:
                 raise ValueError(f"alpha의 길이({alpha_tensor.shape[0]})가 클래스 개수({num_classes})와 일치해야 합니다.")
            # target_indices를 사용하여 각 샘플의 alpha 값 선택
            alpha_t = alpha_tensor.gather(0, target_indices)

        else: # alpha 가중치 없음
            alpha_t = 1.0

        # 6. Focal Loss 계산
        focal_loss_component = alpha_t * (1 - pt)**self.gamma * ce_loss_per_sample

        # 7. reduction 적용
        if self.reduction == 'mean':
            return focal_loss_component.mean()
        elif self.reduction == 'sum':
            return focal_loss_component.sum()
        else: # 'none'
            return focal_loss_component

from transformers import Trainer

class CustomTrainerWithFocalLoss(Trainer):
    """
    Focal Loss (소프트 레이블 지원)를 사용하도록 손실 계산을 오버라이드한 커스텀 Trainer.
    """
    def __init__(self, *args, focal_loss_alpha=0.25, focal_loss_gamma=2.0, num_classes=4, **kwargs):
        super().__init__(*args, **kwargs)
        self.num_classes = num_classes # 클래스 개수 저장

        # alpha 처리: float이면 num_classes 길이의 리스트로 변환
        effective_alpha = focal_loss_alpha
        if isinstance(focal_loss_alpha, (float, int)):
            effective_alpha = [focal_loss_alpha] * self.num_classes
            print(f"Focal Loss: 단일 alpha 값({focal_loss_alpha})을 모든 {self.num_classes}개 클래스에 적용합니다.")
        elif isinstance(focal_loss_alpha, list) and len(focal_loss_alpha) != self.num_classes:
             raise ValueError(f"focal_loss_alpha 리스트의 길이({len(focal_loss_alpha)})가 num_classes({self.num_classes})와 일치해야 합니다.")
        else:
             print(f"Focal Loss: 클래스별 alpha 값({effective_alpha})을 사용합니다.")

        # 최적화된 FocalLossSoftLabels 사용
        self.focal_loss_func = FocalLossSoftLabels(alpha=effective_alpha, gamma=focal_loss_gamma)

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        """
        Trainer가 손실을 계산하는 방식 오버라이드.
        Label Smoothing이 적용된 경우 targets는 (N, C) 형태의 소프트 레이블이 됩니다.
        """
        # labels 키가 있는지 확인 (Trainer 내부적으로 label smoothing 적용 후 전달)
        if "labels" not in inputs:
             raise ValueError("모델 입력에 'labels' 키가 없습니다. 데이터셋 형식을 확인하세요.")

        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.get('logits')

        # FocalLossSoftLabels 함수 호출 (소프트/하드 레이블 모두 처리 가능)
        loss = self.focal_loss_func(logits, labels)

        return (loss, outputs) if return_outputs else loss

from transformers import EarlyStoppingCallback # <--- 이 줄 추가

# 모델 저장 여부 및 wandb 사용 여부 설정
SAVE_MODEL = True
USE_WANDB = True

print("\n" + "=" * 50)
print("모델 훈련 시작")
print("=" * 50)

# 훈련 파라미터 설정
training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=NUM_EPOCHS,
    per_device_train_batch_size=BATCH_SIZE_TRAIN,
    per_device_eval_batch_size=BATCH_SIZE_EVAL,
    warmup_steps=WARMUP_STEPS,
    weight_decay=WEIGHT_DECAY,
    learning_rate=LEARNING_RATE,
    logging_steps=100,
    eval_strategy="epoch",
    save_strategy="epoch" if SAVE_MODEL else "no",
    load_best_model_at_end=SAVE_MODEL,
    metric_for_best_model="accuracy" if SAVE_MODEL else None,
    greater_is_better=True,
    save_total_limit=2 if SAVE_MODEL else 0,
    label_smoothing_factor=0.1, # 10% 레이블 스무딩 적용
    report_to="wandb" if USE_WANDB else "none",  # wandb 로깅 조건부 활성화
    run_name="bert-movie-review-classification" if USE_WANDB else None,
    seed=RANDOM_STATE,
    fp16=torch.cuda.is_available(),
    dataloader_num_workers=2,
    remove_unused_columns=False,
    push_to_hub=False,
    gradient_accumulation_steps=1,
    logging_first_step=True,
    save_safetensors=SAVE_MODEL,
    # === 스케줄러 타입 변경 ===
    lr_scheduler_type="cosine", # <--- 'linear'(기본값) 대신 'cosine'으로 설정
    # === Early Stopping 관련 파라미터 추가 ===
    # early_stopping_patience=3,      # <--- 추가: 이 횟수만큼 평가 점수가 개선되지 않으면 중단 (예: 3 에폭)
    # early_stopping_threshold=0.001, # <--- (선택 사항) 개선으로 간주하기 위한 최소 변화량
    
)
early_stopping_callback = EarlyStoppingCallback(early_stopping_patience=3) # <--- 원하는 patience 값
# Trainer 초기화
trainer = CustomTrainerWithFocalLoss(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    tokenizer=tokenizer,
    num_classes=NUM_CLASSES,
    focal_loss_alpha=[0.11, 0.45, 0.12, 0.32], # 예시 시작점
    focal_loss_gamma=2.0,
    callbacks=[early_stopping_callback], # <--- 콜백 리스트 전달
    data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
    compute_metrics=compute_metrics,
)

# 훈련 정보 출력
print(f"훈련 샘플: {len(train_dataset):,}개")
print(f"검증 샘플: {len(val_dataset):,}개")
print(f"훈련 에포크: {training_args.num_train_epochs}회")
print(f"배치 크기: {BATCH_SIZE_TRAIN} (훈련) / {BATCH_SIZE_EVAL} (검증)")
print(f"학습률: {LEARNING_RATE}")
print(f"시드값: {RANDOM_STATE}")
print(f"디바이스: {device}")
print(f"wandb 사용: {USE_WANDB}")

# 훈련 실행
try:
    training_results = trainer.train()
    print("\n훈련 완료")
    print(f"최종 훈련 손실: {training_results.training_loss:.4f}")

    # 훈련 로그 정보 출력
    if hasattr(training_results, "log_history"):
        print(f"총 훈련 스텝: {training_results.global_step}")

except KeyboardInterrupt:
    print("\n사용자에 의해 훈련이 중단되었습니다.")
    raise
except Exception as e:
    print(f"\n훈련 중 오류 발생: {str(e)}")
    raise