# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/05a_classification.explorer.ipynb.

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

# %% ../../nbs/05a_classification.explorer.ipynb 2
#TODO: optimize imports
from fastcore.all import *

import numpy as np
import pandas as pd
import polvo as pv

from PIL import Image

# %% ../../nbs/05a_classification.explorer.ipynb 7
class DataExplorer:
    def __init__(self, label_ids, idx2image, id2label, label2id):
        store_attr()
        self.grouped_y = self._group_y()
        
    def _group_y(self):
        label_df = pd.DataFrame(self.label_ids, columns=["label_id"])
        label_df['label'] = label_df['label_id'].map(self.id2label) #
        label_df["item_idx"] = label_df.index
        
        label_df = label_df.groupby("label_id").agg({"label": "first", "item_idx": list}).reset_index()
        label_df["n_samples"] = label_df.item_idx.str.len()
        label_df = label_df.sort_values("n_samples", ascending=False).reset_index(drop=True)
        return label_df

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

    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)):
        item_idxs = self.grouped_y.iloc[idx].item_idx
        item_idxs = item_idxs[slice_]
        
        images = [self.idx2image(i) for i in item_idxs]
        xmax, ymax = np.max([image.size for image in images], axis=0)
        
        return pv.image_grid(images, xmax=xmax, ymax=ymax)
