#!/usr/bin python3
""" Command specific tabs of Display Frame of the Faceswap GUI """
import datetime
import gettext
import logging
import os
import sys
import tkinter as tk

from tkinter import ttk
from typing import cast, Dict, Optional, Tuple, TYPE_CHECKING

from .display_graph import TrainingGraph
from .display_page import DisplayOptionalPage
from .custom_widgets import Tooltip
from .analysis import Calculations, Session
from .control_helper import set_slider_rounding
from .utils import FileHandler, get_config, get_images, preview_trigger

if sys.version_info < (3, 8):
    from typing_extensions import get_args, Literal
else:
    from typing import get_args, Literal

if TYPE_CHECKING:
    from PIL import Image

logger = logging.getLogger(__name__)  # pylint: disable=invalid-name

# LOCALES
_LANG = gettext.translation("gui.tooltips", localedir="locales", fallback=True)
_ = _LANG.gettext


class PreviewExtract(DisplayOptionalPage):  # pylint: disable=too-many-ancestors
    """ Tab to display output preview images for extract and convert """

    def display_item_set(self) -> None:
        """ Load the latest preview if available """
        logger.trace("Loading latest preview")  # type:ignore
        size = 256 if self.command == "convert" else 128
        get_images().load_latest_preview(thumbnail_size=int(size * get_config().scaling_factor),
                                         frame_dims=(self.winfo_width(), self.winfo_height()))
        self.display_item = get_images().previewoutput

    def display_item_process(self) -> None:
        """ Display the preview """
        logger.trace("Displaying preview")  # type:ignore
        if not self.subnotebook.children:
            self.add_child()
        else:
            self.update_child()

    def add_child(self) -> None:
        """ Add the preview label child """
        logger.debug("Adding child")
        preview = self.subnotebook_add_page(self.tabname, widget=None)
        lblpreview = ttk.Label(preview, image=get_images().previewoutput[1])
        lblpreview.pack(side=tk.TOP, anchor=tk.NW)
        Tooltip(lblpreview, text=self.helptext, wrap_length=200)

    def update_child(self) -> None:
        """ Update the preview image on the label """
        logger.trace("Updating preview")  # type:ignore
        for widget in self.subnotebook_get_widgets():
            widget.configure(image=get_images().previewoutput[1])

    def save_items(self) -> None:
        """ Open save dialogue and save preview """
        location = FileHandler("dir", None).return_file
        if not location:
            return
        filename = "extract_convert_preview"
        now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = os.path.join(location, f"{filename}_{now}.png")
        get_images().previewoutput[0].save(filename)
        logger.debug("Saved preview to %s", filename)
        print(f"Saved preview to {filename}")


