import logging

import os

import sys

import tkinter as tk

from tkinter import IntVar, Radiobutton, filedialog, messagebox, ttk

from typing import Dict

import numpy as np

import Pmw

from skimage.color import label2rgb

from skimage.measure import label

from dosma.cli import SUPPORTED_QUANTITATIVE_VALUES, SUPPORTED_SCAN_TYPES, parse_args

from dosma.core.io import format_io_utils as fio_utils

from dosma.core.orientation import AXIAL, CORONAL, SAGITTAL

from dosma.gui.dosma_gui import ScanReader

from dosma.gui.gui_utils.filedialog_reader import FileDialogReader

from dosma.gui.im_viewer import IndexTracker

from dosma.gui.preferences_viewer import PreferencesManager

from dosma.msk import knee

import matplotlib

import matplotlib.pyplot as plt

from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk

matplotlib.use("TkAgg")

LARGE_FONT = ("Verdana", 12)

_logger = logging.getLogger(__name__)

class DosmaViewer(tk.Tk):

    def __init__(self, *args, **kwargs):

        tk.Tk.__init__(self, *args, **kwargs)

        container = tk.Frame(self)

        container.pack(side="top", fill="both", expand=True)

        container.grid_rowconfigure(0, weight=1)

        container.grid_columnconfigure(0, weight=1)

        self.frames = {}

        self.protocol("WM_DELETE_WINDOW", self.on_closing)

        for F in (StartPage, DosmaFrame, PageThree, AnalysisFrame):

            frame = F(container, self)

            self.frames[F] = frame

            frame.grid(row=0, column=0, sticky="nsew")

        self.show_frame(StartPage)

        self.pref = PreferencesManager()

    def on_closing(self):

        if messagebox.askokcancel("Quit", "Do you want to quit?"):

            sys.exit()

    def show_frame(self, cont):

        frame = self.frames[cont]

        frame.tkraise()

    def show_preferences(self):

        self.pref.show_window(self)

class StartPage(tk.Frame):

    def __init__(self, parent, controller):

        tk.Frame.__init__(self, parent)

        # photo = tk.PhotoImage(file="./defaults/skel-rotate.gif")

        # label1 = tk.Label(image=photo)

        # label1.pack()

        label = tk.Label(self, text="Start Page", font=LARGE_FONT)

        label.pack(pady=10, padx=10)

        button2 = ttk.Button(self, text="Scan", command=lambda: controller.show_frame(DosmaFrame))

        button2.pack()

        button3 = ttk.Button(

            self, text="Knee Analysis", command=lambda: controller.show_frame(AnalysisFrame)

        )

        button3.pack()

        button3 = ttk.Button(

            self, text="Image Viewer", command=lambda: controller.show_frame(PageThree)

        )

        button3.pack()

        button3 = ttk.Button(

            self, text="Preferences", command=lambda: controller.show_preferences()

        )

        button3.pack()

