import torch

class RankSort(torch.autograd.Function):
	@staticmethod
	def forward(ctx, logits, targets, delta_RS=0.50, eps=1e-10): 

		classification_grads=torch.zeros(logits.shape).cuda()
		
		#Filter fg logits
		fg_labels = (targets > 0.)
		fg_logits = logits[fg_labels]
		fg_targets = targets[fg_labels]
		fg_num = len(fg_logits)

		#Do not use bg with scores less than minimum fg logit
		#since changing its score does not have an effect on precision
		threshold_logit = torch.min(fg_logits)-delta_RS
		relevant_bg_labels=((targets==0) & (logits>=threshold_logit))
		
		relevant_bg_logits = logits[relevant_bg_labels] 
		relevant_bg_grad=torch.zeros(len(relevant_bg_logits)).cuda()
		sorting_error=torch.zeros(fg_num).cuda()
		ranking_error=torch.zeros(fg_num).cuda()
		fg_grad=torch.zeros(fg_num).cuda()
		
		#sort the fg logits
		order=torch.argsort(fg_logits)
		#Loops over each positive following the order
		for ii in order:
			# Difference Transforms (x_ij)
			fg_relations=fg_logits-fg_logits[ii] 
			bg_relations=relevant_bg_logits-fg_logits[ii]

			if delta_RS > 0:
				fg_relations=torch.clamp(fg_relations/(2*delta_RS)+0.5,min=0,max=1)
				bg_relations=torch.clamp(bg_relations/(2*delta_RS)+0.5,min=0,max=1)
			else:
				fg_relations = (fg_relations >= 0).float()
				bg_relations = (bg_relations >= 0).float()

			# Rank of ii among pos and false positive number (bg with larger scores)
			rank_pos=torch.sum(fg_relations)
			FP_num=torch.sum(bg_relations)

			# Rank of ii among all examples
			rank=rank_pos+FP_num
							
			# Ranking error of example ii. target_ranking_error is always 0. (Eq. 7)
			ranking_error[ii]=FP_num/rank	  

			# Current sorting error of example ii. (Eq. 7)
			current_sorting_error = torch.sum(fg_relations*(1-fg_targets))/rank_pos

			#Find examples in the target sorted order for example ii		 
			iou_relations = (fg_targets >= fg_targets[ii])
			target_sorted_order = iou_relations * fg_relations

			#The rank of ii among positives in sorted order
			rank_pos_target = torch.sum(target_sorted_order)

			#Compute target sorting error. (Eq. 8)
			#Since target ranking error is 0, this is also total target error 
			target_sorting_error= torch.sum(target_sorted_order*(1-fg_targets))/rank_pos_target

			#Compute sorting error on example ii
			sorting_error[ii] = current_sorting_error - target_sorting_error
  
			#Identity Update for Ranking Error 
			if FP_num > eps:
				#For ii the update is the ranking error
				fg_grad[ii] -= ranking_error[ii]
				#For negatives, distribute error via ranking pmf (i.e. bg_relations/FP_num)
				relevant_bg_grad += (bg_relations*(ranking_error[ii]/FP_num))

			#Find the positives that are misranked (the cause of the error)
			#These are the ones with smaller IoU but larger logits
			missorted_examples = (~ iou_relations) * fg_relations

			#Denominotor of sorting pmf 
			sorting_pmf_denom = torch.sum(missorted_examples)

			#Identity Update for Sorting Error 
			if sorting_pmf_denom > eps:
				#For ii the update is the sorting error
				fg_grad[ii] -= sorting_error[ii]
				#For positives, distribute error via sorting pmf (i.e. missorted_examples/sorting_pmf_denom)
				fg_grad += (missorted_examples*(sorting_error[ii]/sorting_pmf_denom))

		#Normalize gradients by number of positives 
		classification_grads[fg_labels]= (fg_grad/fg_num)
		classification_grads[relevant_bg_labels]= (relevant_bg_grad/fg_num)

		ctx.save_for_backward(classification_grads)

		return ranking_error.mean(), sorting_error.mean()

	@staticmethod
	def backward(ctx, out_grad1, out_grad2):
		g1, =ctx.saved_tensors
		return g1*out_grad1, None, None, None

