# -----------------------------------------------------------
# Stacked Cross Attention Network implementation based on
# https://arxiv.org/abs/1803.08024.
# "Stacked Cross Attention for Image-Text Matching"
# Kuang-Huei Lee, Xi Chen, Gang Hua, Houdong Hu, Xiaodong He
#
# Writen by Kuang-Huei Lee, 2018
# ---------------------------------------------------------------
"""SCAN model"""

import torch
import torch.nn as nn
import torch.nn.init
import torchvision.models as models
from torch.autograd import Variable
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
from torch.nn.utils.weight_norm import weight_norm
import torch.backends.cudnn as cudnn
from torch.nn.utils.clip_grad import clip_grad_norm_ as clip_grad_norm
import numpy as np
from collections import OrderedDict
from TextModels import BagOfWords
from OT_torch_ import cost_matrix_batch_torch, IPOT_distance_torch_batch_uniform, GW_distance_uniform
import sys


def l1norm(X, dim, eps=1e-8):
	"""L1-normalize columns of X
	"""
	norm = torch.abs(X).sum(dim=dim, keepdim=True) + eps
	X = torch.div(X, norm)
	return X


def l2norm(X, dim, eps=1e-8):
	"""L2-normalize columns of X
	"""
	norm = torch.pow(X, 2).sum(dim=dim, keepdim=True).sqrt() + eps
	X = torch.div(X, norm)
	return X


def EncoderImage(data_name, img_dim, embed_size, precomp_enc_type='basic',
                 no_imgnorm=False):
	"""A wrapper to image encoders. Chooses between an different encoders
	that uses precomputed image features.
	"""
	if precomp_enc_type == 'basic':
		img_enc = EncoderImagePrecomp(
			img_dim, embed_size, no_imgnorm)
	elif precomp_enc_type == 'weight_norm':
		img_enc = EncoderImageWeightNormPrecomp(
			img_dim, embed_size, no_imgnorm)
	else:
		raise ValueError("Unknown precomp_enc_type: {}".format(precomp_enc_type))

	return img_enc


def Image_Embedding(img_dim, embed_size, num_layers, precomp_enc_type='basic',
                    no_imgnorm=False):
	"""A wrapper to image encoders. Encode image features into probability vectors
	"""
	if precomp_enc_type == 'basic':
		img_enc = EncoderImagePrecomp_emb(
			img_dim, embed_size, no_imgnorm)
	elif precomp_enc_type == 'weight_norm':
		img_enc = EncoderImageWeightNormPrecomp_emb(
			img_dim, embed_size, no_imgnorm)
	else:
		raise ValueError("Unknown precomp_enc_type: {}".format(precomp_enc_type))

	return img_enc


def InitImgEmbMap(embed_size, num_img_features):
	return torch.Variable(torch.randn(embed_size, num_img_features))


class EncoderImagePrecomp(nn.Module):

	def __init__(self, img_dim, embed_size, no_imgnorm=False):
		super(EncoderImagePrecomp, self).__init__()
		self.embed_size = embed_size
		self.no_imgnorm = no_imgnorm
		self.fc = nn.Linear(img_dim, embed_size)

		self.init_weights()

	def init_weights(self):
		"""Xavier initialization for the fully connected layer
		"""
		r = np.sqrt(6.) / np.sqrt(self.fc.in_features +
		                          self.fc.out_features)
		self.fc.weight.data.uniform_(-r, r)
		self.fc.bias.data.fill_(0)

	def forward(self, images):
		"""Extract image feature vectors."""
		# assuming that the precomputed features are already l2-normalized

		features = self.fc(images)

		# normalize in the joint embedding space
		if not self.no_imgnorm:
			features = l2norm(features, dim=-1)

		return features

	def load_state_dict(self, state_dict):
		"""Copies parameters. overwritting the default one to
		accept state_dict from Full model
		"""
		own_state = self.state_dict()
		new_state = OrderedDict()
		for name, param in state_dict.items():
			if name in own_state:
				new_state[name] = param

		super(EncoderImagePrecomp, self).load_state_dict(new_state)


