# AUTOGENERATED! DO NOT EDIT! File to edit: ../src/metrics.ipynb.

# %% auto 0
__all__ = ['Metric', 'Mean', 'Accuracy']

# %% ../src/metrics.ipynb 3
from .imports import *
from .util import virtual

# %% ../src/metrics.ipynb 4
class Metric:
    """Base metric class."""
    def __init__(self, name: str, device): 
        self.name   = name
        self.device = torch.device(device)
        self.reset_state() 
    def __repr__(self): return f"{self.name}={self.result()}"

    def update_state(self, inp, tar=None): self.empty=False
    def reset_state(self): self.empty=True        
    
    @virtual
    def _eval(self, inp, tar): pass
    @virtual
    def result(self): pass 

# %% ../src/metrics.ipynb 5
class Mean(Metric):
    """Mean metric, used for loss .."""
    def __init__(self, name: str, device): super().__init__(name, device)   
    @torch.inference_mode()
    def update_state(self, inp: torch.Tensor, tar: torch.Tensor=None, weight: float=1):
        super().update_state(inp, tar)         
        val = self._eval(inp, tar)        
        self.weighted_sum += torch.sum(val * weight)
        self.weight       += weight * torch.numel(val)   
    @torch.inference_mode()
    def reset_state(self): 
        super().reset_state()
        self.weighted_sum = torch.tensor(0.0, device=self.device)
        self.weight       = torch.tensor(0.0, device=self.device)   
    def _eval(self, inp, tar): return inp  
    @torch.inference_mode()
    def result(self): return (self.weighted_sum/self.weight).cpu()

# %% ../src/metrics.ipynb 6
class Accuracy(Mean):   
    """Accuracy metric."""
    @torch.inference_mode()
    def _eval(self, inp, tar): return (inp==tar).float().mean()
