import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F  # for Peak Score logic
from tqdm import tqdm
import numpy as np

# Added 'device' argument to match your pipeline.py call
def train_model(model, train_loader, epochs=10, device=None):
    # Use passed device, or fall back to auto-detection
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    model.to(device)
    optimizer = optim.AdamW(model.parameters(), lr=1e-4)
    criterion = nn.MSELoss()

    for epoch in range(epochs):
        model.train()
        total_loss = 0
        for images, _ in tqdm(train_loader, desc=f"Epoch {epoch+1}"):
            images = images.to(device)
            optimizer.zero_grad()
            
            reconstructed = model(images)
            loss = criterion(reconstructed, images)
            
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
            
        print(f"Epoch {epoch+1} complete. Avg Loss: {total_loss/len(train_loader):.6f}")

# Added 'device' argument to match your pipeline.py call
def evaluate_anomaly(model, test_loader, device=None):
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
    model.eval()
    errors = []
    labels = []
    
    print("Computing anomaly scores (Peak Score Method)...")
    with torch.no_grad():
        for images, label in tqdm(test_loader):
            images = images.to(device)
            reconstructed = model(images)
            
            # 1. Calculate Squared Error map [Batch, C, H, W]
            diff = (images - reconstructed)**2
            
            # 2. Convert to Grayscale (Mean across channels) [Batch, 1, H, W]
            diff_map = torch.mean(diff, dim=1, keepdim=True)
            
            # 3. Peak Score Logic: Apply a 15x15 blur (Average Pooling) 
            # Aggregates the 116p defect signal and reduces single-pixel noise.
            smoothed_diff = F.avg_pool2d(diff_map, kernel_size=15, stride=1, padding=7)
            
            # 4. Get the Max value for each image in the batch (Hotspot detection)
            batch_error, _ = torch.max(smoothed_diff.view(images.size(0), -1), dim=1)
            
            errors.extend(batch_error.cpu().numpy())
            labels.extend(label.numpy())
            
    return errors, labels