"""Width512/depth20 Indra-family CTC recognizer; native class IDs are retained."""
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.checkpoint import checkpoint
class Block(nn.Module):
 def __init__(self,width=512,identity=False):
  super().__init__();self.dw=nn.Conv1d(width,width,17,padding=8,groups=width);self.expand=nn.Conv1d(width,2*width,1);self.fc1=nn.Conv1d(2*width,width//2,1);self.fc2=nn.Conv1d(width//2,2*width,1);self.project=nn.Conv1d(2*width,width,1)
  if identity:nn.init.zeros_(self.project.weight);nn.init.zeros_(self.project.bias)
 def forward(self,x,mask):
  z=F.hardswish(self.expand(self.dw(x)));pool=(z*mask).sum(2,keepdim=True)/mask.sum(2,keepdim=True).clamp_min(1);return(x+self.project(z*torch.sigmoid(self.fc2(F.relu(self.fc1(pool))))))*mask
class Recognizer(nn.Module):
 input_height=64
 def __init__(self,classes=1357,width=512,depth=20,gradient_checkpointing=True):
  super().__init__();assert width==512 and depth==20
  self.stem=nn.Conv2d(1,width,(64,8),stride=(64,4));self.bn=nn.BatchNorm1d(width);self.blocks=nn.ModuleList(Block(width,identity=i%2==1)for i in range(depth));self.head=nn.Conv1d(width,classes,1);self.gradient_checkpointing=gradient_checkpointing;self.bn.eval()
 def train(self,mode=True):super().train(mode);self.bn.eval();return self
 def forward(self,x,widths):
  if x.ndim!=4 or x.shape[1]!=1 or x.shape[2]!=64:raise ValueError('Require Bx1x64xW without height truncation')
  if widths.ndim!=1 or len(widths)!=len(x)or bool((widths<1).any())or bool((widths>x.shape[-1]).any()):raise ValueError('Invalid original crop widths')
  widths=widths.to(device=x.device,dtype=torch.long)
  if x.shape[-1]%4:x=F.pad(x,(0,(-x.shape[-1])%4))
  z=self.stem(F.pad(x,(0,4,0,0))).squeeze(2);lengths=(widths+3)//4;assert bool((lengths<=z.shape[-1]).all());mask=(torch.arange(z.shape[-1],device=z.device)[None,:]<lengths[:,None])[:,None,:].to(z.dtype);z=F.hardswish(self.bn(z))*mask
  for block in self.blocks:
   z=checkpoint(block,z,mask,use_reentrant=False)if self.training and self.gradient_checkpointing and torch.is_grad_enabled()else block(z,mask)
  return self.head(z).transpose(1,2),lengths
 def lift_from_highres(self,source):
  """Net2Wider copies old128 features fourfold; new interleaved blocks start as identities."""
  state=source.state_dict()if isinstance(source,nn.Module)else source
  assert state['stem.weight'].shape==(128,1,64,8)and state['blocks.9.dw.weight'].shape==(128,1,17)
  with torch.no_grad():
   self.stem.weight.copy_(state['stem.weight'].repeat(4,1,1,1));self.stem.bias.copy_(state['stem.bias'].repeat(4))
   for name in ['weight','bias','running_mean','running_var']:getattr(self.bn,name).copy_(state['bn.'+name].repeat(4))
   self.bn.num_batches_tracked.copy_(state['bn.num_batches_tracked'])
   for i in range(10):
    b=self.blocks[2*i];prefix='blocks.'+str(i)+'.';b.dw.weight.copy_(state[prefix+'dw.weight'].repeat(4,1,1));b.dw.bias.copy_(state[prefix+'dw.bias'].repeat(4))
    for name in ['expand','fc1','fc2','project']:
     layer=getattr(b,name);layer.weight.copy_(state[prefix+name+'.weight'].repeat(4,4,1)/4);layer.bias.copy_(state[prefix+name+'.bias'].repeat(4))
   n=state['head.weight'].shape[0];assert n<=self.head.out_channels
   self.head.weight[:n].copy_(state['head.weight'].repeat(1,4,1)/4);self.head.bias[:n].copy_(state['head.bias'])
   if n<self.head.out_channels:self.head.weight[n:].zero_();self.head.bias[n:].fill_(-8.)