"""Definition of metrics used during training."""

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

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

# %% ../../src/pipeline/metrics.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ..imports import *

# %% ../../src/pipeline/metrics.ipynb #36032308-bd0e-4409-9db0-9d89fc258e5a
class Metric(abc.ABC):
    """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        
    
    @abc.abstractmethod
    def _eval(self, inp, tar): pass
        
    @abc.abstractmethod
    def result(self): pass 

# %% ../../src/pipeline/metrics.ipynb #c0402572-bebf-4a04-94b3-4fcedb46bddd
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/pipeline/metrics.ipynb #a057e486-924a-4b5e-adc4-3bcaf575ed51
class Accuracy(Mean):   
    """Accuracy metric."""
    @torch.inference_mode()
    def _eval(self, inp, tar): return (inp==tar).float().mean()
