"""
inference.py
Model loading and inference for the GlaucomaAI backend.

Architecture from [IDSC]_D4.ipynb:
  - EfficientNet-B4 feature extractor (1792-dim embeddings)
  - CDR features (4-dim: vertical, horizontal, area, mean)
  - Quality Score (1-dim)
  - Fused feature vector (1797-dim) → XGBoost or MLP or Ensemble

Feature fusion: [CNN(1792) | CDR(4) | QS(1)] = 1797 dims
Ensemble: 0.5 * XGBoost + 0.5 * MLP (from ensemble_config.json)
"""

import os
import json
import pickle
import torch
import torch.nn as nn
import torchvision.models as models
import numpy as np
import pandas as pd
import xgboost as xgb
from typing import Optional

# ─── Paths ────────────────────────────────────────────────────────────────────
BASE_DIR   = os.path.dirname(os.path.abspath(__file__))
MODEL_DIR  = os.path.join(BASE_DIR, 'model')

MLP_PATH      = os.path.join(MODEL_DIR, 'mlp_model.pth')
XGB_PATH      = os.path.join(MODEL_DIR, 'xgb_model.json')
CONFIG_PATH   = os.path.join(MODEL_DIR, 'ensemble_config.json')
SCALER_PATH   = os.path.join(MODEL_DIR, 'scaler.pkl')
COLUMNS_PATH  = os.path.join(MODEL_DIR, 'feature_columns.pkl')

CNN_FEATURE_DIM = 1792
CDR_DIM         = 4
QS_DIM          = 1
INPUT_DIM       = CNN_FEATURE_DIM + CDR_DIM + QS_DIM  # 1797

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# ─── Singleton model holders ──────────────────────────────────────────────────
_efficientnet = None
_mlp_model    = None
_xgb_model    = None
_ensemble_cfg = None
_scaler       = None
_scaler_loaded = False 
_feature_cols = None
_cols_loaded  = False


# ─── MLP Classifier (exact architecture from notebook Cell 8) ─────────────────

class MLPClassifier(nn.Module):
    def __init__(self, input_dim: int = INPUT_DIM):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(0.4),

            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Dropout(0.3),

            nn.Linear(128, 32),
            nn.BatchNorm1d(32),
            nn.ReLU(),
            nn.Dropout(0.2),

            nn.Linear(32, 1),
            nn.Sigmoid()
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.network(x).squeeze(1)


# ─── Model Loaders ────────────────────────────────────────────────────────────

def load_efficientnet() -> nn.Module:
    global _efficientnet
    if _efficientnet is None:
        print("[InferenceEngine] Loading EfficientNet-B4 (ImageNet pretrained)...")
        net = models.efficientnet_b4(weights='DEFAULT')
        net.classifier = nn.Identity()  # output 1792-dim embeddings
        net = net.to(DEVICE)
        net.eval()
        _efficientnet = net
    return _efficientnet


def load_mlp() -> MLPClassifier:
    global _mlp_model
    if _mlp_model is None:
        print(f"[InferenceEngine] Loading MLP from {MLP_PATH}...")
        model = MLPClassifier(INPUT_DIM)
        if os.path.exists(MLP_PATH):
            state = torch.load(MLP_PATH, map_location=DEVICE)
            model.load_state_dict(state)
            print("  ✓ MLP weights loaded")
        else:
            print("  ⚠ MLP file not found — using random weights (demo mode)")
        model = model.to(DEVICE)
        _mlp_model = model
    return _mlp_model


def load_xgboost() -> Optional[xgb.XGBClassifier]:
    global _xgb_model
    if _xgb_model is None:
        if os.path.exists(XGB_PATH):
            print(f"[InferenceEngine] Loading XGBoost from {XGB_PATH}...")
            _xgb_model = xgb.XGBClassifier()
            _xgb_model.load_model(XGB_PATH)
            print("  ✓ XGBoost loaded")
        else:
            print("  ⚠ XGBoost file not found — demo mode")
            _xgb_model = None
    return _xgb_model


def load_ensemble_config() -> dict:
    global _ensemble_cfg
    if _ensemble_cfg is None:
        if os.path.exists(CONFIG_PATH):
            with open(CONFIG_PATH, 'r') as f:
                _ensemble_cfg = json.load(f)
        else:
            _ensemble_cfg = {'xgb_weight': 0.5, 'mlp_weight': 0.5, 'best_threshold': 0.5}
    return _ensemble_cfg


def load_scaler():
    """Load StandardScaler from model/scaler.pkl."""
    global _scaler, _scaler_loaded
    if not _scaler_loaded:
        _scaler_loaded = True
        if os.path.exists(SCALER_PATH):
            with open(SCALER_PATH, 'rb') as f:
                _scaler = pickle.load(f)
            print(f"[InferenceEngine] ✓ StandardScaler loaded from {SCALER_PATH}")
        else:
            print(f"[InferenceEngine] ⚠ scaler.pkl not found at {SCALER_PATH}")
    return _scaler


