import torch
import torch.nn as nn
import torch.nn.functional as F
import timm

# --- PART 1: The Stem Class ---
class InceptionStem(nn.Module):
    def __init__(self, in_channels=3, out_channels=64):
        super().__init__()
        self.branch1 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels//2, kernel_size=1, stride=1, padding=0),
            nn.ReLU(inplace=True)
        )
        self.branch3 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels//2, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True)
        )
        self.pool_branch = nn.Sequential(
            nn.MaxPool2d(3, stride=1, padding=1),
            nn.Conv2d(in_channels, out_channels//2, kernel_size=1),
            nn.ReLU(inplace=True)
        )
        self.project = nn.Conv2d(out_channels + out_channels//2, out_channels, kernel_size=1)
        self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        b1 = self.branch1(x)
        b3 = self.branch3(x)
        p = self.pool_branch(x)
        cat = torch.cat([b1, b3, p], dim=1)
        out = F.relu(self.bn(self.project(cat)))
        return out

# --- PART 2: The Main Model Class ---
class InceptionViT(nn.Module):
    def __init__(self, vit_model_name='vit_base_patch16_224', pretrained=False, stem_out_channels=64, num_classes=2, dropout=0.3):
        # NOTE: We set pretrained=False here because we are loading YOUR weights,
        # so we don't need to download the ImageNet weights again.
        super().__init__()
        self.stem = InceptionStem(in_channels=3, out_channels=stem_out_channels)

        self.vit = timm.create_model(vit_model_name, pretrained=False, num_classes=0, global_pool='')
        self.embed_dim = self.vit.num_features
        self.stem_pool = nn.AdaptiveAvgPool2d(1)
        self.stem_proj = nn.Linear(stem_out_channels, self.embed_dim)

        self.classifier = nn.Sequential(
            nn.LayerNorm(self.embed_dim),
            nn.Dropout(dropout),
            nn.Linear(self.embed_dim, num_classes)
        )

    def forward(self, x):
        stem_feat_map = self.stem(x)
        
        stem_vec = self.stem_pool(stem_feat_map).view(x.size(0), -1)
        stem_emb = self.stem_proj(stem_vec)

        B = x.size(0)
        x_vit = self.vit.patch_embed(x)
        cls_tokens = self.vit.cls_token.expand(B, -1, -1)
        x_vit = torch.cat((cls_tokens, x_vit), dim=1)
        x_vit = x_vit + self.vit.pos_embed[:, : x_vit.size(1), :].to(x.device)
        x_vit = self.vit.pos_drop(x_vit)

        for blk in self.vit.blocks:
            x_vit = blk(x_vit)
        x_vit = self.vit.norm(x_vit)

        x_vit[:, 0, :] = x_vit[:, 0, :] + stem_emb

        cls_emb = x_vit[:, 0, :]
        out = self.classifier(cls_emb)
        return out

# --- PART 3: The Loader Function ---
def load_model(model_path='InceptionViT_best_model.pth'):
    print(f"Loading InceptionViT from {model_path}...")
    
    # 1. Initialize the empty model structure
    # CRITICAL: We must match the arguments you used in training!
    # Looking at your code: num_classes=?, dropout=0.4
    # Since I don't know your exact num_classes, I will inspect the weights to find it dynamically.
    
    # Load weights first to check dimensions
    state_dict = torch.load(model_path, map_location='cpu')
    
    # Auto-detect number of classes from the last layer weight
    # The key is usually "classifier.2.weight" based on your Sequential block
    if "classifier.2.weight" in state_dict:
        num_classes = state_dict["classifier.2.weight"].shape[0]
        print(f"Auto-detected {num_classes} classes.")
    else:
        num_classes = 2 # Default fallback
        print("Could not auto-detect classes, defaulting to 2.")

    # Create the model
    model = InceptionViT(num_classes=num_classes, dropout=0.4)
    
    # Load the weights
    model.load_state_dict(state_dict)
    model.eval()
    
    return model