import os
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim

# Training
train_file_path = "Train1.csv"
train_data = pd.read_csv(train_file_path)
train_data = train_data.dropna(subset=['price_range'])

valid_values = {0.0, 1.0, 2.0, 3.0}
assert set(train_data['price_range'].unique()) <= valid_values, "Unexpected values in price_range"

class SimpleNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return self.softmax(x)

input_size = len(train_data.columns) - 2  
hidden_size = 50
output_size = len(valid_values)

model = SimpleNN(input_size, hidden_size, output_size)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

epochs = 10
for epoch in range(epochs):
    inputs = torch.tensor(train_data.drop(['price_range', 'ID'], axis=1).values, dtype=torch.float32)
    labels = torch.tensor(train_data['price_range'].values, dtype=torch.long)

    optimizer.zero_grad()

    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

    print(f"Epoch {epoch+1}/{epochs}, Loss: {loss.item()}")

save_path = "model.pth"
torch.save(model.state_dict(), save_path)

# Testing
model.load_state_dict(torch.load("model.pth"))
model.eval() 

test_file_path = "Test1.csv"
test_data = pd.read_csv(test_file_path)

inputs = torch.tensor(test_data.drop(['price_range', 'ID'], axis=1).values, dtype=torch.float32)
with torch.no_grad():
    predictions = model(inputs)
predicted_classes = torch.argmax(predictions, dim=1)

predicted_classes_df = pd.DataFrame(predicted_classes.numpy(), columns=['Predicted_Price_Range'])
predicted_classes_df['Actual_Price_Range'] = test_data['price_range'].values

output_path = 'predictions.csv'
predicted_classes_df.to_csv(output_path, index=False)