import sklearn
from sklearn.datasets import make_circles

# Make a 1000 samples
n_samples = 1000

X, y = make_circles(n_samples, noise=0.03, random_state=42)
len(X), len(y)

X[:5], y[:5]

# Make DataFrame of circle data
import pandas as pd
circles = pd.DataFrame({
    "X1": X[:, 0],
    "X2": X[:, 1],
    "label": y
})
circles.head(10)

# Visualize, visualize, visualize
import matplotlib.pyplot as plt
plt.scatter(
    x=X[:, 0],
    y=X[:, 1],
    c=y,
    cmap=plt.cm.RdYlBu
)

circles.label.value_counts()

X.shape, y.shape

# View the first example of features and labels
X_sample = X[0]
y_sample = y[0]

print(f"Values for one sample of X: {X_sample} and the same for y: {y_sample}")
print(f"Shapes for one sample of X: {X_sample.shape} and the same for y: {y_sample.shape}")

# Turn data into tensors
import torch

X = torch.from_numpy(X).type(torch.float)
y = torch.from_numpy(y).type(torch.float)

X[:5], y[:5]

# Split data into train and test sets
from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

len(X_train), len(X_test), len(y_train), len(y_test)

# Import Pytorch and nn
import torch
from torch import nn

# Make device agnostic code
device = "cuda" if torch.cuda.is_available() else "cpu"
device

# 1. Construc a model that subclasses nn.Module

class CircleModelV1(nn.Module):
    def __init__(self):
        super().__init__()
        # 2. Create 2 nn.Linear layers
        # self.layer_1 = nn.Linear(in_features=2, out_features=5) # takes 2 features and upscales to 5 features
        # self.layer_2 = nn.Linear(in_features=5, out_features=1) # takes 5 features from previous layer and outputs a single feature (same shape as y)
        self.two_linear_layers = nn.Sequential(
            nn.Linear(in_features=2, out_features=5),
            nn.Linear(in_features=5, out_features=1)
        )
    # 3. Define forward method
    def forward(self, x):
        return self.two_linear_layers(x) # x-> layer_1 -> layer_2 -> output

# 4. Instantiate an instance of our model class and send it to target device
model_0 = CircleModelV1().to(device)
model_0

device

next(model_0.parameters()).device

# Let's replicate the model abovee using nn.Sequential()
model_0 = nn.Sequential(
    nn.Linear(in_features=2, out_features=5),
    nn.Linear(in_features=5, out_features=1)
).to(device)

model_0

model_0.state_dict()

# Make predictions
with torch.inference_mode():
    untrained_preds = model_0(X_test.to(device))
    print(f"Length of predictions {len(untrained_preds)}, Shape: {untrained_preds.shape}")
    print(f"Length of test samples: {len(X_test)}, Shape: {X_test.shape}")
    print(f"First 10 predictions: {torch.round(untrained_preds[:10])}")
    print(f"First 10 labels: {y_test[:10]}")

## 2.1 Setup loss function and optimizer
loss_fn = nn.BCEWithLogitsLoss() # BCEWithLogitsLoss has sigmoid activation build in
optimizer = torch.optim.SGD(params=model_0.parameters(), lr=0.1)

# Calculate accuracy - what percentage model predicted right
def accuracy_fn(y_true, y_pred):
    correct = torch.eq(y_true, y_pred).sum().item()
    acc = (correct/len(y_pred)) * 100
    return acc

model_0.eval()
with torch.inference_mode():
    y_logits = model_0(X_test.to(device))[:5]
y_logits

y_test[:5]

# Use the sigmoid activation function on our model logits to turn them into prediction probabilities
y_pred_probs = torch.sigmoid(y_logits)
y_pred_probs

# if probability more than or equal 0.5, then label 1, if less 0.5 than 0
y_preds = torch.round(y_pred_probs)
y_preds

# In full (logits -> prediction probobalities -> prediction labels)
y_pred_labels = torch.round(torch.sigmoid(model_0(X_test.to(device)[:5])))
y_pred_labels

# Check for equality
print(torch.eq(y_preds.squeeze(), y_pred_labels.squeeze()))

# Get rid of extra dimension
y_preds.squeeze()

### Building testing and training loop
torch.manual_seed(42)
torch.cuda.manual_seed(42)

# Set the number of epochs
epochs = 100

# Put data to the target device
X_train, y_train = X_train.to(device), y_train.to(device)
X_test, y_test = X_test.to(device), y_test.to(device)

