import torch
from neurologic.model import Model

class Inference:
    def __init__(self, model_path, device=None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        
        # 1. Load the data (trusting our own source)
        checkpoint = torch.load(model_path, map_location=self.device, weights_only=False)
        
        # 2. Check if it's a dictionary (weights) or the full object
        if isinstance(checkpoint, dict):
            # If it's a dict, we need to know the model structure to load it
            # This is why saving the full object is easier for "dead-simple" use
            raise TypeError("The model file is a dictionary. For 'dead-simple' inference, save the whole model object.")
        else:
            # It's the full model object
            self.model = checkpoint
            
        self.model.to(self.device)
        self.model.eval()