import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from resnet_custom import *
import pdb
import math
from pixelcnn import MaskCNN

device=torch.device("cuda" if torch.cuda.is_available() else "cpu")



def initialize_weights(module):
	"""
	args:
		module: any pytorch module with trainable parameters
	"""

	for m in module.modules():
		if isinstance(m, nn.Conv2d):
			nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
			if m.bias is not None:
				m.bias.data.zero_()

		# if isinstance(m, nn.Linear):
		# 	nn.init.xavier_normal_(m.weight)
		# 	m.bias.data.zero_()

		if isinstance(m, nn.Linear):
			nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
			m.bias.data.zero_()
		
		elif isinstance(m, nn.BatchNorm2d):
			nn.init.constant_(m.weight, 1)
			nn.init.constant_(m.bias, 0)


class CPC_model(nn.Module):
	def __init__(self, input_size = 1024, hidden_size = 128, k = 3, ln = False):
		"""
		args:
			input_size: input size to autoregresser (encoding size)
			hidden_size: number of hidden units in MaskedCNN
			num_layers: number of hidden layers in MaskedCNN
			k: prediction length
		"""
		super(CPC_model, self).__init__()
		
		### Settings
		self.seq_len = 49 # 7 x 7 grid of overlapping 64 x 64 patches extracted from each 256 x 256 image
		self.k = k 
		self.input_size = input_size
		self.hidden_size=hidden_size


		### Networks
		if ln:
			self.encoder = resnet50_ln(pretrained=False)
		else:
			self.encoder = resnet50(pretrained=False)
		self.reg = MaskCNN(n_channel=self.input_size, h=self.hidden_size)
		network_pred = [nn.Linear(input_size, input_size) for i in range(self.k)] #use an indepdent linear layer to predict each future row	
		self.network_pred= nn.ModuleList(network_pred)
		
		# initialize linear network and context network
		initialize_weights(self.network_pred)
		initialize_weights(self.reg)


		### Activation functions
		self.softmax  = nn.Softmax(dim=1)
		self.lsoftmax = nn.LogSoftmax(dim=1)

	def forward(self, x):
		# input = [bs * 7 * 7, 3, 64, 64]
		
		# compute batch_size
		bs = x.size(0) // (self.seq_len)

		rows = int(math.sqrt(self.seq_len))
		cols = int(math.sqrt(self.seq_len))

		# compute latent representation for each patch
		z = self.encoder(x)
		# z.shape: [bs * 7 * 7, 1024]
		
		# reshape z into feature grid: [bs, 7, 7, 1024]
		z = z.contiguous().view(bs, rows, cols, self.input_size)
		
		device = z.device

		#randomly draw a row to predict what is k rows below it, using information in current row and above
		if self.training:
			pred_id = torch.randint(rows - self.k, size=(1,)).long() #low is 0, high is 3 (predicts row 4, 5, 6)
		
		else:
			pred_id = torch.tensor([3]).long()
		
		# feature predictions for the next k rows  e.g.  pred[i] is [bs * cols, 1024] for i in k
		pred = [torch.empty(bs * cols, self.input_size).float().to(device) for i in range(self.k)]

		# ground truth encodings for the next k rows e.g. encode_samples[i] is [bs * cols, 1024] for i in k
		encode_samples = [torch.empty(bs * cols, self.input_size).float().to(device) for i in range(self.k)]

		for i in np.arange(self.k):
			# add ground truth encodings
			start_row = pred_id.item()+i+1
			encode_samples[i] = z[:,start_row, :, :].contiguous().view(bs * cols, self.input_size)

		# reshape feature grid to channel first (required by Pytorch convolution convention)
		z = z.permute(0, 3, 1, 2) 
		# z.shape: from [bs, 7, 7, 1024] --> [bs, 1024, 7, 7]
		
		# apply aggregation to compute context 
		output = self.reg(z)
		# reg is fully convolutional --> output size is [bs, 1024, 7, 7]

		output = output.permute(0, 2, 3, 1) # reshape back to feature grid
		# output.shape: [bs, row, col, 1024]
		
		# context for each patch in the row
		c_t = output[:,pred_id + 1,:, :]
		# c_t.shape: [bs, 1, 7, 1024]

		# reshape for linear classification:
		c_t = c_t.contiguous().view(bs * cols, self.input_size)
		# c_t.shape: [bs * cols, 1024]

		# linear prediction: Wk*c_t
		for i in np.arange(0, self.k):
			if type(self.network_pred) == nn.DataParallel:
				pred[i] = self.network_pred.module[i](c_t)
			
			else:
				pred[i] = self.network_pred[i](c_t)  #e.g. size [bs * cols, 1024]

		nce = 0 # average over prediction length, cols, and batch 
		accuracy = np.zeros((self.k,))
		
		for i in np.arange(0, self.k):
			"""
			goal: can network correctly match predicted features with ground truth features among negative targets 
			i.e. match z_i+k,j with W_k * c_i,j
			postivie target: patch with the correct groundtruth encoding
			negative targets: patches with wrong groundtruth encodings (sampled from other patches in the same image, or other images in the minibatch)

			1) dot product for each k to obtain raw prediction logits 
			total = (a_ij) = [bs * col, bs * col], where a_ij is the logit of ith patch prediction matching jth patch encoding
			
			2) apply softmax along each row to get probability that ith patch prediction matches jth patch encoding 
			we want ith patch prediction to correctly match ith patch encoding, therefore target has 1s along diagnol, and 0s off diagnol

			3) we take the argmax along softmaxed rows to get the patch prediction for the ith patch, this value should be i

			4) compute nce loss as the cross-entropy of classifying the positive sample correctly (sum of logsoftmax along diagnol)

			5) normalize loss by batchsize and k and number of patches in a row
			
			"""
			total = torch.mm(pred[i], torch.transpose(encode_samples[i],0,1)) # e.g. size [bs * col, bs * col]

			accuracy[i] = torch.sum(torch.eq(torch.argmax(self.softmax(total), dim=1), torch.arange(0, bs * cols).to(device))).item() 
			accuracy[i] /= 1. * (bs * cols) 
			
			nce += torch.sum(torch.diag(self.lsoftmax(total))) # nce is a tensor
		
		nce /= -1. * bs * cols * self.k
		# accuracy = 1.*correct.item() / (bs * cols * self.k)
		
		return nce, np.array(accuracy)


# crop data into 64 by 64 with 32 overlap 
def cropdata(data, num_channels=3, kernel_size = 64, stride = 32):
	if len(data.shape) == 3:
		data = data.unsqueeze(0)

	data = data.unfold(2, kernel_size, stride).unfold(3, kernel_size, stride)
	data = data.permute(0,2,3,1,4,5)
	data = data.contiguous().view(-1, num_channels, kernel_size, kernel_size)
	return data

if __name__ == '__main__':
	torch.set_printoptions(threshold=1e6)
	x = torch.rand(2, 3, 256, 256)
	x = cropdata(x)
	print(x.shape)
	model = CPC_model(1024, 256)
	nce, accuracy = model(x)