class EncoderImagePrecomp_emb(nn.Module):

	def __init__(self, img_dim, embed_size, no_imgnorm=False):
		super(EncoderImagePrecomp_emb, self).__init__()
		self.embed_size = embed_size
		self.no_imgnorm = no_imgnorm
		self.fc = nn.Linear(img_dim, embed_size)

		self.init_weights()
		self.softmax = nn.Softmax(dim=-1)

	def init_weights(self):
		"""Xavier initialization for the fully connected layer
		"""
		r = np.sqrt(6.) / np.sqrt(self.fc.in_features +
		                          self.fc.out_features)
		self.fc.weight.data.uniform_(-r, r)
		self.fc.bias.data.fill_(0)

	def forward(self, images):
		"""Extract image feature vectors."""
		# assuming that the precomputed features are already l2-normalized

		features = self.fc(images)

		# normalize in the joint embeddi
		# ng space
		if not self.no_imgnorm:
			features = l2norm(features, dim=-1)
		features = self.softmax(features)
		return features

	def load_state_dict(self, state_dict):
		"""Copies parameters. overwritting the default one to
		accept state_dict from Full model
		"""
		own_state = self.state_dict()
		new_state = OrderedDict()
		for name, param in state_dict.items():
			if name in own_state:
				new_state[name] = param

		super(EncoderImagePrecomp, self).load_state_dict(new_state)


class EncoderImageWeightNormPrecomp(nn.Module):

	def __init__(self, img_dim, embed_size, no_imgnorm=False):
		super(EncoderImageWeightNormPrecomp, self).__init__()
		self.embed_size = embed_size
		self.no_imgnorm = no_imgnorm
		self.fc = weight_norm(nn.Linear(img_dim, embed_size), dim=None)

	def forward(self, images):
		"""Extract image feature vectors."""
		# assuming that the precomputed features are already l2-normalized

		features = self.fc(images)

		# normalize in the joint embedding space
		if not self.no_imgnorm:
			features = l2norm(features, dim=-1)

		return features

	def load_state_dict(self, state_dict):
		"""Copies parameters. overwritting the default one to
		accept state_dict from Full model
		"""
		own_state = self.state_dict()
		new_state = OrderedDict()
		for name, param in state_dict.items():
			if name in own_state:
				new_state[name] = param

		super(EncoderImageWeightNormPrecomp, self).load_state_dict(new_state)


class EncoderImageWeightNormPrecomp_emb(nn.Module):

	def __init__(self, img_dim, embed_size, no_imgnorm=False):
		super(EncoderImageWeightNormPrecomp_emb, self).__init__()
		self.embed_size = embed_size
		self.no_imgnorm = no_imgnorm
		self.fc = weight_norm(nn.Linear(img_dim, embed_size), dim=None)
		self.softmax = nn.Softmax(dim=-1)

	def forward(self, images):
		"""Extract image feature vectors."""
		# assuming that the precomputed features are already l2-normalized

		features = self.fc(images)

		# normalize in the joint embedding space
		if not self.no_imgnorm:
			features = l2norm(features, dim=-1)
		features = self.softmax(features)
		return features

	def load_state_dict(self, state_dict):
		"""Copies parameters. overwritting the default one to
		accept state_dict from Full model
		"""
		own_state = self.state_dict()
		new_state = OrderedDict()
		for name, param in state_dict.items():
			if name in own_state:
				new_state[name] = param

		super(EncoderImageWeightNormPrecomp, self).load_state_dict(new_state)


