utils/model_loader.py

import os
import sys
import logging
from typing import Optional
from functools import lru_cache
from pathlib import Path

# -------------------------------------------------
# Logging setup
# -------------------------------------------------
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# -------------------------------------------------
# TensorFlow import (lazy + fault-tolerant)
# -------------------------------------------------
try:
    import tensorflow as tf
    TF_AVAILABLE = True
    logger.info("TensorFlow successfully imported.")
except ImportError as e:
    TF_AVAILABLE = False
    tf = None
    logger.warning("TensorFlow import failed: %s", e)

# -------------------------------------------------
# Helper: Resolve absolute model directory
# -------------------------------------------------
def _get_models_dir() -> str:
    """Return absolute path to the models directory."""
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    models_dir = os.path.join(base_dir, "models")
    
    # Log diagnostic information
    logger.info("=" * 70)
    logger.info("MODEL DIRECTORY RESOLUTION")
    logger.info(f"Current working directory: {os.getcwd()}")
    logger.info(f"Script location: {__file__}")
    logger.info(f"Base directory: {base_dir}")
    logger.info(f"Models directory: {models_dir}")
    logger.info(f"Models directory exists: {os.path.exists(models_dir)}")
    
    if os.path.exists(models_dir):
        logger.info(f"Models directory contents: {os.listdir(models_dir)}")
    else:
        logger.error(f"❌ Models directory NOT FOUND: {models_dir}")
        logger.error(f"Current directory contents: {os.listdir('.')}")
        # Try alternative paths
        alt_paths = ["models", "./models", "/app/models"]
        for alt_path in alt_paths:
            if os.path.exists(alt_path):
                logger.info(f"✅ Found models at alternative path: {alt_path}")
                logger.info(f"Contents: {os.listdir(alt_path)}")
                models_dir = os.path.abspath(alt_path)
                break
    
    logger.info("=" * 70)
    return models_dir

# -------------------------------------------------
# Helper: Construct model path safely
# -------------------------------------------------
def _resolve_model_path(model_env_key: str, default_filename: str) -> str:
    """Resolve model path from env var or fallback to /models directory."""
    logger.info(f"Resolving model path for: {default_filename}")
    
    # Check environment variable first
    env_path = os.getenv(model_env_key)
    if env_path and os.path.exists(env_path):
        logger.info(f"✅ Using model from env var {model_env_key}: {env_path}")
        return env_path

    # Fallback to models directory
    fallback_path = os.path.join(_get_models_dir(), default_filename)
    
    if os.path.exists(fallback_path):
        file_size = os.path.getsize(fallback_path)
        logger.info(f"✅ Model file found: {fallback_path}")
        logger.info(f"📊 File size: {file_size:,} bytes ({file_size/1024:.2f} KB)")
    else:
        logger.error(f"❌ Model file missing: {fallback_path}")
        
        # Try to find the file in current directory as last resort
        if os.path.exists(default_filename):
            logger.info(f"✅ Found model in current directory: {default_filename}")
            return os.path.abspath(default_filename)
        
        models_in_cwd = os.path.join("models", default_filename)
        if os.path.exists(models_in_cwd):
            logger.info(f"✅ Found model in ./models/: {models_in_cwd}")
            return os.path.abspath(models_in_cwd)
    
    return fallback_path

# -------------------------------------------------
# Generic model loader with caching
# -------------------------------------------------
@lru_cache(maxsize=4)
def _load_model(model_path: str) -> Optional[object]:
    """Load a Keras model once and cache it."""
    logger.info("=" * 70)
    logger.info(f"📥 LOADING MODEL: {model_path}")
    
    if not TF_AVAILABLE:
        logger.error("❌ TensorFlow not available. Cannot load model: %s", model_path)
        logger.info("=" * 70)
        return None

    if not os.path.exists(model_path):
        logger.error(f"❌ Model file not found: {model_path}")
        logger.error(f"Current working directory: {os.getcwd()}")
        logger.error(f"Absolute path: {os.path.abspath(model_path)}")
        logger.info("=" * 70)
        return None

    try:
        file_size = os.path.getsize(model_path)
        logger.info(f"📊 Model file size: {file_size:,} bytes ({file_size/1024/1024:.2f} MB)")
        
        model = tf.keras.models.load_model(model_path)
        logger.info(f"✅ Model loaded successfully: {model_path}")
        logger.info(f"Model input shape: {model.input_shape}")
        logger.info(f"Model output shape: {model.output_shape}")
        logger.info("=" * 70)
        return model
    except Exception as e:
        logger.error(f"❌ Failed to load model: {model_path}")
        logger.exception("Exception details: %s", e)
        logger.info("=" * 70)
        return None

# -------------------------------------------------
# Public APIs (for routes)
# -------------------------------------------------
def load_xception_model() -> Optional[object]:
    """Load and return XceptionNet model."""
    logger.info("🔧 Loading XceptionNet model...")
    path = _resolve_model_path("MODEL_PATH_XCEPTION", "xceptionnet.keras")
    model = _load_model(path)
    if model is None:
        logger.error("❌ XceptionNet model failed to load!")
    return model

def load_mesonet_model() -> Optional[object]:
    """Load and return MesoNet model."""
    logger.info("🔧 Loading MesoNet model...")
    path = _resolve_model_path("MODEL_PATH_MESONET", "mesonet.keras")
    model = _load_model(path)
    if model is None:
        logger.error("❌ MesoNet model failed to load!")
    return model

def clear_model_cache():
    """Clear in-memory cache (useful for memory control or testing)."""
    _load_model.cache_clear()
    logger.info("Model cache cleared.")

# -------------------------------------------------
# Health check utility
# -------------------------------------------------
def verify_models_ready() -> bool:
    """Verify both models are loaded successfully."""
    xception_ok = load_xception_model() is not None
    mesonet_ok = load_mesonet_model() is not None
    return xception_ok and mesonet_ok