import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import ImageClassifierOutput

from .configuration_mobilenet import MobileNetV1Config

class FP32LayerNorm2d(nn.GroupNorm):
    def __init__(self, num_channels):
        super().__init__(1, num_channels)

    def forward(self, x):
        input_dtype = x.dtype
        with torch.autocast(device_type=x.device.type, enabled=False):
            normalized = super().forward(x.float())
        return normalized.to(dtype=input_dtype)

class DepthwiseSeparableConv(nn.Module):
    def __init__(self, in_channels, out_channels, stride, dropout=0.0):
        super().__init__()
        self.use_residual = (stride == 1 and in_channels == out_channels)
        
        self.dw = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=stride, padding=1, groups=in_channels, bias=False)
        self.dw_norm = FP32LayerNorm2d(in_channels)
        self.dw_act = nn.SiLU(inplace=True)
        
        self.pw = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False)
        self.pw_norm = FP32LayerNorm2d(out_channels)
        self.pw_act = nn.SiLU(inplace=True)
        
        self.dropout = nn.Dropout2d(p=dropout) if dropout > 0 else nn.Identity()

        if self.use_residual:
            self.residual_scale = nn.Parameter(torch.tensor(0.1))

    def forward(self, x):
        identity = x
        
        out = self.dw(x)
        out = self.dw_norm(out)
        out = self.dw_act(out)
        
        out = self.pw(out)
        out = self.pw_norm(out)
        out = self.pw_act(out)
        
        out = self.dropout(out)
        
        if self.use_residual:
            return identity + self.residual_scale * out
        return out


class MobileNetV1ForImageClassification(PreTrainedModel):
    config_class = MobileNetV1Config
    
    def __init__(self, config: MobileNetV1Config):
        super().__init__(config)
        self.num_labels = config.num_classes
        
        self.stem = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1, bias=False),
            FP32LayerNorm2d(32),
            nn.SiLU(inplace=True)
        )
        
        self.blocks = nn.Sequential(
            DepthwiseSeparableConv(32, 64, stride=1, dropout=config.block_dropout),
            DepthwiseSeparableConv(64, 128, stride=2, dropout=config.block_dropout),
            DepthwiseSeparableConv(128, 128, stride=1, dropout=config.block_dropout),
            DepthwiseSeparableConv(128, 256, stride=2, dropout=config.block_dropout),
            DepthwiseSeparableConv(256, 256, stride=1, dropout=config.block_dropout),
            DepthwiseSeparableConv(256, 512, stride=2, dropout=config.block_dropout),
            DepthwiseSeparableConv(512, 512, stride=1, dropout=config.block_dropout),
            DepthwiseSeparableConv(512, 512, stride=1, dropout=config.block_dropout),
            DepthwiseSeparableConv(512, 512, stride=1, dropout=config.block_dropout),
            DepthwiseSeparableConv(512, 512, stride=1, dropout=config.block_dropout),
            DepthwiseSeparableConv(512, 512, stride=1, dropout=config.block_dropout),
            DepthwiseSeparableConv(512, 1024, stride=2, dropout=config.block_dropout),
            DepthwiseSeparableConv(1024, 1024, stride=1, dropout=config.block_dropout),
        )
        
        self.gap = nn.AdaptiveAvgPool2d((1, 1))
        self.dropout = nn.Dropout(p=config.final_dropout)
        self.classifier = nn.Linear(1024, config.num_classes)

        self.post_init()

    def _init_weights(self, module):
        if isinstance(module, nn.Conv2d):
            nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.GroupNorm):
            nn.init.ones_(module.weight)
            nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.001)
            if module.bias is not None:
                nn.init.zeros_(module.bias)

    def forward(self, pixel_values=None, labels=None, return_dict=None):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        x = self.stem(pixel_values)
        x = self.blocks(x)
        x = self.gap(x)
        x = torch.flatten(x, 1)
        x = self.dropout(x)
        logits = self.classifier(x)

        loss = None
        if labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))

        if not return_dict:
            output = (logits,)
            return ((loss,) + output) if loss is not None else output

        return ImageClassifierOutput(
            loss=loss,
            logits=logits,
        )