# RNN Based Language Model
class EncoderText(nn.Module):

	def __init__(self, vocab_size, word_dim, embed_size, num_layers,
	             use_bi_gru=False, no_txtnorm=False):
		super(EncoderText, self).__init__()
		self.embed_size = embed_size
		self.no_txtnorm = no_txtnorm

		# word embedding
		self.embed = nn.Embedding(vocab_size, word_dim)

		# caption embedding
		self.use_bi_gru = use_bi_gru
		# if self.use_bi_gru:
		#     embed_size = int(embed_size/2)
		self.rnn = nn.GRU(word_dim, embed_size, num_layers, batch_first=True, bidirectional=use_bi_gru)

		self.init_weights()

	def init_weights(self):
		self.embed.weight.data.uniform_(-0.1, 0.1)

	def forward(self, x, lengths):
		"""Handles variable size captions
		"""
		# Embed word ids to vectors
		x = self.embed(x)
		packed = pack_padded_sequence(x, lengths, batch_first=True)

		# Forward propagate RNN
		out, _ = self.rnn(packed)

		# Reshape *final* output to (batch_size, hidden_size)
		padded = pad_packed_sequence(out, batch_first=True)
		cap_emb, cap_len = padded

		if self.use_bi_gru:
			# print(cap_emb.size(2))
			cap_emb = (cap_emb[:, :, :int(cap_emb.size(2) / 2)] + cap_emb[:, :, int(cap_emb.size(2) / 2):]) / 2

		# normalization in the joint embedding space
		if not self.no_txtnorm:
			cap_emb = l2norm(cap_emb, dim=-1)

		return cap_emb, cap_len, x


# Word Embedding Based Language Model
class WordEmbedding(nn.Module):

	def __init__(self, vocab_size, word_dim, embed_size, num_layers,
	             use_bi_gru=False, no_txtnorm=False):
		super(WordEmbedding, self).__init__()
		self.embed_size = embed_size
		self.no_txtnorm = no_txtnorm

		# word embedding
		self.embed = nn.Embedding(vocab_size, word_dim)

		# # caption embedding
		# self.use_bi_gru = use_bi_gru
		# if self.use_bi_gru:
		#     embed_size = int(embed_size/2)
		# self.rnn = nn.GRU(word_dim, embed_size, num_layers, batch_first=True, bidirectional=use_bi_gru)

		self.init_weights()

	def init_weights(self):
		self.embed.weight.data.uniform_(-0.1, 0.1)

	def forward(self, x, lengths):
		"""Handles variable size captions
		"""
		# Embed word ids to vectors
		x = self.embed(x)
		# packed = pack_padded_sequence(x, lengths, batch_first=True)
		#
		# # Forward propagate RNN
		# out, _ = self.rnn(packed)
		#
		# # Reshape *final* output to (batch_size, hidden_size)
		# padded = pad_packed_sequence(out, batch_first=True)
		# cap_emb, cap_len = padded
		#
		# # if self.use_bi_gru:
		# #     print(cap_emb.size(2))
		# #     cap_emb = (cap_emb[:,:,:cap_emb.size(2)/2] + cap_emb[:,:,cap_emb.size(2)/2:])/2
		#
		# # normalization in the joint embedding space
		# if not self.no_txtnorm:
		#     cap_emb = l2norm(cap_emb, dim=-1)

		cap_emb = x
		cap_len = lengths

		return cap_emb, cap_len, x


