#@title Login to AIcrowd
!pip install -U aicrowd-cli > /dev/null
!aicrowd login 2> /dev/null

#@title Magic Box ⬛ { vertical-output: true, display-mode: "form" }
try:
  import os
  if first_run and os.path.exists("/content/data-purchasing-challenge-2022-starter-kit/data/training"):
    first_run = False
except:
  first_run = True

if first_run:
  %cd /content/
  !git clone http://gitlab.aicrowd.com/zew/data-purchasing-challenge-2022-starter-kit.git > /dev/null
  %cd data-purchasing-challenge-2022-starter-kit
  !aicrowd dataset list -c data-purchasing-challenge-2022
  !aicrowd dataset download -c data-purchasing-challenge-2022
  !mkdir -p data/
  !mv *.tar.gz data/ && cd data && echo "Extracting dataset" && ls *.tar.gz | xargs -n1 -I{} bash -c "tar -xvf {} > /dev/null"


def run_pre_training_phase():
  from run import ZEWDPCBaseRun
  run = ZEWDPCBaseRun()
  run.pre_training_phase = pre_training_phase
  run.pre_training_phase(self=run, training_dataset=training_dataset)
  # NOTE:It is critical that the checkpointing works in a self-contained way
  #      As, the evaluators might choose to run the different phases separately.
  run.save_checkpoint("/tmp/pretrainig_phase_checkpoint.pickle")

def run_purchase_phase():
  from run import ZEWDPCBaseRun
  run = ZEWDPCBaseRun()
  run.pre_training_phase = pre_training_phase
  run.purchase_phase = purchase_phase
  run.load_checkpoint("/tmp/pretrainig_phase_checkpoint.pickle")
  # Hacky way to make it work in notebook
  unlabelled_dataset.purchases = set()
  run.purchase_phase(self=run, unlabelled_dataset=unlabelled_dataset, training_dataset=training_dataset, budget=3000)
  run.save_checkpoint("/tmp/purchase_phase_checkpoint.pickle")
  del run

def run_prediction_phase():
  from run import ZEWDPCBaseRun
  run = ZEWDPCBaseRun()
  run.pre_training_phase = pre_training_phase
  run.purchase_phase = purchase_phase
  run.prediction_phase = prediction_phase
  run.load_checkpoint("/tmp/purchase_phase_checkpoint.pickle")
  run.prediction_phase(self=run, test_dataset=val_dataset)
  del run

#@title a) Runtime Packages<br/><small>Important: Add the packages required by your code here. (space separated)</small> { run: "auto", display-mode: "form" }
apt_packages = "build-essential vim" #@param {type:"string"}
pip_packages = "scikit-image pandas timeout-decorator==0.5.0 numpy wandb" #@param {type:"string"}

!apt install -y $apt_packages git-lfs
!pip install $pip_packages

from evaluator.dataset import ZEWDPCBaseDataset, ZEWDPCProtectedDataset
DATASET_SHUFFLE_SEED = 1022022

# Instantiate Training Dataset
training_dataset = ZEWDPCBaseDataset(
    images_dir="./data/training/images",
    labels_path="./data/training/labels.csv",
    shuffle_seed=DATASET_SHUFFLE_SEED,
)
# Instantiate Unlabelled Dataset
unlabelled_dataset = ZEWDPCProtectedDataset(
    images_dir="./data/unlabelled/images",
    labels_path="./data/unlabelled/labels.csv",
    budget=3000,  # Configurable Parameter
    shuffle_seed=DATASET_SHUFFLE_SEED,
)
# Instantiate Validation Dataset
val_dataset = ZEWDPCBaseDataset(
    images_dir="./data/validation/images",
    labels_path="./data/validation/labels.csv",
    drop_labels=True,
    shuffle_seed=DATASET_SHUFFLE_SEED,
)
val_dataset_gt = ZEWDPCBaseDataset(
    images_dir="./data/validation/images",
    labels_path="./data/validation/labels.csv",
    drop_labels=False,
    shuffle_seed=DATASET_SHUFFLE_SEED,
)

import torch
from torch import nn
from torchvision import models
from torch.optim import Adam, SGD, lr_scheduler
from torchvision import transforms
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import numpy as np
import abc
import datetime
from tqdm import tqdm
from sklearn.metrics import accuracy_score
from sklearn.metrics import hamming_loss
import wandb

from evaluator.dataset import ZEWDPCBaseDataset, ZEWDPCProtectedDataset

