import torch
from models.level1_gatekeeper import GatekeeperModel
from models.level2_router import DiseaseRouterModel

try:
    print("Testing Level 1 Gatekeeper...")
    l1_model = GatekeeperModel(pretrained=False, freeze_backbone=False)
    ckpt = torch.load('weights/level1.pth', map_location='cpu')
    if 'model_state_dict' in ckpt: ckpt = ckpt['model_state_dict']
    l1_model.load_state_dict(ckpt)
    print("Level 1 loaded successfully!")
except Exception as e:
    print(f"Level 1 failed: {e}")

try:
    print("Testing Level 2 Router...")
    l2_model = DiseaseRouterModel(pretrained=False, freeze_backbone=False)
    ckpt = torch.load('weights/level2.pth', map_location='cpu')
    if 'model_state_dict' in ckpt: ckpt = ckpt['model_state_dict']
    l2_model.load_state_dict(ckpt)
    print("Level 2 loaded successfully!")
except Exception as e:
    print(f"Level 2 failed: {e}")