"""
diagnose_models.py – Comprehensive backend model diagnostic
Run: python diagnose_models.py
"""
import os, sys, json, pickle, traceback
import numpy as np

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

PASS = "[PASS]"
FAIL = "[FAIL]"
WARN = "[WARN]"

results = {}

def section(title):
    print(f"\n{'='*60}")
    print(f"  {title}")
    print(f"{'='*60}")

# ─── 1. Files ────────────────────────────────────────────────────────────────
section("1. File Checksizes")
files = {
    'mlp_model.pth':         {'min_kb': 500,  'required': True},
    'xgb_model.json':        {'min_kb': 50,   'required': True},
    'ensemble_config.json':  {'min_kb': 0.05, 'required': True},
    'best_gradcam_model.pth':{'min_kb': 10000,'required': True},
    'scaler.pkl':            {'min_kb': 0.1,  'required': False},
}
for fname, spec in files.items():
    path = os.path.join(MODEL_DIR, fname)
    if os.path.exists(path):
        kb = os.path.getsize(path) / 1024
        ok = kb >= spec['min_kb']
        tag = PASS if ok else WARN
        print(f"  {tag} {fname:35s} {kb:>10.1f} KB")
        results[fname] = ok
    else:
        tag = FAIL if spec['required'] else WARN
        print(f"  {tag} {fname:35s} --- NOT FOUND ---")
        results[fname] = not spec['required']  # WARN for optional

# ─── 2. ensemble_config.json ─────────────────────────────────────────────────
section("2. ensemble_config.json")
try:
    cfg_path = os.path.join(MODEL_DIR, 'ensemble_config.json')
    with open(cfg_path) as f:
        cfg = json.load(f)
    print(f"  {PASS} Loaded OK")
    print(f"       xgb_weight    : {cfg.get('xgb_weight')}")
    print(f"       mlp_weight    : {cfg.get('mlp_weight')}")
    print(f"       best_threshold: {cfg.get('best_threshold')}")
    results['config_ok'] = True
except Exception as e:
    print(f"  {FAIL} {e}")
    results['config_ok'] = False

# ─── 3. scaler.pkl ───────────────────────────────────────────────────────────
section("3. scaler.pkl (StandardScaler)")
scaler_path = os.path.join(MODEL_DIR, 'scaler.pkl')
if os.path.exists(scaler_path):
    try:
        with open(scaler_path, 'rb') as f:
            scaler = pickle.load(f)
        n_feat = scaler.n_features_in_
        ok = n_feat == 1797
        tag = PASS if ok else FAIL
        print(f"  {tag} Loaded OK — n_features_in = {n_feat} (expect 1797)")
        print(f"       mean[0:3]  : {scaler.mean_[:3].round(4)}")
        print(f"       scale[0:3] : {scaler.scale_[:3].round(4)}")
        results['scaler_ok'] = ok
    except Exception as e:
        print(f"  {FAIL} {e}")
        results['scaler_ok'] = False
else:
    print(f"  {WARN} scaler.pkl not found — predictions WILL be biased!")
    print(f"       Export from notebook Step 5.7 and place here.")
    results['scaler_ok'] = None

# ─── 4. XGBoost ──────────────────────────────────────────────────────────────
section("4. XGBoost Model")
try:
    import xgboost as xgb
    model = xgb.XGBClassifier()
    model.load_model(os.path.join(MODEL_DIR, 'xgb_model.json'))
    print(f"  {PASS} Loaded OK")
    print(f"       n_estimators  : {model.n_estimators}")
    # Test inference with dummy 1797-dim input
    dummy = np.random.randn(1, 1797).astype(np.float32)
    prob = model.predict_proba(dummy)[0, 1]
    print(f"  {PASS} Inference OK — dummy prob: {prob:.4f}")
    results['xgb_ok'] = True
except Exception as e:
    print(f"  {FAIL} {traceback.format_exc()}")
    results['xgb_ok'] = False

# ─── 5. MLP ──────────────────────────────────────────────────────────────────
section("5. MLP Model")
try:
    import torch, torch.nn as nn
    class MLPClassifier(nn.Module):
        def __init__(self, input_dim=1797):
            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): return self.network(x).squeeze(1)

    mlp = MLPClassifier()
    state = torch.load(os.path.join(MODEL_DIR, 'mlp_model.pth'), map_location='cpu')
    mlp.load_state_dict(state)
    mlp.eval()
    print(f"  {PASS} Loaded OK")

    dummy = torch.randn(1, 1797)
    with torch.no_grad():
        prob = mlp(dummy).item()
    print(f"  {PASS} Inference OK — dummy prob: {prob:.4f}")
    results['mlp_ok'] = True
except Exception as e:
    print(f"  {FAIL} {traceback.format_exc()}")
    results['mlp_ok'] = False

# ─── 6. EfficientNet Feature Extractor ───────────────────────────────────────
section("6. EfficientNet-B4 Feature Extractor")
try:
    import torch, torch.nn as nn
    import torchvision.models as models
    net = models.efficientnet_b4(weights=None)
    net.classifier = nn.Identity()
    net.eval()
    dummy_img = torch.randn(1, 3, 380, 380)
    with torch.no_grad():
        feats = net(dummy_img)
    print(f"  {PASS} Loaded OK (random weights — no pretrained weights needed for extractor)")
    print(f"       Output shape  : {feats.shape}  (expect [1, 1792])")
    ok = feats.shape == (1, 1792)
    if not ok:
        print(f"  {FAIL} Feature dim mismatch!")
    results['efficientnet_ok'] = ok
