import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models
from typing import Optional, Union

from glam_module import GLAM
from swin_module import SwinWindowAttention
from transformers import EfficientNetModel


class GLAMEfficientNetConfig:
    """Hugging Face-style configuration for 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(nn.Module):
    """EfficientNet (torchvision) + GLAM + Swin Architecture for Classification."""
    def __init__(self, config: GLAMEfficientNetConfig, glam_module_cls, swin_module_cls):
        super().__init__()

        # ✅ 1) Torchvision EfficientNet Backbone
        efficientnet = models.efficientnet_b0(pretrained=False)  # No Hugging Face!
        self.feature_extractor = EfficientNetModel.from_pretrained("google/efficientnet-b0")


        # ✅ 1x1 conv for channel adjustment
        self.conv1x1 = nn.Conv2d(1280, config.embed_dim, kernel_size=1)

        # ✅ 2) Swin Attention Block
        self.swin_attn = swin_module_cls(
            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_module_cls(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):
        """Perform forward pass."""
        # ✅ 1) EfficientNet Backbone
        feats = self.features(pixel_values)              # [B, 1280, H', W']
        feats = self.conv1x1(feats)                      # [B, embed_dim, H', W']
        B, C, H, W = feats.shape

        # ✅ 2) Transformer Branch
        x_perm = feats.permute(0, 2, 3, 1).contiguous()  # [B, H', W', C]
        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)                    # [B, C, H', W']
        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

        # ✅ 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}