def load_feature_columns():
    """Load list of 1797 feature names from model/feature_columns.pkl."""
    global _feature_cols, _cols_loaded
    if not _cols_loaded:
        _cols_loaded = True
        if os.path.exists(COLUMNS_PATH):
            with open(COLUMNS_PATH, 'rb') as f:
                _feature_cols = pickle.load(f)
            print(f"[InferenceEngine] ✓ Feature columns loaded ({len(_feature_cols)})")
        else:
            print(f"[InferenceEngine] ⚠ feature_columns.pkl not found!")
    return _feature_cols


def scale_features(fused: np.ndarray) -> np.ndarray:
    """
    Apply StandardScaler to fused features (1797-dim).
    Uses pandas DataFrame with explicit column names to match notebook.
    """
    scaler = load_scaler()
    cols = load_feature_columns()
    
    if scaler is not None and cols is not None:
        # Create DF to ensure consistency with notebook
        X = pd.DataFrame([fused], columns=cols)
        scaled = scaler.transform(X)
        return scaled.flatten().astype(np.float32)
    
    # Fallback to numpy scaling if cols not available
    if scaler is not None:
        return scaler.transform(fused.reshape(1, -1)).flatten().astype(np.float32)
        
    return fused.astype(np.float32)


# ─── Feature Extraction ───────────────────────────────────────────────────────

def extract_cnn_features(img_tensor: torch.Tensor) -> np.ndarray:
    """
    Extract 1792-dim features from EfficientNet-B4.
    img_tensor: shape (C, H, W) float32, ImageNet normalized.
    """
    net = load_efficientnet()
    with torch.no_grad():
        inp = img_tensor.unsqueeze(0).to(DEVICE)  # (1, C, H, W)
        feats = net(inp)  # (1, 1792)
    return feats.cpu().numpy().flatten()  # (1792,)


# ─── Feature Fusion (from notebook Section 5.6) ───────────────────────────────

def fuse_features(cnn_feats: np.ndarray,
                  cdr: dict,
                  quality_score: float) -> np.ndarray:
    """
    Fuse CNN embedding + CDR features + Quality Score.
    Exact fusion: np.hstack([cnn, cdr, qs]) → 1797-dim
    """
    cdr_vec = np.array([
        cdr['vertical_cdr'],
        cdr['horizontal_cdr'],
        cdr['area_cdr'],
        cdr['mean_cdr'],
    ], dtype=np.float32)

    qs_vec = np.array([quality_score], dtype=np.float32)

    fused = np.hstack([
        cnn_feats.astype(np.float32),
        cdr_vec,
        qs_vec,
    ])  # shape (1797,)

    return fused


# ─── Inference Functions ──────────────────────────────────────────────────────

def predict_xgboost(fused_features: np.ndarray) -> float:
    """XGBoost probability prediction using DataFrame with column names."""
    xgb_model = load_xgboost()
    cols = load_feature_columns()
    
    if xgb_model is None:
        return float(np.random.uniform(0.3, 0.7))
    
    if cols is not None:
        # Use DataFrame for XGBoost to ensure feature name matching
        X = pd.DataFrame([fused_features], columns=cols)
        prob = xgb_model.predict_proba(X)[0, 1]
    else:
        # Fallback to numpy
        X = fused_features.reshape(1, -1)
        prob = xgb_model.predict_proba(X)[0, 1]
        
    return float(prob)


def predict_mlp(fused_features: np.ndarray) -> float:
    """MLP deterministic prediction (dropout disabled)."""
    mlp = load_mlp()
    mlp.eval()
    with torch.no_grad():
        x = torch.FloatTensor(fused_features).unsqueeze(0).to(DEVICE)
        prob = mlp(x).cpu().item()
    return float(prob)


def get_mlp_tensor(fused_features: np.ndarray) -> torch.Tensor:
    """Return fused features as tensor for MC Dropout."""
    return torch.FloatTensor(fused_features).unsqueeze(0)


def predict_ensemble(xgb_prob: float, mlp_prob: float) -> tuple:
    """
    Soft voting: ensemble_prob = w_xgb * xgb_prob + w_mlp * mlp_prob
    Uses weights from ensemble_config.json.
    """
    cfg = load_ensemble_config()
    w_xgb = cfg.get('xgb_weight', 0.5)
    w_mlp = cfg.get('mlp_weight', 0.5)
    threshold = cfg.get('best_threshold', 0.5)
    ens_prob = w_xgb * xgb_prob + w_mlp * mlp_prob
    return float(ens_prob), float(threshold)


def classify_probability(probability: float, threshold: float = 0.5) -> dict:
    """Convert probability to clinical classification."""
    predicted_label = 'GLAUCOMA' if probability >= threshold else 'NORMAL'
    confidence = probability if predicted_label == 'GLAUCOMA' else (1.0 - probability)
    return {
        'label': predicted_label,
        'is_glaucoma': predicted_label == 'GLAUCOMA',
        'probability': round(probability, 4),
        'confidence': round(confidence, 4),
        'threshold': round(threshold, 4),
    }