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


"""Custom DenseNet Backbone"""
class DenseBlock(nn.Module):
    """
    Basic DenseNet block
    """
    def __init__(self, in_channels, growth_rate, num_layers):
        super(DenseBlock, self).__init__()
        self.layers = nn.ModuleList()
        for i in range(num_layers):
            self.layers.append(self._make_layer(in_channels + i * growth_rate, growth_rate))
    
    def _make_layer(self, in_channels, growth_rate):
        layer = nn.Sequential(
            nn.BatchNorm2d(in_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels, 4 * growth_rate, kernel_size=1, bias=False),
            nn.BatchNorm2d(4 * growth_rate),
            nn.ReLU(inplace=True),
            nn.Conv2d(4 * growth_rate, growth_rate, kernel_size=3, padding=1, bias=False)
        )
        return layer
    
    def forward(self, x):
        features = [x]
        for layer in self.layers:
            new_feature = layer(torch.cat(features, dim=1))
            features.append(new_feature)
        return torch.cat(features, dim=1)
    

class TransitionLayer(nn.Module):
    """
    Transition layer between DenseBlocks
    """
    def __init__(self, in_channels, out_channels):
        super(TransitionLayer, self).__init__()
        self.transition = nn.Sequential(
            nn.BatchNorm2d(in_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
            nn.AvgPool2d(kernel_size=2, stride=2)
        )
    
    def forward(self, x):
        return self.transition(x)
    
    
class DenseNetBackbone(nn.Module):
    """
    DenseNet backbone for CAN
    """
    def __init__(self, growth_rate=32, block_config=(6, 12, 24, 16), num_init_features=64):
        super(DenseNetBackbone, self).__init__()
        
        # Initial layer
        self.features = nn.Sequential(
            nn.Conv2d(1, num_init_features, kernel_size=7, stride=2, padding=3, bias=False),
            nn.BatchNorm2d(num_init_features),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )
        
        # DenseBlocks
        num_features = num_init_features
        for i, num_layers in enumerate(block_config):
            block = DenseBlock(num_features, growth_rate, num_layers)
            self.features.add_module(f'denseblock{i+1}', block)
            num_features = num_features + growth_rate * num_layers
            if i != len(block_config) - 1:
                trans = TransitionLayer(num_features, num_features // 2)
                self.features.add_module(f'transition{i+1}', trans)
                num_features = num_features // 2
        
        # Final processing
        self.features.add_module('norm5', nn.BatchNorm2d(num_features))
        self.features.add_module('relu5', nn.ReLU(inplace=True))
        
        self.out_channels = num_features  # 684 (with default configuration)
    
    def forward(self, x):
        return self.features(x)
    

"""Pretrained DenseNet"""
class DenseNetFeatureExtractor(nn.Module):
    def __init__(self, densenet_model, out_channels=684):
        super().__init__()
        # Change input conv to 1 channel
        self.conv0 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        # Copy pretrained weights (average over RGB channels)
        self.conv0.weight.data = densenet_model.features.conv0.weight.data.mean(dim=1, keepdim=True)
        self.features = densenet_model.features
        self.out_channels = out_channels
        # Add a 1x1 conv to match your expected output channels if needed
        self.final_conv = nn.Conv2d(1024, out_channels, kernel_size=1)
        self.final_bn = nn.BatchNorm2d(out_channels)
        self.final_relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv0(x)
        x = self.features.norm0(x)
        x = self.features.relu0(x)
        x = self.features.pool0(x)
        x = self.features.denseblock1(x)
        x = self.features.transition1(x)
        x = self.features.denseblock2(x)
        x = self.features.transition2(x)
        x = self.features.denseblock3(x)
        x = self.features.transition3(x)
        x = self.features.denseblock4(x)
        x = self.features.norm5(x)
        x = self.final_conv(x)
        x = self.final_bn(x)
        x = self.final_relu(x)
        return x
    
    
"""Custom ResNet Backbone"""    
class BasicBlock(nn.Module):
    """
    Basic ResNet block
    """
    expansion = 1
    
    def __init__(self, in_channels, out_channels, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        
        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels * self.expansion:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels * self.expansion, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels * self.expansion)
            )
    
    def forward(self, x):
        identity = x
        
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        
        out = self.conv2(out)
        out = self.bn2(out)
        
        out += self.shortcut(identity)
        out = self.relu(out)
        
        return out
    
    
class Bottleneck(nn.Module):
    """
    Bottleneck ResNet block
    """
    expansion = 4
    
    def __init__(self, in_channels, out_channels, stride=1):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.conv3 = nn.Conv2d(out_channels, out_channels * self.expansion, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        
        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels * self.expansion:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels * self.expansion, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels * self.expansion)
            )
    
    def forward(self, x):
        identity = x
        
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        
        out = self.conv3(out)
        out = self.bn3(out)
        
        out += self.shortcut(identity)
        out = self.relu(out)
        
        return out
    

