# -*- coding: utf-8 -*-
"""cifar100logreg.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1nHegIWH7_eqyO2zWB6chRhJhLQDidmSF
"""

import torch
from torch.autograd import Variable
import torchvision.transforms as transforms
import os
from torchvision.datasets import CIFAR100
import numpy as np
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE

from torch.utils.data import TensorDataset

class LogisticRegression(torch.nn.Module):
    def __init__(self, input_dim, output_dim):
        super(LogisticRegression, self).__init__()
        self.linear = torch.nn.Linear(input_dim, output_dim)

    def forward(self, x):
        outputs = self.linear(x)
        return outputs

train_dataset = CIFAR100(root='./data', train=True, transform=transforms.Compose([
                                      transforms.Resize(32),
                                      transforms.ToTensor(),
                                      transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                                    ]), download=True)
test_dataset = CIFAR100(root='./data', train=False, transform=transforms.Compose([
                                      transforms.Resize(32),
                                      transforms.ToTensor(),
                                      transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                                    ]))

batch_size = 64
n_iters = 1000000
num_epochs = n_iters / (len(train_dataset) / batch_size)
num_epochs = int(num_epochs)

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, 
                                           batch_size=len(train_dataset), 
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset, 
                                          batch_size=len(test_dataset), 
                                          shuffle=False)

# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# model.to(device)
# input_dim = 3*32*32
# output_dim = 100

# model = LogisticRegression(input_dim, output_dim)

for images, labels in train_loader:  
    print('Image batch dimensions:', images.shape)
    print('Image label dimensions:', labels.shape)
    break

X_train = np.array(images.reshape(50000, 3*32*32))
y_labels = np.array(labels)

for images, labels in test_loader:  
    print('Image batch dimensions:', images.shape)
    print('Image label dimensions:', labels.shape)
    break

X_test = np.array(images.reshape(10000,3*32*32))
y_test_labs = np.array(labels)


ncomps = 256


pca = PCA(n_components=ncomps)  
X_train_pca = pca.fit_transform(X_train)

#TSNE
#X_train_pca = TSNE(n_components=3).fit_transform(X_train_pca)

X_train_pca = torch.from_numpy(np.float32(X_train_pca))
X_train_pca = torch.tensor(X_train_pca)
y_train = torch.tensor(y_labels, dtype=torch.long)


X_test_pca = pca.transform(X_test)

#TSNE
#X_test_pca = TSNE(n_components=3).fit_transform(X_test_pca)


X_test_pca = torch.from_numpy(np.float32(X_test_pca))
X_test_pca = torch.tensor(X_test_pca)
y_test = torch.tensor(y_test_labs, dtype=torch.long)



train_ds = TensorDataset(X_train_pca, y_train)
train_dl = torch.utils.data.DataLoader(train_ds, batch_size=batch_size)

test_ds = TensorDataset(X_test_pca, y_test)
test_dl = torch.utils.data.DataLoader(test_ds, batch_size=batch_size)

for images, labels in train_dl:
  print('Image batch dimensions:', images.shape)
  print('Image label dimensions:', labels.shape)
  break


for images, labels in test_dl:
  print('Image batch dimensions:', images.shape)
  print('Image label dimensions:', labels.shape)
  break

input_dim = ncomps
output_dim = 100

model = LogisticRegression(input_dim, output_dim)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

criterion = torch.nn.CrossEntropyLoss()
learning_rate = 0.001
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

train_dataloader = train_dl
test_dataloader = test_dl

iter = 0
for epoch in range(num_epochs):
  print("Epoch: ", epoch)
  for i, (images, labels) in enumerate(train_dataloader):

      #######################
      #  USE GPU FOR MODEL  #
      #######################
      images = images.view(-1, ncomps).requires_grad_().to(device)
      labels = labels.to(device)

      # Clear gradients w.r.t. parameters
      optimizer.zero_grad()

      # Forward pass to get output/logits
      outputs = model(images)

      # Calculate Loss: softmax --> cross entropy loss
      loss = criterion(outputs, labels)

      # Getting gradients w.r.t. parameters
      loss.backward()

      # Updating parameters
      optimizer.step()

      iter += 1

      if iter % 500 == 0:
          # Calculate Accuracy         
          correct = 0
          total = 0
          #Iterate through train dataset
          for images, labels in train_dataloader:
              #######################
              #  USE GPU FOR MODEL  #
              #######################
              images = images.view(-1, ncomps).to(device)

              # Forward pass only to get logits/output
              outputs = model(images)

              # Get predictions from the maximum value
              _, predicted = torch.max(outputs.data, 1)

              # Total number of labels
              total += labels.size(0)

              #######################
              #  USE GPU FOR MODEL  #
              #######################
              # Total correct predictions
              if torch.cuda.is_available():
                  correct += (predicted.cpu() == labels.cpu()).sum()
              else:
                  correct += (predicted == labels).sum()

          accuracy = 100 * correct.item() / total

          # Print Loss
          print('Training: Iteration: {}. Loss: {}. Accuracy: {}'.format(iter, loss.item(), accuracy))
          correct = 0
          total = 0
          # Iterate through test dataset
          for images, labels in test_dataloader:
              #######################
              #  USE GPU FOR MODEL  #
              #######################
              images = images.view(-1, ncomps).to(device)

              # Forward pass only to get logits/output
              outputs = model(images)

              # Get predictions from the maximum value
              _, predicted = torch.max(outputs.data, 1)

              # Total number of labels
              total += labels.size(0)

              #######################
              #  USE GPU FOR MODEL  #
              #######################
              # Total correct predictions
              if torch.cuda.is_available():
                  correct += (predicted.cpu() == labels.cpu()).sum()
              else:
                  correct += (predicted == labels).sum()

          accuracy = 100 * correct.item() / total

          # Print Loss
          print('Testing: Iteration: {}. Loss: {}. Accuracy: {}'.format(iter, loss.item(), accuracy))