class ZEWDPCBaseRun:

    def __init__(self):
        self.evaluation_state = {}
        # Model parameters
        self.BATCH_SIZE = 32
        self.NUM_WORKERS = 2
        self.LEARNING_RATE = 0.001
        self.NUM_CLASSES = 4
        self.TOPK= 3
        self.THRESHOLD = 0.5
        self.NUM_EPOCS = 50
        self.EVAL_FREQ = 5

        self.model = models.efficientnet_b1(num_classes = self.NUM_CLASSES)
        ## CHANGE CPU CUDA HERE
        self.model.cuda()
        # self.model.cpu()

        self.trainable_parameters = filter(lambda param: param.requires_grad, self.model.parameters())
        self.optimizer = Adam(self.trainable_parameters, lr=self.LEARNING_RATE)
        
        self.epoch = 0

        self.lr_scheduler_ = lr_scheduler.ReduceLROnPlateau(
            self.optimizer, mode='max', patience=2, verbose=True
        )
        self.criterion = nn.BCEWithLogitsLoss()

        # WandB setup
        model_name = "EfficientNet-B1"
        wandb_config={
            "model": model_name,
            "learning_rate": self.LEARNING_RATE,
            "samples_per_gpu": self.BATCH_SIZE,
            "workers_per_gpu": self.NUM_WORKERS,
            "augmentations": "NO",
            "description": "Pre-Training. No augmentation"
        }
        wandb.init(project="<PROJECT_NAME>", entity="<USERNAME>", name=model_name, config=wandb_config)
        

    def pre_training_phase(
        self, training_dataset: ZEWDPCBaseDataset, register_progress=lambda x: False
    ):
        print("\n================> Pre-Training Phase\n")
        # Creating transformations
        train_transform = transforms.Compose([
            transforms.ToTensor(),
        ])
        training_dataset.set_transform(train_transform)
        train_loader = DataLoader(
            dataset=training_dataset,
            batch_size=self.BATCH_SIZE,
            shuffle=False,
            num_workers=self.NUM_WORKERS,
        )

        def run_epoch():
            for _, batch in enumerate(train_loader):

                ## CHANGE CPU CUDA HERE
                x, y = batch["image"].cuda(), batch["label"]
                # x, y = batch["image"].cpu(), batch["label"]

                pred_y = self.model(x)
                # Change the shape of true labels here. Because for last batch the no. of images can be less
                y = torch.cat(y, dim=0).reshape(
                    self.NUM_CLASSES, pred_y.shape[0]
                    ).T.type(torch.FloatTensor)
                
                ## CHANGE CPU CUDA HERE. Comment for CPU
                y = y.cuda()
                loss = self.criterion(pred_y, y)
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

                # 416 = BATCH_SIZE*13
                if self.global_step % 416 == 0:
                    wandb.log({
                        "Step": (self.epoch*5000)+self.global_step,
                        "Loss": loss,
                    })
                    print("[{}] Training [epoch {}, step {}], loss: {:4f}".format(
                        datetime.datetime.now(), self.epoch, self.global_step, loss))
                self.global_step += self.BATCH_SIZE
        
        epoch_range = tqdm(range(self.epoch, self.NUM_EPOCS))
        for i in epoch_range:
            epoch_range.set_description(f"Epoch: {i}")
            self.global_step = 0
            run_epoch()
            register_progress(i) # Epoch as progress
            if (i+1)%self.EVAL_FREQ == 0:
                predictions = self.prediction_phase(val_dataset)
                self.evaluation(predictions)
            self.epoch += 1
        print("Execution Complete of Training Phase.")


    def purchase_phase(
        self,
        unlabelled_dataset: ZEWDPCProtectedDataset,
        training_dataset: ZEWDPCBaseDataset,
        budget=1000,
        register_progress=lambda x: False,
    ):
        """
        # Purchase Phase
        -------------------------
        In this phase of the competition, you have access to
        the unlabelled_dataset (an instance of `ZEWDPCProtectedDataset`)
        and the training_dataset (an instance of `ZEWDPCBaseDataset`)
        {see datasets.py for more details}, and a purchase budget.

        You can iterate over both the datasets and access the images without restrictions.
        However, you can probe the labels of the unlabelled_dataset only until you
        run out of the label purchasing budget.

        PARTICIPANT_TODO: Add your code here
        """
        print("\n================> Purchase Phase | Budget = {}\n".format(budget))

        register_progress(0.0) #Register Progress
        
        for sample in tqdm(unlabelled_dataset):
            idx = sample["idx"]
            
            # image = unlabelled_dataset.__getitem__(idx)
            # print(image)

            # Budgeting & Purchasing Labels
            if budget > 0:
                label = unlabelled_dataset.purchase_label(idx)

            budget -= 1
        register_progress(1.0) #Register Progress
        print("Execution Complete of Purchase Phase.")

    def prediction_phase(
        self,
        test_dataset: ZEWDPCBaseDataset,
        register_progress=lambda x: False,
    ):
        """
        # Prediction Phase
        -------------------------
        In this phase of the competition, you have access to the test dataset, and you
        are supposed to make predictions using your trained models.

        Returns:
            np.ndarray of shape (n, 4)
                where n is the number of samples in the test set
                and 4 refers to the 4 labels to be predicted for each sample
                for the multi-label classification problem.

        PARTICIPANT_TODO: Add your code here
        """
        print(
            "\n================> Prediction Phase : - on {} images\n".format(
                len(test_dataset)
            )
        )
        
        test_transform = transforms.Compose([
            transforms.ToTensor(),
        ])
        test_dataset.set_transform(test_transform)
        test_loader = DataLoader(
            dataset=test_dataset,
            batch_size=self.BATCH_SIZE,
            shuffle=False,
            num_workers=self.NUM_WORKERS,
        )
    
        def convert_to_label(preds):
            return np.array((torch.sigmoid(preds) > 0.5), dtype=int).tolist()

        predictions = []
        self.model.eval()
        with torch.no_grad():
            for _, batch in enumerate(test_loader):
                ## CHANGE CPU CUDA HERE
                # X= batch['image'].cpu()
                X = batch['image'].cuda()

                pred_y = self.model(X)

                # Convert to labels
                pred_y_labels = []
                for arr in pred_y:
                    ## CHANGE CPU CUDA HERE
                    pred_y_labels.append(convert_to_label(arr.cpu())) # For CUDA
                    # pred_y_labels.append(convert_to_label(arr)) # For CPU

                # Save the results
                predictions.extend(pred_y_labels)
            register_progress(1.0)

        predictions = np.array(predictions) # random predictions
        print("Execution Complete of Purchase Phase.")
        return predictions

    def evaluation(self, predictions):
        from evaluator.evaluation_metrics import accuracy_score, hamming_loss, exact_match_ratio

        y_true = val_dataset_gt._get_all_labels()
        y_pred = predictions

        accuracy_score = accuracy_score(y_true, y_pred)
        hamming_loss_score = hamming_loss(y_true, y_pred)
        exact_match_ratio_score = exact_match_ratio(y_true, y_pred)

        wandb.log({
            "Epoch": self.epoch+1,
            "Accuracy": accuracy_score,
            "Hamming Loss": hamming_loss_score,
            "Match ratio": exact_match_ratio_score
        })

        print("Accuracy Score : ", accuracy_score)
        print("Hamming Loss : ", hamming_loss_score)
        print("Exact Match Ratio : ", exact_match_ratio_score)

    def save_checkpoint(self, checkpoint_path):
        """
        Saves the checkpoint in the checkpoint_path directory. Each checkpoint will be saved for epoch_x
        """
        save_dict = {
            'epoch': self.epoch + 1,
            'model_state_dict': self.model.state_dict(),
            'optim_state_dict': self.optimizer.state_dict(),
        }
        torch.save(save_dict, checkpoint_path)
        print(f"Checkpont epoch:{self.epoch} Model saved at {checkpoint_path}")

    def load_checkpoint(self, checkpoint_path):
        """
        Load the latest checkpoint from the experiment
        """
        ## CHANGE CPU CUDA HERE
        checkpoint_model = torch.load(checkpoint_path, map_location="cuda:0")
        # checkpoint_model = torch.load(checkpoint_path, map_location="cpu")
        
        self.latest_epoch = checkpoint_model['epoch']
        self.model.load_state_dict(checkpoint_model['model_state_dict'])
        self.optimizer.load_state_dict(checkpoint_model['optim_state_dict'])
        print('loading checkpoint success (epoch {})'.format(self.latest_epoch))

