import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
from sklearn.model_selection import train_test_split

import pandas as pd 
import numpy as np

train_df = pd.read_csv('digit recognizor.csv')

train_df.info()

train_df.head()

train_x = train_df.drop(['label'], axis=1)
train_y = train_df['label']

train_x.values.dtype

X_train, X_test, y_train, y_test = train_test_split( train_x, train_y, test_size=0.2, random_state=42)

X_train.shape

X_train = torch.tensor(X_train.values,dtype=torch.float)
y_train = torch.LongTensor(y_train.values)

X_test = torch.tensor(X_test.values,dtype=torch.float)
y_test = torch.LongTensor(y_test.values)

# X_train = torch.tensor(X_train.values)
# y_train = torch.tensor(y_train.values)

# X_test = torch.tensor(X_test.values)
# y_test = torch.tensor(y_test.values)

X_train.dtype

y_train.dtype

# Get cpu or gpu device for training.
device = "cpu"
print(f"Using {device} device")

# Define model
class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

model = NeuralNetwork().to(device)
print(model)

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

def train(X,y, model, loss_fn, optimizer):
    batch = 64
    size = len(X)
    model.train()
    X, y = X.to(device), y.to(device)
    pred = model(X)
    loss = loss_fn(pred, y)
    # Backpropagation
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if batch % 100 == 0:
        loss, current = loss.item(), batch * len(X)
        print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

def test(X,y, model, loss_fn):
    size = len(X) 
    num_batches = 64
    model.eval()
    test_loss, correct = 0, 0
    
    with torch.no_grad():
        X, y = X.to(device), y.to(device)
        pred = model(X)
        test_loss += loss_fn(pred, y).item()
        correct += (pred.argmax(1) == y).type(torch.float).sum().item()
        
    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

epochs = 10
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train(X_train,y_train, model, loss_fn, optimizer)
    test(X_test,y_test, model, loss_fn)
print("Done!")

for name, param in model.named_parameters():
    print(name,'-->',param.type(),'-->',param.dtype,'-->',param.shape)