def func_attention(query, context, opt, smooth, eps=1e-8):
	"""
	query: (n_context, queryL, d)
	context: (n_context, sourceL, d)
	"""
	batch_size_q, queryL = query.size(0), query.size(1)
	batch_size, sourceL = context.size(0), context.size(1)

	# Get attention
	# --> (batch, d, queryL)
	queryT = torch.transpose(query, 1, 2)

	# (batch, sourceL, d)(batch, d, queryL)
	# --> (batch, sourceL, queryL)
	attn = torch.bmm(context, queryT)
	if opt.raw_feature_norm == "softmax":
		# --> (batch*sourceL, queryL)
		attn = attn.view(batch_size * sourceL, queryL)
		attn = nn.Softmax(dim=1)(attn)
		# --> (batch, sourceL, queryL)
		attn = attn.view(batch_size, sourceL, queryL)
	elif opt.raw_feature_norm == "l2norm":
		attn = l2norm(attn, 2)
	elif opt.raw_feature_norm == "clipped_l2norm":
		attn = nn.LeakyReLU(0.1)(attn)
		attn = l2norm(attn, 2)
	elif opt.raw_feature_norm == "l1norm":
		attn = l1norm_d(attn, 2)
	elif opt.raw_feature_norm == "clipped_l1norm":
		attn = nn.LeakyReLU(0.1)(attn)
		attn = l1norm_d(attn, 2)
	elif opt.raw_feature_norm == "clipped":
		attn = nn.LeakyReLU(0.1)(attn)
	elif opt.raw_feature_norm == "no_norm":
		pass
	else:
		raise ValueError("unknown first norm type:", opt.raw_feature_norm)
	# --> (batch, queryL, sourceL)
	attn = torch.transpose(attn, 1, 2).contiguous()
	# --> (batch*queryL, sourceL)
	attn = attn.view(batch_size * queryL, sourceL)
	attn = nn.Softmax(dim=1)(attn * smooth)
	# --> (batch, queryL, sourceL)
	attn = attn.view(batch_size, queryL, sourceL)
	# --> (batch, sourceL, queryL)
	attnT = torch.transpose(attn, 1, 2).contiguous()

	# --> (batch, d, sourceL)
	contextT = torch.transpose(context, 1, 2)
	# (batch x d x sourceL)(batch x sourceL x queryL)
	# --> (batch, d, queryL)
	weightedContext = torch.bmm(contextT, attnT)
	# --> (batch, queryL, d)
	weightedContext = torch.transpose(weightedContext, 1, 2)

	return weightedContext, attnT


def cosine_similarity(x1, x2, dim=1, eps=1e-8):
	"""Returns cosine similarity between x1 and x2, computed along dim."""
	w12 = torch.sum(x1 * x2, dim)
	w1 = torch.norm(x1, 2, dim)
	w2 = torch.norm(x2, 2, dim)
	return (w12 / (w1 * w2).clamp(min=eps)).squeeze()



def computeMatchmap_batch(I, A):
	assert (I.dim() == 3)
	assert (A.dim() == 3)
	assert (I.size(2) == A.size(2))
	bs = I.size(0)
	D = I.size(2)
	# Ir = I.view(bs, -1, D)
	# Ar= A.view(bs, D, -1)
	matchmap = torch.bmm(I, A.transpose(2, 1))
	return matchmap


def matchmapSim_batch(M, simtype):
	assert (M.dim() == 3)
	if simtype == 'SISA':
		return M.mean(dim=2).mean(dim=1, keepdim=True)
	elif simtype == 'MISA':
		M_maxH, _ = M.max(1)
		# M_maxHW, _ = M_maxH.max(1)
		return M_maxH.mean(dim=1, keepdim=True)
	elif simtype == 'SIMA':
		M_maxT, _ = M.max(2)
		return M_maxT.mean(dim=1, keepdim=True)  # .mean(dim=1)
	else:
		raise ValueError


