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

# %% auto 0
__all__ = ['image_grid', 'group_by_count', 'ClassificationExplorer', 'image_size', 'image_sizes_hist', 'one_batch_with_idxs',
           'show_tfms']

# %% ../nbs/00_core.ipynb 3
#TODO: optimize imports
from fastai.vision.all import *

# %% ../nbs/00_core.ipynb 15
def image_grid(
    shows: Sequence[Union[PILBase, TensorImageBase]],
    ncols=3,
    nrows=None,
    pad=0.0,
    figsize=None,
    show=True,
    **kwargs
):
    from mpl_toolkits.axes_grid1 import ImageGrid

    nrows = nrows or int(np.ceil(len(shows) / ncols))
    figsize = figsize or (14, 4 * nrows)

    fig = plt.figure(figsize=figsize)
    grid = ImageGrid(fig, 111, (nrows, ncols), axes_pad=pad, **kwargs)

    for ax, show in zip(grid, shows):
        show(ctx=ax)

    if show:
        plt.show()

    return grid

# %% ../nbs/00_core.ipynb 16
def group_by_count(label_df):
    label_df["idx"] = label_df.index
#     label_df = label_df.groupby("label_id")["idx"].apply(list).reset_index()
    label_df = label_df.groupby("label_id").agg({"label": "first", "idx": list}).reset_index()
    label_df["n_samples"] = label_df.idx.str.len()
    label_df = label_df.sort_values("n_samples", ascending=False).reset_index(drop=True)
    return label_df

# %% ../nbs/00_core.ipynb 17
class ClassificationExplorer:
    def __init__(self, x_tl, y_tl):
        self.x_tl = x_tl
        self.y_tl = y_tl
        
        self.id2label = y_tl.vocab
        self.label2id = {label: i for i, label in enumerate(self.id2label)}

        self.label_df = pd.DataFrame(y_tl, columns=["label_id"])
        self.label_df['label'] = self.label_df['label_id'].map(y_tl.vocab.__getitem__)
        
        self.grouped_y = group_by_count(self.label_df)

    def most_common(self):
        "Return a dataframe ordered by most common classes."
        return self.grouped_y

    def show_label_id(self, id, slice_=slice(None)):
        # TODO: can we also pick the class name to show?
        idx = self.grouped_y[self.grouped_y["label_id"] == id].index[0]
        return self.show_most_common(idx=idx, slice_=slice_)
    
    def show_label(self, label, slice_=slice(None)):
        return self.show_label_id(self.label2id[label], slice_=slice_)

    def show_most_common(self, idx: int, slice_=slice(None)):
        idxs = self.grouped_y.iloc[idx].idx
        idxs = idxs[slice_]
        shows = [self.x_tl[i].show for i in idxs]
        return image_grid(shows)
    
    @classmethod
    def from_datasets(cls, dss):
        return cls(dss.tls[0], dss.tls[1])

# %% ../nbs/00_core.ipynb 21
def image_size(image_file):
    with Image.open(str(image_file)) as image:
        return image.size

# %% ../nbs/00_core.ipynb 22
def image_sizes_hist(
    image_files # Sequence of image filepaths.
):
    "Plot a histogram of widths and heights."
    sizes = parallel(image_size, image_files, progress=progress_bar)
    widths, heights = zip(*sizes)
    plt.hist(widths, label='width')
    plt.hist(heights, label='heights')
    plt.legend()

# %% ../nbs/00_core.ipynb 25
def one_batch_with_idxs(
    dl, # A `DataLoader` instance. Commonly `dls.train` or `dls.valid`
    unique_idx=None, # If specified, show only images with index `unique_idx`.
):
    "Show original and transformed version of images."
    old_get_idxs = dl.get_idxs

    idxs = dl.get_idxs() if unique_idx is None else [unique_idx] * len(dl.items)
    dl.get_idxs = lambda: idxs
    x, y = dl.one_batch()

    dl.get_idxs = old_get_idxs
    
    return (x, y), idxs

# %% ../nbs/00_core.ipynb 26
@delegates(image_grid)
def show_tfms(
    dl, # A `DataLoader` instance. Commonly `dls.train` or `dls.valid`
    unique_idx=None, # If specified, show only images with index `unique_idx`.
    max_n=5, # Maximum number of samples.
    **kwargs,
):
    "Show original and transformed version of images."
    (xs, ys), idxs = one_batch_with_idxs(dl, unique_idx=unique_idx)
    
    shows = []
    for i, x, idx in zip(range(max_n), xs, idxs):
        def _s(ctx, idx=idx):
            return show_at(dl.dataset, idx, ctx=show_title(idx, ctx=ctx))
        shows.append(_s)
        shows.append(x.show)
        
    ncols = kwargs.pop('ncols', None) or 2
    return image_grid(shows, ncols=ncols, **kwargs)
