import math
import torch
import torch.nn.functional as F

def color_loss(pred,target):
    """Fit one coherent target; distribution matching discourages gray averages."""
    pred=F.adaptive_avg_pool2d(pred,(64,64)).float();target=target.float()
    pixel=F.l1_loss(pred/20,target/20)
    p=F.avg_pool2d(pred,2).flatten(2);t=F.avg_pool2d(target,2).flatten(2)
    angle=torch.arange(8,device=p.device,dtype=p.dtype)*math.pi/8
    directions=torch.stack([angle.cos(),angle.sin()],1)
    a=torch.einsum('kc,bcn->bkn',directions,p).sort(-1).values
    b=torch.einsum('kc,bcn->bkn',directions,t).sort(-1).values
    distribution=F.l1_loss(a/20,b/20)
    # Match boundaries in the teacher target, never minimize gradients toward zero.
    edge=(F.l1_loss((pred[:,:,:,1:]-pred[:,:,:,:-1])/20,(target[:,:,:,1:]-target[:,:,:,:-1])/20)+F.l1_loss((pred[:,:,1:]-pred[:,:,:-1])/20,(target[:,:,1:]-target[:,:,:-1])/20))/2
    return pixel+.25*distribution+.1*edge,{'pixel':pixel,'distribution':distribution,'edge':edge}