"""
AI-Generated Image Detector - Inference Script
Detects whether an image is real or AI-generated using frequency analysis + deep learning.

Usage:
    python inference.py --image path/to/image.jpg
    python inference.py --image https://example.com/image.png
    python inference.py --image_dir path/to/folder/
"""
import os
import io
import math
import argparse
import json
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
from pathlib import Path


# Import model architecture from train.py
from train import FrequencyAwareDetector


def load_model(model_dir=".", device="cuda" if torch.cuda.is_available() else "cpu"):
    """Load trained FrequencyAwareDetector model."""
    config_path = os.path.join(model_dir, "detector_config.json")
    weights_path = os.path.join(model_dir, "model_state_dict.pt")
    
    if os.path.exists(config_path):
        with open(config_path) as f:
            config = json.load(f)
    else:
        config = {
            "backbone_name": "microsoft/swinv2-tiny-patch4-window8-256",
            "num_labels": 2, "dct_patch_size": 32,
            "num_freq_bands": 8, "fft_bins": 32,
        }
    
    model = FrequencyAwareDetector(
        backbone_name=config["backbone_name"],
        num_labels=config["num_labels"],
        dct_patch_size=config["dct_patch_size"],
        num_freq_bands=config["num_freq_bands"],
        fft_bins=config["fft_bins"],
    )
    
    if os.path.exists(weights_path):
        state_dict = torch.load(weights_path, map_location=device)
        model.load_state_dict(state_dict)
        print(f"✓ Loaded weights from {weights_path}")
    else:
        print("⚠ No weights found, using random initialization")
    
    model.to(device)
    model.eval()
    return model, config


def get_transform(image_size=256):
    """Standard evaluation transform."""
    return Compose([
        Resize((image_size + 32, image_size + 32)),
        CenterCrop((image_size, image_size)),
        ToTensor(),
        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])


def predict_single(model, image_path, transform, device="cpu"):
    """Predict whether a single image is real or AI-generated."""
    if image_path.startswith("http"):
        import requests
        response = requests.get(image_path)
        img = Image.open(io.BytesIO(response.content)).convert("RGB")
    else:
        img = Image.open(image_path).convert("RGB")
    
    pixel_values = transform(img).unsqueeze(0).to(device)
    
    with torch.no_grad():
        output = model(pixel_values=pixel_values)
        logits = output["logits"]
        probs = torch.softmax(logits, dim=1)
        pred = probs.argmax(dim=1).item()
        confidence = probs[0][pred].item()
    
    labels = {0: "Real", 1: "AI-Generated"}
    return {
        "prediction": labels[pred],
        "confidence": confidence,
        "real_probability": probs[0][0].item(),
        "ai_generated_probability": probs[0][1].item(),
    }


def main():
    parser = argparse.ArgumentParser(description="Detect AI-generated images")
    parser.add_argument("--image", type=str, help="Path or URL to single image")
    parser.add_argument("--image_dir", type=str, help="Directory of images to analyze")
    parser.add_argument("--model_dir", type=str, default=".", help="Directory containing model weights")
    parser.add_argument("--image_size", type=int, default=256)
    parser.add_argument("--device", type=str, default="auto")
    args = parser.parse_args()
    
    if args.device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    else:
        device = args.device
    
    print(f"Device: {device}")
    model, config = load_model(args.model_dir, device)
    transform = get_transform(args.image_size)
    
    if args.image:
        result = predict_single(model, args.image, transform, device)
        print(f"\n{'='*50}")
        print(f"Image: {args.image}")
        print(f"Prediction: {result['prediction']}")
        print(f"Confidence: {result['confidence']:.2%}")
        print(f"  Real probability: {result['real_probability']:.4f}")
        print(f"  AI-generated probability: {result['ai_generated_probability']:.4f}")
        print(f"{'='*50}")
    
    elif args.image_dir:
        extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.tiff'}
        image_files = [
            f for f in Path(args.image_dir).iterdir()
            if f.suffix.lower() in extensions
        ]
        
        print(f"\nAnalyzing {len(image_files)} images from {args.image_dir}...\n")
        
        results = []
        for img_path in sorted(image_files):
            try:
                result = predict_single(model, str(img_path), transform, device)
                results.append(result)
                status = "🤖" if result["prediction"] == "AI-Generated" else "📷"
                print(f"  {status} {img_path.name}: {result['prediction']} ({result['confidence']:.1%})")
            except Exception as e:
                print(f"  ❌ {img_path.name}: Error - {e}")
        
        real_count = sum(1 for r in results if r["prediction"] == "Real")
        ai_count = sum(1 for r in results if r["prediction"] == "AI-Generated")
        print(f"\nSummary: {real_count} Real, {ai_count} AI-Generated out of {len(results)} images")
    
    else:
        parser.print_help()


if __name__ == "__main__":
    main()