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

__all__ = ['ImageLabeler', 'SingleClassImageLabeler', 'MultiClassImageLabeler']

# Cell
import json
import pandas as pd
import logging

from pathlib import Path
from forgebox.files import file_detail
from forgebox.html import DOM
from typing import List, Dict
from tqdm.notebook import tqdm

# to avoid confusions for 2 Image class
from PIL import Image as PILImage
from ipywidgets import Image as ImageWidget

from ipywidgets import interact, interact_manual, Button, SelectMultiple, \
    Output, HBox, VBox

# Cell
class ImageLabeler:
    default_file = "unpackai_labeler.json"

    def __init__(self,
                 image_folder: Path,
                 formats: List[str] = ["jpg", "jpeg", "png", "bmp"],
                 ):
        """
        path: Path, a folder full of images
        formats: a list of allowed formats
        """
        self.image_folder = str(image_folder)

        self.file_df = file_detail(image_folder)
        self.file_df["path"] = self.file_df["path"]\
            .apply(lambda x: str(Path(x).resolve()))
        self.filter_image(formats)
        progress_path = Path(self.image_folder)/self.default_file
        if (progress_path).exists():
            try:
                print(f"load from progress {progress_path}")
                with open(progress_path, "r") as f:
                    self.progress = json.loads(f.read())
            except:
                print("loading progress error")
        self.output = Output()

    def __repr__(self):
        return f"{self.__class__.__name__} on [{self.image_folder}({len(self.image_df)})], see labeler.image_df"

    def filter_image(
        self,
        formats: List[str] = ["jpg", "jpeg", "png", "bmp"]
    ) -> pd.DataFrame:
        """
        Filter the file dataframe to image only files
        assign image_df attribute to the object
        """
        formats += list(map(lambda x: x.upper(), formats))
        self.image_df = self.file_df[self.file_df.file_type.isin(
            formats)].reset_index(drop=True)
        return self.image_df

    def __call__(self, *args, **kwargs):
        raise NotImplementedError(
            f"Please use SingleClassImageLabeler, or MultiClassImageLabeler")

    @property
    def identifier(self):
        return self.progress['meta']['identifier']

    def save_progress(
        self,
        location: Path = None,
        filename=None,
    ):
        """
        Save the progress to location/filename
        default save to current directory ./unpackai_imglbl.json
        """
        filename = self.default_file if filename is None else filename
        location = Path(
            self.image_folder) if location is None else Path(location)
        with open(location/filename, "w") as f:
            f.write(json.dumps(self.progress))
        logging.info(f"Progess Saved to {location/filename}")

    @classmethod
    def load_saved(cls, filepath="./unpackai_imglbl.json"):
        """
        Load saved labeler's progress
        """
        with open(filepath, "r") as f:
            progress = json.loads(f.read())
        image_folder = progress['meta']['image_folder']
        obj = cls(image_folder)
        obj.progress = progress
        return obj

    def new_progress(self, labels: List[str], identifier: str = "path"):
        """
        create a new progress
        keys:
        - meta: meta data
        - data: a dictionary
            - key: image path
            - value: labeled result
        """
        self.progress = dict(
            meta=dict(
                image_folder=self.image_folder,
                labels=labels,
                identifier=identifier,
            ),
            data=dict((str(k), None) for k in list(self.image_df[identifier]))
        )

    def __call__(self, labels: List[str] = ["pos", "neg"]):
        """
        Start the labeling
        labels: a list of labels, more then one string,
            but not in hundreds or more, please
        """
        self.labels = labels
        if hasattr(self, "progress") == False:
            self.new_progress(labels)

        # iteration with progress bar
        for k, v in tqdm(self.progress['data'].items(), leave=False):
            if v is None:
                yield k

    def __getitem__(self, key: str):
        """
        render a page according to key
        """
        row = self.get_row_data(key)
        self.output.clear_output()
        with self.output:
            # display subject image
            with PILImage.open(
                    row[self.identifier]).resize((512, 512)) as img:
                display(img)
            label_btns = self.create_label_btns(row)

            # the key, usually the image path
            key = row[self.identifier]

            # current labeled label
            current = self.progress['data'][key]
            if current is not None:
                DOM(f"Current Label:{current}", "h5")()

            # navigation buttons
            nav_btns = list(btn for btn in [self.create_show_last_btn(key),
                                            self.create_show_next_btn(key),
                                            self.create_save_btn(),
                                            self.create_save_to_csv(),
                                            ] if btn is not None)
            display(VBox([label_btns,
                          HBox(nav_btns)
                          ]))

    def get_row_data(self, key: str):
        identifier = self.identifier
        row = dict(self.image_df.query(
            f"{identifier}=='{key}'").to_dict(orient='records')[0])
        return row

    def render_page(self):
        """
        Render a new page
        """
        try:
            key = next(self.gen)
        except StopIteration:
            self.save_progress()
            self.done_page()
            return
        self[key]

    def create_show_last_btn(self, key: str) -> Button:
        """
        create a show last button
            with current key
        """
        keys = list(self.progress["data"].keys())
        idx = keys.index(str(key))
        if idx == 0:
            return None
        last_key = keys[idx-1]

        def show_last_click():
            self[last_key]
        btn = Button(description="Last", icon="arrow-left")
        btn.click = show_last_click
        return btn

    def create_show_next_btn(self, key: str) -> Button:
        """
        create a show next button
            with current key
        """
        keys = list(self.progress["data"].keys())
        idx = keys.index(str(key))
        if idx >= len(self.progress["data"])-1:
            return None
        next_key = keys[idx+1]

        def show_next_click():
            self[next_key]
        btn = Button(description="Next", icon="arrow-right")
        btn.click = show_next_click
        return btn

    def create_save_btn(self) -> Button:
        """
        create a save JSON button
        """
        btn = Button(description="Save JSON", icon='save')
        btn.click = self.save_progress
        return btn

    @staticmethod
    def csv_val_processing(v):
        return v

    def save_to_csv(self):
        DOM("Please name a filepath for csv file like ./progress.csv", "div")()

        # ask interactively for the csv saving path
        @interact_manual
        def save_csv(path="./progress.csv"):
            if len(self.progress['data']) == 0:
                DOM("Nothing to save", "div")()

            # the labeled results, filter out the empety progress
            keys, vals = zip(*list(
                (k, self.csv_val_processing(v))
                for k, v in self.progress["data"].items() if v is not None))
            pd.DataFrame({"path": keys, "label": vals}
                         ).to_csv(path, index=False)
            DOM(f"Progress saved to: '{path}'", "div")()

    def create_save_to_csv(self) -> Button:
        """
        create a save to csv button
        """
        btn = Button(description="CSV", icon='save')
        btn.click = self.save_to_csv
        return btn

    @property
    def get_y(self):
        mapping = dict(
            (str(k), v) for k, v in self.progress["data"].items()
            if v is not None)
        print(f"Creating dataset with {len(mapping)} labels")

        def get_y_func(x):
            return str(mapping.get(str(x)))
        return get_y_func


