import torch
import torch.nn as nn
from torch.utils.data import Dataset


# PyTorch Dataset
class RealEstateDataset(Dataset):
    def __init__(self, df, categorical_columns, numerical_columns):
        self.X_cat = torch.tensor(df[categorical_columns].values, dtype=torch.long)
        self.X_num = torch.tensor(df[numerical_columns].values, dtype=torch.float32)
        self.y_price = torch.tensor(df['sale_price'].values, dtype=torch.float32)
        self.y_date = torch.tensor(df['sale_date'].values, dtype=torch.long)

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

    def __getitem__(self, idx):
        return self.X_cat[idx], self.X_num[idx], self.y_price[idx], self.y_date[idx]


# Model Definition
class RealEstateModel(nn.Module):
    def __init__(self, num_numerical, embedding_dims, num_date_classes):
        super(RealEstateModel, self).__init__()

        # Embeddings for categorical variables
        self.embeddings = nn.ModuleList([
            nn.Embedding(num_classes, dim)
            for num_classes, dim in embedding_dims
        ])

        # Fully connected layers
        embedding_dim_sum = sum(dim for _, dim in embedding_dims)
        self.fc1 = nn.Linear(embedding_dim_sum + num_numerical, 256)
        self.fc2 = nn.Linear(256, 128)

        # Output layers
        self.price_head = nn.Linear(128, 1)
        self.date_head = nn.Linear(128, num_date_classes)

    def forward(self, x_cat, x_num):
        # Embedding layers
        x = [emb(x_cat[:, i]) for i, emb in enumerate(self.embeddings)]
        x = torch.cat(x, dim=1)

        # Concatenate with numerical features
        x = torch.cat([x, x_num], dim=1)

        # Fully connected layers
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))

        # Output heads
        price = self.price_head(x)
        sale_date = self.date_head(x)

        return price, sale_date