import torch


def compute_errors(gt, pred, crop=True):
    """
    计算深度指标
    Args:
        gt:    ground truth [B,128,416]
        pred:  prediction   [B,1,128,416]
        crop:
    Returns:
        abs_rel, sq_rel, rms, log_rms, a1, a2, a3
    """
    abs_rel, sq_rel, rms, log_rms, a1, a2, a3 = 0, 0, 0, 0, 0, 0, 0
    batch_size = gt.size(0)

    mask_1 = gt > 0
    mask_2 = gt <= 80
    mask = mask_1 * mask_2
    pred_masked = pred.squeeze() * mask

    for current_gt, current_pred in zip(gt, pred_masked):
        # 对当前gt做均值方差归一化
        valid_gt = ((current_gt - current_gt.min()) / (current_gt.max() - current_gt.min())).clamp(1e-3, 1)
        valid_pred = ((current_pred - current_pred.min()) / (current_pred.max() - current_pred.min())).clamp(1e-3, 1)
        valid_pred = valid_pred * torch.median(valid_gt) / torch.median(valid_pred)  # 不明白
        # abs_diff += torch.mean(torch.abs(valid_gt - valid_pred))
        # Relative Absolute Error，绝对相对误差
        abs_rel += torch.mean(torch.abs(valid_gt - valid_pred) / valid_gt)
        # Relative Squared Error，平方相对误差
        sq_rel += torch.mean(((valid_gt - valid_pred) ** 2) / valid_gt)
        # Root Mean Squared Error，均方根误差
        rms += torch.sqrt(torch.mean((valid_gt - valid_pred) ** 2))
        # Log Root Mean Squared Error，对数均方根误差
        log_rms += torch.sqrt(torch.mean((torch.log(valid_gt)-torch.log(valid_pred))**2))
        # 阈值精确度
        thresh = torch.max((valid_gt / valid_pred), (valid_pred / valid_gt))
        a1 += (thresh < 1.25).float().mean()
        a2 += (thresh < 1.25 ** 2).float().mean()
        a3 += (thresh < 1.25 ** 3).float().mean()

    return [metric.item() / batch_size for metric in [abs_rel, sq_rel, rms, log_rms, a1, a2, a3]]