import torch
import torch.nn.functional as F
import json
import os
from huggingface_hub import hf_hub_download
from model import PhysicsGuidedBNN

class PGBNNPredictor:
    """
    Inference and Uncertainty Quantification engine.
    """
    def __init__(self, model: PhysicsGuidedBNN, config: dict):
        self.model = model
        self.config = config
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        self.model.eval()

    @classmethod
    def from_pretrained(cls, repo_id: str):
        config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
        weights_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")

        with open(config_path, "r") as f:
            config = json.load(f)

        model = PhysicsGuidedBNN(config)
        model.load_state_dict(torch.load(weights_path, map_location="cpu"))
        return cls(model, config)

    def predict(self, x: torch.Tensor, num_samples: int = 30):
        """
        Performs Monte Carlo forward passes to compute predictive mean and epistemic uncertainty.
        """
        x = x.to(self.device)
        probs_list = []

        with torch.no_grad():
            for _ in range(num_samples):
                logits = self.model(x)
                probs = F.softmax(logits, dim=-1)[:, 1] # Probability of fault (class 1)
                probs_list.append(probs.unsqueeze(0))

        probs_tensor = torch.cat(probs_list, dim=0) # [num_samples, batch_size]
        
        # Mean prediction across MC samples
        mean_prob = torch.mean(probs_tensor, dim=0)
        
        # Epistemic Uncertainty (Standard Deviation across MC samples)
        epistemic_uncertainty = torch.std(probs_tensor, dim=0)

        return mean_prob, epistemic_uncertainty