class PreviewTrain(DisplayOptionalPage):  # pylint: disable=too-many-ancestors
    """ Training preview image(s) """
    def __init__(self, *args, **kwargs) -> None:
        self.update_preview = get_config().tk_vars.update_preview
        super().__init__(*args, **kwargs)

    def add_options(self) -> None:
        """ Add the additional options """
        self._add_option_refresh()
        self._add_option_mask_toggle()
        super().add_options()

    def _add_option_refresh(self) -> None:
        """ Add refresh button to refresh preview immediately """
        logger.debug("Adding refresh option")
        btnrefresh = ttk.Button(self.optsframe,
                                image=get_images().icons["reload"],
                                command=lambda x="update": preview_trigger().set(x))  # type:ignore
        btnrefresh.pack(padx=2, side=tk.RIGHT)
        Tooltip(btnrefresh,
                text=_("Preview updates at every model save. Click to refresh now."),
                wrap_length=200)
        logger.debug("Added refresh option")

    def _add_option_mask_toggle(self) -> None:
        """ Add button to toggle mask display on and off """
        logger.debug("Adding mask toggle option")
        btntoggle = ttk.Button(
            self.optsframe,
            image=get_images().icons["mask2"],
            command=lambda x="mask_toggle": preview_trigger().set(x))  # type:ignore
        btntoggle.pack(padx=2, side=tk.RIGHT)
        Tooltip(btntoggle,
                text=_("Click to toggle mask overlay on and off."),
                wrap_length=200)
        logger.debug("Added mask toggle option")

    def display_item_set(self) -> None:
        """ Load the latest preview if available """
        logger.trace("Loading latest preview")  # type:ignore
        if not self.update_preview.get():
            logger.trace("Preview not updated")  # type:ignore
            return
        get_images().load_training_preview()
        self.display_item = get_images().previewtrain

    def display_item_process(self) -> None:
        """ Display the preview(s) resized as appropriate """
        logger.trace("Displaying preview")  # type:ignore
        sortednames = sorted(list(get_images().previewtrain.keys()))
        existing = self.subnotebook_get_titles_ids()
        should_update = self.update_preview.get()

        for name in sortednames:
            if name not in existing:
                self.add_child(name)
            elif should_update:
                tab_id = existing[name]
                self.update_child(tab_id, name)

        if should_update:
            self.update_preview.set(False)

    def add_child(self, name: str) -> None:
        """ Add the preview canvas child

        Parameters
        ----------
        name: str
            The name of the notebook tab to add
        """
        logger.debug("Adding child")
        preview = PreviewTrainCanvas(self.subnotebook, name)
        preview = self.subnotebook_add_page(name, widget=preview)
        Tooltip(preview, text=self.helptext, wrap_length=200)
        self.vars["modified"].set(get_images().previewtrain[name][2])

    def update_child(self, tab_id: int, name: str) -> None:
        """ Update the preview canvas

        Parameters
        ----------
        tab_id: int
            The index of the tab to update
        name: str
            The name of the tab to update
        """
        logger.debug("Updating preview")
        if self.vars["modified"].get() != get_images().previewtrain[name][2]:
            self.vars["modified"].set(get_images().previewtrain[name][2])
            widget = self.subnotebook_page_from_id(tab_id)
            widget.reload()

    def save_items(self) -> None:
        """ Open save dialogue and save preview """
        location = FileHandler("dir", None).return_file
        if not location:
            return
        for preview in self.subnotebook.children.values():
            preview.save_preview(location)


class PreviewTrainCanvas(ttk.Frame):  # pylint: disable=too-many-ancestors
    """ Canvas to hold a training preview image

    Parameters
    ----------
    parent: :class:`tkinter.ttk.Notebook`
        The notebook that the training image canvas belongs to
    previewname: str
        The name of the preview image displayed in the canvas
    """
    def __init__(self, parent: ttk.Notebook, previewname: str) -> None:
        logger.debug("Initializing %s: (previewname: '%s')", self.__class__.__name__, previewname)
        ttk.Frame.__init__(self, parent)

        self.name = previewname
        get_images().resize_image(self.name, None)
        self.previewimage = get_images().previewtrain[self.name][1]

        self.canvas = tk.Canvas(self, bd=0, highlightthickness=0)
        self.canvas.pack(side=tk.TOP, fill=tk.BOTH, expand=True)
        self.imgcanvas = self.canvas.create_image(0,
                                                  0,
                                                  image=self.previewimage,
                                                  anchor=tk.NW)
        self.bind("<Configure>", self.resize)
        logger.debug("Initialized %s:", self.__class__.__name__)

    def resize(self, event: tk.Event) -> None:
        """  Resize the image to fit the frame, maintaining aspect ratio

        Parameters
        ----------
        event: :class:`tkinter.Event`
            The resize event object
        """
        logger.trace("Resizing preview image")  # type:ignore
        framesize: Optional[Tuple[int, int]] = (event.width, event.height)
        # Sometimes image is resized before frame is drawn
        framesize = None if framesize == (1, 1) else framesize
        get_images().resize_image(self.name, framesize)
        self.reload()

    def reload(self) -> None:
        """ Reload the preview image """
        logger.trace("Reloading preview image")  # type:ignore
        self.previewimage = get_images().previewtrain[self.name][1]
        self.canvas.itemconfig(self.imgcanvas, image=self.previewimage)

    def save_preview(self, location: str) -> None:
        """ Save the figure to file.

        Parameters
        ----------
        location: str
            The full path to the location to save the preview image
        """
        filename = self.name
        now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = os.path.join(location, f"{filename}_{now}.png")
        cast("Image.Image", get_images().previewtrain[self.name][0]).save(filename)
        logger.debug("Saved preview to %s", filename)
        print(f"Saved preview to {filename}")


