"""
utils/model_loader.py

Loads all 4 trained models from the models/ directory.
Models are expected to be PyTorch .pt or .pth files.
"""

import torch
import torchvision.models as models
import os


def build_model(num_classes: int) -> torch.nn.Module:
    """
    Build a ResNet-50 model with a custom final classification layer.
    
    We use ResNet-50 as the backbone — a powerful, proven architecture.
    The final fully-connected layer is replaced to match our number of classes.

    Args:
        num_classes: How many categories this model predicts.

    Returns:
        A PyTorch model (not yet loaded with weights).
    """
    model = models.resnet50(weights=None)  # No pretrained weights — we'll load ours
    # Replace the final layer to match our class count
    in_features = model.fc.in_features
    model.fc = torch.nn.Linear(in_features, num_classes)
    return model


def load_model(path: str, num_classes: int, device: torch.device) -> torch.nn.Module:
    """
    Load a trained model from a .pt or .pth file.

    Args:
        path: File path to the saved model weights.
        num_classes: Number of output classes for this model.
        device: CPU or GPU.

    Returns:
        Loaded model ready for inference.
    """
    if not os.path.exists(path):
        raise FileNotFoundError(
            f"Model file not found: '{path}'\n"
            f"Make sure your trained model is saved at this location."
        )

    model = build_model(num_classes)
    # Load the saved weights into the model
    state_dict = torch.load(path, map_location=device)
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()  # Set to evaluation mode (disables dropout, batchnorm training behavior)
    return model


def load_all_models() -> dict:
    """
    Load all 4 models and return them in a dictionary.

    Model files expected in models/ folder:
        - tree_vs_nontree.pt  (2 classes: Tree, Non-Tree)
        - species.pt          (2 classes: Mango, White Gum)
        - mango_stage.pt      (4 classes: Seedling, Sapling, Mature, Overmature)
        - gum_stage.pt        (4 classes: Seedling, Sapling, Mature, Overmature)

    Returns:
        Dictionary with keys: tree_vs_nontree, species, mango_stage, gum_stage
    """
    # Use GPU if available, otherwise CPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Define model configurations: (file_path, number_of_classes)
    model_configs = {
        "tree_vs_nontree": ("models/tree_vs_nontree.pt", 2),
        "species":         ("models/species.pt",         2),
        "mango_stage":     ("models/mango_stage.pt",     4),
        "gum_stage":       ("models/gum_stage.pt",       4),
    }

    loaded_models = {}
    for name, (path, num_classes) in model_configs.items():
        loaded_models[name] = load_model(path, num_classes, device)

    return loaded_models