import torch
torch.autograd.set_detect_anomaly(True)
import torch.nn as nn
import pytorch_lightning as pl
import models.ViLT.vilt.modules.vision_transformer as vit
import pdb
from transformers.models.bert.modeling_bert import BertConfig, BertEmbeddings
from models.ViLT.vilt.modules import heads, objectives
# from models.ViLT.vilt.modules import vilt_utils


class ViLTransformerSS(pl.LightningModule):
	def __init__(self, config):
		super().__init__()
		self.save_hyperparameters()

		self.bert_config = BertConfig(
			vocab_size=config["vocab_size"],
			hidden_size=config["hidden_size"],
			num_hidden_layers=config["num_layers"],
			num_attention_heads=config["num_heads"],
			intermediate_size=config["hidden_size"] * config["mlp_ratio"],
			max_position_embeddings=config["max_text_len"],
			hidden_dropout_prob=config["drop_rate"],
			attention_probs_dropout_prob=config["drop_rate"],
		)

		self.text_embeddings = BertEmbeddings(self.bert_config)
		self.text_embeddings.apply(objectives.init_weights)

		self.token_type_embeddings = nn.Embedding(2, config["hidden_size"])
		self.token_type_embeddings.apply(objectives.init_weights)

		if self.hparams.config["checkpoint"] == "":
			self.transformer = getattr(vit, self.hparams.config["vit"])(
				pretrained=True, config=self.hparams.config
			)
		else:
			self.transformer = getattr(vit, self.hparams.config["vit"])(
				pretrained=False, config=self.hparams.config
			)

		self.pooler = heads.Pooler(config["hidden_size"])
		self.pooler.apply(objectives.init_weights)

		if config["loss_names"]["mlm"] > 0:
			self.mlm_score = heads.MLMHead(self.bert_config)
			self.mlm_score.apply(objectives.init_weights)

		if config["loss_names"]["itm"] > 0:
			self.itm_score = heads.ITMHead(config["hidden_size"])
			self.itm_score.apply(objectives.init_weights)

		if config["loss_names"]["mpp"] > 0:
			self.mpp_score = heads.MPPHead(self.bert_config)
			self.mpp_score.apply(objectives.init_weights)

		# ===================== Downstream ===================== #

		hs = self.hparams.config["hidden_size"]

		if self.hparams.config["loss_names"]["vqa"] > 0:
			vs = self.hparams.config["vqav2_label_size"]
			self.vqa_classifier = nn.Sequential(
				nn.Linear(hs, hs * 2),
				nn.LayerNorm(hs * 2),
				nn.GELU(),
				nn.Linear(hs * 2, vs),
			)
			self.vqa_classifier.apply(objectives.init_weights)

		if self.hparams.config["loss_names"]["nlvr2"] > 0:
			self.nlvr2_classifier = nn.Sequential(
				nn.Linear(hs * 2, hs * 2),
				nn.LayerNorm(hs * 2),
				nn.GELU(),
				nn.Linear(hs * 2, 2),
			)
			self.nlvr2_classifier.apply(objectives.init_weights)
			emb_data = self.token_type_embeddings.weight.data
			self.token_type_embeddings = nn.Embedding(3, hs)
			self.token_type_embeddings.apply(objectives.init_weights)
			self.token_type_embeddings.weight.data[0, :] = emb_data[0, :]
			self.token_type_embeddings.weight.data[1, :] = emb_data[1, :]
			self.token_type_embeddings.weight.data[2, :] = emb_data[1, :]

		if self.hparams.config["loss_names"]["irtr"] > 0:
			self.rank_output = nn.Linear(hs, 1)
			self.rank_output.weight.data = self.itm_score.fc.weight.data[1:, :]
			self.rank_output.bias.data = self.itm_score.fc.bias.data[1:]
			self.margin = 0.2
			for p in self.itm_score.parameters():
				p.requires_grad = False

		# vilt_utils.set_metrics(self)
		self.current_tasks = list()

		# ===================== load downstream (test_only) ======================

		if self.hparams.config["checkpoint"] != "":
			ckpt = torch.load(self.hparams.config["checkpoint"], map_location="cpu")
			state_dict = ckpt['state_dict']
			# interpolate position embedding
			org_embeds = state_dict['text_embeddings.position_embeddings.weight']
			org_tokens_len = org_embeds.shape[0]
			new_tokens_len = self.text_embeddings.position_embeddings.weight.shape[0] 
			if org_tokens_len > new_tokens_len:
				state_dict['text_embeddings.position_embeddings.weight'] = org_embeds[:new_tokens_len,:]
			else:
				extra_tokens = new_tokens_len - org_tokens_len
				extra_tokens_embeds = nn.Embedding(extra_tokens, self.hparams.config["hidden_size"])
				extra_tokens_embeds_weight = extra_tokens_embeds.weight
				state_dict['text_embeddings.position_embeddings.weight'] = torch.cat([org_embeds,extra_tokens_embeds_weight],dim=0)
			state_dict['text_embeddings.position_ids'] = self.text_embeddings.position_ids
			msg = self.load_state_dict(state_dict, strict=False)
			print(msg)
	
	def infer(
		self,
		batch,
		mask_text=False,
		mask_image=False,
		image_token_type_idx=1,
		image_embeds=None,
		image_masks=None,
	):
		
		imgkey = "image"

		do_mlm = "_mlm" if mask_text else ""
		text_ids = batch[f"text_ids{do_mlm}"]
		# text_labels = batch[f"text_labels{do_mlm}"]
		text_masks = batch[f"text_masks"]
		text_embeds = self.text_embeddings(text_ids)
		query_ids = batch[f"query_ids"]
		query_embeds = self.text_embeddings(query_ids)
		query_masks = batch[f"query_masks"]
        
		check = torch.isnan(query_embeds)
		# print(-1,torch.isnan(check).any())  
		# print((check == True).nonzero(as_tuple=True))
        
        
		for i, blk in enumerate(self.transformer.blocks):
			query_embeds, _ = blk(query_embeds, mask=query_masks)       
		query_embeds = self.transformer.norm(query_embeds)
        

		if image_embeds is None and image_masks is None:
			img = batch[imgkey]
			(
				image_embeds,
				image_masks,
				patch_index,
				image_labels,
			) = self.transformer.visual_embed(
				img,
				max_image_len=self.hparams.config["max_image_len"],
				mask_it=mask_image,
			)
		else:
			patch_index, image_labels = (
				None,
				None,
			)

		text_embeds, image_embeds = (
			text_embeds + self.token_type_embeddings(torch.zeros_like(text_masks)),
			image_embeds + self.token_type_embeddings(torch.full_like(image_masks, image_token_type_idx)),
		)

		# co_embeds = text_embeds
		# co_masks = text_masks
		co_embeds = torch.cat([text_embeds, image_embeds], dim=1)
		co_masks = torch.cat([text_masks, image_masks], dim=1)

		x = co_embeds

		for i, blk in enumerate(self.transformer.blocks):
			x, _attn = blk(x, mask=co_masks)

		x = self.transformer.norm(x)
		text_feats, image_feats = (
			x[:, : text_embeds.shape[1]],
			x[:, text_embeds.shape[1] :],
		)
		# text_feats, image_feats = x, None
		cls_feats = self.pooler(x)

		ret = {
			"query_embeds": query_embeds,
			"query_atts": query_masks,  
			"query_cls": query_embeds[:,0,:],         
			"text_feats": text_feats,
			"image_feats": image_feats,
			"multi_cls": cls_feats,
			"raw_cls_feats": x[:, 0],
			"image_labels": image_labels,
			"image_atts": image_masks,
			"multi_embeds": x,
            "multi_atts": co_masks,
			# "text_labels": text_labels,
			# "text_ids": text_ids,
			"text_masks": text_masks,
			"patch_index": patch_index,
		}

		return ret
	
	@torch.no_grad()
	def output_multi(self,batch,image_token_type_idx=1):
				
		imgkey = "image"

		text_ids = batch["text_ids"]
		text_masks = batch["text_masks"]
		text_embeds = self.text_embeddings(text_ids)
		
		img = batch[imgkey]
		
		(
			image_embeds,
			image_masks,
			patch_index,
			image_labels,
		) = self.transformer.visual_embed(
			img,
			max_image_len=self.hparams.config["max_image_len"],
			mask_it=False,
		)

		text_embeds, image_embeds = (
			text_embeds + self.token_type_embeddings(torch.zeros_like(text_masks)),
			image_embeds
			+ self.token_type_embeddings(
				torch.full_like(image_masks, image_token_type_idx)
			),
		)
        
		# co_embeds = text_embeds
		# co_masks = text_masks
		co_embeds = torch.cat([text_embeds, image_embeds], dim=1)
		co_masks = torch.cat([text_masks, image_masks], dim=1)

		x = co_embeds

		for i, blk in enumerate(self.transformer.blocks):
			x, _attn = blk(x, mask=co_masks)

		x = self.transformer.norm(x)
		text_feats, image_feats = (
			x[:, : text_embeds.shape[1]],
			x[:, text_embeds.shape[1] :],
		)
		# text_feats, image_feats = x, None
		cls_feats = self.pooler(x)

		ret = {
			"text_feats": text_feats,
			"image_feats": image_feats,
			"multi_cls": cls_feats,
			"raw_cls_feats": x[:, 0],
			"image_labels": image_labels,
			"image_atts": image_masks,
			"multi_embeds": x,
            "multi_atts": co_masks,
			# "text_labels": text_labels,
			# "text_ids": text_ids,
			"text_masks": text_masks,
			"patch_index": patch_index,
		}

		return ret




	def forward(self, batch):
		ret = dict()
		if len(self.current_tasks) == 0:
			ret.update(self.infer(batch))
			return ret

		# Masked Language Modeling
		if "mlm" in self.current_tasks:
			ret.update(objectives.compute_mlm(self, batch))

		# Masked Patch Prediction
		if "mpp" in self.current_tasks:
			ret.update(objectives.compute_mpp(self, batch))

		# Image Text Matching
		if "itm" in self.current_tasks:
			ret.update(objectives.compute_itm_wpa(self, batch))

		# Visual Question Answering
		if "vqa" in self.current_tasks:
			ret.update(objectives.compute_vqa(self, batch))

		# Natural Language for Visual Reasoning 2
		if "nlvr2" in self.current_tasks:
			ret.update(objectives.compute_nlvr2(self, batch))

		# Image Retrieval and Text Retrieval
		if "irtr" in self.current_tasks:
			ret.update(objectives.compute_irtr(self, batch))

		return ret

	def training_step(self, batch, batch_idx):
		# vilt_utils.set_task(self)
		output = self(batch)
		total_loss = sum([v for k, v in output.items() if "loss" in k])

		return total_loss

	def training_epoch_end(self, outs):
		pass
		# vilt_utils.epoch_wrapup(self)

	def validation_step(self, batch, batch_idx):
		# vilt_utils.set_task(self)
		output = self(batch)

	def validation_epoch_end(self, outs):
		pass
		# vilt_utils.epoch_wrapup(self)

	def test_step(self, batch, batch_idx):
		# vilt_utils.set_task(self)
		output = self(batch)
		ret = dict()

		if self.hparams.config["loss_names"]["vqa"] > 0:
			ret.update(objectives.vqa_test_step(self, batch, output))

		return ret

	def test_epoch_end(self, outs):
		model_name = self.hparams.config["checkpoint"].split("/")[-1][:-5]

		if self.hparams.config["loss_names"]["vqa"] > 0:
			objectives.vqa_test_wrapup(outs, model_name)
		# vilt_utils.epoch_wrapup(self)

	def configure_optimizers(self):
		pass
		# return vilt_utils.set_schedule(self)