class GraphDisplay(DisplayOptionalPage):  # pylint: disable=too-many-ancestors
    """ The Graph Tab of the Display section """
    def __init__(self,
                 parent: ttk.Notebook,
                 tab_name: str,
                 helptext: str,
                 wait_time: int,
                 command: Optional[str] = None) -> None:
        self._trace_vars: Dict[Literal["smoothgraph", "display_iterations"],
                               Tuple[tk.BooleanVar, str]] = {}
        super().__init__(parent, tab_name, helptext, wait_time, command)

    def set_vars(self) -> None:
        """ Add graphing specific variables to the default variables.

        Overrides original method.

        Returns
        -------
        dict
            The variable names with their corresponding tkinter variable
        """
        tk_vars = super().set_vars()

        smoothgraph = tk.DoubleVar()
        smoothgraph.set(0.900)
        tk_vars["smoothgraph"] = smoothgraph

        raw_var = tk.BooleanVar()
        raw_var.set(True)
        tk_vars["raw_data"] = raw_var

        smooth_var = tk.BooleanVar()
        smooth_var.set(True)
        tk_vars["smooth_data"] = smooth_var

        iterations_var = tk.IntVar()
        iterations_var.set(10000)
        tk_vars["display_iterations"] = iterations_var

        logger.debug(tk_vars)
        return tk_vars

    def on_tab_select(self) -> None:
        """ Callback for when the graph tab is selected.

        Pull latest data and run the tab's update code when the tab is selected.
        """
        logger.debug("Callback received for '%s' tab", self.tabname)
        if self.display_item is not None:
            get_config().tk_vars.refresh_graph.set(True)
        self._update_page()

    def add_options(self) -> None:
        """ Add the additional options """
        self._add_option_refresh()
        super().add_options()
        self._add_option_raw()
        self._add_option_smoothed()
        self._add_option_smoothing()
        self._add_option_iterations()

    def _add_option_refresh(self) -> None:
        """ Add refresh button to refresh graph immediately """
        logger.debug("Adding refresh option")
        tk_var = get_config().tk_vars.refresh_graph
        btnrefresh = ttk.Button(self.optsframe,
                                image=get_images().icons["reload"],
                                command=lambda: tk_var.set(True))
        btnrefresh.pack(padx=2, side=tk.RIGHT)
        Tooltip(btnrefresh,
                text=_("Graph updates at every model save. Click to refresh now."),
                wrap_length=200)
        logger.debug("Added refresh option")

    def _add_option_raw(self) -> None:
        """ Add check-button to hide/display raw data """
        logger.debug("Adding display raw option")
        tk_var = self.vars["raw_data"]
        chkbtn = ttk.Checkbutton(
            self.optsframe,
            variable=tk_var,
            text="Raw",
            command=lambda v=tk_var: self._display_data_callback("raw", v))  # type:ignore
        chkbtn.pack(side=tk.RIGHT, padx=5, anchor=tk.W)
        Tooltip(chkbtn, text=_("Display the raw loss data"), wrap_length=200)

    def _add_option_smoothed(self) -> None:
        """ Add check-button to hide/display smoothed data """
        logger.debug("Adding display smoothed option")
        tk_var = self.vars["smooth_data"]
        chkbtn = ttk.Checkbutton(
            self.optsframe,
            variable=tk_var,
            text="Smoothed",
            command=lambda v=tk_var: self._display_data_callback("smoothed", v))  # type:ignore
        chkbtn.pack(side=tk.RIGHT, padx=5, anchor=tk.W)
        Tooltip(chkbtn, text=_("Display the smoothed loss data"), wrap_length=200)

    def _add_option_smoothing(self) -> None:
        """ Add a slider to adjust the smoothing amount """
        logger.debug("Adding Smoothing Slider")
        tk_var = self.vars["smoothgraph"]
        min_max = (0, 0.999)
        hlp = _("Set the smoothing amount. 0 is no smoothing, 0.99 is maximum smoothing.")

        ctl_frame = ttk.Frame(self.optsframe)
        ctl_frame.pack(padx=2, side=tk.RIGHT)

        lbl = ttk.Label(ctl_frame, text="Smoothing:", anchor=tk.W)
        lbl.pack(pady=5, side=tk.LEFT, anchor=tk.N, expand=True)

        tbox = ttk.Entry(ctl_frame, width=6, textvariable=tk_var, justify=tk.RIGHT)
        tbox.pack(padx=(0, 5), side=tk.RIGHT)

        ctl = ttk.Scale(
            ctl_frame,
            variable=tk_var,
            command=lambda val, var=tk_var, dt=float, rn=3, mm=min_max:  # type:ignore
            set_slider_rounding(val, var, dt, rn, mm))
        ctl["from_"] = min_max[0]
        ctl["to"] = min_max[1]
        ctl.pack(padx=5, pady=5, fill=tk.X, expand=True)
        for item in (tbox, ctl):
            Tooltip(item,
                    text=hlp,
                    wrap_length=200)
        logger.debug("Added Smoothing Slider")

    def _add_option_iterations(self) -> None:
        """ Add a slider to adjust the amount if iterations to display """
        logger.debug("Adding Iterations Slider")
        tk_var = self.vars["display_iterations"]
        min_max = (0, 100000)
        hlp = _("Set the number of iterations to display. 0 displays the full session.")

        ctl_frame = ttk.Frame(self.optsframe)
        ctl_frame.pack(padx=2, side=tk.RIGHT)

        lbl = ttk.Label(ctl_frame, text="Iterations:", anchor=tk.W)
        lbl.pack(pady=5, side=tk.LEFT, anchor=tk.N, expand=True)

        tbox = ttk.Entry(ctl_frame, width=6, textvariable=tk_var, justify=tk.RIGHT)
        tbox.pack(padx=(0, 5), side=tk.RIGHT)

        ctl = ttk.Scale(
            ctl_frame,
            variable=tk_var,
            command=lambda val, var=tk_var, dt=int, rn=1000, mm=min_max:  # type:ignore
            set_slider_rounding(val, var, dt, rn, mm))
        ctl["from_"] = min_max[0]
        ctl["to"] = min_max[1]
        ctl.pack(padx=5, pady=5, fill=tk.X, expand=True)
        for item in (tbox, ctl):
            Tooltip(item,
                    text=hlp,
                    wrap_length=200)
        logger.debug("Added Iterations Slider")

    def display_item_set(self) -> None:
        """ Load the graph(s) if available """
        if Session.is_training and Session.logging_disabled:
            logger.trace("Logs disabled. Hiding graph")  # type:ignore
            self.set_info("Graph is disabled as 'no-logs' has been selected")
            self.display_item = None
            self._clear_trace_variables()
        elif Session.is_training and self.display_item is None:
            logger.trace("Loading graph")  # type:ignore
            self.display_item = Session
            self._add_trace_variables()
        elif Session.is_training and self.display_item is not None:
            logger.trace("Graph already displayed. Nothing to do.")  # type:ignore
        else:
            logger.trace("Clearing graph")  # type:ignore
            self.display_item = None
            self._clear_trace_variables()

    def display_item_process(self) -> None:
        """ Add a single graph to the graph window """
        if not Session.is_training:
            logger.debug("Waiting for Session Data to become available to graph")
            self.after(1000, self.display_item_process)
            return

        logger.debug("Adding graph")
        existing = list(self.subnotebook_get_titles_ids().keys())

        loss_keys = self.display_item.get_loss_keys(Session.session_ids[-1])
        if not loss_keys:
            # Reload if we attempt to get loss keys before data is written
            logger.debug("Waiting for Session Data to become available to graph")
            self.after(1000, self.display_item_process)
            return

        loss_keys = [key for key in loss_keys if key != "total"]
        display_tabs = sorted(set(key[:-1].rstrip("_") for key in loss_keys))

        for loss_key in display_tabs:
            tabname = loss_key.replace("_", " ").title()
            if tabname in existing:
                continue

            display_keys = [key for key in loss_keys if key.startswith(loss_key)]
            data = Calculations(session_id=Session.session_ids[-1],
                                display="loss",
                                loss_keys=display_keys,
                                selections=["raw", "smoothed"],
                                smooth_amount=self.vars["smoothgraph"].get())
            self.add_child(tabname, data)

    def _smooth_amount_callback(self, *args) -> None:
        """ Update each graph's smooth amount on variable change """
        try:
            smooth_amount = self.vars["smoothgraph"].get()
        except tk.TclError:
            # Don't update when there is no value in the variable
            return
        logger.debug("Updating graph smooth_amount: (new_value: %s, args: %s)",
                     smooth_amount, args)
        for graph in self.subnotebook.children.values():
            graph.calcs.set_smooth_amount(smooth_amount)

    def _iteration_limit_callback(self, *args) -> None:
        """ Limit the amount of data displayed in the live graph on a iteration slider
        variable change. """
        try:
            limit = self.vars["display_iterations"].get()
        except tk.TclError:
            # Don't update when there is no value in the variable
            return
        logger.debug("Updating graph iteration limit: (new_value: %s, args: %s)",
                     limit, args)
        for graph in self.subnotebook.children.values():
            graph.calcs.set_iterations_limit(limit)

    def _display_data_callback(self, line: str, variable: tk.BooleanVar) -> None:
        """ Update the displayed graph lines based on option check button selection.

        Parameters
        ----------
        line: str
            The line to hide or display
        variable: :class:`tkinter.BooleanVar`
            The tkinter variable containing the ``True`` or ``False`` data for this display item
        """
        var = variable.get()
        logger.debug("Updating display %s to %s", line, var)
        for graph in self.subnotebook.children.values():
            graph.calcs.update_selections(line, var)

    def add_child(self, name: str, data: Calculations) -> None:
        """ Add the graph for the selected keys.

        Parameters
        ----------
        name: str
            The name of the graph to add to the notebook
        data: :class:`~lib.gui.analysis.stats.Calculations`
            The object holding the data to be graphed
        """
        logger.debug("Adding child: %s", name)
        graph = TrainingGraph(self.subnotebook, data, "Loss")
        graph.build()
        graph = self.subnotebook_add_page(name, widget=graph)
        Tooltip(graph, text=self.helptext, wrap_length=200)

    def save_items(self) -> None:
        """ Open save dialogue and save graphs """
        graphlocation = FileHandler("dir", None).return_file
        if not graphlocation:
            return
        for graph in self.subnotebook.children.values():
            graph.save_fig(graphlocation)

    def _add_trace_variables(self) -> None:
        """ Add tracing for when the option sliders are updated, for updating the graph. """
        for name, action in zip(get_args(Literal["smoothgraph", "display_iterations"]),
                                (self._smooth_amount_callback, self._iteration_limit_callback)):
            var = self.vars[name]
            if name not in self._trace_vars:
                self._trace_vars[name] = (var, var.trace("w", action))

    def _clear_trace_variables(self) -> None:
        """ Clear all of the trace variables from :attr:`_trace_vars` and reset the dictionary. """
        if self._trace_vars:
            for name, (var, trace) in self._trace_vars.items():
                logger.debug("Clearing trace from variable: %s", name)
                var.trace_vdelete("w", trace)
            self._trace_vars = {}

    def close(self) -> None:
        """ Clear the plots from RAM """
        self._clear_trace_variables()
        if self.subnotebook is None:
            logger.debug("No graphs to clear. Returning")
            return

        for name, graph in self.subnotebook.children.items():
            logger.debug("Clearing: %s", name)
            graph.clear()
        super().close()