%matplotlib inline
import matplotlib.pyplot as plt
import seaborn as sns

import numpy as np
import pandas as pd

from sklearn.utils import shuffle

#Define the model 
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

from tqdm import tqdm 

from sklearn.datasets import load_boston
boston = load_boston()
print( "Type of boston dataset:", type(Boston))
boston_df = pd.DataFrame(boston['data'] )
boston_df.columns = boston['feature_names']
boston_df['PRICE']= boston['target']

sns.distplot(boston_df['PRICE']);
sns.pairplot(boston_df);

from sklearn.model_selection  import train_test_split

y = boston_df['PRICE']
X = boston_df.iloc[:,0:13]
X_train, X_test, y_train, y_test = train_test_split(X, y, 
                                                    test_size=0.3, 
                                                    random_state=0)
X_train=X_train.values
y_train=y_train.values
X_test=X_test.values
y_test=y_test.values

batch_size = 50
num_epochs = 200
learning_rate = 0.01
size_hidden= 100

batch_no = len(X_train) // batch_size  #batches
cols=X_train.shape[1] #Number of columns in input matrix
n_output=1

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

print("Executing the model on :",device)
class Net(torch.nn.Module):
    def __init__(self, n_feature, size_hidden, n_output):
        super(Net, self).__init__()
        self.hidden = torch.nn.Linear(cols, size_hidden)
        self.predict = torch.nn.Linear(size_hidden, n_output)

    def forward(self, x):
        x = F.relu(self.hidden(x))
        x = self.predict(x) 
        return x

model = Net(cols, size_hidden, n_output)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
criterion = torch.nn.MSELoss(reduction='sum')

loss_train_history = []
loss_test_history = []
loss_test_records = []

for epoch in tqdm(range(num_epochs)):
    running_loss = 0.0
    X_train, y_train = shuffle(X_train, y_train)
    for i in range(batch_no):
        start = i * batch_size
        end = start + batch_size
        inputs = Variable(torch.FloatTensor(X_train[start:end]))
        labels = Variable(torch.FloatTensor(y_train[start:end]))
        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, torch.unsqueeze(labels,dim=1))
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        
    loss_train_history.append(running_loss)
    inputs = Variable(torch.FloatTensor(X_test))
    outputs = model(inputs)
    labels = Variable(torch.FloatTensor(y_test))
    loss = criterion(outputs, torch.unsqueeze(labels,dim=1))
    loss_test_history.append(loss.item())
    if min(loss_test_history) == loss_test_history[-1]:
      model_path = 'model.pth'
      torch.save(model.state_dict(), model_path)
      loss_test_records.append(epoch)

model.load_state_dict(torch.load(model_path))
fig, ax = plt.subplots()
ax.plot(loss_train_history, color="blue")
ax.plot(loss_test_history, color="red")
plt.show()
fig, ax = plt.subplots()
ax.plot(loss_train_history, color="blue")
ax.plot(loss_test_history, color="red")
plt.ylim(0, 8000)
plt.show()
plt.plot(loss_test_records)
plt.show()

from sklearn.metrics import r2_score

X = Variable(torch.FloatTensor(X_train)) 
result = model(X)
pred=result.data[:,0].numpy()
print(len(pred),len(y_train))
print(r2_score(pred,y_train))

X = Variable(torch.FloatTensor(X_test)) 
result = model(X)
pred=result.data[:,0].numpy()
print(len(pred),len(y_test))
print(r2_score(pred,y_test))