import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig
from transformers import EfficientNetModel
from typing import Optional, Union


# --------------------------------------------------
# Import your GLAM, SwinWindowAttention blocks here
# --------------------------------------------------
# from .glam_module import GLAM
# from .swin_module import SwinWindowAttention


class GLAMEfficientNetConfig(PretrainedConfig):
    """Hugging Face-style configuration for GLAM EfficientNet."""
    model_type = "glam_efficientnet"

    def __init__(self,
                 num_classes: int = 3,
                 embed_dim: int = 512,
                 num_heads: int = 8,
                 window_size: int = 7,
                 reduction_ratio: int = 8,
                 dropout: float = 0.5,
                 **kwargs):
        super().__init__(**kwargs)
        self.num_classes = num_classes
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.window_size = window_size
        self.reduction_ratio = reduction_ratio
        self.dropout = dropout


class GLAMEfficientNetForClassification(PreTrainedModel):
    """Hugging Face-style Model for EfficientNet + GLAM + Swin Architecture."""

    config_class = GLAMEfficientNetConfig

    def __init__(self, config: GLAMEfficientNetConfig):
        super().__init__(config)

        # 1) EfficientNet Backbone
        self.features = EfficientNetModel.from_pretrained("google/efficientnet-b0").features
        self.conv1x1 = nn.Conv2d(1280, config.embed_dim, kernel_size=1)

        # 2) Swin Attention Block
        self.swin_attn = SwinWindowAttention(
            embed_dim=config.embed_dim,
            window_size=config.window_size,
            num_heads=config.num_heads,
            dropout=config.dropout
        )
        self.pre_attn_norm = nn.LayerNorm(config.embed_dim)
        self.post_attn_norm = nn.LayerNorm(config.embed_dim)

        # 3) GLAM Block
        self.glam = GLAM(in_channels=config.embed_dim, reduction_ratio=config.reduction_ratio)

        # 4) Self-Adaptive Gating
        self.gate_fc = nn.Linear(config.embed_dim, 1)

        # Final classification
        self.dropout = nn.Dropout(config.dropout)
        self.classifier = nn.Linear(config.embed_dim, config.num_classes)

    def forward(self, pixel_values, labels=None, **kwargs):
        # 1) Extract EfficientNet Features
        feats = self.features(pixel_values).last_hidden_state
        feats = self.conv1x1(feats)

        B, C, H, W = feats.shape

        # 2) Transformer Branch
        x_perm = feats.permute(0, 2, 3, 1).contiguous()
        x_norm = self.pre_attn_norm(x_perm).permute(0, 3, 1, 2).contiguous()
        x_norm = self.dropout(x_norm)

        T_out = self.swin_attn(x_norm)

        T_out = self.post_attn_norm(T_out.permute(0, 2, 3, 1).contiguous())
        T_out = T_out.permute(0, 3, 1, 2).contiguous()

        # 3) GLAM Branch
        G_out = self.glam(feats)

        # 4) Self-Adaptive Gating
        gap_feats = F.adaptive_avg_pool2d(feats, (1, 1)).view(B, C)
        g = torch.sigmoid(self.gate_fc(gap_feats)).view(B, 1, 1, 1)

        F_out = g * T_out + (1 - g) * G_out

        # 5) Final Pooling + Classifier
        pooled = F.adaptive_avg_pool2d(F_out, (1, 1)).view(B, -1)
        logits = self.classifier(self.dropout(pooled))

        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits, labels)

        return {"loss": loss, "logits": logits}