for epoch in range(epochs):
    model_0.train()
    
    # forward pass
    y_logits = model_0(X_train).squeeze()
    y_pred = torch.round(torch.sigmoid(y_logits))
    
    # calculate the loss and accuracy
    loss = loss_fn(y_logits, y_train)
    acc = accuracy_fn(y_true=y_train, y_pred=y_pred)
    
    optimizer.zero_grad()
    
    # backpropogation
    loss.backward()
    
    # gradient descent
    optimizer.step()
    
    ### Testing
    model_0.eval()
    with torch.inference_mode():
        test_logits = model_0(X_test).squeeze()
        test_pred = torch.round(torch.sigmoid(test_logits))
        
        # calculate test loss and accuracy
        test_loss = loss_fn(test_logits, y_test)
        test_acc = accuracy_fn(y_true=y_test, y_pred=test_pred)
    
    # print out what happening
    if epoch % 10 == 0:
        print(f"Epoch: {epoch} | Loss: {loss:.5f}, Acc: {acc:.2f}%| Test loss: {test_loss:.5f}, test acc: {test_acc:.2f}%")

import requests
from pathlib import Path

# Download helper functions from Learn Pytorch repo

if Path("helper_functions.py").is_file():
    print("already exist")
else:
    request = requests.get("https://raw.githubusercontent.com/mrdbourke/pytorch-deep-learning/main/helper_functions.py")
    with open("helper_functions.py", "wb") as f:
        f.write(request.content)
        
from helper_functions import plot_predictions, plot_decision_boundary

# Plot decision boundary of the model
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.title("Train")
plot_decision_boundary(model_0, X_train, y_train)
plt.subplot(1, 2, 2)
plt.title("Test")
plot_decision_boundary(model_0, X_test, y_test)