class ResNetBackbone(nn.Module):
    """
    ResNet backbone for CAN model, designed to output similar dimensions as DenseNet
    """
    def __init__(self, block_type='bottleneck', layers=[3, 4, 6, 3], num_init_features=64):
        super(ResNetBackbone, self).__init__()
        
        # Initial layer
        self.conv1 = nn.Conv2d(1, num_init_features, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(num_init_features)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        
        # Define block type
        if block_type == 'basic':
            block = BasicBlock
            expansion = 1
        elif block_type == 'bottleneck':
            block = Bottleneck
            expansion = 4
        else:
            raise ValueError(f"Unknown block type: {block_type}")
        
        # Create layers
        self.layer1 = self._make_layer(block, num_init_features, 64, layers[0], stride=1)
        self.layer2 = self._make_layer(block, 64 * expansion, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 128 * expansion, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 256 * expansion, 512, layers[3], stride=2)
        
        # Final processing to match DenseNet output channels
        self.final_conv = nn.Conv2d(512 * expansion, 684, kernel_size=1)
        self.final_bn = nn.BatchNorm2d(684)
        self.final_relu = nn.ReLU(inplace=True)
        
        self.out_channels = 684  # Match DenseNet output channels
        
        # Initialize weights
        self._initialize_weights()
    
    def _make_layer(self, block, in_channels, out_channels, num_blocks, stride):
        layers = []
        layers.append(block(in_channels, out_channels, stride))
        for _ in range(1, num_blocks):
            layers.append(block(out_channels * block.expansion, out_channels))
        return nn.Sequential(*layers)
    
    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        
        x = self.final_conv(x)
        x = self.final_bn(x)
        x = self.final_relu(x)
        
        return x

    

"""Pretrained ResNet"""
class ResNetFeatureExtractor(nn.Module):
    def __init__(self, resnet_model, out_channels=684):
        super().__init__()
        # Change input conv to 1 channel
        self.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.conv1.weight.data = resnet_model.conv1.weight.data.sum(dim=1, keepdim=True)  # average weights if needed
        self.bn1 = resnet_model.bn1
        self.relu = resnet_model.relu
        self.maxpool = resnet_model.maxpool
        self.layer1 = resnet_model.layer1
        self.layer2 = resnet_model.layer2
        self.layer3 = resnet_model.layer3
        self.layer4 = resnet_model.layer4
        # Add a 1x1 conv to match DenseNet output channels if needed
        self.final_conv = nn.Conv2d(2048, out_channels, kernel_size=1)
        self.final_bn = nn.BatchNorm2d(out_channels)
        self.final_relu = nn.ReLU(inplace=True)
        self.out_channels = out_channels

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.final_conv(x)
        x = self.final_bn(x)
        x = self.final_relu(x)
        return x
    
    
"""Channel Attention"""
class ChannelAttention(nn.Module):
    """
    Channel-wise attention mechanism
    """
    def __init__(self, in_channels, ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        
        self.fc = nn.Sequential(
            nn.Conv2d(in_channels, in_channels // ratio, kernel_size=1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels // ratio, in_channels, kernel_size=1, bias=False)
        )
        self.sigmoid = nn.Sigmoid()
    
    def forward(self, x):
        avg_out = self.fc(self.avg_pool(x))
        max_out = self.fc(self.max_pool(x))
        out = avg_out + max_out
        return self.sigmoid(out)
    
    
"""Multi-scale Couting Module"""
class MSCM(nn.Module):
    """
    Multi-Scale Counting Module
    """
    def __init__(self, in_channels, num_classes):
        super(MSCM, self).__init__()
        
        # Branch 1: 3x3 kernel
        self.branch1 = nn.Sequential(
            nn.Conv2d(in_channels, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Dropout2d(p=0.2)
        )
        self.attention1 = ChannelAttention(256)
        
        # Branch 2: 5x5 kernel
        self.branch2 = nn.Sequential(
            nn.Conv2d(in_channels, 256, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.Dropout2d(p=0.2)
        )
        self.attention2 = ChannelAttention(256)
        
        # 1x1 Conv layer to reduce channels and create counting map
        self.conv_reduce = nn.Conv2d(512, num_classes, kernel_size=1)
        self.sigmoid = nn.Sigmoid()
    
    def forward(self, x):
        # Process branch 1
        out1 = self.branch1(x)
        out1 = out1 * self.attention1(out1)
        
        # Process branch 2
        out2 = self.branch2(x)
        out2 = out2 * self.attention2(out2)
        
        # Concatenate features from both branches
        concat_features = torch.cat([out1, out2], dim=1)  # Shape: B x 512 x H x W
        
        # Create counting map
        count_map = self.sigmoid(self.conv_reduce(concat_features))  # Shape: B x C x H x W
        
        # Apply sum-pooling to create 1D counting vector
        # Sum over the entire feature map along height and width
        count_vector = torch.sum(count_map, dim=(2, 3))  # Shape: B x C
        
        return count_map, count_vector
    

"""Positional Encoding"""
class PositionalEncoding(nn.Module):
    """
    Positional encoding for attention decoder
    """
    def __init__(self, d_model, max_seq_len=1024):
        super(PositionalEncoding, self).__init__()
        self.d_model = d_model
        
        # Create positional encoding matrix
        pe = torch.zeros(max_seq_len, d_model)
        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        # x shape: B x H x W x d_model
        b, h, w, _ = x.shape
        
        # Ensure we have enough positional encodings for the feature map size
        if h*w > self.pe.size(0):       #type: ignore
            # Dynamically extend positional encodings if needed
            device = self.pe.device
            extended_pe = torch.zeros(h*w, self.d_model, device=device)                                                             #type: ignore
            position = torch.arange(0, h*w, dtype=torch.float, device=device).unsqueeze(1)                                          #type: ignore
            div_term = torch.exp(torch.arange(0, self.d_model, 2, device=device).float() * (-math.log(10000.0) / self.d_model))     #type: ignore
            
            extended_pe[:, 0::2] = torch.sin(position * div_term)
            extended_pe[:, 1::2] = torch.cos(position * div_term)
            
            pos_encoding = extended_pe.view(h, w, -1)
        else:
            # Use pre-computed positional encodings
            pos_encoding = self.pe[:h*w].view(h, w, -1)                                                                             #type: ignore
              
        pos_encoding = pos_encoding.unsqueeze(0).expand(b, -1, -1, -1)  # B x H x W x d_model
        return pos_encoding
    
    
"""Counting-combined Attentional Decoder"""
class CCAD(nn.Module):
    """
    Counting-Combined Attentional Decoder
    """
    def __init__(self, input_channels, hidden_size, embedding_dim, num_classes, use_coverage=True):
        super(CCAD, self).__init__()
        
        self.hidden_size = hidden_size
        self.embedding_dim = embedding_dim
        self.use_coverage = use_coverage
        
        # Input layer to reduce feature map
        self.feature_proj = nn.Conv2d(input_channels, hidden_size * 2, kernel_size=1)
        
        # Positional encoding
        self.pos_encoder = PositionalEncoding(hidden_size * 2)
        
        # Embedding layer for output symbols
        self.embedding = nn.Embedding(num_classes, embedding_dim)
        
        # GRU cell
        self.gru = nn.GRUCell(embedding_dim + hidden_size + num_classes, hidden_size)
        
        # Attention
        self.attention_w = nn.Linear(hidden_size * 2, hidden_size)
        self.attention_v = nn.Linear(hidden_size, 1)
        if use_coverage:
            self.coverage_proj = nn.Linear(1, hidden_size)
        
        # Output layer
        self.out = nn.Linear(hidden_size + hidden_size + num_classes, num_classes)
        self.dropout = nn.Dropout(p=0.3)
        
    def forward(self, feature_map, count_vector, target=None, teacher_forcing_ratio=0.5, max_len=200):
        batch_size = feature_map.size(0)
        device = feature_map.device
        
        # Transform feature map
        projected_features = self.feature_proj(feature_map)  # B x 2*hidden_size x H x W
        H, W = projected_features.size(2), projected_features.size(3)
        
        # Reshape feature map to B x H*W x 2*hidden_size
        projected_features = projected_features.permute(0, 2, 3, 1).contiguous()  # B x H x W x 2*hidden_size
        
        # Add positional encoding
        pos_encoding = self.pos_encoder(projected_features)  # B x H x W x 2*hidden_size
        projected_features = projected_features + pos_encoding
        
        # Reshape for attention processing
        projected_features = projected_features.view(batch_size, H*W, -1)  # B x H*W x 2*hidden_size
        
        # Initialize initial hidden state
        h_t = torch.zeros(batch_size, self.hidden_size, device=device)
        
        # Initialize coverage attention if used
        if self.use_coverage:
            coverage = torch.zeros(batch_size, H*W, 1, device=device)
        
        # First <SOS> token
        y_t_1 = torch.ones(batch_size, dtype=torch.long, device=device)
        
        # Prepare target sequence if provided
        if target is not None:
            max_len = target.size(1)
        
        # Array to store predictions
        outputs = torch.zeros(batch_size, max_len, self.embedding.num_embeddings, device=device)
        
        for t in range(max_len):
            # Apply embedding to the previous symbol
            embedded = self.embedding(y_t_1)  # B x embedding_dim
            
            # Compute attention
            attention_input = self.attention_w(projected_features)  # B x H*W x hidden_size
            
            # Add coverage attention if used
            if self.use_coverage:
                coverage_input = self.coverage_proj(coverage.float())                        #type: ignore
                attention_input = attention_input + coverage_input
            
            # Add hidden state to attention
            h_expanded = h_t.unsqueeze(1).expand(-1, H*W, -1)  # B x H*W x hidden_size
            attention_input = torch.tanh(attention_input + h_expanded)
            
            # Compute attention weights
            e_t = self.attention_v(attention_input).squeeze(-1)  # B x H*W
            alpha_t = F.softmax(e_t, dim=1)  # B x H*W
            
            # Update coverage if used
            if self.use_coverage:
                coverage = coverage + alpha_t.unsqueeze(-1)                                  #type: ignore
            
            # Compute context vector
            alpha_t = alpha_t.unsqueeze(1)  # B x 1 x H*W
            context = torch.bmm(alpha_t, projected_features).squeeze(1)  # B x 2*hidden_size
            context = context[:, :self.hidden_size]  # Take the first half as context vector
            
            # Combine embedding, context vector, and count vector
            gru_input = torch.cat([embedded, context, count_vector], dim=1)
            
            # Update hidden state
            h_t = self.gru(gru_input, h_t)
            
            # Predict output symbol
            output = self.out(torch.cat([h_t, context, count_vector], dim=1))
            outputs[:, t] = output
            
            # Decide the next input symbol
            if target is not None and torch.rand(1).item() < teacher_forcing_ratio:
                y_t_1 = target[:, t]
            else:
                # Greedy decoding
                _, y_t_1 = output.max(1)
        
        return outputs
    

"""Full model CAN (Counting-Aware Network)"""
class CAN(nn.Module):
    """
    Counting-Aware Network for handwritten mathematical expression recognition
    """
    def __init__(self, num_classes, backbone=None, hidden_size=256, embedding_dim=256, use_coverage=True):
        super(CAN, self).__init__()
        
        # Backbone
        if backbone is None:
            self.backbone = DenseNetBackbone()
        else:
            self.backbone = backbone
        backbone_channels = self.backbone.out_channels
        
        # Multi-Scale Counting Module
        self.mscm = MSCM(backbone_channels, num_classes)
        
        # Counting-Combined Attentional Decoder
        self.decoder = CCAD(
            input_channels=backbone_channels,
            hidden_size=hidden_size,
            embedding_dim=embedding_dim,
            num_classes=num_classes,
            use_coverage=use_coverage
        )
        
        # Save parameters for later use
        self.hidden_size = hidden_size
        self.embedding_dim = embedding_dim
        self.num_classes = num_classes
        self.use_coverage = use_coverage
    
    def init_hidden_state(self, visual_features):
        """
        Initialize hidden state and cell state for LSTM
        
        Args:
            visual_features: Visual features from backbone
            
        Returns:
            h, c: Initial hidden and cell states
        """
        batch_size = visual_features.size(0)
        device = visual_features.device
        
        # Initialize hidden state with zeros
        h = torch.zeros(1, batch_size, self.hidden_size, device=device)
        c = torch.zeros(1, batch_size, self.hidden_size, device=device)
        
        return h, c

    def forward(self, x, target=None, teacher_forcing_ratio=0.5):
        # Extract features from backbone
        features = self.backbone(x)
        
        # Compute count map and count vector from MSCM
        count_map, count_vector = self.mscm(features)
        
        # Decode with CCAD
        outputs = self.decoder(features, count_vector, target, teacher_forcing_ratio)
        
        return outputs, count_vector
    
    def calculate_loss(self, outputs, targets, count_vectors, count_targets, lambda_count=0.01):
        """
        Compute the combined loss function for CAN
        
        Args:
            outputs: Predicted output sequence from decoder
            targets: Actual target sequence
            count_vectors: Predicted count vector
            count_targets: Actual target count vector
            lambda_count: Weight for counting loss
        
        Returns:
            Total loss: L = L_cls + λ * L_counting
        """
        # Loss for decoder (cross entropy)
        L_cls = F.cross_entropy(outputs.view(-1, outputs.size(-1)), targets.view(-1))
        
        # Loss for counting (MSE)
        L_counting = F.smooth_l1_loss(count_vectors / self.num_classes, count_targets / self.num_classes)
        
        # Total loss
        total_loss = L_cls + lambda_count * L_counting
        
        return total_loss, L_cls, L_counting

    def recognize(self, images, max_length=150, start_token=None, end_token=None, beam_width=5):
        """
        Recognize the handwritten expression using beam search (batch_size=1 only).
        
        Args:
            images: Input image tensor, shape (1, channels, height, width)
            max_length: Maximum length of the output sequence
            start_token: Start token index
            end_token: End token index
            beam_width: Beam width for beam search
            
        Returns:
            best_sequence: List of token indices
            attention_weights: List of attention weights for visualization
        """
        if images.size(0) != 1:
            raise ValueError("Beam search is implemented only for batch_size=1")
        
        device = images.device
        
        # Encode the image
        visual_features = self.backbone(images)
        
        # Get count vector
        _, count_vector = self.mscm(visual_features)
        
        # Prepare feature map for decoder
        projected_features = self.decoder.feature_proj(visual_features)  # (1, 2*hidden_size, H, W)
        H, W = projected_features.size(2), projected_features.size(3)
        projected_features = projected_features.permute(0, 2, 3, 1).contiguous()  # (1, H, W, 2*hidden_size)
        pos_encoding = self.decoder.pos_encoder(projected_features)  # (1, H, W, 2*hidden_size)
        projected_features = projected_features + pos_encoding  # (1, H, W, 2*hidden_size)
        projected_features = projected_features.view(1, H*W, -1)  # (1, H*W, 2*hidden_size)
        
        # Initialize beams
        beam_sequences = [torch.tensor([start_token], device=device)] * beam_width  # List of (seq_len) tensors
        beam_scores = torch.zeros(beam_width, device=device)  # (beam_width)
        h_t = torch.zeros(beam_width, self.hidden_size, device=device)  # (beam_width, hidden_size)
        if self.use_coverage:
            coverage = torch.zeros(beam_width, H*W, device=device)  # (beam_width, H*W)
        
        all_attention_weights = []
        
        for step in range(max_length):
            # Get current tokens for all beams
            current_tokens = torch.tensor([seq[-1] for seq in beam_sequences], device=device)  # (beam_width)
            
            # Apply embedding
            embedded = self.decoder.embedding(current_tokens)  # (beam_width, embedding_dim)
            
            # Compute attention for each beam
            attention_input = self.decoder.attention_w(projected_features.expand(beam_width, -1, -1))  # (beam_width, H*W, hidden_size)
            if self.use_coverage:
                coverage_input = self.decoder.coverage_proj(coverage.unsqueeze(-1))  # (beam_width, H*W, hidden_size)            #type: ignore
                attention_input = attention_input + coverage_input
            
            h_expanded = h_t.unsqueeze(1).expand(-1, H*W, -1)  # (beam_width, H*W, hidden_size)
            attention_input = torch.tanh(attention_input + h_expanded)
            
            e_t = self.decoder.attention_v(attention_input).squeeze(-1)  # (beam_width, H*W)
            alpha_t = F.softmax(e_t, dim=1)  # (beam_width, H*W)
            
            all_attention_weights.append(alpha_t.detach())
            
            if self.use_coverage:
                coverage = coverage + alpha_t            #type: ignore
            
            context = torch.bmm(alpha_t.unsqueeze(1), projected_features.expand(beam_width, -1, -1)).squeeze(1)  # (beam_width, 2*hidden_size)
            context = context[:, :self.hidden_size]  # (beam_width, hidden_size)
            
            # Expand count_vector to (beam_width, num_classes)
            count_vector_expanded = count_vector.expand(beam_width, -1)  # (beam_width, num_classes)
            
            gru_input = torch.cat([embedded, context, count_vector_expanded], dim=1)  # (beam_width, embedding_dim + hidden_size + num_classes)
            
            h_t = self.decoder.gru(gru_input, h_t)  # (beam_width, hidden_size)
            
            output = self.decoder.out(torch.cat([h_t, context, count_vector_expanded], dim=1))  # (beam_width, num_classes)
            scores = F.log_softmax(output, dim=1)  # (beam_width, num_classes)
            
            # Compute new scores for all beam-token combinations
            new_beam_scores = beam_scores.unsqueeze(1) + scores  # (beam_width, num_classes)
            new_beam_scores_flat = new_beam_scores.view(-1)  # (beam_width * num_classes)
            
            # Select top beam_width scores and indices
            topk_scores, topk_indices = new_beam_scores_flat.topk(beam_width)
            
            # Determine which beam and token each top score corresponds to
            beam_indices = topk_indices // self.num_classes  # (beam_width)
            token_indices = topk_indices % self.num_classes  # (beam_width)
            
            # Create new beam sequences and states
            new_beam_sequences = []
            new_h_t = []
            if self.use_coverage:
                new_coverage = []
            for i in range(beam_width):
                prev_beam_idx = beam_indices[i].item()
                token = token_indices[i].item()
                new_seq = torch.cat([beam_sequences[prev_beam_idx], torch.tensor([token], device=device)])           #type: ignore
                new_beam_sequences.append(new_seq)
                new_h_t.append(h_t[prev_beam_idx])
                if self.use_coverage:
                    new_coverage.append(coverage[prev_beam_idx])           #type: ignore
            
            # Update beams
            beam_sequences = new_beam_sequences
            beam_scores = topk_scores
            h_t = torch.stack(new_h_t)
            if self.use_coverage:
                coverage = torch.stack(new_coverage)           #type: ignore
        
        # Select the sequence with the highest score
        best_idx = beam_scores.argmax()
        best_sequence = beam_sequences[best_idx].tolist()
        
        # Remove <start> and stop at <end>
        if best_sequence[0] == start_token:
            best_sequence = best_sequence[1:]
        if end_token in best_sequence:
            end_idx = best_sequence.index(end_token)
            best_sequence = best_sequence[:end_idx]
        
        return best_sequence, all_attention_weights


def create_can_model(num_classes, hidden_size=256, embedding_dim=256, use_coverage=True, pretrained_backbone=False, backbone_type='densenet'):
    """
    Create CAN model with either DenseNet or ResNet backbone
    
    Args:
        num_classes: Number of symbol classes
        pretrained_backbone: Whether to use a pretrained backbone
        backbone_type: Type of backbone to use ('densenet' or 'resnet')
    
    Returns:
        CAN model
    """
    # Create backbone
    if backbone_type == 'densenet':
        if pretrained_backbone:
            densenet = models.densenet121(pretrained=True)
            backbone = DenseNetFeatureExtractor(densenet, out_channels=684)
        else:
            backbone = DenseNetBackbone()
    elif backbone_type == 'resnet':
        if pretrained_backbone:
            resnet = models.resnet50(pretrained=True)
            backbone = ResNetFeatureExtractor(resnet, out_channels=684)
        else:
            backbone = ResNetBackbone(block_type='bottleneck', layers=[3, 4, 6, 3])
    else:
        raise ValueError(f"Unknown backbone type: {backbone_type}")
    
    # Create model
    model = CAN(
        num_classes=num_classes,
        backbone=backbone,
        hidden_size=hidden_size,
        embedding_dim=embedding_dim,
        use_coverage=use_coverage
    )
    
    return model


# # Example usage
# if __name__ == "__main__":
#     # Create CAN model with 101 symbol classes (example)
#     num_classes = 101  # Number of symbol classes + special tokens like <SOS>, <EOS>
#     model = create_can_model(num_classes)
    
#     # Create dummy input data
#     batch_size = 4
#     input_image = torch.randn(batch_size, 1, 128, 384)  # B x C x H x W
#     target = torch.randint(0, num_classes, (batch_size, 50))  # B x max_len
    
#     # Forward pass
#     outputs, count_vectors = model(input_image, target)
    
#     # Print output shapes
#     print(f"Outputs shape: {outputs.shape}")  # B x max_len x num_classes
#     print(f"Count vectors shape: {count_vectors.shape}")  # B x num_classes