#!/usr/bin/env python3
"""
SE-AlexNet: Unified Inference Script
=====================================

Minimal inference demo for loading any SE-AlexNet variant and running a forward pass.
Copy-paste ready. Under 20 lines of core logic.

Usage:
  >>> from inference import SEModelPipeline
  >>> pipe = SEModelPipeline('se-location3/facebased/squeeze-32')
  >>> probs = pipe.predict('path/to/face_image.jpg')  # → torch.Tensor, shape (1, num_classes)

CLI:
  python inference.py se-location3/facebased/squeeze-32 --image face.jpg
"""

import os
import sys
import json
import torch
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
from typing import Union, Optional

# Import model classes from local modeling.py
from modeling import (
    AlexNet,
    AlexNetWithSE_L1,
    AlexNetWithSE_L2,
    AlexNetWithSE_L3,
    VGG16,
    MODEL_REGISTRY,
    load_model_from_config,
)


# ── Image Preprocessing ─────────────────────────────────────────────────────
# Standard ImageNet-style preprocessing (used for all variants)
INPUT_SIZE = 224

_preprocess = transforms.Compose([
    transforms.Resize((INPUT_SIZE, INPUT_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])


# ── Model Pipeline ──────────────────────────────────────────────────────────
class SEModelPipeline:
    """
    Minimal pipeline to load and run inference with any SE-AlexNet variant.

    Parameters
    ----------
    model_dir : str
        Path relative to the repo root, e.g. 'se-location3/facebased/squeeze-32'.
    device : str
        Torch device, defaults to 'cpu'.
    """
    def __init__(self, model_dir: str, device: str = 'cpu'):
        # Resolve repo root (parent of this file)
        repo_root = os.path.dirname(os.path.abspath(__file__))
        full_dir = os.path.join(repo_root, model_dir)

        # Load config
        config_path = os.path.join(full_dir, 'config.json')
        if not os.path.exists(config_path):
            raise FileNotFoundError(f'config.json not found in {full_dir}')
        with open(config_path, 'r') as f:
            self.config = json.load(f)

        # Find weights (prefer .safetensors, fallback to .pth)
        safetensors_path = os.path.join(full_dir, 'model.safetensors')
        pth_path = os.path.join(full_dir, 'model.pth')

        weights_path = None
        if os.path.exists(safetensors_path):
            weights_path = safetensors_path
        elif os.path.exists(pth_path):
            weights_path = pth_path
        else:
            raise FileNotFoundError(
                f'No weights found (model.safetensors or model.pth) in {full_dir}'
            )

        # Build model and load weights
        self.device = device
        self.model = load_model_from_config(self.config, weights_path, device)
        self.model.eval()

        # Store metadata
        self.num_classes = self.config['num_classes']
        self.model_type = self.config['model_type']
        self.pretraining = self.config['pretraining']

        print(f'Loaded {self.model_type} | {self.pretraining} | '
              f'{self.num_classes} classes → {device}')

    def predict(self, image: Union[str, Image.Image, torch.Tensor]) -> torch.Tensor:
        """
        Run inference on a single image.

        Args:
            image: Path to image, PIL Image, or preprocessed tensor (3×224×224).

        Returns:
            torch.Tensor of shape (1, num_classes) with class probabilities.
        """
        if isinstance(image, str):
            image = Image.open(image).convert('RGB')
        if isinstance(image, Image.Image):
            image = _preprocess(image)
        # At this point image is a tensor (3, H, W) or (1, 3, H, W)
        if image.dim() == 3:
            image = image.unsqueeze(0)  # add batch dim

        image = image.to(self.device)

        with torch.no_grad():
            logits = self.model(image)
            probs = F.softmax(logits, dim=1)

        return probs

    def predict_topk(self, image: Union[str, Image.Image, torch.Tensor],
                     k: int = 3):
        """Return top-k class indices and probabilities."""
        probs = self.predict(image)
        topk_probs, topk_indices = torch.topk(probs, k, dim=1)
        return topk_indices[0].cpu().numpy(), topk_probs[0].cpu().numpy()


# ── CLI ─────────────────────────────────────────────────────────────────────
if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(
        description='SE-AlexNet Inference — forward pass on a single image'
    )
    parser.add_argument('model_dir',
                        help='Model directory (e.g. se-location3/facebased/squeeze-32)')
    parser.add_argument('--image', '-i', required=True,
                        help='Path to input image')
    parser.add_argument('--device', default='cpu',
                        help='Device (cpu, cuda, mps)')
    parser.add_argument('--topk', type=int, default=3,
                        help='Show top-k predictions')

    args = parser.parse_args()

    pipe = SEModelPipeline(args.model_dir, device=args.device)
    indices, probs = pipe.predict_topk(args.image, k=args.topk)

    print(f'\nTop-{args.topk} predictions:')
    for idx, prob in zip(indices, probs):
        print(f'  Class {idx}: {prob:.4f} ({prob*100:.1f}%)')