import torch
import numpy as np
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from IPython import display
display.set_matplotlib_formats('svg')

# https://playground.tensorflow.org/

# generate data for class A and class B

A = [1, 1]
B = [5, 1]

#x, y
a = [A[0] + np.random.randn(100), A[1] + np.random.randn(100)]
b = [B[0] + np.random.randn(100), B[1] + np.random.randn(100)]

#y
ya = np.zeros((100, 1))
yb = np.ones((100, 1))
labels = np.vstack( (ya, yb) ) 
data = np.hstack( (a, b) ).T 

x = torch.tensor(data).float()
y = torch.tensor(labels).float()

plt.plot(x[np.where(y==0)[0], 0], x[np.where(y==0)[0], 1],'bs')
plt.plot(x[np.where(y==1)[0], 0], x[np.where(y==1)[0], 1],'ks')
plt.show()

# there are many possible architectures of nn
# in this case we create a classification network in the form of
# input - linear - relu - linear - sigmoid

#instead of this part we use model = nn.Sequential():
#####
    #def forward(X, w, b):
    #    return torch.matmul(X, w) + b 

    #w = torch.randn(2, 1, requires_grad=True) 
    #b = torch.randn(1, requires_grad=True)
#####

model = nn.Sequential(
    nn.Linear(2, 1),
    nn.ReLU(),
    nn.Linear(1, 1),
    nn.Sigmoid()
)

print("model", model)
print("model.weight", model[0].weight)
print("model.bias", model[0].bias)
print("model.weight", model[2].weight)
print("model.bias", model[2].bias)

loss_fn = nn.BCELoss()
optimizer = optim.SGD(model.parameters(),lr=0.05)


epochs=1000
train_losses = torch.zeros(epochs)
for epoch in range(epochs):
    
    pred = model(x)
    
    loss = loss_fn(pred, y)
    train_losses[epoch] = loss
    
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

plt.plot(train_losses.detach(),'o',markerfacecolor='w',linewidth=.1)
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.show()

predictions = model(x)
print(predictions)
p_labels = predictions > 0.5
print(p_labels)