import torch
import torch.nn as nn

def _sigmoid(x):
	y = torch.clamp(x.sigmoid_(), min=1e-4, max=1-1e-4)
	return y

def _sigmoid12(x):
	y = torch.clamp(x.sigmoid_(), 1e-12)
	return y

def _gather_feat(feat, ind):
	dim = feat.size(2)
	ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)
	feat = feat.gather(1, ind)
	return feat

def _tranpose_and_gather_feat(feat, ind):
	feat = feat.permute(0, 2, 3, 1).contiguous()
	feat = feat.view(feat.size(0), -1, feat.size(3))
	feat = _gather_feat(feat, ind)
	return feat

def flip_tensor(x):
	return torch.flip(x, [3])
	# tmp = x.detach().cpu().numpy()[..., ::-1].copy()
	# return torch.from_numpy(tmp).to(x.device)

def flip_lr(x, flip_idx):
	tmp = x.detach().cpu().numpy()[..., ::-1].copy()
	shape = tmp.shape
	for e in flip_idx:
		tmp[:, e[0], ...], tmp[:, e[1], ...] = \
			tmp[:, e[1], ...].copy(), tmp[:, e[0], ...].copy()
	return torch.from_numpy(tmp.reshape(shape)).to(x.device)

def flip_lr_off(x, flip_idx):
	tmp = x.detach().cpu().numpy()[..., ::-1].copy()
	shape = tmp.shape
	tmp = tmp.reshape(tmp.shape[0], 17, 2, 
										tmp.shape[2], tmp.shape[3])
	tmp[:, :, 0, :, :] *= -1
	for e in flip_idx:
		tmp[:, e[0], ...], tmp[:, e[1], ...] = \
			tmp[:, e[1], ...].copy(), tmp[:, e[0], ...].copy()
	return torch.from_numpy(tmp.reshape(shape)).to(x.device)

def _nms(heat, kernel=3):
	pad = (kernel - 1) // 2

	hmax = nn.functional.max_pool2d(
			heat, (kernel, kernel), stride=1, padding=pad)
	keep = (hmax == heat).float()
	return heat * keep

def _topk_channel(scores, K=100):
	batch, cat, height, width = scores.size()
	
	topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), K)

	topk_inds = topk_inds % (height * width)
	topk_ys   = (topk_inds / width).int().float()
	topk_xs   = (topk_inds % width).int().float()

	return topk_scores, topk_inds, topk_ys, topk_xs

def _topk(scores, K=100):
	batch, cat, height, width = scores.size()
		
	topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), K)

	topk_inds = topk_inds % (height * width)
	topk_ys   = (topk_inds / width).int().float()
	topk_xs   = (topk_inds % width).int().float()
		
	topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), K)
	topk_clses = (topk_ind / K).int()
	topk_inds = _gather_feat(
			topk_inds.view(batch, -1, 1), topk_ind).view(batch, K)
	topk_ys = _gather_feat(topk_ys.view(batch, -1, 1), topk_ind).view(batch, K)
	topk_xs = _gather_feat(topk_xs.view(batch, -1, 1), topk_ind).view(batch, K)

	return topk_score, topk_inds, topk_clses, topk_ys, topk_xs
