import torch
from torch.utils.data import Dataset
import numpy as np

class FashionReviewDataset(Dataset):
    def __init__(self, df, tokenizer, max_len=128):
        self.df = df
        self.tokenizer = tokenizer
        self.max_len = max_len
        
        # 데이터프레임에서 필요한 컬럼 추출
        self.texts = df['input_text'].astype(str).tolist()
        self.ratings = df['rating'].astype(float).tolist()
        # Trust Score가 이미 계산되어 있다면 가져오고, 없다면 기본값 1.0
        self.weights = df.get('trust_score', pd.Series([1.0]*len(df))).astype(float).tolist()

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        text = self.texts[idx]
        rating = self.ratings[idx]
        weight = self.weights[idx]

        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_len,
            padding='max_length',
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt',
        )

        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'rating': torch.tensor(rating, dtype=torch.float),
            'weight': torch.tensor(weight, dtype=torch.float)
        }