# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/50_loss.ipynb.

# %% auto 0
__all__ = ['GeodesicLoss', 'AutoEncoderLoss', 'OptimalTransportLoss', 'DensityLoss', 'MMDLoss']

# %% ../nbs/50_loss.ipynb 6
#| export


# %% ../nbs/50_loss.ipynb 8
from types import FunctionType
from typing import Self, Tuple

# %% ../nbs/50_loss.ipynb 10
#| export


# %% ../nbs/50_loss.ipynb 12
#| export


# %% ../nbs/50_loss.ipynb 14
try: import numpy as np
except ImportError: ...

try: import pandas as pd
except ImportError: ...

try: import ot
except: ...

# %% ../nbs/50_loss.ipynb 16
try: import torch, torch.nn as nn
except ImportError: ...

# %% ../nbs/50_loss.ipynb 18
#| export


# %% ../nbs/50_loss.ipynb 20
from quac import tensor
from etrc import OTMethod, DiffusionMethod

# %% ../nbs/50_loss.ipynb 22
from .atyp import LossFunction, DiffusionFunction
from .kwds import DiffusionKeywords
from .init import init_loss
from .diff import Diffusion

# %% ../nbs/50_loss.ipynb 25
class GeodesicLoss(nn.Module, Diffusion):
    def __init__(
        self: Self,
        kind: DiffusionMethod = DiffusionMethod.PHATEDistance,
        *args, **kwargs: DiffusionKeywords
    ) -> None:
        super(Diffusion, GeodesicLoss).__init__()
        
        geocls: DiffusionFunction = DiffusionMethod.get(kind)
        self.fn = geocls(*args, **kwargs)
        self.mse = nn.MSELoss()
        
    def __call__(self, source: tensor, target: tensor):
        dist_source = self.fn.fit(source)
        # TODO: we should check if this is square regardless of what is in original code
        dist_target = torch.cdist(target, target) ** 2
        # dist_target = self.fn.fit(target)
        return self.mse(dist_target, dist_source)

# %% ../nbs/50_loss.ipynb 27
class AutoEncoderLoss(nn.Module):
    def __init__(
        self,
        useemb: bool = True,
        usedis: bool = True,
        userec: bool = True,
        embloss: LossFunction | None = None,
        disloss: LossFunction | None = None,
        recloss: LossFunction | None = None,
        **kwargs
    ):
        self.useemb = useemb
        self.usedis = usedis
        self.userec = userec
        
        self.embloss = embloss # embedding loss
        self.disloss = disloss # discriminator loss
        self.recloss = recloss # reconstruction loss

    def __call__(self, sample: tensor, target: tensor, sample_p: tensor, target_p: tensor) -> tuple[float, float, float]:
        loss_e, loss_d, loss_r = init_loss(sample), init_loss(sample), init_loss(sample)
        
        if self.usedis: loss_d = self.disloss(sample, target_p)
        if self.useemb: loss_e = self.embloss(target_p, target)
        if self.userec: loss_e = self.recloss(sample_p, sample)
        return loss_d, loss_e, loss_r

# %% ../nbs/50_loss.ipynb 29
class OptimalTransportLoss(nn.Module):
    def __init__(self, method: OTMethod = OTMethod.emd):
        super(OptimalTransportLoss, self).__init__()
        self.method: OTMethod = OTMethod.safe(method)
        self.fn: FunctionType | None = OTMethod.imp(key=method)
        
    def __call__(self, source: tensor, target: tensor) -> tensor:
        loss = OTMethod.loss(source, target, method=self.method, func=self.fn, )
        return loss

# %% ../nbs/50_loss.ipynb 31
class DensityLoss(nn.Module):
    def __init__(self, hinge_value: float = 0.01, top_k: int = 5, dim: int = 2):
        super(DensityLoss, self).__init__()
        self.hinge_value = hinge_value
        self.top_k = top_k
        self.dim = dim
        
    def calc_cdist_across_one_t(self, source, target):
        c_dist = torch.stack([torch.cdist(source, target)])
        return c_dist
    
    def calc_cdist_across_time(self, source, target):
        n_samples = source.size(0)
        # NOTE: previously t_start=1, changing here because torchdyn returns all t-steps, 0 included
        t_start = 0
        c_dist = torch.stack([torch.cdist(source[i], target[i]) for i in range(t_start, n_samples)])
        return c_dist
    
    def _check_is_local(self, source) -> bool:
        n_dims = source.dim()
        if n_dims < 3: return True
        return False
    
    def _check_k(self, source):
        n_last = source.size(-1)
        if n_last < self.top_k: return n_last
        return self.top_k
        
            
    def __call__(self, source, target):
        is_local = self._check_is_local(source)
        top_k = self._check_k(source)
        
        if is_local:
             c_dist = self.calc_cdist_across_one_t(source, target)
        else:
            c_dist = self.calc_cdist_across_time(source, target)
            
        values, _ = torch.topk(c_dist, top_k, dim=self.dim, largest=False, sorted=False)
        values -= self.hinge_value
        values[values < 0] = 0
        loss = torch.mean(values)
        return loss

# %% ../nbs/50_loss.ipynb 33
class MMDLoss(nn.Module):
    '''https://github.com/ZongxianLee/MMD_Loss.Pytorch/blob/master/mmd_loss.py'''
    def __init__(self, kernel_mul = 2.0, kernel_num = 5):
        super(MMDLoss, self).__init__()
        self.kernel_num = kernel_num
        self.kernel_mul = kernel_mul
        self.fix_sigma = None
        return
    
    def guassian_kernel(self, source, target, kernel_mul=2.0, kernel_num=5, fix_sigma=None):
        n_samples = int(source.size(0)) + int(target.size(0))
        total = torch.cat([source, target], dim=0)
        total0 = total.unsqueeze(0).expand(int(total.size(0)), int(total.size(0)), int(total.size(1)))
        total1 = total.unsqueeze(1).expand(int(total.size(0)), int(total.size(0)), int(total.size(1)))
        L2_distance = ((total0-total1)**2).sum(2) 
        if fix_sigma:
            bandwidth = fix_sigma
        else:
            bandwidth = torch.sum(L2_distance.data) / (n_samples**2-n_samples)
        bandwidth /= kernel_mul ** (kernel_num // 2)
        bandwidth_list = [bandwidth * (kernel_mul**i) for i in range(kernel_num)]
        kernel_val = [torch.exp(-L2_distance / bandwidth_temp) for bandwidth_temp in bandwidth_list]
        return sum(kernel_val)

    def forward(self, source, target):
        batch_size = int(source.size()[0])
        kernels = self.guassian_kernel(source, target, kernel_mul=self.kernel_mul, kernel_num=self.kernel_num, fix_sigma=self.fix_sigma)
        XX = kernels[:batch_size, :batch_size]
        YY = kernels[batch_size:, batch_size:]
        XY = kernels[:batch_size, batch_size:]
        YX = kernels[batch_size:, :batch_size]
        loss = torch.mean(XX + YY - XY -YX)
        return loss
