# Base Dependencies
# -----------------
import numpy as np
from typing import Optional

# 3rd-Party Dependencies
# -----------------
import torch
from torch import Tensor
from torchmetrics import Metric
from torchmetrics.utilities.checks import _input_format_classification
from sklearn.metrics import precision_score, recall_score, f1_score


def _make_binary(preds: torch.Tensor, target: torch.Tensor):

    # obtain decimal values from one-hot encoding
    preds2 = preds.argmax(axis=1).int()
    target2 = target.argmax(axis=1).int()
    
    # replace positive classes by 1
    preds2[preds2 != 0] = 1
    target2[target2 != 0] = 1

    return preds2, target2


class DetectionF1Score(Metric):
    def __init__(self, ) -> None:
        super().__init__()
        self.add_state("y_true", default=torch.Tensor([]).int())
        self.add_state("y_pred", default=torch.Tensor([]).int())

    def update(self, preds: torch.Tensor, target: torch.Tensor): 
        preds, target, datatype = _input_format_classification(preds, target)
        p, t = _make_binary(preds, target)
        self.y_pred = torch.cat((self.y_pred, p))
        self.y_true = torch.cat([self.y_true, t])

    def compute(self) -> torch.Tensor:
        """Computes f-beta over state."""
        score =  f1_score(y_true=self.y_true.cpu().numpy(), y_pred=self.y_pred.cpu().numpy(), average="binary")
        return torch.tensor(score)


class DetectionPrecision(Metric):
    def __init__(self, ) -> None:
        super().__init__()
        self.add_state("y_true", default=torch.Tensor([]).int())
        self.add_state("y_pred", default=torch.Tensor([]).int())

    def update(self, preds: torch.Tensor, target: torch.Tensor): 
        preds, target, datatype = _input_format_classification(preds, target)
        p, t = _make_binary(preds, target)
        self.y_pred = torch.cat((self.y_pred, p))
        self.y_true = torch.cat([self.y_true, t])

    def compute(self) -> torch.Tensor: 
        score = precision_score(y_true=self.y_true.cpu().numpy(), y_pred=self.y_pred.cpu().numpy(), average="binary")
        return torch.tensor(score)


class DetectionRecall(Metric):
    def __init__(self, ) -> None:
        super().__init__()
        self.add_state("y_true", default=torch.Tensor([]).int())
        self.add_state("y_pred", default=torch.Tensor([]).int())

    def update(self, preds: torch.Tensor, target: torch.Tensor): 
        preds, target, datatype = _input_format_classification(preds, target)
        p, t = _make_binary(preds, target)
        self.y_pred = torch.cat((self.y_pred, p))
        self.y_true = torch.cat([self.y_true, t])

    def compute(self) -> torch.Tensor:
        score = recall_score(y_true=self.y_true.cpu().numpy(), y_pred=self.y_pred.cpu().numpy(), average="binary")
        return torch.tensor(score)