class aLRPLoss(torch.autograd.Function):
	@staticmethod
	def forward(ctx, logits, targets, regression_losses, delta=1., eps=1e-5): 
		classification_grads=torch.zeros(logits.shape).cuda()
		
		#Filter fg logits
		fg_labels = (targets == 1)
		fg_logits = logits[fg_labels]
		fg_num = len(fg_logits)

		#Do not use bg with scores less than minimum fg logit
		#since changing its score does not have an effect on precision
		threshold_logit = torch.min(fg_logits)-delta

		#Get valid bg logits
		relevant_bg_labels=((targets==0)&(logits>=threshold_logit))
		relevant_bg_logits=logits[relevant_bg_labels] 
		relevant_bg_grad=torch.zeros(len(relevant_bg_logits)).cuda()
		rank=torch.zeros(fg_num).cuda()
		prec=torch.zeros(fg_num).cuda()
		fg_grad=torch.zeros(fg_num).cuda()
		
		max_prec=0										   
		#sort the fg logits
		order=torch.argsort(fg_logits)
		#Loops over each positive following the order
		for ii in order:
			#x_ij s as score differences with fgs
			fg_relations=fg_logits-fg_logits[ii] 
			#Apply piecewise linear function and determine relations with fgs
			fg_relations=torch.clamp(fg_relations/(2*delta)+0.5,min=0,max=1)
			#Discard i=j in the summation in rank_pos
			fg_relations[ii]=0

			#x_ij s as score differences with bgs
			bg_relations=relevant_bg_logits-fg_logits[ii]
			#Apply piecewise linear function and determine relations with bgs
			bg_relations=torch.clamp(bg_relations/(2*delta)+0.5,min=0,max=1)

			#Compute the rank of the example within fgs and number of bgs with larger scores
			rank_pos=1+torch.sum(fg_relations)
			FP_num=torch.sum(bg_relations)
			#Store the total since it is normalizer also for aLRP Regression error
			rank[ii]=rank_pos+FP_num
							
			#Compute precision for this example to compute classification loss 
			prec[ii]=rank_pos/rank[ii]				
			#For stability, set eps to a infinitesmall value (e.g. 1e-6), then compute grads
			if FP_num > eps:   
				fg_grad[ii] = -(torch.sum(fg_relations*regression_losses)+FP_num)/rank[ii]
				relevant_bg_grad += (bg_relations*(-fg_grad[ii]/FP_num))   
					
		#aLRP with grad formulation fg gradient
		classification_grads[fg_labels]= fg_grad
		#aLRP with grad formulation bg gradient
		classification_grads[relevant_bg_labels]= relevant_bg_grad 
 
		classification_grads /= (fg_num)
	
		cls_loss=1-prec.mean()
		ctx.save_for_backward(classification_grads)

		return cls_loss, rank, order

	@staticmethod
	def backward(ctx, out_grad1, out_grad2, out_grad3):
		g1, =ctx.saved_tensors
		return g1*out_grad1, None, None, None, None
	
	
