import torch

from torch.nn import Module

from torch import Tensor

from typing import Tuple

from tqdm import tqdm

def get_epistemic_uncertainty(model: Module, x: Tensor, n: int = 10) -> Tuple[Tensor, Tensor]:

    '''

    Estimates epistemic uncertainty with n monte carlo predictions of model on x.

    Returns:

        standard deviation uncertainty, mean prediction

    '''

    model = model.train()

    with torch.no_grad():

        uncertain_preds = [model(x).detach().cpu() for _ in tqdm(range(n), leave=False)]

    model = model.eval()

    uncertain_preds_tensor = torch.stack(uncertain_preds)

    epistemic_uncertainty = uncertain_preds_tensor.std(dim=0)

    mean_prediction = uncertain_preds_tensor.mean(dim=0)

    return epistemic_uncertainty, mean_prediction