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

# %% auto 0
__all__ = ['ImagePair', 'Sameness', 'Pairs']

# %% ../nbs/pairs.ipynb 2
import random
from typing import Union

from fastai.vision.all import *
from fastprogress.fastprogress import *

import fastai_datasets.patches
from .utils import *

# %% ../nbs/pairs.ipynb 3
class ImagePair(fastuple):
    """Adds showing functionality to fastai's `fastuple`"""
    @classmethod
    def create(cls, fns): return cls(tuple(PILImage.create(f) for f in fns))
    
    def show(self, ctx=None, **kwargs):
        t1,t2 = ToTensor()(self)
        assert type(t1) == type(t2)
        line = t1.new_zeros((t1.shape[0], t1.shape[1], 1))
        joint_image = torch.cat([t1,line,t2], dim=2).as_subclass(type(t1))
        return joint_image.show(ctx=ctx, **kwargs)

class Sameness(Categorize):
    """Converts boolean/binary labels into "Not Same"/"Same" labels"""
    def __init__(self):
        super().__init__(['Not Same', 'Same'], sort=False)

    def encodes(self, o: Union[bool, int]):
        return super().encodes(self.vocab[o])

# %% ../nbs/pairs.ipynb 5
@typedispatch
def show_batch(x:ImagePair, y, samples, ctxs=None, max_n=9, nrows=None, ncols=3, figsize=None, **kwargs):
    if figsize is None: figsize = (ncols*4, max_n//ncols * 2)
    if ctxs is None: ctxs = get_grid(min(len(samples), max_n), nrows=nrows, ncols=ncols, figsize=figsize)
    ctxs = show_batch[TensorImage](x, y, samples, ctxs=ctxs, max_n=max_n, **kwargs)
    return ctxs

@typedispatch
def show_results(x:ImagePair, y:TensorCategory, samples, outs, ctxs=None, max_n=9, nrows=None, ncols=3, figsize=None, **kwargs):
    if figsize is None: figsize = (ncols*4, max_n//ncols * 2.5)
    if ctxs is None: ctxs = get_grid(min(len(samples), max_n), nrows=nrows, ncols=ncols, figsize=figsize)
    ctxs = show_results[TensorImage, TensorCategory](x, y, samples, outs, ctxs=ctxs, max_n=max_n, **kwargs)

# %% ../nbs/pairs.ipynb 6
def _pairs_for_split(singles: DataLoaders, split_idx: int, factor: int):
    assert singles.n_inp == 1

    indices = L(range_of(singles))[singles.splits[split_idx]]
    num = int(len(indices) * factor)

    class_map = defaultdict(list)
    for i, c in progress_bar(indices.zipwith(singles.i2t.subset(split_idx)), leave=False):
        class_map[singles.vocab[c]].append(i)

    @return_list
    def _positive_pairs():
        multi_item_class_map = {k: v for k, v in class_map.items() if len(v)>1}
        for _ in progress_bar(range(num//2), leave=False):
            c, idxs = random.choice(list(multi_item_class_map.items()))
            yield tuple(random.sample(idxs, 2))

    @return_list
    def _negative_pairs():
        for _ in progress_bar(range(num//2), leave=False):
            (c1, idxs1), (c2, idxs2) = random.sample(list(class_map.items()), 2)
            yield (random.choice(idxs1), random.choice(idxs2))

    return _positive_pairs() + _negative_pairs()


def Pairs(singles: Datasets,  # Used to construct pairs
          factor=10,  # ratio between numer of pairs and number of single items, for each split
          tuple_type=ImagePair  # To add `show` logic for pairs
          ) -> Datasets:
    """Fixed dataset of randomly-sampled pairs """
    assert singles.n_inp == 1
    pairs_for_splits = [_pairs_for_split(singles, i, factor) for i in range(singles.n_subsets)]

    splits = [slice(len(pairs_for_splits[0]))]
    for pairs in pairs_for_splits[1:]:
        start = splits[-1].stop
        splits.append(slice(start, start+len(pairs)))

    return Datasets(
        [L(x) for items in pairs_for_splits for x in items],
        tfms=[
            lambda o: tuple_type(singles.tls[0][o]),
            [lambda o: bool(singles.i2t[o[0]] == singles.i2t[o[1]]), Sameness()]
        ],
        splits=splits
    )
