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

# %% auto 0
__all__ = ['ThresholdSiamese']

# %% ../nbs/pair_matching.ipynb 3
import torch
from torch import nn
import numpy as np
from fastai.vision.all import *
from fastprogress.fastprogress import *

from .siamese import *
from .utils import *


class ThresholdSiamese(nn.Module):
    """Combines a deep similarity metric with a threshold to produce a binary classifier."""
    @delegates(DistanceSiamese)
    def __init__(self, *args, **kwargs):
        super().__init__()
        self.distance = DistanceSiamese(*args, **kwargs)
        self.threshold = Threshold()

    def forward(self, x):
        return self.threshold(self.distance(x))

    @delegates(DistanceSiamese.plot_distance_histogram)
    def plot_distance_histogram(self, *args, **kwargs):
        self.distance.plot_distance_histogram(*args, **kwargs)
        plt.axvline(self.threshold.t.item(), linestyle='--', color='grey', label='Threshold')

# %% ../nbs/pair_matching.ipynb 10
@patch
def fit_threshold(self: ThresholdSiamese, train_dl: DataLoader):
    """Picks a threshold that maximizes the accuracy on a dataloader"""
    self.eval().to(train_dl.device)
    with torch.no_grad():
        distances, targets = zip(*((self.distance(x), y) for x, y in progress_bar(train_dl, leave=False)))
        distances, targets = torch.cat(distances).as_subclass(Tensor), torch.cat(targets).as_subclass(Tensor)

        return self.threshold.fit(distances, targets) 