def score(images, captions, cap_lens, x, opt):
	"""
	Images: (n_image, n_regions, d) matrix of images
	Captions: (n_caption, max_n_word, d) matrix of captions
	CapLens: (n_caption) array of caption lengths
	"""
	similarities = []
	similarities_OT = []
	n_image = images.size(0)
	n_caption = captions.size(0)
	for i in range(n_caption):
		# Get the i-th text description
		if cap_lens is not None:
			n_word = cap_lens[i]
			cap_i = captions[i, :n_word, :].unsqueeze(0).contiguous()
		else:
			cap_i = captions[i].unsqueeze(0).contiguous()
		# x_i = x[i, :n_word, :].unsqueeze(0).contiguous()
		# x_i_expand = x_i.repeat(n_image, 1, 1)
		# --> (n_image, n_word, d)
		cap_i_expand = cap_i.repeat(n_image, 1, 1)
		"""
			word(query): (n_image, n_word, d)
			image(context): (n_image, n_regions, d)
			weiContext: (n_image, n_word, d)
			attn: (n_image, n_region, n_word)
		"""

		row_sim = matchmapSim_batch(computeMatchmap_batch(images, cap_i_expand), opt.sim_type_dot)

		#
		# weiContext, attn = func_attention(cap_i_expand, images, opt, smooth=opt.lambda_softmax)
		# cap_i_expand = cap_i_expand.contiguous()
		# weiContext = weiContext.contiguous()
		# x_i_expand = x_i_expand.contiguous()
		# # (n_image, n_word)
		#
		#
		# row_sim = cosine_similarity(cap_i_expand, weiContext, dim=2)
		# if opt.agg_func == 'LogSumExp':
		#     row_sim.mul_(opt.lambda_lse).exp_()
		#     row_sim = row_sim.sum(dim=1, keepdim=True)
		#     row_sim = torch.log(row_sim) / opt.lambda_lse
		# elif opt.agg_func == 'Max':
		#     row_sim = row_sim.max(dim=1, keepdim=True)[0]
		# elif opt.agg_func == 'Sum':
		#     row_sim = row_sim.sum(dim=1, keepdim=True)
		# elif opt.agg_func == 'Mean':
		#     row_sim = row_sim.mean(dim=1, keepdim=True)
		# else:
		#     raise ValueError("unknown aggfunc: {}".format(opt.agg_func))

		cos_distance = cost_matrix_batch_torch(cap_i_expand.transpose(2, 1), images.transpose(2, 1))
		# if opt.raw_feature_norm == "clipped_l2norm":
		# 	cos_distance = 1-nn.LeakyReLU(0.1)(cos_distance)
		# elif opt.raw_feature_norm == "abs":
		# 	cos_distance = 1-torch.abs(cos_distance)
		if cap_lens is not None:
			OT_sim = IPOT_distance_torch_batch_uniform(cos_distance, n_image, images.size(1), n_word, opt.OT_iteration)
		else:
			OT_sim = IPOT_distance_torch_batch_uniform(cos_distance, n_image, images.size(1), images.size(1), opt.OT_iteration)

		similarities.append(row_sim)
		similarities_OT.append(OT_sim)

	# (n_image, n_caption)
	similarities = torch.cat(similarities, 1)
	similarities_OT = torch.cat(similarities_OT, 1)

	return similarities, similarities_OT


class ImgEncoder(nn.Module):

	def __init__(self, img_dim, embed_size, no_imgnorm=False):
		super(ImgEncoder, self).__init__()
		self.embed_size = embed_size
		self.no_imgnorm = no_imgnorm
		self.fc = nn.Linear(img_dim, embed_size)
		self.init_weights()

	def init_weights(self):
		"""Xavier initialization for the fully connected layer
		"""
		r = np.sqrt(6.) / np.sqrt(self.fc.in_features +
		                          self.fc.out_features)
		self.fc.weight.data.uniform_(-r, r)
		self.fc.bias.data.fill_(0)

	def forward(self, images):
		"""Extract image feature vectors."""
		# assuming that the precomputed features are already l2-normalized

		features = self.fc(images)

		# normalize in the joint embedding space
		if not self.no_imgnorm:
			features = l2norm(features, dim=-1)

		return features

	def load_state_dict(self, state_dict):
		"""Copies parameters. overwritting the default one to
		accept state_dict from Full model
		"""
		own_state = self.state_dict()
		new_state = OrderedDict()
		for name, param in state_dict.items():
			if name in own_state:
				new_state[name] = param

		super(ImgEncoder, self).load_state_dict(new_state)