class AnalysisFrame(tk.Frame):

    __TISSUES_KEY = "Tissues"

    __QUANTITATIVE_VALUES_KEY = "Quantitative values"

    __LOAD_PATH_KEY = "Load data"

    __PID_KEY = "pid"

    __MEDIAL_TO_LATERAL_ORIENTATION_KEY = "ml"

    def __init__(self, parent, controller):

        tk.Frame.__init__(self, parent)

        self.manager: Dict = {}

        self.gui_manager: Dict = {}

        self.balloon = Pmw.Balloon()

        self.__init_manager()

        self.__base_gui()

        self.preferences = PreferencesManager()

        self.file_dialog_reader = FileDialogReader()

        self.scan_reader = ScanReader(self)

        button1 = ttk.Button(self, text="Home", command=lambda: controller.show_frame(StartPage))

        button1.pack(anchor="se", side="right")

        button1 = ttk.Button(self, text="Run", command=lambda: self.execute())

        button1.pack(anchor="sw", side="left")

    def execute(self):

        try:

            load_path = self.manager[self.__LOAD_PATH_KEY].get()

            if not load_path:

                raise ValueError("Load path not defined")

            preferences_str = self.preferences.get_cmd_line_str().strip()

            tissue_str = ""

            for c, t in enumerate(self.manager[self.__TISSUES_KEY]):

                if t.get():

                    tissue_str += "--%s " % knee.SUPPORTED_TISSUES[c].STR_ID

            tissue_str = tissue_str.strip()

            if not tissue_str:

                raise ValueError("No tissues selected")

            qv_str = ""

            for c, qv in enumerate(self.manager[self.__QUANTITATIVE_VALUES_KEY]):

                if qv.get():

                    qv_str += "--%s " % SUPPORTED_QUANTITATIVE_VALUES[c].name.lower()

            qv_str = qv_str.strip()

            if not qv_str:

                raise ValueError("No quantitative values selected")

            pid = self.manager[self.__PID_KEY].get()

            medial_to_lateral = self.manager[self.__MEDIAL_TO_LATERAL_ORIENTATION_KEY].get()

            if not pid:

                raise ValueError("No PID was provided")

            # analysis string

            str_f = "--l %s %s knee %s --pid %s %s %s" % (

                load_path,

                preferences_str,

                tissue_str,

                pid,

                "--ml" if medial_to_lateral else "",

                qv_str,

            )

            str_f = str_f.strip()

            parse_args(str_f.split())

        except Exception as e:

            tk.messagebox.showerror(str(type(e)), e.__str__())

    def __init_manager(self):

        self.manager[self.__LOAD_PATH_KEY] = tk.StringVar()

        self.manager[self.__TISSUES_KEY] = [

            tk.BooleanVar() for i in range(len(knee.SUPPORTED_TISSUES))

        ]

        self.manager[self.__QUANTITATIVE_VALUES_KEY] = [

            tk.BooleanVar() for i in range(len(SUPPORTED_QUANTITATIVE_VALUES))

        ]

        self.manager[self.__PID_KEY] = tk.StringVar()

        self.manager[self.__MEDIAL_TO_LATERAL_ORIENTATION_KEY] = tk.BooleanVar()

    def __display_pid_info(self):

        hb = tk.Frame(self)

        hb.pack(side="top", anchor="nw")

        _label = tk.Label(hb, text=self.__PID_KEY.upper())

        _label.pack(side="left", anchor="w", pady=10)

        t = tk.Entry(hb, textvariable=self.manager[self.__PID_KEY])

        t.pack(side="left", anchor="w", pady=10)

        self.balloon.bind(_label, "Patient id")

    def __display_data_loader(self):

        hb = tk.Frame(self)

        filedialog = FileDialogReader(self.manager[self.__LOAD_PATH_KEY])

        b = tk.Button(

            hb,

            text=self.__LOAD_PATH_KEY,

            command=lambda fd=filedialog: self.manager[self.__LOAD_PATH_KEY].set(

                fd.get_save_dirpath()

            ),

        )

        b.pack(side="left", anchor="nw", pady=10)

        _label = tk.Label(hb, textvariable=self.manager[self.__LOAD_PATH_KEY])

        _label.pack(side="left", anchor="nw", pady=10)

        hb.pack(side="top", anchor="nw")

    def __display_multi_option(self, label, options_list, boolvar_list):

        hb = tk.Frame(self)

        _label = tk.Label(hb, text="%s:" % label)

        _label.pack(side="left", anchor="w")

        hb.pack(side="top", anchor="nw")

        frames = [tk.Frame(hb)] * (len(options_list) // 3 + 1)

        for ind, option in enumerate(options_list):

            f = frames[ind // 3]

            b = tk.Checkbutton(f, text=option, variable=boolvar_list[ind])

            b.pack(side="top", anchor="nw", pady=5)

        for f in frames:

            f.pack(side="left", anchor="nw")

        return hb

    def __display_tissues(self):

        tissue_names = [x.FULL_NAME for x in knee.SUPPORTED_TISSUES]

        _label = self.__display_multi_option(

            self.__TISSUES_KEY, tissue_names, self.manager[self.__TISSUES_KEY]

        )

        self.balloon.bind(_label, "Tissues to analyze")

    def __display_quant_vals(self):

        quantitative_value_names = [x.name for x in SUPPORTED_QUANTITATIVE_VALUES]

        _label = self.__display_multi_option(

            self.__QUANTITATIVE_VALUES_KEY,

            quantitative_value_names,

            self.manager[self.__QUANTITATIVE_VALUES_KEY],

        )

        self.balloon.bind(_label, "Quantitative values to analyze")

    def __display_knee_info(self):

        hb = tk.Frame(self)

        hb.pack(side="top", anchor="nw")

        _label = tk.Label(hb, text="Medial -> Lateral: ")

        _label.pack(side="left", anchor="w", pady=10)

        t = tk.Checkbutton(hb, variable=self.manager[self.__MEDIAL_TO_LATERAL_ORIENTATION_KEY])

        t.pack(side="left", anchor="w", pady=10)

        self.balloon.bind(_label, "Select if Dicoms proceed in medial->lateral direction")

    def __base_gui(self):

        self.__display_data_loader()

        self.__display_pid_info()

        self.__display_tissues()

        self.__display_knee_info()

        self.__display_quant_vals()

class DosmaFrame(tk.Frame):

    __SCAN_KEY = "Scan"

    __TISSUES_KEY = "Tissues"

    __DICOM_PATH_KEY = "Read dicoms"

    __LOAD_PATH_KEY = "Load data"

    __SAVE_PATH_KEY = "Save path"

    __DATA_KEY = "data"  # Track option menu for dicom/load path

    __DATA_PATH_KEY = "datapath"  # Track filepath associated with option menu

    __IGNORE_EXTENSION_KEY = "Ignore extension"

    def __init__(self, parent, controller):

        tk.Frame.__init__(self, parent)

        self.file_dialog_reader = FileDialogReader()

        self.manager: Dict = {}

        self.gui_manager: Dict = {}

        self.balloon = Pmw.Balloon()

        self.__init_manager()

        self.__base_gui()

        self.preferences = PreferencesManager()

        self.scan_reader = ScanReader(self)

        button1 = ttk.Button(self, text="Home", command=lambda: controller.show_frame(StartPage))

        button1.pack(anchor="se", side="right")

        button1 = ttk.Button(self, text="Run", command=lambda: self.execute())

        button1.pack(anchor="sw", side="left")

        self.InitUI()

    def execute(self):

        try:

            save_path = self.manager[self.__SAVE_PATH_KEY].get()

            if not save_path:

                raise ValueError("Save path not defined")

            action_str = self.scan_reader.get_cmd_line_str().strip()

            if not action_str:

                raise ValueError("No action selected")

            preferences_str = self.preferences.get_cmd_line_str().strip()

            source = "d"

            if self.manager[self.__DATA_KEY].get() == self.__LOAD_PATH_KEY:

                source = "l"

            tissue_str = ""

            for c, t in enumerate(self.manager[self.__TISSUES_KEY]):

                if t.get():

                    tissue_str += "--%s " % knee.SUPPORTED_TISSUES[c].STR_ID

            tissue_str = tissue_str.strip()

            if not tissue_str:

                raise ValueError("No tissues selected")

            ignore_ext = self.manager[self.__IGNORE_EXTENSION_KEY].get()

            str_f = "--%s %s --s %s %s %s %s %s %s" % (

                source,

                self.manager[self.__DATA_PATH_KEY].get(),

                save_path,

                preferences_str,

                "--ignore_ext" if ignore_ext else "",

                self.manager[self.__SCAN_KEY].get(),

                tissue_str,

                action_str,

            )

            _logger.info("CMD LINE INPUT: %s" % str_f)

            parse_args(str_f.split())

        except Exception as e:

            tk.messagebox.showerror(str(type(e)), e.__str__())

    def __init_manager(self):

        self.manager[self.__SCAN_KEY] = tk.StringVar()

        self.manager[self.__TISSUES_KEY] = [

            tk.BooleanVar() for i in range(len(knee.SUPPORTED_TISSUES))

        ]

        self.manager[self.__DATA_KEY] = tk.StringVar()

        self.manager[self.__DATA_PATH_KEY] = tk.StringVar()

        self.manager[self.__SCAN_KEY].trace_add("write", self.__on_scan_change)

        self.manager[self.__SAVE_PATH_KEY] = tk.StringVar()

        self.manager[self.__IGNORE_EXTENSION_KEY] = tk.BooleanVar()

    def __on_scan_change(self, *args):

        scan_id = self.manager[self.__SCAN_KEY].get()

        scan = None

        for x in SUPPORTED_SCAN_TYPES:

            if x.NAME == scan_id:

                scan = x

        self.scan_reader.load_scan(scan)

        assert scan is not None, "No scan selected"

    def __update_svar(self, *args):

        svar = self.manager[self.__DATA_PATH_KEY]

        selected_option = self.manager[self.__DATA_KEY].get()

        if selected_option == self.__DICOM_PATH_KEY:

            fp = self.file_dialog_reader.get_volume_filepath(

                selected_option, im_type=fio_utils.ImageDataFormat.dicom

            )

        elif selected_option == self.__LOAD_PATH_KEY:

            fp = self.file_dialog_reader.get_dirpath(selected_option)

        else:

            raise ValueError("%s key not found" % self.__DATA_KEY)

        if not fp:

            svar.set("")

            return

        svar.set(fp)

        if selected_option == self.__LOAD_PATH_KEY:

            self.manager[self.__SAVE_PATH_KEY].set(fp)

    def __display_data_loader(self):

        s_var = self.manager[self.__DATA_PATH_KEY]

        hb = tk.Frame(self)

        label = tk.Label(hb, text="Data source: ")

        label.pack(side="left", anchor="nw", pady=10)

        options = [self.__DICOM_PATH_KEY, self.__LOAD_PATH_KEY]

        menu = tk.OptionMenu(

            hb, self.manager[self.__DATA_KEY], *options, command=self.__update_svar

        )

        menu.pack(side="left", anchor="nw", pady=10)

        label = tk.Label(hb, textvariable=s_var)

        label.pack(side="left", anchor="nw", pady=10)

        hb.pack(side="top", anchor="nw")

        self.balloon.bind(hb, "Read dicoms or load data")

        hb = tk.Frame(self)

        # filedialog = FileDialogReader(self.manager[self.__SAVE_PATH_KEY])

        b = tk.Button(

            hb,

            text=self.__SAVE_PATH_KEY,

            command=lambda fd=self.file_dialog_reader: self.manager[self.__SAVE_PATH_KEY].set(

                fd.get_save_dirpath()

            ),

        )

        b.pack(side="left", anchor="nw", pady=10)

        label = tk.Label(hb, textvariable=self.manager[self.__SAVE_PATH_KEY])

        label.pack(side="left", anchor="nw", pady=10)

        hb.pack(side="top", anchor="nw")

        hb = tk.Frame(self)

        b = tk.Checkbutton(

            hb, text=self.__IGNORE_EXTENSION_KEY, variable=self.manager[self.__IGNORE_EXTENSION_KEY]

        )

        b.pack(side="left", anchor="nw", pady=10)

        self.balloon.bind(b, "Ignore '.dcm' extension when loading dicoms")

        hb.pack(side="top", anchor="nw")

    def __display_tissues(self):

        hb = tk.Frame(self)

        _label = tk.Label(hb, text="Tissues:")

        _label.pack(side="left", anchor="w")

        hb.pack(side="top", anchor="nw")

        frames = [tk.Frame(hb)] * (len(knee.SUPPORTED_TISSUES) // 3 + 1)

        for ind, tissue in enumerate(knee.SUPPORTED_TISSUES):

            f = frames[ind // 3]

            b = tk.Checkbutton(

                f, text=tissue.FULL_NAME, variable=self.manager[self.__TISSUES_KEY][ind]

            )

            b.pack(side="top", anchor="nw", pady=5)

        for f in frames:

            f.pack(side="left", anchor="nw")

        self.balloon.bind(_label, "Tissues to analyze")

    def __base_gui(self):

        self.__display_data_loader()

        self.__display_tissues()

        hb = tk.Frame(self)

        scan_label = tk.Label(hb, text="Scan:")

        scan_label.pack(side="left", anchor="nw", pady=10)

        options = [x.NAME for x in SUPPORTED_SCAN_TYPES]

        scan_dropdown = tk.OptionMenu(hb, self.manager[self.__SCAN_KEY], *options)

        scan_dropdown.pack(side="left", anchor="nw", pady=10)

        hb.pack(side="top", anchor="nw")

    def InitUI(self):

        self.text_box = tk.Text(self, wrap="word", height=11, width=50)

        self.text_box.pack(anchor="s", side="bottom")

class PageThree(tk.Frame):

    SUPPORTED_FORMATS = (("nifti files", "*.nii\.gz"), ("dicom files", "*.dcm"))  # noqa: W605

    __base_filepath = "../"

    _ORIENTATIONS = [("sagittal", SAGITTAL), ("coronal", CORONAL), ("axial", AXIAL)]

    def __init__(self, parent, controller):

        tk.Frame.__init__(self, parent)

        self._im_display = None

        self.binding_vars: Dict = {}

        fig, ax = plt.subplots(1, 1)

        X = np.random.rand(20, 20, 40)

        self.tracker = IndexTracker(ax, X)

        canvas = FigureCanvasTkAgg(fig, self)

        canvas.draw()

        canvas.get_tk_widget().pack(side=tk.TOP, fill=tk.BOTH, expand=True)

        canvas.mpl_connect("scroll_event", self.tracker.onscroll)

        toolbar = NavigationToolbar2Tk(canvas, self)

        toolbar.update()

        canvas._tkcanvas.pack(side=tk.TOP, fill=tk.BOTH, expand=True)

        self.im = None

        self.mask = None

        self._im_display = None

        button1 = ttk.Button(

            self, text="Back to Home", command=lambda: controller.show_frame(StartPage)

        )

        button1.pack(side=tk.BOTTOM, anchor="sw")

        button2 = ttk.Button(self, text="Load main image", command=self.load_volume_callback)

        button2.pack()

        button3 = ttk.Button(self, text="Load mask", command=self.load_mask_callback)

        button3.pack()

        self.init_reformat_display()

    def __reformat_callback(self, *args):

        self.im_update()

    def init_reformat_display(self):

        orientation_var = IntVar(0)

        orientation_var.trace_add("write", self.__reformat_callback)

        count = 0

        for text, _value in self._ORIENTATIONS:

            b = Radiobutton(self, text=text, variable=orientation_var, value=count)

            b.pack(side=tk.TOP, anchor="w")

            count += 1

        self._orientation = orientation_var

    def load_volume_callback(self):

        im = self.load_volume()

        if not im:

            return

        self.im = im

        self.mask = None

        self.im_update()

    def load_mask_callback(self):

        if not self.im:

            messagebox.showerror("Loading mask failed", "Main image must be loaded prior to mask")

            return

        mask = self.load_volume("Load mask")

        mask.reformat(self.im.orientation, inplace=True)

        try:

            self.__verify_mask_size(self.im.volume, mask.volume)

        except Exception as e:

            messagebox.showerror("Loading mask failed", str(e))

            return

        self.mask = mask

        self.im_update()

    def __verify_mask_size(self, im: np.ndarray, mask: np.ndarray):

        if mask.ndim != 3:

            raise ValueError("Dimension mismatch. Mask must be 3D")

        if im.shape != mask.shape:

            raise ValueError(

                "Dimension mismatch. Image of shape %s, but mask of shape %s"

                % (str(im.shape), str(mask.shape))

            )

    def im_update(self):

        orientation = self.orientation

        self.im.reformat(orientation, inplace=True)

        im = self.im.volume

        im = im / np.max(im)

        if self.mask:

            self.mask.reformat(orientation, inplace=True)

            label_image = label(self.mask.volume)

            im = self.__labeltorgb_3d__(im, label_image, 0.3)

        self.im_display = im

    def __labeltorgb_3d__(self, im: np.ndarray, labels: np.ndarray, alpha: float = 0.3):

        im_rgb = np.zeros(im.shape + (3,))  # rgb channel

        for s in range(im.shape[2]):

            im_slice = im[..., s]

            labels_slice = labels[..., s]

            im_rgb[..., s, :] = label2rgb(labels_slice, image=im_slice, bg_label=0, alpha=alpha)

        return im_rgb

    def load_volume(self, title="Select volume file(s)"):

        files = filedialog.askopenfilenames(initialdir=self.__base_filepath, title=title)

        if len(files) == 0:

            return

        filepath = files[0]

        self.__base_filepath = os.path.dirname(filepath)

        if filepath.endswith(".dcm"):

            filepath = os.path.dirname(filepath)

        im = fio_utils.generic_load(filepath, 1)

        return im

    @property

    def orientation(self):

        ind = self._orientation.get()

        return self._ORIENTATIONS[ind][1]

    @property

    def im_display(self):

        return self._im_display

    @im_display.setter

    def im_display(self, value):

        self._im_display = value

        self.tracker.x = self._im_display