import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.fft
import torchvision.models as models
from torch.utils.data import Dataset
from PIL import Image

def extract_frequency_features(image_tensor):
    # Convert RGB to Grayscale
    gray_image = 0.2989 * image_tensor[0] + 0.5870 * image_tensor[1] + 0.1140 * image_tensor[2]
    # 2D Fast Fourier Transform to look at noise grids
    fft_result = torch.fft.fft2(gray_image)
    fft_shifted = torch.fft.fftshift(fft_result)
    magnitude_spectrum = torch.log(torch.abs(fft_shifted) + 1e-8)
    return magnitude_spectrum.unsqueeze(0) 

class BluemarkDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        self.data_dir = data_dir
        self.transform = transform
        self.image_paths = []
        self.labels = []
        self.classes = {'clean': 0, 'marked': 1, 'scrubbed': 2}
        
        for class_name, label in self.classes.items():
            class_dir = os.path.join(data_dir, class_name)
            if os.path.isdir(class_dir):
                for img_name in os.listdir(class_dir):
                    if img_name.lower().endswith(('.png', '.jpg', '.jpeg')):
                        self.image_paths.append(os.path.join(class_dir, img_name))
                        self.labels.append(label)

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        image = Image.open(img_path).convert('RGB')
        label = self.labels[idx]
        if self.transform:
            image = self.transform(image)
        freq_features = extract_frequency_features(image)
        return freq_features, label

class ResNet50FrequencyDetector(nn.Module):
    def __init__(self, num_classes=3):
        super().__init__()
        # Load the massive ResNet-50 architecture
        self.resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
        
        # SURGICAL MODIFICATION: 
        # ResNet expects 3 RGB channels. We change the first layer to accept our 1-channel FFT map.
        original_conv1 = self.resnet.conv1
        self.resnet.conv1 = nn.Conv2d(
            1, # Changed from 3 to 1
            original_conv1.out_channels, 
            kernel_size=original_conv1.kernel_size, 
            stride=original_conv1.stride, 
            padding=original_conv1.padding, 
            bias=False
        )
        
        # Modify the final classification layer to output our 3 classes instead of 1000
        num_ftrs = self.resnet.fc.in_features
        self.resnet.fc = nn.Sequential(
            nn.Dropout(0.5), # Keep dropout to prevent memorizing the dataset
            nn.Linear(num_ftrs, num_classes)
        )

    def forward(self, x):
        return self.resnet(x)