class ContrastiveLoss(nn.Module):
	"""
	Compute contrastive loss
	"""

	def __init__(self, opt, margin=0, max_violation=False):
		super(ContrastiveLoss, self).__init__()
		self.opt = opt
		self.margin = margin
		self.max_violation = max_violation

	def forward(self, im, s, s_l, x):
		# x is the word embedding of sentence
		# compute image-sentence score matrix
		# if self.opt.cross_attn == 't2i':
		#     scores, scores_OT = xattn_score_t2i_OT(im, s, s_l, x, self.opt)
		# elif self.opt.cross_attn == 'i2t':
		#     scores, scores_OT = xattn_score_i2t_OT(im, s, s_l, x, self.opt)
		# else:
		#     raise ValueError("unknown first norm type:", self.opt.raw_feature_norm)

		scores, scores_OT = score(im, s, s_l, x, self.opt)
		scores = scores + self.opt.alpha * scores_OT
		diagonal = scores.diag().view(im.size(0), 1)
		d1 = diagonal.expand_as(scores)
		d2 = diagonal.t().expand_as(scores)

		diagonal_OT = scores_OT.diag().view(im.size(0), 1)
		d1_OT = diagonal_OT.expand_as(scores_OT)
		d2_OT = diagonal_OT.t().expand_as(scores_OT)

		# compare every diagonal score to scores in its column
		# caption retrieval
		cost_s = (self.margin + scores - d1).clamp(min=0)
		cost_s_OT = (self.margin + scores_OT - d1_OT).clamp(min=0)
		# compare every diagonal score to scores in its row
		# image retrieval
		cost_im = (self.margin + scores - d2).clamp(min=0)
		cost_im_OT = (self.margin + scores_OT - d2_OT).clamp(min=0)

		# clear diagonals
		mask = torch.eye(scores.size(0)) > .5
		I = Variable(mask)
		if torch.cuda.is_available():
			I = I.cuda()
		cost_s = cost_s.masked_fill_(I, 0)
		cost_im = cost_im.masked_fill_(I, 0)

		cost_s_OT = cost_s_OT.masked_fill_(I, 0)
		cost_im_OT = cost_im_OT.masked_fill_(I, 0)

		alpha = self.opt.alpha  # .cuda()
		# cost_s = cost_s + alpha * cost_s_OT
		# cost_im = cost_im + alpha * cost_im_OT

		# keep the maximum violating negative for each query
		if self.max_violation:
			cost_s = cost_s.max(1)[0]
			cost_im = cost_im.max(0)[0]
			cost_s_OT = cost_s_OT.max(1)[0]
			cost_im_OT = cost_im_OT.max(0)[0]

		return cost_s.sum() + cost_im.sum(), cost_s_OT.sum() + cost_im_OT.sum()



class Model(object):
	"""
	dot product model plus Optimal Transport
	"""

	def __init__(self, opt):
		# Build Models
		self.grad_clip = opt.grad_clip
		self.img_enc = EncoderImage(opt.data_name, opt.img_dim, opt.embed_size,
		                            precomp_enc_type=opt.precomp_enc_type,
		                            no_imgnorm=opt.no_imgnorm)
		self.txt_enc = EncoderText(opt.vocab_size, opt.word_dim,
		                           opt.embed_size, opt.num_layers,
		                           use_bi_gru=opt.bi_gru,
		                           no_txtnorm=opt.no_txtnorm)
		if torch.cuda.is_available():
			self.img_enc.cuda()
			self.txt_enc.cuda()
			cudnn.benchmark = True

		# Loss and Optimizer
		self.criterion = ContrastiveLoss(opt=opt,
		                                        margin=opt.margin,
		                                        max_violation=opt.max_violation)
		params = list(self.txt_enc.parameters())
		params += list(self.img_enc.fc.parameters())

		self.params = params

		self.optimizer = torch.optim.Adam(params, lr=opt.learning_rate)

		self.Eiters = 0

	def state_dict(self):
		state_dict = [self.img_enc.state_dict(), self.txt_enc.state_dict()]
		return state_dict

	def load_state_dict(self, state_dict):
		self.img_enc.load_state_dict(state_dict[0])
		self.txt_enc.load_state_dict(state_dict[1])

	def train_start(self):
		"""switch to train mode
		"""
		self.img_enc.train()
		self.txt_enc.train()

	def val_start(self):
		"""switch to evaluate mode
		"""
		self.img_enc.eval()
		self.txt_enc.eval()

	def forward_emb(self, images, captions, lengths, volatile=False):
		"""Compute the image and caption embeddings
		"""
		# Set mini-batch dataset
		images = Variable(images, volatile=volatile)
		captions = Variable(captions, volatile=volatile)
		if torch.cuda.is_available():
			images = images.cuda()
			captions = captions.cuda()

		# Forward
		img_emb = self.img_enc(images)

		# cap_emb (tensor), cap_lens (list), embedded_sentence
		cap_emb, cap_lens, x = self.txt_enc(captions, lengths)
		return img_emb, cap_emb, cap_lens, x

	def forward_loss(self, img_emb, cap_emb, cap_len, x, **kwargs):
		"""Compute the loss given pairs of image and caption embeddings
		"""
		loss, loss_OT = self.criterion(img_emb, cap_emb, cap_len, x)
		self.logger.update('Le', loss.item(), img_emb.size(0))
		self.logger.update('Le_OT', loss_OT.item(), img_emb.size(0))
		return loss

	def train_emb(self, images, captions, lengths, ids=None, *args):
		"""One training step given images and captions.
		image: bs * 36 * 2048
		"""
		self.Eiters += 1
		self.logger.update('Eit', self.Eiters)
		self.logger.update('lr', self.optimizer.param_groups[0]['lr'])

		# compute the embeddings
		img_emb, cap_emb, cap_lens, x = self.forward_emb(images, captions, lengths)

		# measure accuracy and record loss
		self.optimizer.zero_grad()
		loss = self.forward_loss(img_emb, cap_emb, cap_lens, x)

		# compute gradient and do SGD step
		loss.backward()
		if self.grad_clip > 0:
			clip_grad_norm(self.params, self.grad_clip)
		self.optimizer.step()

