model_manager.py

"""
Model Manager - Handles multiple CNN models for Alzheimer's (MRI) and
Parkinson's (DaTscan) classification, with graceful fallback when models
are not available.

PD imaging = DaTscan ONLY:
  - densenet121_parkinsonsDATSCAN.keras  (Keras, 2-class)
  - parkinsons_densenet169DATSCAN.keras  (Keras, 2-class)
  - parkinsons_densenet201DATSCAN.keras  (Keras, 2-class)
  - parkinsons_3dcnnDATSCAN.pth         (PyTorch 3D CNN, 2-class)

AD imaging = MRI:
  - alzheimers_densenet121.pth
  - alzheimers_densenet169.pth
  - alzheimers_densenet201.pth
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Dict, Optional, Tuple, List
import numpy as np
import torch
import torch.nn as nn
import torchvision.models as tv
from PIL import Image
import torchvision.transforms as transforms
from io import BytesIO

logger = logging.getLogger("app.models.model_manager")

# ── AD MRI Model configurations ───────────────────────────────────────────────
MODEL_CONFIGS = {
    # Alzheimer's MRI models (PyTorch .pth)
    "ad_dn121": {
        "name": "Alzheimer's DenseNet121 (MRI)",
        "condition": "alzheimers",
        "imaging_type": "mri",
        "architecture": "densenet121",
        "framework": "pytorch",
        "num_classes": 4,
        "filename": "alzheimers_densenet121.pth",
        "class_names": ["Mild Demented", "Moderate Demented", "Non Demented", "Very Mild Demented"],
    },
    "ad_dn169": {
        "name": "Alzheimer's DenseNet169 (MRI)",
        "condition": "alzheimers",
        "imaging_type": "mri",
        "architecture": "densenet169",
        "framework": "pytorch",
        "num_classes": 4,
        "filename": "alzheimers_densenet169.pth",
        "class_names": ["Mild Demented", "Moderate Demented", "Non Demented", "Very Mild Demented"],
    },
    "ad_dn201": {
        "name": "Alzheimer's DenseNet201 (MRI)",
        "condition": "alzheimers",
        "imaging_type": "mri",
        "architecture": "densenet201",
        "framework": "pytorch",
        "num_classes": 4,
        "filename": "alzheimers_densenet201.pth",
        "class_names": ["Mild Demented", "Moderate Demented", "Non Demented", "Very Mild Demented"],
    },

    # Parkinson's DaTscan models — prefer retrained .pth, fall back to .keras
    "pd_datscan_dn121": {
        "name": "Parkinson's DaTscan DenseNet121",
        "condition": "parkinsons",
        "imaging_type": "datscan",
        "architecture": "densenet121",
        "framework": "pytorch",
        "num_classes": 2,
        # Retrained .pth takes priority; .keras kept as fallback filename
        "filename": "parkinsons_densenet121.pth",
        "filename_fallback": "densenet121_parkinsonsDATSCAN.keras",
        "class_names": ["No Parkinson's", "Parkinson's Disease"],
    },
    "pd_datscan_dn169": {
        "name": "Parkinson's DaTscan DenseNet169",
        "condition": "parkinsons",
        "imaging_type": "datscan",
        "architecture": "densenet169",
        "framework": "pytorch",
        "num_classes": 2,
        "filename": "parkinsons_densenet169.pth",
        "filename_fallback": "parkinsons_densenet169DATSCAN.keras",
        "class_names": ["No Parkinson's", "Parkinson's Disease"],
    },
    "pd_datscan_dn201": {
        "name": "Parkinson's DaTscan DenseNet201",
        "condition": "parkinsons",
        "imaging_type": "datscan",
        "architecture": "densenet201",
        "framework": "pytorch",
        "num_classes": 2,
        "filename": "parkinsons_densenet201.pth",
        "filename_fallback": "parkinsons_densenet201DATSCAN.keras",
        "class_names": ["No Parkinson's", "Parkinson's Disease"],
    },
    "pd_datscan_3dcnn": {
        "name": "Parkinson's DaTscan 3D CNN",
        "condition": "parkinsons",
        "imaging_type": "datscan",
        "architecture": "3dcnn",
        "framework": "pytorch",
        "num_classes": 2,
        "filename": "parkinsons_3dcnnDATSCAN.pth",
        "class_names": ["No Parkinson's", "Parkinson's Disease"],
        "input_3d": True,
    },
}

# ── Ensemble configurations ────────────────────────────────────────────────────
ENSEMBLE_CONFIGS = {
    "ad_homogeneous": {
        "name": "Alzheimer's MRI Homogeneous Ensemble (DenseNet 121+169+201)",
        "condition": "alzheimers",
        "imaging_type": "mri",
        "models": ["ad_dn121", "ad_dn169", "ad_dn201"],
        "weights": [0.4, 0.3, 0.3],
    },
    "pd_datscan_ensemble": {
        "name": "Parkinson's DaTscan Ensemble (DenseNet 121+169+201)",
        "condition": "parkinsons",
        "imaging_type": "datscan",
        "models": ["pd_datscan_dn121", "pd_datscan_dn169", "pd_datscan_dn201"],
        "weights": [0.4, 0.3, 0.3],
    },
}

# Accepted DaTscan file extensions
DATSCAN_EXTENSIONS = {".nii", ".gz", ".dcm", ".png", ".jpg", ".jpeg"}


class ModelManager:
    """Manages MRI (AD) and DaTscan (PD) models with graceful fallback."""

    def __init__(self):
        self.models_dir = Path(__file__).resolve().parent.parent.parent / "saved_models"
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.loaded_models: Dict[str, Optional[object]] = {}
        self.model_status: Dict[str, str] = {}
        self.image_transform = self._get_image_transform()
        self._initialize_models()

    def _get_image_transform(self):
        return transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
        ])

    # ── PyTorch model builder ──────────────────────────────────────────────────

    def _build_pytorch_model(self, config: dict) -> nn.Module:
        arch = config["architecture"]
        num_classes = config["num_classes"]

        if arch == "densenet121":
            model = tv.densenet121(weights=None)
            in_features = model.classifier.in_features
            model.classifier = nn.Sequential(
                nn.Linear(in_features, 256), nn.ReLU(), nn.Dropout(0.3),
                nn.Linear(256, num_classes),
            )
        elif arch == "densenet169":
            model = tv.densenet169(weights=None)
            in_features = model.classifier.in_features
            model.classifier = nn.Sequential(
                nn.Linear(in_features, 256), nn.ReLU(), nn.Dropout(0.3),
                nn.Linear(256, num_classes),
            )
        elif arch == "densenet201":
            model = tv.densenet201(weights=None)
            in_features = model.classifier.in_features
            model.classifier = nn.Sequential(
                nn.Linear(in_features, 256), nn.ReLU(), nn.Dropout(0.3),
                nn.Linear(256, num_classes),
            )
        elif arch == "3dcnn":
            model = self._build_3dcnn(num_classes)
        else:
            raise ValueError(f"Unsupported pytorch architecture: {arch}")
        return model

    def _build_3dcnn(self, num_classes: int = 2) -> nn.Module:
        """Simple 3D CNN for DaTscan volumetric input."""
        class Simple3DCNN(nn.Module):
            def __init__(self, n_classes):
                super().__init__()
                self.features = nn.Sequential(
                    nn.Conv3d(1, 32, 3, padding=1), nn.BatchNorm3d(32), nn.ReLU(),
                    nn.MaxPool3d(2),
                    nn.Conv3d(32, 64, 3, padding=1), nn.BatchNorm3d(64), nn.ReLU(),
                    nn.MaxPool3d(2),
                    nn.Conv3d(64, 128, 3, padding=1), nn.BatchNorm3d(128), nn.ReLU(),
                    nn.AdaptiveAvgPool3d((4, 4, 4)),
                )
                self.classifier = nn.Sequential(
                    nn.Flatten(),
                    nn.Linear(128 * 4 * 4 * 4, 256), nn.ReLU(), nn.Dropout(0.4),
                    nn.Linear(256, n_classes),
                )
            def forward(self, x):
                return self.classifier(self.features(x))
        return Simple3DCNN(num_classes)

    # ── Keras model loader ─────────────────────────────────────────────────────

    def _load_keras_model(self, model_key: str) -> Tuple[Optional[object], str]:
        config = MODEL_CONFIGS[model_key]
        model_path = self.models_dir / config["filename"]
        if not model_path.exists():
            logger.warning("Keras model file not found: %s", model_path)
            return None, "Model file not found"
        try:
            import os
            os.environ["TF_USE_LEGACY_KERAS"] = "1"
            import tensorflow as tf
            model = tf.keras.models.load_model(str(model_path), compile=False)
            logger.info("Loaded Keras model: %s", config["name"])
            return model, "Active"
        except Exception as e:
            logger.error("Failed to load Keras model %s: %s", config["name"], e)
            return None, f"Load error: {str(e)}"

    # ── PyTorch model loader ───────────────────────────────────────────────────

    def _load_pytorch_model(self, model_key: str) -> Tuple[Optional[nn.Module], str]:
        config = MODEL_CONFIGS[model_key]
        model_path = self.models_dir / config["filename"]

        # If primary .pth not found, try fallback (old .keras → skip, just report missing)
        if not model_path.exists():
            fallback = config.get("filename_fallback")
            if fallback:
                fallback_path = self.models_dir / fallback
                if fallback_path.exists() and fallback_path.suffix in (".keras", ".h5"):
                    # Keras fallback — delegate to keras loader
                    return self._load_keras_model_from_path(config, fallback_path)
            logger.warning("Model file not found: %s", model_path)
            return None, "Model file not found"
        try:
            model = self._build_pytorch_model(config)
            state_dict = torch.load(str(model_path), map_location=self.device)
            if isinstance(state_dict, dict) and "model_state_dict" in state_dict:
                state_dict = state_dict["model_state_dict"]
            # strict=False allows loading models whose classifier head differs slightly
            model.load_state_dict(state_dict, strict=False)
            model.to(self.device)
            model.eval()
            logger.info("Loaded PyTorch model: %s", config["name"])
            return model, "Active"
        except Exception as e:
            logger.error("Failed to load PyTorch model %s: %s", config["name"], e)
            return None, f"Load error: {str(e)}"

    def _load_keras_model_from_path(self, config: dict, model_path: Path) -> Tuple[Optional[object], str]:
        try:
            import os
            os.environ["TF_USE_LEGACY_KERAS"] = "1"
            import tensorflow as tf
            model = tf.keras.models.load_model(str(model_path), compile=False)
            logger.info("Loaded Keras fallback model: %s", config["name"])
            return model, "Active (Keras fallback)"
        except Exception as e:
            logger.error("Failed to load Keras fallback %s: %s", config["name"], e)
            return None, f"Load error: {str(e)}"

    def _initialize_models(self):
        logger.info("Initializing model manager (AD-MRI + PD-DaTscan)...")
        for model_key, config in MODEL_CONFIGS.items():
            # All models now use PyTorch; Keras fallback handled inside _load_pytorch_model
            model, status = self._load_pytorch_model(model_key)
            self.loaded_models[model_key] = model
            self.model_status[model_key] = status
        logger.info("Model initialization complete")

    # ── Availability ───────────────────────────────────────────────────────────

    def get_available_models(self, condition: str = None) -> List[dict]:
        available = []
        for model_key, config in MODEL_CONFIGS.items():
            if condition and config["condition"] != condition:
                continue
            if self.model_status.get(model_key) == "Active":
                available.append({
                    "key": model_key,
                    "name": config["name"],
                    "condition": config["condition"],
                    "imaging_type": config.get("imaging_type", "mri"),
                    "architecture": config["architecture"],
                    "framework": config.get("framework", "pytorch"),
                    "status": "Active",
                })
        return available

    def get_model_status(self) -> Dict[str, str]:
        return self.model_status.copy()

    # ── PyTorch image prediction (AD MRI) ─────────────────────────────────────

    def predict_image(self, model_key: str, image_bytes: bytes, filename: str = "") -> dict:
        """Make prediction using a PyTorch model on standard image bytes."""
        if model_key not in MODEL_CONFIGS:
            return {"error": f"Unknown model: {model_key}"}

        config = MODEL_CONFIGS[model_key]

        # Validate DaTscan extensions
        if config.get("imaging_type") == "datscan" and filename:
            ext = Path(filename).suffix.lower()
            # .nii.gz has compound suffix
            if filename.endswith(".nii.gz"):
                ext = ".nii.gz"
            if ext not in DATSCAN_EXTENSIONS and ext != ".nii.gz":
                return {"error": f"Invalid file type '{ext}' for DaTscan analysis. Accepted: .nii, .nii.gz, .dcm, .png, .jpg"}

        # 3D CNN needs special handling
        if config.get("input_3d"):
            return self.predict_3dcnn(model_key, image_bytes, filename)

        model = self.loaded_models.get(model_key)
        if model is None:
            return {"error": f"Model {model_key} is not available"}

        try:
            image = Image.open(BytesIO(image_bytes)).convert("RGB")
            inputs = self.image_transform(image).unsqueeze(0).to(self.device)

            with torch.no_grad():
                outputs = model(inputs)
                probs = torch.softmax(outputs, dim=1)
                pred_class = torch.argmax(probs, dim=1).item()
                confidence = float(probs[0][pred_class].item())

            return {
                "model_key": model_key,
                "model_name": config["name"],
                "condition": config["condition"],
                "imaging_type": config.get("imaging_type", "mri"),
                "prediction": pred_class,
                "confidence": confidence,
                "class_name": config["class_names"][pred_class],
                "all_probabilities": {
                    cn: float(p) for cn, p in zip(config["class_names"], probs[0].cpu().numpy())
                },
                "status": "success",
            }
        except Exception as e:
            logger.error("PyTorch prediction failed for %s: %s", model_key, e)
            return {"error": f"Prediction failed: {str(e)}"}

    # ── Keras image prediction (PD DaTscan DenseNet) ──────────────────────────

    def predict_keras_image(self, model_key: str, image_bytes: bytes, filename: str = "") -> dict:
        """Run a Keras DaTscan model on 2D image/slice bytes."""
        if model_key not in MODEL_CONFIGS:
            return {"error": f"Unknown model: {model_key}"}

        config = MODEL_CONFIGS[model_key]

        # Extension check
        if filename:
            ext = Path(filename).suffix.lower()
            fname_lower = filename.lower()
            if fname_lower.endswith(".nii.gz"):
                ext = ".nii.gz"
            if ext not in DATSCAN_EXTENSIONS:
                return {"error": f"Invalid file type '{ext}' for DaTscan. Accepted: .nii, .nii.gz, .dcm, .png, .jpg"}

        model = self.loaded_models.get(model_key)
        if model is None:
            return {"error": f"Keras model {model_key} is not available"}

        try:
            from app.preprocessing.datscan_preprocessor import DaTscanPreprocessor
            preprocessor = DaTscanPreprocessor()
            img_array = preprocessor.preprocess_2d(image_bytes, filename)  # (224, 224, 3) float32

            import numpy as _np
            batch = _np.expand_dims(img_array, 0)  # (1, 224, 224, 3)
            preds = model.predict(batch, verbose=0)  # (1, num_classes)
            probs = preds[0]
            pred_class = int(_np.argmax(probs))
            confidence = float(probs[pred_class])

            return {
                "model_key": model_key,
                "model_name": config["name"],
                "condition": config["condition"],
                "imaging_type": "datscan",
                "prediction": pred_class,
                "confidence": confidence,
                "class_name": config["class_names"][pred_class],
                "all_probabilities": {
                    cn: float(p) for cn, p in zip(config["class_names"], probs)
                },
                "status": "success",
            }
        except Exception as e:
            logger.error("Keras DaTscan prediction failed for %s: %s", model_key, e)
            return {"error": f"DaTscan prediction failed: {str(e)}"}

    # ── 3D CNN prediction (PD DaTscan volumetric) ─────────────────────────────

    def predict_3dcnn(self, model_key: str, volume_bytes: bytes, filename: str = "") -> dict:
        """Run the 3D CNN on a NIfTI volume (.nii or .nii.gz required)."""
        if model_key not in MODEL_CONFIGS:
            return {"error": f"Unknown model: {model_key}"}

        config = MODEL_CONFIGS[model_key]
        fname_lower = (filename or "").lower()
        if not (fname_lower.endswith(".nii") or fname_lower.endswith(".nii.gz")):
            return {"error": "3D CNN requires a NIfTI file (.nii or .nii.gz)."}

        model = self.loaded_models.get(model_key)
        if model is None:
            return {"error": f"3D CNN model {model_key} is not available"}

        try:
            from app.preprocessing.datscan_preprocessor import DaTscanPreprocessor
            preprocessor = DaTscanPreprocessor()
            volume_tensor = preprocessor.preprocess_3d(volume_bytes, filename)  # (1, 1, D, H, W)
            volume_tensor = volume_tensor.to(self.device)

            with torch.no_grad():
                outputs = model(volume_tensor)
                probs = torch.softmax(outputs, dim=1)
                pred_class = int(torch.argmax(probs, dim=1).item())
                confidence = float(probs[0][pred_class].item())

            return {
                "model_key": model_key,
                "model_name": config["name"],
                "condition": config["condition"],
                "imaging_type": "datscan",
                "prediction": pred_class,
                "confidence": confidence,
                "class_name": config["class_names"][pred_class],
                "all_probabilities": {
                    cn: float(p) for cn, p in zip(config["class_names"], probs[0].cpu().numpy())
                },
                "status": "success",
            }
        except Exception as e:
            logger.error("3D CNN prediction failed for %s: %s", model_key, e)
            return {"error": f"3D CNN prediction failed: {str(e)}"}

    # ── Ensemble prediction ────────────────────────────────────────────────────

    def predict_ensemble(self, ensemble_key: str, image_bytes: bytes, filename: str = "") -> dict:
        if ensemble_key not in ENSEMBLE_CONFIGS:
            return {"error": f"Unknown ensemble: {ensemble_key}"}

        ensemble_config = ENSEMBLE_CONFIGS[ensemble_key]
        model_predictions = []
        weights = ensemble_config["weights"]

        for model_key in ensemble_config["models"]:
            result = self.predict_image(model_key, image_bytes, filename)
            if "error" not in result:
                model_predictions.append(result)

        if not model_predictions:
            return {"error": "No models available in ensemble"}

        if len(weights) != len(model_predictions):
            weights = [1.0 / len(model_predictions)] * len(model_predictions)

        combined_probs: Dict[str, float] = {}
        total_weight = 0.0
        for pred, weight in zip(model_predictions, weights):
            total_weight += weight
            for class_name, prob in pred["all_probabilities"].items():
                combined_probs[class_name] = combined_probs.get(class_name, 0) + prob * weight

        for cn in combined_probs:
            combined_probs[cn] /= total_weight

        final_class = max(combined_probs, key=lambda x: combined_probs[x])
        final_confidence = combined_probs[final_class]
        first_config = MODEL_CONFIGS[ensemble_config["models"][0]]

        return {
            "ensemble_key": ensemble_key,
            "ensemble_name": ensemble_config["name"],
            "condition": ensemble_config["condition"],
            "imaging_type": ensemble_config.get("imaging_type", "mri"),
            "prediction": first_config["class_names"].index(final_class),
            "confidence": final_confidence,
            "class_name": final_class,
            "all_probabilities": combined_probs,
            "model_contributions": [
                {"model": p["model_name"], "weight": w, "confidence": p["confidence"]}
                for p, w in zip(model_predictions, weights)
            ],
            "status": "success",
        }


# ── Singleton ──────────────────────────────────────────────────────────────────
_model_manager: Optional[ModelManager] = None


def get_model_manager() -> ModelManager:
    global _model_manager
    if _model_manager is None:
        _model_manager = ModelManager()
    return _model_manager