class APLoss(torch.autograd.Function):
	@staticmethod
	def forward(ctx, logits, targets, delta=1.): 
		classification_grads=torch.zeros(logits.shape).cuda()
		
		#Filter fg logits
		fg_labels = (targets == 1)
		fg_logits = logits[fg_labels]
		fg_num = len(fg_logits)

		#Do not use bg with scores less than minimum fg logit
		#since changing its score does not have an effect on precision
		threshold_logit = torch.min(fg_logits)-delta

		#Get valid bg logits
		relevant_bg_labels=((targets==0)&(logits>=threshold_logit))
		relevant_bg_logits=logits[relevant_bg_labels] 
		relevant_bg_grad=torch.zeros(len(relevant_bg_logits)).cuda()
		rank=torch.zeros(fg_num).cuda()
		prec=torch.zeros(fg_num).cuda()
		fg_grad=torch.zeros(fg_num).cuda()
		
		max_prec=0										   
		#sort the fg logits
		order=torch.argsort(fg_logits)
		#Loops over each positive following the order
		for ii in order:
			#x_ij s as score differences with fgs
			fg_relations=fg_logits-fg_logits[ii] 
			#Apply piecewise linear function and determine relations with fgs
			fg_relations=torch.clamp(fg_relations/(2*delta)+0.5,min=0,max=1)
			#Discard i=j in the summation in rank_pos
			fg_relations[ii]=0

			#x_ij s as score differences with bgs
			bg_relations=relevant_bg_logits-fg_logits[ii]
			#Apply piecewise linear function and determine relations with bgs
			bg_relations=torch.clamp(bg_relations/(2*delta)+0.5,min=0,max=1)

			#Compute the rank of the example within fgs and number of bgs with larger scores
			rank_pos=1+torch.sum(fg_relations)
			FP_num=torch.sum(bg_relations)
			#Store the total since it is normalizer also for aLRP Regression error
			rank[ii]=rank_pos+FP_num
							
			#Compute precision for this example 
			current_prec=rank_pos/rank[ii]
			
			#Compute interpolated AP and store gradients for relevant bg examples
			if (max_prec<=current_prec):
				max_prec=current_prec
				relevant_bg_grad += (bg_relations/rank[ii])
			else:
				relevant_bg_grad += (bg_relations/rank[ii])*(((1-max_prec)/(1-current_prec)))
			
			#Store fg gradients
			fg_grad[ii]=-(1-max_prec)
			prec[ii]=max_prec 

		#aLRP with grad formulation fg gradient
		classification_grads[fg_labels]= fg_grad
		#aLRP with grad formulation bg gradient
		classification_grads[relevant_bg_labels]= relevant_bg_grad 
 
		classification_grads /= fg_num
	
		cls_loss=1-prec.mean()
		ctx.save_for_backward(classification_grads)

		return cls_loss

	@staticmethod
	def backward(ctx, out_grad1):
		g1, =ctx.saved_tensors
		return g1*out_grad1, None, None

