import torch
import torch.nn as nn #multi later nueral network
import torch.optim as optim #optimization
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
#tensor is multi dimensional structure for deep leanring models

# 1-D data
scaler = np.array(3.14)
scaler

array(3.14)

# 2-D data
matrix = np.arange(12).reshape(3,4)
matrix

array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

plt.figure(figsize=(3,3))
sns.heatmap(matrix, square=True, annot=True)
plt.show()

# multi-dimensional data
tensor_01 = torch.rand(2,3,4)
print(tensor_01.shape)
print(tensor_01)

torch.Size([2, 3, 4])
tensor([[[0.8282, 0.5561, 0.7492, 0.7471],
         [0.2671, 0.9685, 0.7803, 0.9106],
         [0.1462, 0.5651, 0.4328, 0.5159]],

        [[0.1904, 0.4941, 0.0195, 0.2649],
         [0.8705, 0.2303, 0.7087, 0.8983],
         [0.2587, 0.8834, 0.0723, 0.1490]]])

#toy data (two classes), manually defined data
np.random.seed(0)
c_0 = np.random.randn(100,2) + np.array([-2,0]) #create random numbers
c_1 = np.random.randn(100,2) + np.array([2,0]) #create random numbers
X = np.vstack([c_0,c_1]) #stack the data sets horizontally
Y = np.hstack([np.zeros(100), np.ones(100)]) #stack the data sets vertically

#visualize the data
plt.figure(figsize=(6,4))
plt.scatter(X[:,0], X[:, 1], c=Y, cmap='coolwarm', alpha = 0.5)
plt.show()

#turn data into tensors (multi dimensional format)
X_t = torch.tensor(X, dtype=torch.float32)
Y_t = torch.tensor(Y, dtype=torch.long)

#build a tiny model neural network,
model = nn. Sequential (
    nn.Linear(2, 8), nn.ReLU(),
    nn.Linear(8,8), nn.ReLU(),
    nn.Linear(8,2)
)

#model training, optimizing to make less mistakes /be refined
optimizer = optim.Adam(model.parameters(), lr=0.03)
loss_fc= nn.CrossEntropyLoss()

for epoch in range(300):
  optimizer.zero_grad()
  y_pred = model(X_t)
  loss = loss_fc(y_pred, Y_t)
  loss.backward()
  optimizer.step()
  if epoch % 50 ==0:
    print(f"Epoch {epoch}, loss={loss.item():.3f}")

Epoch 0, loss=0.727
Epoch 50, loss=0.062
Epoch 100, loss=0.049
Epoch 150, loss=0.045
Epoch 200, loss=0.040
Epoch 250, loss=0.037

#visualize the decision boundry
xx, yy = np.meshgrid(np.linspace(-5,5,200), np.linspace(-3,3,200))
grid = torch.tensor(np.c_[xx.ravel(), yy.ravel()], dtype=torch.float32)

with torch.no_grad():
  probs= torch.softmax(model(grid), dim=1)[:,1].reshape(xx.shape)

plt.contour(xx, yy, probs, levels=20, cmap='coolwarm', alpha=0.5)
plt.scatter(X[Y==0, 0], X[Y==0,1], label='Class 0')
plt.scatter(X[Y==1,0], X[Y==1,1], label='Class 1')
plt.legend()
plt.show()