"""Algebraic pointwise-Conv -> Linear lowering; unchanged masked SE semantics."""
import copy
import torch
from torch import nn
from torch.nn import functional as F

def linear(conv):
    layer = nn.Linear(conv.in_channels, conv.out_channels)
    with torch.no_grad():
        layer.weight.copy_(conv.weight[..., 0])
        layer.bias.copy_(conv.bias)
    return layer

class PortableBlock(nn.Module):
    def __init__(self, block):
        super().__init__()
        self.dw = copy.deepcopy(block.dw)
        for name in ['expand', 'fc1', 'fc2', 'project']:
            setattr(self, name, linear(getattr(block, name)))

    def forward(self, x, mask):
        z = F.hardswish(self.expand(self.dw(x.transpose(1,2)).transpose(1,2)))
        pool = (z * mask).sum(1, keepdim=True) / mask.sum(1, keepdim=True).clamp_min(1.)
        gate = torch.sigmoid(self.fc2(F.relu(self.fc1(pool))))
        return (x + self.project(z * gate)) * mask

class PortableRecognizer(nn.Module):
    """Input widths are true unpadded widths; canvas must be a multiple of eight."""
    def __init__(self, source):
        super().__init__()
        self.index_dtype = torch.int64
        self.stem = copy.deepcopy(source.stem)
        self.bn = copy.deepcopy(source.bn)
        self.blocks = nn.ModuleList(PortableBlock(b) for b in source.blocks)
        self.head = linear(source.head)

    def forward(self, pixels, widths):
        z = self.stem(F.pad(pixels, (0,4,0,0))).squeeze(2)
        lengths = (widths + 3) // 4
        mask = (torch.arange(z.shape[-1], device=z.device, dtype=self.index_dtype)[None,:] < lengths[:,None])[:,:,None].to(z.dtype)
        z = F.hardswish(self.bn(z)).transpose(1,2) * mask
        for block in self.blocks:
            z = block(z, mask)
        return self.head(z), lengths

class Compact(nn.Module):
    def __init__(self, source):
        super().__init__()
        self.model = source
    def forward(self, pixels, widths):
        logits, lengths = self.model(pixels, widths)
        confidence, tokens = logits.softmax(-1).max(-1)
        return tokens.to(torch.int32), confidence, lengths.to(torch.int32)