import torch
import pandas as pd
import matplotlib.pyplot as plt

file = './RegressionModels/PredictSallary/Salary_dataset.csv'
df = pd.read_csv(file)

X = torch.tensor(df["YearsExperience"].values, dtype=torch.float32).unsqueeze(1)
y = torch.tensor(df['Salary'].values, dtype=torch.float32).unsqueeze(1) 

model = torch.nn.Sequential(
    torch.nn.Linear(1,1)
)

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

# 5. Training loop
for epoch in range(1000):
    y_pred = model(X)
    loss = loss_fn(y_pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if epoch % 100 == 99:
        print(f'Epoch {epoch+1}, Loss: {loss.item():.2f}')

# 6. Plot the results
plt.scatter(X.numpy(), y.numpy(), label='Actual data')
plt.plot(X.numpy(), model(X).detach().numpy(), color='red', label='Model prediction')
plt.xlabel('Years of Experience')
plt.ylabel('Salary')
plt.legend()
plt.show()