import torch
from torch.utils.data import Dataset
import pandas as pd
import joblib

class HeartDiseaseDataset(Dataset):
    def __init__(self, csv_path):
        self.df = pd.read_csv(csv_path)
        # Load metadata to ensure we use the correct columns
        metadata = joblib.load('assets/model_metadata.joblib')
        self.cat_cols = metadata['cat_cols']
        self.num_cols = metadata['num_cols']
        self.target = metadata['target']

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

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        
        # Categorical features must be Long tensors for nn.Embedding
        x_cat = torch.tensor(row[self.cat_cols].values.astype(int), dtype=torch.long)
        
        # Numerical features must be Float tensors
        x_num = torch.tensor(row[self.num_cols].values.astype(float), dtype=torch.float)
        
        # Target
        y = torch.tensor(row[self.target], dtype=torch.float).unsqueeze(0)
        
        return x_cat, x_num, y