import numpy as np
import re
import os

from PyQt5 import QtWidgets as QW
from PyQt5.QtCore import Qt, QEvent
from PyQt5.QtGui import QImage, QPixmap

from df3d.core import Core
import glob


def main():
    """Main entry point to run the GUI."""
    import sys

    cli_args = parse_cli_args(sys.argv)

    assert len(
        glob.glob(
            os.path.join(cli_args["input_folder"], cli_args["output_subfolder"])
            + "/df3d_result*.pkl"
        )
    ), f"Before running the GUI, run df3d-cli on folder {cli_args['input_folder']} first and generate df3d_result file"

    app = QW.QApplication([])
    window = DeepflyGUI()
    window.setup(**cli_args)
    window.set_width(app.desktop().size().width())
    window.show()
    app.exec_()


def parse_cli_args(argv):
    """Parses the argument string argv.

    Returns:
    A simple namespace with the parsed arguments values.
    """
    args = {}
    args["output_subfolder"] = "df3d"
    try:
        args["input_folder"] = argv[1]
        args["num_images_max"] = int(argv[2])
        args["output_subfolder"] = argv[2]
    except (IndexError, ValueError):
        pass
    return args


class DeepflyGUI(QW.QWidget):
    """Graphical User Interface Widget for DeepFly."""

    def __init__(self):
        QW.QWidget.__init__(self)
        self.img_id = 0
        self.core = None

    def setup(self, input_folder=None, output_subfolder=None, num_images_max=None):
        """Configures the interface and prompts user for missing data."""

        if not input_folder:
            input_folder = self.prompt_for_directory()

        if not output_subfolder:
            output_subfolder = self.prompt_output_subdirectory_name()

        self.core = Core(input_folder, output_subfolder, num_images_max, None)
        self.setup_layout()
        self.onclick_image_mode()

    def set_width(self, width):
        """Sets the GUI's window's width and resizes its height accordingly."""

        hw_ratio = self.core.image_shape[0] * 1.2 / self.core.image_shape[1]
        height = int(width / hw_ratio)
        self.resize(width, height)

    def setup_layout(self):
        """Creates the GUI layout (buttons and stuff.)"""

        # --- Create checkboxes ---
        # self.checkbox_solve_bp = QW.QCheckBox("Auto-correct", self)
        # self.checkbox_solve_bp.setChecked(True)
        # self.checkbox_solve_bp.stateChanged.connect(self.update_frame)

        # --- Create buttons ---
        def mb(text, onClick):
            b = QW.QPushButton(text, self)
            b.setMaximumWidth(b.fontMetrics().boundingRect(text).width() + 27)
            b.clicked.connect(onClick)
            return b

        self.button_prev = mb("<", self.onclick_prev_image)
        self.button_next = mb(">", self.onclick_next_image)
        self.button_last = mb(">>", self.onclick_last_image)
        self.button_first = mb("<<", self.onclick_first_image)
        #self.button_prev_err = mb("< previous error", self.onclick_prev_error)
        #self.button_next_err = mb("next error >", self.onclick_next_error)
        #self.button_pose_save = mb("Save", self.onclick_save_pose)
        self.button_pose_mode = mb("Pose", self.onclick_pose_mode)
        self.button_image_mode = mb("Image", self.onclick_image_mode)
        #self.button_heatmap_mode = mb("Prob. Map", self.onclick_heatmap_mode)
        button_textbox_img_id_go = mb("Go",  self.onclick_goto_img)
        #self.button_calibrate_calc = mb("Calibration", self.onclick_calibrate)

        self.button_correction_mode = \
            mb("Correction", self.onclick_correction_mode)
        #self.button_camera_order = \
        #    mb("Camera ordering", self.onclick_camera_order)
        #self.button_pose_estimate = \
        #    mb("2D Pose Estimation", self.onclick_pose2d_estimation)

        self.button_image_mode.setCheckable(True)
        self.button_correction_mode.setCheckable(True)
        self.button_pose_mode.setCheckable(True)
        # self.button_heatmap_mode.setCheckable(True)

        self.textbox_img_id = QW.QLineEdit(str(self.img_id), self)
        self.textbox_img_id.setFixedWidth(100)

        self.combo_joint_id = QW.QComboBox(self)
        self.combo_joint_id.addItem("View all joints", [])
        for i in range(self.core.number_of_joints):
            self.combo_joint_id.addItem(f"View joint {i}", [i])
        self.combo_joint_id.activated[str].connect(self.update_frame)

        # --- Create widgets to display images and pose results ---
        def make_image_view(cam_id):
            iv = QW.QLabel()
            iv.setScaledContents(True)
            iv.cam_id = cam_id
            return iv

        top_row = [make_image_view(cam_id) for cam_id in [0, 1, 2]]
        bottom_row = [make_image_view(cam_id) for cam_id in [4, 5, 6]]
        self.image_views = top_row + bottom_row

        for image_view in self.image_views:
            image_view.installEventFilter(self)

        # --- Layouts ---
        layout_h_images = QW.QHBoxLayout()
        layout_h_images.setSpacing(1)
        for image_pose in top_row:
            layout_h_images.addWidget(image_pose)
            image_pose.resize(image_pose.sizeHint())

        layout_h_images_bot = QW.QHBoxLayout()
        layout_h_images_bot.setSpacing(1)
        for image_pose in bottom_row:
            layout_h_images_bot.addWidget(image_pose)
            image_pose.resize(image_pose.sizeHint())

        layout_h_buttons_top = QW.QHBoxLayout()
        layout_h_buttons_top.setSpacing(3)
        layout_h_buttons_top.setAlignment(Qt.AlignRight)
        #layout_h_buttons_top.addWidget(self.button_pose_estimate)
        #layout_h_buttons_top.addWidget(self.button_calibrate_calc)
        #layout_h_buttons_top.addWidget(self.button_camera_order)
        #layout_h_buttons_top.addWidget(self.button_pose_save)
        layout_h_buttons_top.addStretch()
        layout_h_buttons_top.addWidget(self.button_image_mode)
        layout_h_buttons_top.addWidget(self.button_pose_mode)
        layout_h_buttons_top.addWidget(self.button_correction_mode)
        # layout_h_buttons_top.addWidget(self.button_heatmap_mode)

        layout_h_buttons = QW.QHBoxLayout()
        layout_h_buttons.setSpacing(1)
        layout_h_buttons.addWidget(self.button_first)
        layout_h_buttons.addWidget(self.button_prev)
        layout_h_buttons.addWidget(self.button_next)
        layout_h_buttons.addWidget(self.button_last)
        layout_h_buttons.addWidget(self.textbox_img_id)
        layout_h_buttons.addWidget(button_textbox_img_id_go)
        layout_h_buttons.addStretch()
        #layout_h_buttons.addWidget(self.button_prev_err)
        #layout_h_buttons.addWidget(self.button_next_err)
        layout_h_buttons.addStretch()
        # layout_h_buttons.addWidget(self.checkbox_solve_bp)
        layout_h_buttons.addStretch()
        layout_h_buttons.addWidget(self.combo_joint_id)

        layout_v = QW.QVBoxLayout()
        layout_v.addLayout(layout_h_buttons_top)
        layout_v.addLayout(layout_h_images)
        layout_v.addLayout(layout_h_images_bot)
        layout_v.addLayout(layout_h_buttons)
        layout_v.setSpacing(0)

        self.setLayout(layout_v)
        self.setWindowTitle(self.core.input_folder)

    # ------------------------------------------------------------------
    # user input

    def onclick_camera_order(self):
        """Prompts for a new camera ordering."""

        cidread2cid = self.prompt_for_camera_ordering()
        if cidread2cid is None:  # canceled
            return
        if self.core.update_camera_ordering(cidread2cid):
            self.update_frame()
        else:
            self.display_error_message("Ordering not changed (wrong format).")

    def onclick_pose2d_estimation(self):
        """Runs the core's pose2d_estimation routine and switches to correction mode."""
        self.core.pose2d_estimation()
        self.onclick_correction_mode()

    def onclick_first_image(self):
        """Displays the first image."""
        self.display_img(0)

    def onclick_last_image(self):
        """Displays the last image."""
        self.display_img(self.core.max_img_id)

    def onclick_prev_image(self):
        """Displays the previous image if is exists."""
        self.display_img(max(self.img_id - 1, 0))

    def onclick_next_image(self):
        """Displays the next image if it exists."""
        self.display_img(min(self.core.max_img_id, self.img_id + 1))

    def onclick_prev_error(self):
        """Among previous images, displays the last one with an error, if it exists."""
        prev_img = self.core.prev_error(self.img_id)
        if prev_img is not None:
            self.display_img(prev_img)
        else:
            msg = "No error remaining among previous images"
            self.display_error_message(msg)

    def onclick_next_error(self):
        """Among next images, displays the first one with an error, if it exists."""
        next_img = self.core.next_error(self.img_id)
        if next_img is not None:
            self.display_img(next_img)
        else:
            msg = "No error remaining among next images"
            self.display_error_message(msg)

    def onclick_calibrate(self):
        """Prompts for a calibration range and calibrates."""
        rng = self.prompt_for_calibration_range()
        if rng is not None:  # not canceled
            self.core.calibrate_calc(*rng)

    def onclick_save_pose(self):
        """Saves the pose estimation results and manual corrections."""
        self.core.save_pose()
        self.core.save_corrections()

    def onclick_goto_img(self):
        """Displays the image whose id is in the textbox, if valid id."""
        try:
            img_id = int(self.textbox_img_id.text())
            self.display_img(img_id)
            self.setFocus()
        except BaseException:
            msg = "Textbox content should be an integer image id"
            self.display_error_message(msg)
            self.textbox_img_id.setText(str(self.img_id))

    def onclick_image_mode(self):
        """Switches the GUI to image mode."""
        self.uncheck_mode_buttons()
        self.button_image_mode.setChecked(True)
        self.combo_joint_id.setEnabled(False)
        self.correction_controls_enabled(False)
        self.belief_propagation_enabled = lambda: False
        self.display_method = lambda c, i, j: self.core.get_image(c, i)
        self.update_frame()

    def onclick_pose_mode(self):
        """Switches the GUI to pose estimation results mode."""
        if not self.core.has_pose:
            return False
        self.uncheck_mode_buttons()
        self.button_pose_mode.setChecked(True)
        self.combo_joint_id.setEnabled(True)
        self.correction_controls_enabled(False)
        self.belief_propagation_enabled = lambda: False
        self.display_method = lambda c, i, j: self.core.plot_2d(c, i, joints=j)
        self.update_frame()

    def onclick_correction_mode(self):
        """Switches the GUI to manual corrections mode."""
        if not self.core.has_pose:
            return False
        self.uncheck_mode_buttons()
        self.button_correction_mode.setChecked(True)
        self.combo_joint_id.setEnabled(True)
        self.correction_controls_enabled(True)

        # self.belief_propagation_enabled = lambda: \
        #    self.checkbox_solve_bp.isChecked()

        self.display_method = lambda c, i, j: self.core.plot_2d(
            c, i, with_corrections=True, joints=j
        )

        self.update_frame()

    def keyPressEvent(self, event):
        switch = {
            Qt.Key_Return: self.onclick_goto_img,
            Qt.Key_A: self.onclick_prev_image,
            Qt.Key_D: self.onclick_next_image,
            # Qt.Key_H: self.onclick_heatmap_mode,
            Qt.Key_I: self.onclick_image_mode,
            Qt.Key_X: self.onclick_pose_mode,
            Qt.Key_C: self.onclick_correction_mode,
            Qt.Key_T: self.onclick_save_pose,
        }
        default_action = lambda: None
        action = switch.get(event.key(), default_action)
        action()

    # ------------------------------------------------------------------
    # user feedback and prompt

    def prompt_for_directory(self):
        """Prompts for a directory and returns its path.

        Returns:
        String: path to a directory.
        """
        return str(
            QW.QFileDialog.getExistingDirectory(
                self,
                directory="./",
                caption="Select Directory",
                options=QW.QFileDialog.DontUseNativeDialog,
            )
        )

    def prompt_output_subdirectory_name(self):
        """Prompts for the ouput subdirectory name.

        Returns:
        String: the name of the subdirectory in which to write output.
        """
        ok_pressed = False
        while not ok_pressed:
            text, ok_pressed = QW.QInputDialog.getText(
                self,
                "Name of output sub-directory",
                "Name:",
                QW.QLineEdit.Normal,
                "df3d",
            )
        return str(text)

    def prompt_for_camera_ordering(self):
        """Prompts for a camera ordering.

        Returns:
        list[int]: a list of camera ids.
        """
        text, ok_pressed = QW.QInputDialog.getText(
            self, "Rename Images", "Camera order:", QW.QLineEdit.Normal, ""
        )
        if ok_pressed:
            cidread2cid = re.findall(r"[0-9]+", text)
            cidread2cid = [int(x) for x in cidread2cid]
            return cidread2cid

    def prompt_for_calibration_range(self):
        """Prompts for a range of ids to use for calibration.

        Returns:
        pair[int, int]: (min_id, max_id)
        """
        text, okPressed = QW.QInputDialog.getText(
            self,
            "Calibration",
            "Range of images:",
            QW.QLineEdit.Normal,
            f"0-{self.core.max_img_id}",
        )
        if okPressed:
            numbers = re.findall(r"[0-9]+", text)
            numbers = [int(x) for x in numbers]
            if len(numbers) == 2:
                return numbers[0], numbers[1]
            else:
                msg = "Please provide a range such as: 0-10"
                self.display_error_message(msg)

    def display_error_message(self, message):
        """Displays an error message to user."""
        msgBox = QW.QMessageBox()
        msgBox.setText(message)
        msgBox.exec()

    # ------------------------------------------------------------------

    def uncheck_mode_buttons(self):
        self.button_correction_mode.setChecked(False)
        # self.button_heatmap_mode.setChecked(False)
        self.button_image_mode.setChecked(False)
        self.button_pose_mode.setChecked(False)

    def correction_controls_enabled(self, enabled):
        pass
        #self.button_next_err.setEnabled(enabled)
        #self.button_prev_err.setEnabled(enabled)
        #self.checkbox_solve_bp.setEnabled(enabled)


    def display_img(self, img_id):
        self.img_id = img_id
        # print('wtf')
        self.textbox_img_id.setText(str(self.img_id))
        self.update_frame()

    def update_frame(self):
        for iv in self.image_views:
            self.update_image_view(iv)

    def update_image_view(self, iv):
        joints = self.combo_joint_id.currentData()
        image_array = self.display_method(iv.cam_id, self.img_id, joints)
        im = image_array.astype(np.uint8)
        height, width, _ = im.shape
        bytesPerLine = 3 * width
        qIm = QImage(im, width, height, bytesPerLine, QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(qIm)
        pixmap = pixmap.scaledToWidth(400)
        iv.setPixmap(pixmap)

    def eventFilter(self, iv, e):
        """Event filter listening to the image views mouse events
        to handle manual corrections
        """
        MousePress = QEvent.MouseButtonPress
        MouseMove = QEvent.MouseMove
        left_press = e.type() == MousePress and e.buttons() == Qt.LeftButton
        left_move = e.type() == MouseMove and e.buttons() == Qt.LeftButton
        correction_mode = self.button_correction_mode.isChecked()

        if correction_mode and (left_press or left_move):
            frame = iv.frameGeometry()
            x = int(e.x() * self.core.image_shape[0] / frame.width())
            y = int(e.y() * self.core.image_shape[1] / frame.height())

            if left_press:
                joint = self.core.nearest_joint(iv.cam_id, self.img_id, x, y)
                self.joint_being_corrected = joint
                return False

            elif left_move:
                joint = self.joint_being_corrected
                self.core.move_joint(iv.cam_id, self.img_id, joint, x, y)
                self.update_image_view(iv)
                return False

        return super().eventFilter(iv, e)


if __name__ == "__main__":
    main()