except Exception as e:
    print(f"  {FAIL} {traceback.format_exc()}")
    results['efficientnet_ok'] = False

# ─── 7. Grad-CAM Model ───────────────────────────────────────────────────────
section("7. Grad-CAM Model (best_gradcam_model.pth)")
try:
    import torch, torch.nn as nn
    import torchvision.models as models

    class EfficientNetGradCAM(nn.Module):
        def __init__(self):
            super().__init__()
            base = models.efficientnet_b4(weights=None)
            self.features   = base.features
            self.avgpool    = base.avgpool
            self.classifier = nn.Sequential(
                nn.Dropout(0.4), nn.Linear(1792, 128), nn.ReLU(),
                nn.Dropout(0.3), nn.Linear(128, 1), nn.Sigmoid()
            )
            self.gradients = None; self.activations = None
        def save_gradient(self, grad): self.gradients = grad
        def forward(self, x):
            x = self.features(x)
            if x.requires_grad: x.register_hook(self.save_gradient)
            self.activations = x
            x = self.avgpool(x)
            x = torch.flatten(x, 1)
            return self.classifier(x)

    gcam = EfficientNetGradCAM()
    state = torch.load(os.path.join(MODEL_DIR, 'best_gradcam_model.pth'), map_location='cpu')
    gcam.load_state_dict(state)
    gcam.eval()
    print(f"  {PASS} Weights loaded OK")

    dummy_img = torch.randn(1, 3, 380, 380)
    dummy_img.requires_grad_(True)
    out = gcam(dummy_img)
    prob = out.item()
    gcam.zero_grad()
    out.backward()
    print(f"  {PASS} Forward+backward pass OK — dummy prob: {prob:.4f}")
    grad_ok = gcam.gradients is not None
    print(f"  {PASS if grad_ok else FAIL} Gradient hook: {'captured' if grad_ok else 'MISSING'}")
    results['gradcam_ok'] = grad_ok
except Exception as e:
    print(f"  {FAIL} {traceback.format_exc()}")
    results['gradcam_ok'] = False

# ─── 8. Full Pipeline Simulation ─────────────────────────────────────────────
section("8. Full Feature Fusion Simulation")
try:
    cnn_feats = np.random.randn(1792).astype(np.float32)
    cdr = {'vertical_cdr': 0.55, 'horizontal_cdr': 0.50, 'area_cdr': 0.28, 'mean_cdr': 0.52}
    qs  = 4.0

    cdr_vec = np.array([cdr['vertical_cdr'], cdr['horizontal_cdr'], cdr['area_cdr'], cdr['mean_cdr']])
    fused   = np.hstack([cnn_feats, cdr_vec, np.array([qs])]).astype(np.float32)
    print(f"  {PASS} Fused feature shape: {fused.shape}  (expect 1797,)")

    if results.get('scaler_ok'):
        scaled = scaler.transform(fused.reshape(1, -1)).flatten()
        print(f"  {PASS} StandardScaler applied — mean of scaled: {scaled.mean():.4f} (near 0 expected)")

    xgb_prob = model.predict_proba(fused.reshape(1, -1))[0, 1]
    print(f"  {PASS} XGBoost prob: {xgb_prob:.4f}")

    import torch
    mlp.eval()
    with torch.no_grad():
        mlp_prob = mlp(torch.FloatTensor(fused).unsqueeze(0)).item()
    print(f"  {PASS} MLP prob: {mlp_prob:.4f}")

    ens = 0.5 * xgb_prob + 0.5 * mlp_prob
    threshold = cfg.get('best_threshold', 0.5)
    label = 'GLAUCOMA' if ens >= threshold else 'NORMAL'
    print(f"  {PASS} Ensemble prob: {ens:.4f}  threshold: {threshold:.4f}  → {label}")
    results['pipeline_ok'] = True
except Exception as e:
    print(f"  {FAIL} {traceback.format_exc()}")
    results['pipeline_ok'] = False

# ─── Summary ─────────────────────────────────────────────────────────────────
section("SUMMARY")
all_required = ['config_ok', 'xgb_ok', 'mlp_ok', 'efficientnet_ok', 'gradcam_ok', 'pipeline_ok']
for k, v in results.items():
    if v is True:   icon = "✅"
    elif v is None: icon = "⚠️ "
    else:           icon = "❌"
    print(f"  {icon}  {k}")

missing_scaler = results.get('scaler_ok') is None
if missing_scaler:
    print(f"\n  ⚠️  scaler.pkl MISSING — predictions will be BIASED!")
    print(f"     Export from Google Colab notebook (Step 5.7) and upload to model/")

all_pass = all(results.get(k) for k in all_required)
print(f"\n{'✅ ALL REQUIRED MODELS OK' if all_pass else '❌ SOME MODELS FAILED — CHECK ABOVE'}")