class CircleModelV1(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer_1 = nn.Linear(in_features=2, out_features=10)
        self.layer_2 = nn.Linear(in_features=10, out_features=10)
        self.layer_3 = nn.Linear(in_features=10, out_features=1)
        
    def forward(self, x):
        return self.layer_3(self.layer_2(self.layer_1(x)))
    
model_1 = CircleModelV1().to(device)
model_1

loss_fn = torch.nn.BCEWithLogitsLoss()
optimizer = torch.optim.SGD(params=model_1.parameters(), lr=0.1)

### Building testing and training loop
torch.manual_seed(42)
torch.cuda.manual_seed(42)
epochs = 1000

X_train, y_train = X_train.to(device), y_train.to(device)
X_test, y_test = X_test.to(device), y_test.to(device)

for epoch in range(epochs):
    model_1.train()
    
    train_logits = model_1(X_train).squeeze()
    train_preds = torch.round(torch.sigmoid(train_logits))
    
    loss = loss_fn(train_logits, y_train)
    train_acc = accuracy_fn(y_true=y_train, y_pred=train_preds)
    
    optimizer.zero_grad()
    
    loss.backward()
    
    optimizer.step()
    
    model_1.eval()
    with torch.inference_mode():
        test_logits = model_1(X_test).squeeze()
        test_preds = torch.round(torch.sigmoid(test_logits))
        
        test_loss = loss_fn(test_logits, y_test)
        test_acc = accuracy_fn(y_true= y_test, y_pred=test_preds)
    
    if epoch % 100 == 0:
        print(f"Epoch: {epoch} | Loss: {loss}. Accuracy: {train_acc} | Test Loss: {test_loss}. Test accuracy: {test_acc}")

# Plot decision boundary of the model
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.title("Train")
plot_decision_boundary(model_1, X_train, y_train)
plt.subplot(1, 2, 2)
plt.title("Test")
plot_decision_boundary(model_1, X_test, y_test)

### 6.1 Recreating non-linear data(red and blue circles)

import matplotlib.pyplot as plt
from sklearn.datasets import make_circles

n_samples = 1000
X, y = make_circles(n_samples,
                    noise=0.03,
                   random_state=42)
plt.scatter(X[:, 0], X[:, 1], c=y, cmap=plt.cm.RdYlBu)

# Convert data to tensors and train and test splits
import torch
from sklearn.model_selection import train_test_split

X = torch.from_numpy(X).type(torch.float)
y = torch.from_numpy(y).type(torch.float)

# Split into train and test
X_train, X_test, y_train, y_test = train_test_split(X,
                                                   y,
                                                   test_size=0.2,
                                                   random_state=42)

X_train[:5], y_train[:5]

# Build a model with non-linear activation function
from torch import nn

class CircleModelV2(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer_1 = nn.Linear(in_features=2, out_features=10)
        self.layer_2 = nn.Linear(in_features=10, out_features=10)
        self.layer_3 = nn.Linear(in_features=10, out_features=1)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        return self.layer_3(self.relu(self.layer_2(self.relu(self.layer_1(x)))))

model_3 = CircleModelV2().to(device)
model_3

loss_fn = torch.nn.BCEWithLogitsLoss()
optimizer = torch.optim.SGD(params=model_3.parameters(), lr=0.1)

torch.manual_seed(42)
torch.cuda.manual_seed(42)
epochs = 2500

X_train, y_train = X_train.to(device), y_train.to(device)
X_test, y_test = X_test.to(device), y_test.to(device)

for epoch in range(epochs):
    model_3.train()
    y_logits = model_3(X_train).squeeze()
    loss = loss_fn(y_logits, y_train)
    y_pred = torch.round(torch.sigmoid(y_logits))
    acc = accuracy_fn(y_true=y_train, y_pred=y_pred)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    model_3.eval()
    with torch.inference_mode():
        y_test_logits = model_3(X_test).squeeze()
        loss_test = loss_fn(y_test_logits, y_test)
        y_test_pred = torch.round(torch.sigmoid(y_test_logits))
        test_acc = accuracy_fn(y_true=y_test, y_pred=y_test_pred)
        
        if epoch % 100 == 0:
            print(f"epoch: {epoch}. Loss: {loss:.2f}. Accuracy: {acc:.2f}. Test Loss {loss_test:.2f}. Test accuracy: {test_acc:.2f}")

# Plot decision boundary of the model
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.title("Train")
plot_decision_boundary(model_3, X_train, y_train)
plt.subplot(1, 2, 2)
plt.title("Test")
plot_decision_boundary(model_3, X_test, y_test)

# make predictions
X_train, y_train = X_train.to("cpu"), y_train.to("cpu")
X_test, y_test = X_test.to("cpu"), y_test.to("cpu")

model_3.eval()
with torch.inference_mode():
    y_preds = torch.round(torch.sigmoid(model_3(X_test))).squeeze()

y_preds[:10], y_test[:10]

# Create a tensor
A = torch.arange(-10, 10, 1, dtype=torch.float32)
A.dtype

# Visualize the data
plt.plot(A)

plt.plot(torch.relu(A))

def relu(x: torch.Tensor) -> torch.Tensor:
    return torch.maximum(torch.tensor(0), x) # inputs must be tensors

plt.plot(relu(A))

# Let's do same for sigmoid
def sigmoid(x):
    return 1 / (1 + torch.exp(-x))

plt.plot(torch.sigmoid(A))

plt.plot(sigmoid(A))

import torch
import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs
from sklearn.model_selection import train_test_split

# Set hyper parameters for data creation
NUM_CLASSES = 4
NUM_FEATURES = 2
RANDOM_SEED = 42

# 1. Create multi-class data
X_blob, y_blob = make_blobs(n_samples=1000,
                           n_features=NUM_FEATURES,
                           centers=NUM_CLASSES,
                           cluster_std=1.5,
                           random_state=RANDOM_SEED)
# 2. Turn data into tensors
X_blob = torch.from_numpy(X_blob).type(torch.float)
y_blob = torch.from_numpy(y_blob).type(torch.LongTensor)

# 3. Split into train and test
X_blob_train, X_blob_test, y_blob_train, y_blob_test = train_test_split(X_blob, y_blob, test_size=0.2, random_state=RANDOM_SEED)

# 4. Visualize
plt.figure(figsize=(10, 7))
plt.scatter(X_blob[:, 0], X_blob[:, 1], c=y_blob, cmap=plt.cm.RdYlBu)

# Create device agnostic code
device = "cuda" if torch.cuda.is_available() else "cpu"
device

# Build multi-class classification model
class BlobModel(nn.Module):
    def __init__(self, input_features, output_features, hidden_units=8):
        """
        Initializes multi-class classification model.
        
        Args:
        input_features (int): Number of input features to the model
        output_features (int): Number of output features (number of output classes)
        hidden_units (int): Number of hidden units between layers, default 8
        """
        super().__init__()
        
        self.linear_layer_stack = nn.Sequential(
            nn.Linear(in_features=input_features, out_features=hidden_units),
            nn.ReLU(),
            nn.Linear(in_features=hidden_units, out_features=hidden_units),
            nn.ReLU(),
            nn.Linear(in_features=hidden_units, out_features=output_features)
        )
    
    def forward(self, x):
        return self.linear_layer_stack(x)
    
    
# Create an instance of blob model and send to target device
model_4 = BlobModel(input_features=2,
                   output_features=4,
                   hidden_units=8).to(device)
model_4

# Create a loss function for multi-class classification
loss_fn = nn.CrossEntropyLoss()

optimizer = torch.optim.SGD(params=model_4.parameters(), lr=0.1)

model_4.eval()
with torch.inference_mode():
    y_logits = model_4(X_blob_test.to(device))[:5]

y_logits

# Convert our model's logit output to prediction probabilities
y_pred_probs = torch.softmax(y_logits, dim=1)
y_pred_probs[:5]

torch.sum(y_pred_probs[0])

torch.argmax(y_pred_probs[0])

# Conver our model's prediction probabilities to prediction labels
y_preds = torch.argmax(y_pred_probs, dim=1)
y_preds

torch.manual_seed(42)
torch.cuda.manual_seed(42)
epochs = 100

X_blob_train, y_blob_train = X_blob_train.to(device), y_blob_train.to(device)
X_blob_test, y_blob_test = X_blob_test.to(device), y_blob_test.to(device)

for epoch in range(epochs):
    model_4.train()
    
    y_logits = model_4(X_blob_train)
    y_preds = torch.argmax(torch.softmax(y_logits, dim=1), dim=1)
    loss = loss_fn(y_logits, y_blob_train)
    acc = accuracy_fn(y_true=y_blob_train, y_pred=y_preds)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    model_4.eval()
    with torch.inference_mode():
        y_test_logits = model_4(X_blob_test)
        y_test_preds = torch.argmax(torch.softmax(y_test_logits, dim=1), dim=1)
        loss_test = loss_fn(y_test_logits, y_blob_test)
        acc_test = accuracy_fn(y_true=y_blob_test, y_pred=y_test_preds)
    
        if epoch % 10 == 0:
            print(f"Epoch {epoch} | Loss: {loss:.2f} | Accuracy: {acc:.2f} | Test loss: {loss_test:.2f} | Test Accuracy: {acc_test:.2f}")

# Make predictions
model_4.eval()
model_4.to(device)
with torch.inference_mode():
    y_logits = model_4(X_blob_test.to(device))
    y_preds = torch.softmax(y_logits, dim=1).argmax(dim=1)

y_preds[:10], y_blob_test[:10]

# Plot decision boundary of the model
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.title("Train")
plot_decision_boundary(model_4, X_blob_train, y_blob_train)
plt.subplot(1, 2, 2)
plt.title("Test")
plot_decision_boundary(model_4, X_blob_test, y_blob_test)

import torchmetrics

from torchmetrics import Accuracy

# Setup metric
torchmetric_accuracy = Accuracy(task='multiclass',                                           
                                     num_classes=NUM_CLASSES).to(device)

# Calculate accuracy
torchmetric_accuracy(y_preds, y_blob_test)

import torch
from torch import nn
from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split

moons_numpy = make_moons(n_samples=1000, random_state=42)
X, y = moons_numpy[0], moons_numpy[1]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

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

X_train = torch.from_numpy(X_train).to(device).type(torch.float)
X_test = torch.from_numpy(X_test).to(device).type(torch.float)
y_train = torch.from_numpy(y_train).to(device).type(torch.float)
y_test = torch.from_numpy(y_test).to(device).type(torch.float)

len(X_train), len(y_train),len(X_test), len(y_test)

# 2- Visualizing the dataset.
from matplotlib import pyplot as plt

# When the label y is 0, the class is represented with a blue square.
# When the label y is 1, the class is represented with a green triangle.
plt.plot(moons_numpy[0][:, 0][moons_numpy[1]==1], moons_numpy[0][:, 1][moons_numpy[1]==1], "bs")
plt.plot(moons_numpy[0][:, 0][moons_numpy[1]==0], moons_numpy[0][:, 1][moons_numpy[1]==0], "g^")

# X contains two features, x1 and x2
plt.xlabel(r"$x_1$", fontsize=20)
plt.ylabel(r"$x_2$", fontsize=20)

# Simplifying the plot by removing the axis scales.
plt.xticks([])
plt.yticks([])

# Displaying the plot.
plt.show()

class MoonModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(in_features=2, out_features=10),
            nn.ReLU(),
            nn.Linear(in_features=10, out_features=10),
            nn.ReLU(),
            nn.Linear(in_features=10, out_features=1)
        )
    
    def forward(self, x):
        return self.layers(x)
    
moon_model = MoonModel().to(device)
moon_model

loss_fn = nn.BCEWithLogitsLoss()
optimizer = torch.optim.SGD(params = moon_model.parameters(), lr=0.1)

import torchmetrics

torch.manual_seed(42)
torch.cuda.manual_seed(42)

accuracy_fn = Accuracy(task='binary').to(device)

epochs = 810

for epoch in range(epochs):
    moon_model.train()
    
    y_logits = moon_model(X_train).squeeze()
    loss = loss_fn(y_logits, y_train)
    acc = accuracy_fn(torch.round(torch.sigmoid(y_logits)), y_train)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    moon_model.eval()
    y_test_logits = moon_model(X_test).squeeze()
    acc_test = accuracy_fn(torch.round(torch.sigmoid(y_test_logits)), y_test)
    loss_test = loss_fn(y_test_logits, y_test)
    
    if epoch % 100 == 0:
        print(f"Epoch: {epoch} | Loss {loss} | Accuracy {acc}| Test Loss {loss_test} | Test Accuracy {acc_test}")
    

# Plot decision boundary of the model
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.title("Train")
plot_decision_boundary(moon_model, X_train, y_train)
plt.subplot(1, 2, 2)
plt.title("Test")
plot_decision_boundary(moon_model, X_test, y_test)

def tanh(z):
	return (torch.exp(z) - torch.exp(-z)) / (torch.exp(z) + torch.exp(-z))

# Code for creating a spiral dataset from CS231n
import numpy as np
N = 100 # number of points per class
D = 2 # dimensionality
K = 3 # number of classes
X = np.zeros((N*K,D)) # data matrix (each row = single example)
y = np.zeros(N*K, dtype='uint8') # class labels
for j in range(K):
  ix = range(N*j,N*(j+1))
  r = np.linspace(0.0,1,N) # radius
  t = np.linspace(j*4,(j+1)*4,N) + np.random.randn(N)*0.2 # theta
  X[ix] = np.c_[r*np.sin(t), r*np.cos(t)]
  y[ix] = j
# lets visualize the data
plt.scatter(X[:, 0], X[:, 1], c=y, s=40, cmap=plt.cm.Spectral)
plt.show()

len(X), len(y)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
X_train = torch.from_numpy(X_train).type(torch.float).to(device)
X_test = torch.from_numpy(X_test).type(torch.float).to(device)
y_train = torch.from_numpy(y_train).type(torch.torch.LongTensor).to(device)
y_test = torch.from_numpy(y_test).type(torch.torch.LongTensor).to(device)

print(len(X_train), len(y_train),len(X_test), len(y_test))
print(X_train.device, y_train.device, X_test.device, y_test.device)

class SpiralModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(in_features=2, out_features=10),
            nn.ReLU(),
            nn.Linear(in_features=10, out_features=10),
            nn.ReLU(),
            nn.Linear(in_features=10, out_features=3)
        )
    
    def forward(self, x):
        return self.layers(x)
    