def shard_score(images, captions, caplens, x, opt, shard_size=128):
    """
    Computer pairwise t2i image-caption distance with locality sharding
    """
    n_im_shard = (len(images) - 1) // shard_size + 1
    n_cap_shard = (len(captions) - 1) // shard_size + 1

    d = np.zeros((len(images), len(captions)))
    for i in range(n_im_shard):
        im_start, im_end = shard_size * i, min(shard_size * (i + 1), len(images))
        for j in range(n_cap_shard):
            sys.stdout.write('\r>> shard_xattn_t2i batch (%d,%d)' % (i, j))
            cap_start, cap_end = shard_size * j, min(shard_size * (j + 1), len(captions))
            with torch.no_grad():
                im = Variable(torch.from_numpy(images[im_start:im_end])).cuda()
                s = Variable(torch.from_numpy(captions[cap_start:cap_end])).cuda()
                x_ = Variable(torch.from_numpy(x[cap_start:cap_end])).cuda()
            if caplens is not None:
                l = caplens[cap_start:cap_end]
            else:
                l = None
            sim, sim_OT = score(im, s, l, x_, opt)
            if caplens is not None:
                d[im_start:im_end, cap_start:cap_end] = sim.cpu().numpy() + opt.alpha * sim_OT.cpu().numpy()
            else:
                d[im_start:im_end, cap_start:cap_end] = sim_OT.cpu().numpy()
    sys.stdout.write('\n')
    return d

def shard_score_cluster(images, captions, caplens, x, cluster, opt, shard_size=128):
    """
    Computer pairwise t2i image-caption distance with locality sharding
    """
    n_im_shard = (len(images) - 1) // shard_size + 1
    n_cap_shard = (len(captions) - 1) // shard_size + 1

    d = np.zeros((len(images), len(captions)))
    for i in range(n_im_shard):
        im_start, im_end = shard_size * i, min(shard_size * (i + 1), len(images))
        for j in range(n_cap_shard):
            sys.stdout.write('\r>> shard_xattn_t2i batch (%d,%d)' % (i, j))
            cap_start, cap_end = shard_size * j, min(shard_size * (j + 1), len(captions))
            with torch.no_grad():
                im = Variable(torch.from_numpy(images[im_start:im_end])).cuda()
                s = Variable(torch.from_numpy(captions[cap_start:cap_end])).cuda()
                x_ = Variable(torch.from_numpy(x[cap_start:cap_end])).cuda()
            if caplens is not None:
                l = caplens[cap_start:cap_end]
            else:
                l = None
            sim, sim_OT = score(im, s, l, x_, opt)
            if caplens is not None:
                d[im_start:im_end, cap_start:cap_end] = sim.cpu().numpy() + opt.alpha * sim_OT.cpu().numpy()
            else:
                d[im_start:im_end, cap_start:cap_end] = sim_OT.cpu().numpy()
    sys.stdout.write('\n')
    return d