import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from typing import List, Tuple, Optional

# Use Swin-T as feature extractor (best for defect detection in 2025)
from swin_transformer import swin_tiny_patch4_window7_224  # from previous code

class DefectImageTransformer(nn.Module):
    """
    State-of-the-Art Anomaly Detection Transformer for Industrial Images
    Works on MVTec AD, BTAD, VisA, Real-IAD, etc.
    """
    def __init__(
        self,
        pretrained_swin_path: str = "swin_tiny_patch4_window7_224.pth",
        layers_to_extract: List[int] = [1, 2, 3],  # stage 2,3,4
        memory_bank_size: int = 10000,
        feature_dim: int = 768,
        temperature: float = 0.07,
        k_neighbors: int = 9,
    ):
        super().__init__()

        # 1. Load pretrained Swin-T (ImageNet-1K or ImageNet-22K)
        self.backbone = swin_tiny_patch4_window7_224(pretrained=False)
        if pretrained_swin_path:
            state = torch.load(pretrained_swin_path, map_location='cpu')
            self.backbone.load_state_dict(state, strict=False)

        self.layers_to_extract = layers_to_extract
        self.feature_dims = [96, 192, 384, 768]  # Swin-T dims
        self.memory_bank_size = memory_bank_size
        self.k = k_neighbors
        self.temp = temperature

        # 2. Feature pyramid extraction hooks
        self.features = []
        self.hooks = []
        self._register_hooks()

        # 3. Memory bank (will be filled during fit())
        self.register_buffer("memory_bank", torch.zeros(memory_bank_size, feature_dim))
        self.register_buffer("memory_labels", torch.zeros(memory_bank_size))  # 0=good

        # 4. Lightweight decoder heads (for pixel-level map)
        self.decoder_heads = nn.ModuleList([
            nn.Conv2d(768, 384, 3, padding=1),
            nn.Conv2d(384, 192, 3, padding=1),
            nn.Conv2d(192, 1, 3, padding=1),  # anomaly score map
        ])

    def _register_hooks(self):
        def hook_fn(module, input, output, idx):
            self.features.append(output)

        # Hook into Swin stages
        for i, layer in enumerate(self.backbone.layers):
            if i in self.layers_to_extract:
                hook = layer.register_forward_hook(
                    lambda m, i, o, idx=i: hook_fn(m, i, o, idx))
                self.hooks.append(hook)

    def forward_features(self, x):
        self.features.clear()
        _ = self.backbone.forward_features(x)  # triggers hooks
        return self.features  # list of [B, C, H, W] from different stages

    def build_memory_bank(self, train_loader, device='cuda'):
        """Call once on defect-free training data"""
        self.eval()
        memory_features = []

        with torch.no_grad():
            for img, _ in train_loader:
                img = img.to(device)
                feats = self.forward_features(img)
                # Use stage-3 features (best for defects)
                feat = feats[-1]  # [B, 768, H, W]
                feat = F.adaptive_avg_pool2d(feat, (1, 1)).view(img.size(0), -1)  # global pool
                memory_features.append(feat.cpu())

        memory_features = torch.cat(memory_features, dim=0)
        # Subsample to memory_bank_size
        if len(memory_features) > self.memory_bank_size:
            idxs = torch.randperm(len(memory_features))[:self.memory_bank_size]
            memory_features = memory_features[idxs]

        self.memory_bank.data = memory_features
        print(f"Memory bank built: {self.memory_bank.shape}")

    def get_anomaly_map(self, x):
        """Returns pixel-level anomaly heatmap"""
        self.eval()
        with torch.no_grad():
            feats = self.forward_features(x)  # list of multi-scale features

            # Use stage-4 (deepest) + upsample
            deep_feat = feats[-1]  # [B, 768, 7, 7]
            x = F.interpolate(deep_feat, size=x.shape[2:], mode='bilinear', align_corners=False)

            for layer in self.decoder_heads:
                x = layer(x)
                x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)

            anomaly_map = torch.sigmoid(x.squeeze(1))  # [B, H, W]
            return anomaly_map

    def forward(self, x):
        feats = self.forward_features(x)
        test_feat = feats[-1]  # [B, 768, H, W]
        test_feat = F.adaptive_avg_pool2d(test_feat, 1).view(x.size(0), -1)  # [B, 768]

        # Nearest neighbor distance in memory bank
        dist = torch.cdist(test_feat, self.memory_bank)  # [B, memory_size]
        dist, _ = torch.topk(dist, k=self.k, dim=1, largest=False)  # smallest k distances
        score = dist.mean(dim=1)  # image-level anomaly score

        # Pixel-level map
        anomaly_map = self.get_anomaly_map(x)

        return {
            'anomaly_score': score,
            'anomaly_map': anomaly_map,
            'feature': test_feat
        }

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Example: MVTec AD bottle dataset (only good images for training)
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                                  std=[0.229, 0.224, 0.225]),
])

train_dataset = datasets.ImageFolder("mvtec/bottle/train/good", transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=False)

# Initialize model
model = DefectImageTransformer().cuda()

# Build memory bank from defect-free images
model.build_memory_bank(train_loader)

# Now test on defective images
test_img = torch.randn(1, 3, 224, 224).cuda()
result = model(test_img)
print("Anomaly Score:", result['anomaly_score'].item())
# > 0.1 = likely defect

# 1. Use larger patch size for small defects
model = DefectImageTransformer(patch_size=8)  # better for tiny cracks

# 2. Use Swin-B or Swin-L for ultra-high-res (1024px+)
# 3. Add test-time augmentation
# 4. Ensemble with Reverse Distillation (teacher-student)
# 5. Use DINOv2 or CLIP features as extra memory bank → even better