spiral_model = SpiralModel().to(device)
spiral_model

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(params = spiral_model.parameters(), lr=0.01)

import torchmetrics

torch.manual_seed(42)
torch.cuda.manual_seed(42)

accuracy_fn = Accuracy(task='multiclass', num_classes=3).to(device)

epochs = 60001

for epoch in range(epochs):
    moon_model.train()
    
    y_logits = spiral_model(X_train)
    loss = loss_fn(y_logits, y_train.to(device))
    acc = accuracy_fn(torch.softmax(y_logits, dim=1).argmax(dim=1), y_train)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    moon_model.eval()
    y_test_logits = spiral_model(X_test)
    acc_test = accuracy_fn(torch.softmax(y_test_logits, dim=1).argmax(dim=1), y_test)
    loss_test = loss_fn(y_test_logits, y_test)
    
    if epoch % 10000 == 0:
        print(f"Epoch: {epoch} | Loss {loss:.2f} | Accuracy {acc:.2f}| Test Loss {loss_test:.2f} | Test Accuracy {acc_test:.2f}")

# Plot decision boundary of the model
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.title("Train")
plot_decision_boundary(spiral_model, X_train, y_train)
plt.subplot(1, 2, 2)
plt.title("Test")
plot_decision_boundary(spiral_model, X_test, y_test)