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, 'torch.argsort', 'torch.argsort', False, 'import torch\n'), (126, 'torch.argsort', 'torch.argsort', False, 'import torch\n'), (196, 'torch.argsort', 'torch.argsort', False, 'import torch\n'), (17, 'torch.min', 'torch.min', False, 'import torch\n'), (42, 'torch.sum', 'torch.sum', False, 'import torch\n'), (43, 'torch.sum', 'torch.sum', False, 'import torch\n'), (59, 'torch.sum', 'torch.sum', False, 'import torch\n'), (80, 'torch.sum', 'torch.sum', False, 'import torch\n'), (114, 'torch.min', 'torch.min', False, 'import torch\n'), (132, 'torch.clamp', 'torch.clamp', False, 'import torch\n'), (139, 'torch.clamp', 'torch.clamp', False, 'import torch\n'), (143, 'torch.sum', 'torch.sum', False, 'import torch\n'), (184, 'torch.min', 'torch.min', False, 'import torch\n'), (202, 'torch.clamp', 'torch.clamp', False, 'import torch\n'), (209, 'torch.clamp', 'torch.clamp', False, 'import torch\n'), (213, 'torch.sum', 'torch.sum', False, 'import torch\n'), (7, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (22, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (23, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (24, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (35, 'torch.clamp', 'torch.clamp', False, 'import torch\n'), (36, 'torch.clamp', 'torch.clamp', False, 'import torch\n'), (52, 'torch.sum', 'torch.sum', False, 'import torch\n'), (63, 'torch.sum', 'torch.sum', False, 'import torch\n'), (105, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (120, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (121, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (122, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (142, 'torch.sum', 'torch.sum', False, 'import torch\n'), (175, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (190, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (191, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (192, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (212, 'torch.sum', 'torch.sum', False, 'import torch\n'), (151, 'torch.sum', 'torch.sum', False, 'import torch\n')]