# this mounts your Google Drive to the Colab VM.
from google.colab import drive
drive.mount('/content/drive', force_remount=True)

import os
os.mkdir('/home/data/')
%cd '/home/data/'

!pip install --upgrade fastai 
!pip install aicrowd-cli

API_KEY = '52ab6eb031245b7028158e2f3e993174' #Please enter your API Key from [https://www.aicrowd.com/participants/me]
!aicrowd login --api-key $API_KEY

!aicrowd dataset download --challenge f1-team-classification -j 3

!rm -rf data
!mkdir data

!unzip train.zip  -d data/train
!unzip val.zip -d data/val
!unzip test.zip  -d data/test

!mv train.csv data/train.csv
!mv val.csv data/val.csv
!mv sample_submission.csv data/sample_submission.csv

import torch
import torch.nn as nn
from torch.utils.data import Dataset,DataLoader,RandomSampler
from torchvision import transforms
import torch.optim as optim
import torchvision.models as models

import numpy as np
import pandas as pd
from PIL import Image
import matplotlib.pyplot as plt

import time
import os
import copy

class ImageDataset(Dataset):
  def __init__(self,ImageFold,df,lblDict,transforms):
    self.ImageFold=ImageFold
    self.df=df
    self.lblDict=lblDict
    self.trans=transforms

  def __len__(self):
    return len(self.df)

  def __getitem__(self,ind):
    im=self.load_image(self.df.iloc[ind][0])
    lbl=self.lblDict[self.df.iloc[ind][1]]
    im=self.trans(im)
    return im, lbl


  def load_image(self,ind):
    return Image.open(self.ImageFold+str(self.df.iloc[ind][0])+'.jpg')

data_transform = transforms.Compose([
        transforms.Resize(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])

df_train=pd.read_csv('data/train.csv')
df_val=pd.read_csv('data/val.csv')

ds_train=ImageDataset('data/train/',df_train,{'redbull' : 0 , 'mercedes' : 1},transforms=data_transform)
ds_val=ImageDataset('data/val/',df_val,{'redbull' : 0 , 'mercedes' : 1},transforms=data_transform)

dl_train=DataLoader(ds_train,batch_size=64,shuffle=True,num_workers=2)
dl_val=DataLoader(ds_val,batch_size=64,shuffle=True,num_workers=2)

dataloaders_dict={'train':dl_train, 'val':dl_val}

def train_model(model, dataloaders, criterion, optimizer, num_epochs=25, is_inception=False):
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)
                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):

                  outputs = model(inputs)
                  loss = criterion(outputs, labels)

                  _, preds = torch.max(outputs, 1)

                    # backward + optimize only if in training phase
                  if phase == 'train':
                    loss.backward()
                    optimizer.step()

                # statistics
                running_loss += loss.detach().item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
            epoch_loss = running_loss / len(dataloaders[phase].dataset)
            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())


        print()

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))


    model.load_state_dict(best_model_wts)
    return model

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
num_epochs=5

vgg16 = models.vgg16(pretrained=True)
vgg16.classifier[6]=nn.Linear(4096,2)
vgg16.to(device)
optimizer = optim.SGD(vgg16.parameters(), lr=0.001, momentum=0.9)
criterion = nn.CrossEntropyLoss()

# Train and evaluate
model_ft = train_model(vgg16, dataloaders_dict, criterion, optimizer, num_epochs=num_epochs)

optimizer = optim.SGD(vgg16.parameters(), lr=0.0003, momentum=0.9)
model_ft = train_model(vgg16, dataloaders_dict, criterion, optimizer, num_epochs=num_epochs)

data_test_transform = transforms.Compose([
        transforms.Resize(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])

vgg16.eval()

A=[[i for i in range(10000)],['']*10000]
df=pd.DataFrame(A).transpose()
df.columns=['ImageID','bboxes']
i=0
for f in os.listdir('data/test/'):
  im=Image.open('data/test/'+f)
  tens=torch.reshape(data_test_transform(im),(1,3,224,224))
  inputs = tens.to(device)
  outputs = np.argmax(vgg16(inputs).detach().cpu().numpy())
  if outputs == 1:
    df.iloc[int(f.split('.')[0]),1]='mercedes'
  else:
    df.iloc[int(f.split('.')[0]),1]='redbull'

df.to_csv('/content/drive/MyDrive/Colab Notebooks/aicrowd/f1/1 challange/submission.csv',index=False)