import tempfile
checkpoint_path = tempfile.NamedTemporaryFile(delete=False).name

run = ZEWDPCBaseRun()
## Pre - Training process
run.pre_training_phase(training_dataset)
run.save_checkpoint(checkpoint_path)
del run

# ## Purchasing phase
run = ZEWDPCBaseRun()
run.load_checkpoint(checkpoint_path)
run.purchase_phase(unlabelled_dataset, training_dataset, budget=3000)
run.save_checkpoint(checkpoint_path)
del run

## Prediction phase
run = ZEWDPCBaseRun()
run.load_checkpoint(checkpoint_path)
predictions = run.prediction_phase(val_dataset)
assert type(predictions) == np.ndarray
assert predictions.shape == (len(val_dataset), 4)

## Evaluation Phase
from evaluator.evaluation_metrics import accuracy_score, hamming_loss, exact_match_ratio

y_true = val_dataset_gt._get_all_labels()
y_pred = predictions

accuracy_score = accuracy_score(y_true, y_pred)
hamming_loss_score = hamming_loss(y_true, y_pred)
exact_match_ratio_score = exact_match_ratio(y_true, y_pred)

print("Accuracy Score : ", accuracy_score)
print("Hamming Loss : ", hamming_loss_score)
print("Exact Match Ratio : ", exact_match_ratio_score)