# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/00vision_triplet.ipynb (unless otherwise specified).

__all__ = ['SampleEpisode', 'compute_distance_matrix', 'EpisodeDataLoader', 'get_preds', 'distance',
           'remove_duplicates', 'map5', 'accuracy']

# Cell
from fastai.vision.all import *
from torch.utils.data import Sampler

# Cell
class SampleEpisode(Sampler):
    def __init__(self, ys, n_episodes, n_way, k_shot, hard_samples=True,
                 distance_matrix=None):
        self.epoch_size = n_episodes
        self.n_way, self.k_shot = n_way, k_shot
        self.ids = ys
        self.classes = np.array(list(set(self.ids)))
        self.hard_samples = hard_samples
        self.distance_matrix = distance_matrix
        self._epochs = []

    def __len__(self):
        return self.epoch_size*self.n_way*self.k_shot

    def __iter__(self): return iter(self.get_epoch())

    def get_epoch(self):
        """Get indices for one epoch of size epoch_size"""
        idx = []
        for n in range(self.epoch_size):
            idx = [*idx, *self.get_batch()]
        return idx

    def get_batch(self):
        """Get indices for one mini-batch"""
        idx = []
        hard_samples = self.hard_samples if self.distance_matrix is not None else False
        try:
            support_classes = np.random.choice(self.classes, size=self.n_way, replace=False)
        except:
            support_classes = np.random.choice(self.classes, size=self.n_way, replace=True)
        for k in support_classes:
            if hard_samples:
                s = self.get_hard_samples(k, idx)
                idx = [*idx[:-1], *s]
            else:
                s = self.random_sample(k)
                idx = [*idx, *s]
        if hard_samples: idx = idx[:-1]
        return idx

    def random_sample(self, k):
        """Random samples are n-way k-shot"""
        idx = []
        where = np.where(self.ids==k)[0]
        available = len(where)
        replace = True if available < self.k_shot else False
        if replace:
            idx.append(np.random.choice(where, size=available, replace=False))
            idx.append(np.random.choice(where, size=self.k_shot-available, replace=True))
        else:
            idx.append(np.random.choice(where, size=self.k_shot, replace=False))
        return np.concatenate(idx)

    def get_hard_samples(self, k, idx=None):
        """
        Hard samples are selected as the positives with higher distance and
        negatives with smaller distance from the anchor image
        """
        dists = self.distance_matrix
        if len(idx) == 0:
            where = np.where(self.ids==k)[0]
            anchor_idx = np.random.choice(where, size=1, replace=False)
        else:
            anchor_idx = np.array([idx[-1]])
            k = self.ids[self.ids==self.ids[anchor_idx]][0]
            where = np.where(self.ids==k)[0]

        if len(where) > self.k_shot-1:
            hardest_positives = np.array(where[dists[anchor_idx, where]
                                 .sort(descending=True)[1][:self.k_shot-1]])
        else:
            hardest_positives = [*where,
                *np.random.choice(where, size=self.k_shot-len(where)-1, replace=True)]
        ids_in_idx = np.zeros((len(self.ids)))
        if len(idx) != 0:
            classes_in_idx = np.array(list(set(self.ids[idx])))
            for c in classes_in_idx: ids_in_idx[self.ids==c] = 1
        where = np.where((self.ids!=k) & (ids_in_idx==0))[0]
        if len(where)==0: where = np.where(self.ids!=k)[0]
        hardest_negative = np.array([where[dists[anchor_idx, where].sort().indices[0]]])
        return [*anchor_idx, *hardest_positives, *hardest_negative]


# Cell
def compute_distance_matrix(dataloader, model=None):
    with torch.no_grad():
        model.eval()
        embs = []
        t = children_and_parameters(model)[0][0][0].weight.data.type()
        for xb, yb in dataloader:
            emb = model(xb.type(t)).float()
            embs.append(emb)
        embs = torch.cat(embs, dim=0)
        dmat = []
        for e in embs:
            s = (e.view(1,-1) - embs).pow_(2).sum(1).pow_(0.5).unsqueeze(1)
            dmat.append(s)
        dmat = torch.cat(dmat, dim=1).cpu()
        model.train()
        gc.collect()
        return dmat

# Cell
class EpisodeDataLoader(DataLoader):
    _sampler = SampleEpisode
    def __init__(self, fix_dl, dataset=None, model=None, n_episodes=100, n_way=16, k_shot=4, **kwargs):
        super().__init__(dataset=dataset, **kwargs)
        self.fix_dl = fix_dl
        self.n_way, self.k_shot = n_way, k_shot
        ys = np.array([int(np.nanmax(o[0][1].cpu())) for o in self.fix_dl.dataset])
        self.sampler = self._sampler(ys, n_episodes, n_way, k_shot)
        self.model = model
        self.n_episodes = n_episodes

    def before_iter(self):
        self.sampler.distance_matrix = compute_distance_matrix(self.fix_dl, self.model)

    def get_idxs(self):
        return list(self.sampler.__iter__())

    def __len__(self):
        return self.n_episodes

    def create_item(self, s):
        try:o = next(self.it) if s is None else self.dataset[s]
        except:set_trace()
        return o

# Cell
def get_preds(dataloader=None, model=None):
    with torch.no_grad():
        model.eval()
        preds, ys = [], []
        t = children_and_parameters(model)[0][0][0].weight.data.type()
        for xb, yb in dataloader:
            xb, yb = mask2category(xb,yb)
            p = model(xb.type(t)).float()
            preds.append(p)
            ys.append(yb)
        preds = torch.cat(preds, dim=0)
        ys = torch.cat(ys, dim=0)
        model.train()
        gc.collect()
        return preds, ys

def distance(valid_preds, train_preds, train_ys, n=100):
    """Compute distance of each valid sample to the train samples"""
    classes_preds, dists = [], []
    for x in valid_preds:
        d = (x.view(1,-1) - train_preds).pow(2).sum(-1).pow(0.5).sort()
        classes_preds.append(train_ys[d[1][:n]].view(1,-1))
        dists.append(d[0][:n].view(1,-1))
    classes_preds = torch.cat(classes_preds, dim=0)
    dists = torch.cat(dists, dim=0)
    return classes_preds, dists

def remove_duplicates(classes_preds, n=5):
    final_classes = []
    for cl in classes_preds:
        g = []
        for c in cl:
            if len(g) < n and c not in g:
                g.append(c.item())
        final_classes.append(g)
    return torch.tensor(final_classes)

def map5(final_classes, valid_y):
    out = 0
    for i in range(5):
        out+=(final_classes[:, i] == valid_y).float().mean().item()*(1/(i+1))
    return out

def accuracy(x, y, dls=None, model=None, fix_dl_train_idx=2, fix_dl_valid_idx=3):
    train_preds, train_ys = get_preds(dls[fix_dl_train_idx], model)
    valid_preds, valid_ys = get_preds(dls[fix_dl_valid_idx], model)
    classes_preds, dists = distance(valid_preds, train_preds, train_ys)
    final_classes = classes_preds#remove_duplicates(classes_preds)
    accuracy = (final_classes[:, 0] == valid_ys).float().mean().cpu().item()
    return torch.FloatTensor([accuracy])