import torch

from .binary_classification_metrics import get_binary_metrics
from .regression_metrics import get_regression_metrics
from .metrics_utils import check_metric_is_better


def reverse_los(y, los_info):
    return y * los_info["los_std"] + los_info["los_mean"]

def get_all_metrics(preds, labels, task, los_info):
    # convert preds and labels to tensor if they are ndarray type
    if isinstance(preds, torch.Tensor) == False:
        preds = torch.tensor(preds)
    if isinstance(labels, torch.Tensor) == False:
        labels = torch.tensor(labels)
    
    if task == "outcome":
        return get_binary_metrics(preds, labels)
    elif task == "los":
        return get_regression_metrics(reverse_los(preds, los_info), reverse_los(labels[:, 1], los_info))
    elif task == "multitask":
        return get_binary_metrics(preds[:, 0], labels[:, 0]) | get_regression_metrics(reverse_los(preds[:, 1], los_info), reverse_los(labels[:, 1], los_info))
    else:
        raise ValueError("Task not supported")