import torch
import torch.nn as nn
import torch.optim as optim
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# διαβάζουμε το αρχείο δεδομένων για την εκπαίδευση
df = pd.read_csv("mydatafile.csv")
features = [   "RAIN",    "NIGHT",    "LIGHTING",    "VISIBILITY",    "SPEED",    "SURFACE"]
X = df[features].values
y = df["DANGER"].values.reshape(-1, 1)
X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=0.2, random_state=42)
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.float32)

# περιγραφή του Multi-layer perceptron μοντέλου
class SimpleMLP(nn.Module):
    def __init__(self):
        super(SimpleMLP, self).__init__()
        
        self.layer1 = nn.Linear(6, 32)
        self.layer2 = nn.Linear(32, 16)
        self.output = nn.Linear(16, 1)
        
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid() 
        
    def forward(self, x):
        x = self.relu(self.layer1(x))
        x = self.relu(self.layer2(x))
        x = self.sigmoid(self.output(x))
        return x


##### ΚΥΡΙΩΣ ΠΡΟΓΡΑΜΜΑ ########

# Δημιουργία του νευρωνικού δικτύου
model = SimpleMLP()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Εκπαίδευση
epochs = 200

for epoch in range(epochs):

    # Forward pass
    predictions = model(X_train)
    loss = criterion(predictions, y_train)
    
    # Backward pass
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if (epoch+1) % 10 == 0:
        print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}")


# Αξιολόγηση
with torch.no_grad():
    test_predictions = model(X_test)
    test_loss = criterion(test_predictions, y_test)
    print("\nTest Data:", X_test)
    print("\nTargets:", y_test)
    print("\nPredictions:",test_predictions)
    print("\nTest Loss:", test_loss.item())


# Αποθήκευση εκπαιδευμένων weights κ biases σε αρχείο
torch.save(model.state_dict(), "road_danger_mlp_weights.pth")
print("\nModel weights saved to road_danger_mlp_weights.pth")