"""
utils/predictor.py

Handles image preprocessing and running predictions through each model.
Returns: (label, confidence_score, input_tensor)
The input_tensor is passed back so Grad-CAM can use it.
"""

import torch
import torchvision.transforms as transforms
from PIL import Image
import numpy as np


# ──────────────────────────────────────────────────────────
# IMAGE PREPROCESSING
# ──────────────────────────────────────────────────────────
# Standard ImageNet normalization values — used because most 
# deep learning models are pretrained on ImageNet.
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]

def preprocess_image(image: Image.Image) -> torch.Tensor:
    """
    Convert a PIL image into a normalized tensor ready for model input.

    Steps:
    1. Resize to 224x224 (standard input size for ResNet)
    2. Convert to PyTorch tensor (values 0–1)
    3. Normalize using ImageNet mean and std

    Args:
        image: A PIL Image (RGB).

    Returns:
        A tensor of shape [1, 3, 224, 224] — batch of 1 image.
    """
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
    ])
    # Add a batch dimension: [3, 224, 224] → [1, 3, 224, 224]
    tensor = transform(image).unsqueeze(0)
    return tensor


# ──────────────────────────────────────────────────────────
# LABEL DEFINITIONS
# ──────────────────────────────────────────────────────────
TREE_LABELS    = ["Non-Tree", "Tree"]
SPECIES_LABELS = ["Mango", "White Gum"]
STAGE_LABELS   = ["Seedling", "Sapling", "Mature", "Overmature"]


# ──────────────────────────────────────────────────────────
# INFERENCE HELPER
# ──────────────────────────────────────────────────────────
def run_inference(model: torch.nn.Module, tensor: torch.Tensor) -> tuple:
    """
    Run a forward pass through a model and return prediction results.

    Args:
        model: A loaded PyTorch model in eval mode.
        tensor: Input tensor of shape [1, 3, 224, 224].

    Returns:
        (predicted_index, confidence_score)
    """
    device = next(model.parameters()).device
    tensor = tensor.to(device)

    with torch.no_grad():  # No gradient tracking needed for inference
        outputs = model(tensor)           # Raw logits: [1, num_classes]
        probs = torch.softmax(outputs, dim=1)  # Convert to probabilities
        confidence, predicted_idx = torch.max(probs, dim=1)

    return predicted_idx.item(), confidence.item()


# ──────────────────────────────────────────────────────────
# PHASE 1: TREE vs NON-TREE
# ──────────────────────────────────────────────────────────
def predict_tree(image: Image.Image, model: torch.nn.Module) -> tuple:
    """
    Predict whether the image contains a tree or not.

    Args:
        image: PIL Image.
        model: Loaded tree vs non-tree model.

    Returns:
        (label: str, confidence: float, tensor: torch.Tensor)
        e.g. ("Tree", 0.97, tensor)
    """
    tensor = preprocess_image(image)
    idx, confidence = run_inference(model, tensor)
    label = TREE_LABELS[idx]
    return label, confidence, tensor


# ──────────────────────────────────────────────────────────
# PHASE 2: SPECIES DETECTION
# ──────────────────────────────────────────────────────────
def predict_species(image: Image.Image, model: torch.nn.Module) -> tuple:
    """
    Predict the tree species: Mango or White Gum.

    Args:
        image: PIL Image.
        model: Loaded species detection model.

    Returns:
        (label: str, confidence: float, tensor: torch.Tensor)
        e.g. ("Mango", 0.89, tensor)
    """
    tensor = preprocess_image(image)
    idx, confidence = run_inference(model, tensor)
    label = SPECIES_LABELS[idx]
    return label, confidence, tensor


# ──────────────────────────────────────────────────────────
# PHASE 3: STAGE CLASSIFICATION
# ──────────────────────────────────────────────────────────
def predict_stage(image: Image.Image, model: torch.nn.Module, species: str) -> tuple:
    """
    Predict the growth stage of the tree.
    Routes to the appropriate model based on detected species.

    Args:
        image: PIL Image.
        model: The stage model for the detected species (mango or gum).
        species: Species name (used for logging/display only here).

    Returns:
        (label: str, confidence: float, tensor: torch.Tensor)
        e.g. ("Mature", 0.92, tensor)
    """
    tensor = preprocess_image(image)
    idx, confidence = run_inference(model, tensor)
    label = STAGE_LABELS[idx]
    return label, confidence, tensor