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

# %% auto 0
__all__ = ['normalized_squared_euclidean_distance', 'ContrastiveLoss', 'DistanceSiamese']

# %% ../nbs/siamese.ipynb 3
from torch import nn
from torch.nn.functional import normalize

from fastai.vision.all import *


def normalized_squared_euclidean_distance(x1, x2):
    r"""
    Squared Euclidean distance over normalized vectors:
    $$\left\| \frac{x_1}{\|x_1\|}-\frac{x_2}{\|x_2\|} \right\|^2 $$
    """
    assert x1.dim() <= 2
    assert x2.dim() <= 2
    x1 = normalize(x1, dim=-1)
    x2 = normalize(x2, dim=-1)
    return (x1 - x2).pow(2).sum(dim=-1)


class ContrastiveLoss(BaseLoss):
    @delegates(nn.HingeEmbeddingLoss)
    def __init__(self, **kwargs):
        super().__init__(loss_cls=nn.HingeEmbeddingLoss, **kwargs)

    def __call__(self, input, target):
        return super().__call__(input, 2*target-1)  # hinge_embedding_loss expects targets to be 1 or -1


# @dataclass
class DistanceSiamese(Module):
    """Outputs the distance between two inputs in feature space"""
    def __init__(self, 
                 backbone: Module,  # embeds inputs in a feature space
                 distance_metric = normalized_squared_euclidean_distance):
        self.backbone = backbone
        self.distance_metric = distance_metric

    def forward(self, x):
        f1, f2 = L(x).map(self.backbone).map(nn.Flatten())
        return self.distance_metric(f1, f2)


# %% ../nbs/siamese.ipynb 4
from matplotlib.ticker import PercentFormatter
from fastprogress.fastprogress import *
import numpy as np

@patch
def plot_distance_histogram(self: DistanceSiamese, pairs_dls: Union[TfmdDL, Dict[str, TfmdDL]], label='Distance'):
    """Plots a histogram of intra-class and inter-class distances"""    
    if isinstance(pairs_dls, DataLoader):
        pairs_dls = {
            'Intra-Class': pairs_dls.by_target['Same'],
            'Inter-Class': pairs_dls.by_target['Not Same']
        }

    self.eval()

    with torch.no_grad():
        for label, dl in pairs_dls.items():
            self._hist(dl, label)
    
    plt.legend()
    plt.xlabel(label)


@patch
def _hist(self: DistanceSiamese, dl, label):
    self.to(dl.device)
    processed_batches = [self(x) for x, _ in progress_bar(dl, leave=False)]
    distances = torch.cat(processed_batches).cpu().numpy()
    weights = np.ones_like(distances) / len(distances)  # normalization for percentage ticks
    plt.hist(distances, label=label, alpha=.5, edgecolor='black', lw=1, weights=weights)
    plt.gca().yaxis.set_major_formatter(PercentFormatter(1))
