import torch, torch.nn as nn
class DW(nn.Module):  # depthwise-separable bottleneck (MobileFaceNet style)
    def __init__(s,i,o,st,exp=2):
        super().__init__(); h=i*exp; s.res=(st==1 and i==o)
        s.b=nn.Sequential(nn.Conv2d(i,h,1,bias=False),nn.BatchNorm2d(h),nn.PReLU(h),
            nn.Conv2d(h,h,3,st,1,groups=h,bias=False),nn.BatchNorm2d(h),nn.PReLU(h),
            nn.Conv2d(h,o,1,bias=False),nn.BatchNorm2d(o))
    def forward(s,x): return x+s.b(x) if s.res else s.b(x)
class TinyFaceNet(nn.Module):
    """~1.0M params, 112x112 in, 128-d embedding."""
    def __init__(s,emb=128,n_cls=10,w=1.6):
        super().__init__(); c=lambda v:int(v*w)
        s.stem=nn.Sequential(nn.Conv2d(3,c(32),3,2,1,bias=False),nn.BatchNorm2d(c(32)),nn.PReLU(c(32)),
                             nn.Conv2d(c(32),c(32),3,1,1,groups=c(32),bias=False),nn.BatchNorm2d(c(32)),nn.PReLU(c(32)))
        cfg=[(c(32),c(48),2),(c(48),c(48),1),(c(48),c(48),1),(c(48),c(96),2),(c(96),c(96),1),(c(96),c(96),1),(c(96),c(96),1),
             (c(96),c(128),2),(c(128),c(128),1),(c(128),c(128),1)]
        s.blocks=nn.Sequential(*[DW(i,o,st) for i,o,st in cfg])
        s.head=nn.Sequential(nn.Conv2d(c(128),c(256),1,bias=False),nn.BatchNorm2d(c(256)),nn.PReLU(c(256)),
                             nn.Conv2d(c(256),c(256),7,1,0,groups=c(256),bias=False),nn.BatchNorm2d(c(256)),
                             nn.Flatten(),nn.Linear(c(256),emb,bias=False),nn.BatchNorm1d(emb))
        s.distill=nn.Linear(emb,512)        # training-only: map to teacher space
        s.cls=nn.Linear(emb,n_cls)          # training-only
    def forward(s,x): return s.head(s.blocks(s.stem(x)))
if __name__=="__main__":
    m=TinyFaceNet(); core=sum(p.numel() for n,p in m.named_parameters() if not n.startswith(("distill","cls")))
    print("deploy params",core, "total",sum(p.numel() for p in m.parameters())); print(m(torch.zeros(2,3,112,112)).shape)