import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin
import pandas as pd

# Set device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  
print('Device:', device)


# Define model for the model class

torch.manual_seed(42)

model = nn.Sequential(
    nn.Linear(12, 12),
    nn.ReLU(),
    nn.Linear(12, 6),
    nn.ReLU(),
    nn.Linear(6, 1),
    nn.Sigmoid()
)

# Define model class
class MyModel(nn.Module, PyTorchModelHubMixin):

    def __init__(self):
        super().__init__()  # Initialize nn.Module
        self.model = model

    def forward(self, x):
        return self.model(x)  # Assume this model has a defined forward pass

# EndpointHandler class
class EndpointHandler:
    def __init__(self, path=""):
        self.model = MyModel.from_pretrained("damiano216/pay-boo-2")  # Load from Hugging Face
        self.model.to(device)  # Move model to GPU if available
        self.model.eval()  # Set model to evaluation mode

    def __call__(self, data):

        print(f"Payload: {data}")

        # Create a Pandas DataFrame
        payloadDataFrame = pd.DataFrame(data['chargeData'])
        print(payloadDataFrame)

        
        new_data_tensor = torch.tensor(payloadDataFrame.values, dtype=torch.float).to(device)  # Ensure tensor is on device
        print(f"new_data_tensor: {new_data_tensor}")
        
        # Make predictions
        with torch.no_grad():
            predictions = self.model(new_data_tensor)

        # Interpret predictions
        print(f"Predictions: {predictions[0].item()}")

        return predictions[0].item()