"""
This example is largely adapted from https://github.com/pytorch/examples/blob/master/imagenet/main.py
"""
import os
from argparse import ArgumentParser, Namespace
import pickle
from collections import OrderedDict

import numpy as np
from sklearn.metrics import roc_auc_score, accuracy_score
from sklearn.preprocessing import LabelBinarizer


import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision

from pytorch_lightning.core import LightningModule
from pytorch_lightning.trainer import Trainer

from networks import define_net
from datasets import define_dataloader

import pdb
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

class ImageClassifier(LightningModule):

	def __init__(self, opt: Namespace):
		super(ImageClassifier, self).__init__()
		self.opt = opt
		self.net = define_net(opt)
		self.loss_fun = nn.CrossEntropyLoss()
		self.log_pkl = {'step_train_loss':[], 'epoch_val_loss':[], 'epoch_val_score':[], 'epoch_val_label':[]}

	def forward(self, x):
		return self.net.forward(x)

	def training_step(self, batch, batch_idx):
		data, target = batch

		output = self.net.forward(data)
		loss_train = self.loss_fun(output, target).mean()

		self.log('train loss', loss_train)
		#self.log_pkl['step_train_loss'].append(loss_train.detach().cpu().numpy())
		return loss_train

	def validation_step(self, batch, batch_idx):
		data, target = batch

		output = self.net.forward(data)
		loss_val = self.loss_fun(output, target)

		output = OrderedDict({
			'y_pred': output,
			'y_label': target,
			'val_loss': loss_val,
		})

		return output

	def validation_epoch_end(self, outputs):
		tqdm_dict = {}
		y_score_all = torch.vstack([output['y_pred'] for output in outputs]).detach().cpu().numpy()
		y_label_all = torch.hstack([output['y_label'] for output in outputs]).detach().cpu().numpy()
		y_label_all = LabelBinarizer().fit([0,1,2,3,4,5,6,7,8,9]).transform(y_label_all)
		avg_val_loss = torch.stack([output['val_loss'] for output in outputs]).mean()

		for i in range(self.opt.num_class):
			# Per-class AUC 
			y_score_i = np.expand_dims(y_score_all[:,i], axis=1)
			y_label_i = np.expand_dims(y_label_all[:,i], axis=1)
			auc = roc_auc_score(y_label_i[:,0], y_score_i[:,0], average='micro')
			self.log('val auc [Class %d]' % i, auc)

			# Per-class Accuracy
			y_pred_i = np.array(y_score_all.argmax(axis=1) == i, dtype=np.int)
			acc = accuracy_score(np.squeeze(y_label_i), y_pred_i)
			self.log('val acc [Class %d]' % i, acc)

		self.log('val loss', avg_val_loss)

		### Manual Logging
		self.log_pkl['epoch_val_loss'].append(avg_val_loss)
		self.log_pkl['epoch_val_score'].append(y_score_all)
		self.log_pkl['epoch_val_label'].append(y_label_all)
		with open(os.path.join(self.opt.log_path, 'results.pkl'), 'wb') as handle:
			pickle.dump(self.log_pkl, handle, protocol=pickle.HIGHEST_PROTOCOL)



	def configure_optimizers(self):
		optimizer = torch.optim.SGD(self.net.parameters(), lr=self.opt.lr, weight_decay=self.opt.wd, momentum=self.opt.momentum)
		return [optimizer]

	def train_dataloader(self):
		train_dataloader = define_dataloader(split='train', opt=self.opt)
		return train_dataloader

	def val_dataloader(self):
		val_dataloader = define_dataloader(split='val', opt=self.opt)
		return val_dataloader

	@classmethod
	def __accuracy(cls, output, target, topk=(1,)):
		"""Computes the accuracy over the k top predictions for the specified values of k"""
		with torch.no_grad():
			maxk = max(topk)
			batch_size = target.size(0)

			_, pred = output.topk(maxk, 1, True, True)
			pred = pred.t()
			correct = pred.eq(target.view(1, -1).expand_as(pred))

			res = []
			for k in topk:
				correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
				res.append(correct_k.mul_(100.0 / batch_size))

			if len(res) == 1: return res[0]
			return res