[((27, 8, 27, 32), 'torch.argsort', 'torch.argsort', ({(27, 22, 27, 31): 'fg_logits'}, {}), '(fg_logits)', False, 'import torch\n'), ((126, 8, 126, 32), 'torch.argsort', 'torch.argsort', ({(126, 22, 126, 31): 'fg_logits'}, {}), '(fg_logits)', False, 'import torch\n'), ((196, 8, 196, 32), 'torch.argsort', 'torch.argsort', ({(196, 22, 196, 31): 'fg_logits'}, {}), '(fg_logits)', False, 'import torch\n'), ((17, 20, 17, 40), 'torch.min', 'torch.min', ({(17, 30, 17, 39): 'fg_logits'}, {}), '(fg_logits)', False, 'import torch\n'), ((42, 12, 42, 35), 'torch.sum', 'torch.sum', ({(42, 22, 42, 34): 'fg_relations'}, {}), '(fg_relations)', False, 'import torch\n'), ((43, 10, 43, 33), 'torch.sum', 'torch.sum', ({(43, 20, 43, 32): 'bg_relations'}, {}), '(bg_relations)', False, 'import torch\n'), ((59, 21, 59, 51), 'torch.sum', 'torch.sum', ({(59, 31, 59, 50): 'target_sorted_order'}, {}), '(target_sorted_order)', False, 'import torch\n'), ((80, 23, 80, 52), 'torch.sum', 'torch.sum', ({(80, 33, 80, 51): 'missorted_examples'}, {}), '(missorted_examples)', False, 'import torch\n'), ((114, 20, 114, 40), 'torch.min', 'torch.min', ({(114, 30, 114, 39): 'fg_logits'}, {}), '(fg_logits)', False, 'import torch\n'), ((132, 16, 132, 67), 'torch.clamp', 'torch.clamp', (), '', False, 'import torch\n'), ((139, 16, 139, 67), 'torch.clamp', 'torch.clamp', (), '', False, 'import torch\n'), ((143, 10, 143, 33), 'torch.sum', 'torch.sum', ({(143, 20, 143, 32): 'bg_relations'}, {}), '(bg_relations)', False, 'import torch\n'), ((184, 20, 184, 40), 'torch.min', 'torch.min', ({(184, 30, 184, 39): 'fg_logits'}, {}), '(fg_logits)', False, 'import torch\n'), ((202, 16, 202, 67), 'torch.clamp', 'torch.clamp', (), '', False, 'import torch\n'), ((209, 16, 209, 67), 'torch.clamp', 'torch.clamp', (), '', False, 'import torch\n'), ((213, 10, 213, 33), 'torch.sum', 'torch.sum', ({(213, 20, 213, 32): 'bg_relations'}, {}), '(bg_relations)', False, 'import torch\n'), ((7, 23, 7, 48), 'torch.zeros', 'torch.zeros', ({(7, 35, 7, 47): 'logits.shape'}, {}), '(logits.shape)', False, 'import torch\n'), ((22, 16, 22, 35), 'torch.zeros', 'torch.zeros', ({(22, 28, 22, 34): 'fg_num'}, {}), '(fg_num)', False, 'import torch\n'), ((23, 16, 23, 35), 'torch.zeros', 'torch.zeros', ({(23, 28, 23, 34): 'fg_num'}, {}), '(fg_num)', False, 'import torch\n'), ((24, 10, 24, 29), 'torch.zeros', 'torch.zeros', ({(24, 22, 24, 28): 'fg_num'}, {}), '(fg_num)', False, 'import torch\n'), ((35, 17, 35, 71), 'torch.clamp', 'torch.clamp', (), '', False, 'import torch\n'), ((36, 17, 36, 71), 'torch.clamp', 'torch.clamp', (), '', False, 'import torch\n'), ((52, 27, 52, 65), 'torch.sum', 'torch.sum', ({(52, 37, 52, 64): '(fg_relations * (1 - fg_targets))'}, {}), '(fg_relations * (1 - fg_targets))', False, 'import torch\n'), ((63, 25, 63, 70), 'torch.sum', 'torch.sum', ({(63, 35, 63, 69): '(target_sorted_order * (1 - fg_targets))'}, {}), '(target_sorted_order * (1 - fg_targets))', False, 'import torch\n'), ((105, 23, 105, 48), 'torch.zeros', 'torch.zeros', ({(105, 35, 105, 47): 'logits.shape'}, {}), '(logits.shape)', False, 'import torch\n'), ((120, 7, 120, 26), 'torch.zeros', 'torch.zeros', ({(120, 19, 120, 25): 'fg_num'}, {}), '(fg_num)', False, 'import torch\n'), ((121, 7, 121, 26), 'torch.zeros', 'torch.zeros', ({(121, 19, 121, 25): 'fg_num'}, {}), '(fg_num)', False, 'import torch\n'), ((122, 10, 122, 29), 'torch.zeros', 'torch.zeros', ({(122, 22, 122, 28): 'fg_num'}, {}), '(fg_num)', False, 'import torch\n'), ((142, 14, 142, 37), 'torch.sum', 'torch.sum', ({(142, 24, 142, 36): 'fg_relations'}, {}), '(fg_relations)', False, 'import torch\n'), ((175, 23, 175, 48), 'torch.zeros', 'torch.zeros', ({(175, 35, 175, 47): 'logits.shape'}, {}), '(logits.shape)', False, 'import torch\n'), ((190, 7, 190, 26), 'torch.zeros', 'torch.zeros', ({(190, 19, 190, 25): 'fg_num'}, {}), '(fg_num)', False, 'import torch\n'), ((191, 7, 191, 26), 'torch.zeros', 'torch.zeros', ({(191, 19, 191, 25): 'fg_num'}, {}), '(fg_num)', False, 'import torch\n'), ((192, 10, 192, 29), 'torch.zeros', 'torch.zeros', ({(192, 22, 192, 28): 'fg_num'}, {}), '(fg_num)', False, 'import torch\n'), ((212, 14, 212, 37), 'torch.sum', 'torch.sum', ({(212, 24, 212, 36): 'fg_relations'}, {}), '(fg_relations)', False, 'import torch\n'), ((151, 20, 151, 61), 'torch.sum', 'torch.sum', ({(151, 30, 151, 60): '(fg_relations * regression_losses)'}, {}), '(fg_relations * regression_losses)', False, 'import torch\n')]