"""Original compact baseline, NOT a reproduction of MSHNet or DNANet."""
import torch
from torch import nn
from torch.nn import functional as F

class TinyUNet(nn.Module):
    def __init__(self, width=16):
        super().__init__()
        def block(a,b):
            return nn.Sequential(nn.Conv2d(a,b,3,padding=1),nn.ReLU(),
                                 nn.Conv2d(b,b,3,padding=1),nn.ReLU())
        self.enc = block(1,width)
        self.mid = block(width,width*2)
        self.dec = block(width*3,width)
        self.head = nn.Conv2d(width,1,1)
    def forward(self,x):
        a = self.enc(x)
        b = self.mid(F.max_pool2d(a,2,ceil_mode=True))
        b = F.interpolate(b,size=a.shape[-2:],mode='bilinear',align_corners=False)
        return self.head(self.dec(torch.cat((a,b),1)))

class CandidateGate(nn.Module):
    def __init__(self):
        super().__init__()
        self.net=nn.Sequential(nn.Linear(5,16),nn.ReLU(),nn.Linear(16,1))
    def forward(self,x):
        return self.net(x).squeeze(-1)

def segmentation_loss(logits, target, name='bce_dice'):
    p = logits.sigmoid()
    bce = F.binary_cross_entropy_with_logits(logits,target)
    dims=tuple(range(1,p.ndim)); eps=1e-6
    inter=(p*target).sum(dims)
    if name=='bce': return bce
    if name=='focal':
        pt=p*target+(1-p)*(1-target)
        weight=.25*target+.75*(1-target)
        return (weight*(1-pt).pow(2)*F.binary_cross_entropy_with_logits(logits,target,reduction='none')).mean()
    dice=(1-(2*inter+eps)/(p.sum(dims)+target.sum(dims)+eps)).mean()
    if name=='dice': return dice
    if name=='bce_dice': return bce+dice
    if name=='iou': return (1-(inter+eps)/(p.sum(dims)+target.sum(dims)-inter+eps)).mean()
    raise ValueError(name)

def gate_loss(logits, labels, areas, dual, alpha, size_weight=True, tau=1.):
    import math
    pos=labels>0.5; neg=~pos
    weight=areas.clamp_min(1).rsqrt().clamp(.1,1) if size_weight else torch.ones_like(areas)
    positive=(F.softplus(-logits[pos])*weight[pos]).sum()/max(1,int(pos.sum()))
    upper=F.softplus(logits[neg]/tau).sum()/math.log(2)
    return positive+dual*(upper-alpha), upper