import torch


def picp(target, predictions:list, total = True):
    """
    Calculate PICP (prediction interval coverage probability) or simply the % of true
    values in the predicted intervals

    Parameters
    ----------
    target : torch.Tensor
        true values of the target variable
    predictions : list
        - predictions[0] = y_pred_upper, predicted upper limit of the target variable (torch.Tensor)
        - predictions[1] = y_pred_lower, predicted lower limit of the target variable (torch.Tensor)
    total : bool, default = True
        - When total is set to True, return overall PICP
        - When total is set to False, return PICP along the horizon

    Returns
    -------
    torch.Tensor
        The PICP, which depending on the value of 'total' is either a scalar (PICP in %, for
        significance level alpha = 0.05, PICP should >= 95%)
        or 1d-array over the horizon, in which case it is expected to decrease as we move
        along the horizon. Generally, higher is better.
    """

    # coverage_horizon = torch.zeros(targets.shape[1], device= targets.device,requires_grad=True)
    # for i in range(targets.shape[1]):
    #     # for each step in forecast horizon, calcualte the % of true values in the predicted interval
    #     coverage_horizon[i] = (torch.sum((targets[:, i] > y_pred_lower[:, i]) &
    #                             (targets[:, i] <= y_pred_upper[:, i])) / targets.shape[0]) * 100
    assert len(predictions) == 2
    #torch.set_printoptions(precision=5)
    y_pred_upper = predictions[0]
    y_pred_lower = predictions[1]
    coverage_horizon = 100. * (torch.sum((target > y_pred_lower) &
                                         (target <= y_pred_upper), dim=0)) / target.shape[0]

    coverage_total = torch.sum(coverage_horizon) / target.shape[1]
    if total:
        return coverage_total
    else:
        return coverage_horizon