class SingleClassImageLabeler(ImageLabeler):
    default_file = "unpackai_single.json"

    def __init__(self, image_folder: Path):
        """
        path: Path, a folder full of images
        """
        super().__init__(image_folder)

    def __call__(self, labels: List[str] = ["pos", "neg"]):
        self.gen = super().__call__(labels)

        self.render_page()

        display(self.output)

    def create_label_btns(self, row: Dict[str, str]) -> HBox:
        """
        Create labels control
        """
        btns = []

        def create_callback(label):
            def callback():
                k = row[self.identifier]
                self.progress["data"][str(k)] = label
                self.render_page()
            return callback
        for label in self.labels:
            btn = Button(description=label, icon="check-circle")
            btn.click = create_callback(label)
            btns.append(btn)

        return HBox(btns)

    def done_page(self):
        self.output.clear_output()
        with self.output:
            DOM("That's the end of the iteration", "h3")()
            display(HBox([
                self.create_save_btn(),
                self.create_save_to_csv(),
            ]))


class MultiClassImageLabeler(ImageLabeler):
    default_file = "unpackai_multi.json"

    def __init__(self, image_folder: Path):
        """
        path: Path, a folder full of images
        """
        super().__init__(image_folder)

    def __call__(self, labels: List[str] = ["pos", "neg"]):
        self.gen = super().__call__(labels)
        DOM("press Command(mac) or Ctrl(win/linux) to select multiple", "h4")()
        self.render_page()
        display(self.output)

    @staticmethod
    def csv_val_processing(v):
        """
        processing on value before save to csv
        """
        return ",".join(v)

    def create_label_btns(self, row):
        """
        Create labels control for multicategorcial image dataset
        """
        btns = []
        select = SelectMultiple(options=self.labels)
        btn = Button(description="Okay!", icon="check-circle")

        def callback():
            k = row[self.identifier]
            self.progress["data"][str(k)] = list(select.value)
            self.render_page()

        btn.click = callback

        return HBox([select, btn])

    def done_page(self):
        self.output.clear_output()
        with self.output:
            DOM("That's the end of the iteration", "h3")()
            display(HBox([
                self.create_save_btn(),
                self.create_save_to_csv(),
            ]))

    @property
    def get_y(self):
        mapping = dict(
            (str(k), v) for k, v in self.progress["data"].items()
            if v is not None)
        print(f"Creating dataset with {len(mapping)} labels")

        def get_y_func(x):
            target = mapping.get(str(x))
            target = list() if target is None else target
            return target
        return get_y_func

    @staticmethod
    def gety_from_csv(csv_path: Path, path:str='path', label:str='label', sep=','):
        """
        csv_path: path for csv file
        sep: separator, default:',', could be '\t' for tab
        path: name of column contains image path info
        label: name of the column contains label info
        """
        df = pd.read_csv(csv_path, sep = sep)
        df[label] = df[label].fillna("")
        df[label] = df[label].apply(
            lambda x:list(map(
                lambda i:i.strip(),str(x).replace('，',',').split(',')))
            if x !='' else []
        )
        all_labels = []
        for lbls in list(df[label]):
            all_labels += lbls
        all_labels = list(set(all_labels))
        print(f"All possible labels:\t{all_labels}")
        df[path] = df[path].apply(lambda x:str(Path(x).resolve()))
        mapping = dict(df[[path, label]].values)
        def get_y(x):
            y = mapping.get(str(Path(x).resolve()))
            if y is None:
                return []
            return y
        return get_y