import torch
import torch.nn as nn
from dataset import get_data
from model import LinearModel
from matplotlib import pyplot as plt
def train():
    x_train, y_train = get_data()
    model = LinearModel()
    loss_fn = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(),lr=0.01)
    epochs = 1000
    y_axis = []
    for epoch in range(epochs):
        pred = model(x_train)
        loss = loss_fn(pred,y_train)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        y_axis.append(loss.item())
        if epoch % 100 == 0:
            print(f'Epoch {epoch}, Loss: {loss.item()}')
    x_axis = range(epochs)
    plt.plot(x_axis, y_axis)
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.title('Training Loss over Epochs')
    
    
    torch.save(model.state_dict(),'models/model.pt')
    print('Model saved!')
    plt.show()
        
if __name__ == '__main__':
    train()