import pandas as pd
import torch
import torch.nn as nn
from sklearn.preprocessing import StandardScaler

df = pd.read_csv("./RegressionModels/CaliforniaHousing/housing.csv")


df = df.dropna(subset=df.columns[:8])
df['ocean_proximity_encoded'] = df['ocean_proximity'].astype('category').cat.codes     #Encodes text values as numerical ones

# print(df)
# print(df.iloc[:,0:8])

scaler_x = StandardScaler()
scaled_X = scaler_x.fit_transform(df.iloc[:,0:8].join(df["ocean_proximity_encoded"]).values)
X = torch.tensor(scaled_X, dtype=torch.float32)

scaler_y = StandardScaler()
scaled_Y = scaler_y.fit_transform(df["median_house_value"].values.reshape(-1,1))
Y = torch.tensor(scaled_Y, dtype=torch.float32)

model = torch.nn.Sequential(
    torch.nn.Linear(9, 18),
    torch.nn.ReLU(),
    torch.nn.Linear(18, 1)
)

loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-2)


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
X = X.to(device)
Y = Y.to(device)


for epoch in range(3000):
    pred_y = model(X)
    loss = loss_fn(pred_y, Y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if epoch % 99 == 0:
        print('Epoch: ', epoch, f"loss: {loss.item():.2f}")