import torch
torch.cuda.is_available()
# True

"""
Explanation: <img src="../Pierian-Data-Logo.PNG">
<br>
<strong><center>Copyright 2019. Created by Jose Marcial Portilla.</center></strong>
What is CUDA?
Most people confuse CUDA for a language or maybe an API. It is not.
It’s more than that. CUDA is a parallel computing platform and programming model that makes using a GPU for general purpose computing simple and elegant. The developer still programs in the familiar C, C++, Fortran, or an ever expanding list of supported languages, and incorporates extensions of these languages in the form of a few basic keywords.
These keywords let the developer express massive amounts of parallelism and direct the compiler to the portion of the application that maps to the GPU.
How do I install PyTorch for GPU?
Refer to video, its dependent on whether you have an NVIDIA GPU card or not.
How do I know if I have CUDA available?
End of explanation
"""


## Get Id of default device
torch.cuda.current_device()

# 0
torch.cuda.get_device_name(0) # Get name device with ID '0'

# Returns the current GPU memory usage by 
# tensors in bytes for a given device
torch.cuda.memory_allocated()

# Returns the current GPU memory managed by the
# caching allocator in bytes for a given device
torch.cuda.memory_cached()

"""
Explanation: Using GPU and CUDA
We've provided 2 versions of our yml file, a GPU version and a CPU version. To use GPU, you need to either manually create a virtual environment, please watch the video related to this lecture, as not every computer can run GPU, you need CUDA and an NVIDIA GPU.
End of explanation
"""


# CPU
a = torch.FloatTensor([1.,2.])

a

a.device

# GPU
a = torch.FloatTensor([1., 2.]).cuda()

a.device

torch.cuda.memory_allocated()

"""
Explanation: Using CUDA instead of CPU
End of explanation
"""


import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split

import pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline

class Model(nn.Module):
    def __init__(self, in_features=4, h1=8, h2=9, out_features=3):
        super().__init__()
        self.fc1 = nn.Linear(in_features,h1)    # input layer
        self.fc2 = nn.Linear(h1, h2)            # hidden layer
        self.out = nn.Linear(h2, out_features)  # output layer
        
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.out(x)
        return x

torch.manual_seed(32)
model = Model()

# From the discussions here: discuss.pytorch.org/t/how-to-check-if-model-is-on-cuda
next(model.parameters()).is_cuda

gpumodel = model.cuda()

next(gpumodel.parameters()).is_cuda

df = pd.read_csv('../Data/iris.csv')
X = df.drop('target',axis=1).values
y = df['target'].values
X_train, X_test, y_train, y_test = train_test_split(X,y,test_size=0.2,random_state=33)

"""
Explanation: Sending Models to GPU
End of explanation
"""


X_train = torch.FloatTensor(X_train).cuda()
X_test = torch.FloatTensor(X_test).cuda()
y_train = torch.LongTensor(y_train).cuda()
y_test = torch.LongTensor(y_test).cuda()

trainloader = DataLoader(X_train, batch_size=60, shuffle=True)
testloader = DataLoader(X_test, batch_size=60, shuffle=False)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

import time
epochs = 100
losses = []
start = time.time()
for i in range(epochs):
    i+=1
    y_pred = gpumodel.forward(X_train)
    loss = criterion(y_pred, y_train)
    losses.append(loss)
    
    # a neat trick to save screen space:
    if i%10 == 1:
        print(f'epoch: {i:2}  loss: {loss.item():10.8f}')

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
print(f'TOTAL TRAINING TIME: {time.time()-start}')

"""
Explanation: Convert Tensors to .cuda() tensors
End of explanation
"""


