# fair_rec/loss.py (또는 model.py) 내부

import torch
import torch.nn as nn

class TrustWeightedMSELoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.mse = nn.MSELoss(reduction='none')
        
        # [수정됨] Study 1의 J-Curve 결과 반영 (Beta Coefficients)
        # 5점(0.201) > 4점(0.135) > 1점(0.114) > 3점(0.0) > 2점(-0.039)
        self.trust_weights = {
            5: 0.201,   # Decisive Validation (가장 높은 신뢰도)
            4: 0.135,   # Safe Baseline (5점보다 낮음!)
            3: 0.000,   # Reference Point
            2: -0.039,  # Ambiguous Dissatisfaction (신뢰도 가장 낮음)
            1: 0.114    # Decisive Warning (확실한 불만은 오히려 신뢰도 높음)
        }

    def forward(self, predictions, targets, ratings, lengths=None):
        # 1. 기본 MSE 계산
        loss = self.mse(predictions, targets)
        
        # 2. 신뢰도 가중치 매핑
        device = predictions.device
        weights = torch.tensor([self.trust_weights.get(int(r), 1.0) for r in ratings], device=device)
        
        # 3. (옵션) Hollow 5-star Penalty: 5점인데 내용이 너무 짧으면(10단어 미만) 신뢰도 깎기
        if lengths is not None:
            # 5점(ratings==5) 이면서 길이가 짧은(lengths < 10) 경우
            mask_hollow = (ratings == 5) & (lengths < 10)
            # 가중치를 0.201에서 절반 수준으로 낮춤
            weights[mask_hollow] = weights[mask_hollow] * 0.5 

        # 4. 최종 Loss 계산
        weighted_loss = loss * weights
        return weighted_loss.mean()