# Third party modules
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import matplotlib.pyplot as plt
import numpy as np
import tkinter as tk
from tkinter import messagebox

class MaterialParamSelector:
    def __init__(
            self, root, x, y, tth, preselected_materials, label, on_complete):

        self.root = root
        self.root.title('Material Parameter Selection')
        self.on_complete = on_complete  # Completion callback

        # Reference data
        self.ref_data_x = x
        self.ref_data_y = y
        self.ref_data_label = label
        self.tth = tth

        # Materials
        self.materials = []
        self.selected_material = None

        # Create plot
        self.figure, self.ax = plt.subplots()
        self.legend_handles = []
        self.canvas = FigureCanvasTkAgg(self.figure, self.root)
        self.canvas.get_tk_widget().grid(row=0, column=0, rowspan=12)

        # Widgets for material list and parameters
        self.material_listbox = tk.Listbox(root, height=5)
        self.material_listbox.grid(row=0, column=1, rowspan=2)
        self.material_listbox.bind(
            '<<ListboxSelect>>', self.on_material_select)

        self.add_material_button = tk.Button(
            root, text='Add Material', command=self.add_material)
        self.add_material_button.grid(row=0, column=2)

        self.remove_material_button = tk.Button(
            root, text='Remove Material', command=self.remove_material)
        self.remove_material_button.grid(row=1, column=2)

        # Parameter fields
        self.fields = {}
        for i, field in enumerate(
                ['Material Name', 'Space Group',
                 'a', 'b', 'c', 'alpha', 'beta', 'gamma']):
            if i > 1:
                units = 'Angstroms' if i < 5 else 'degrees'
                text = f'{field} ({units})'
            else:
                text = field
            tk.Label(root, text=text).grid(row=i+2, column=1)
            entry = tk.Entry(root)
            entry.grid(row=i+2, column=2)
            self.fields[field] = entry

        self.update_button = tk.Button(
            root, text='Update Material Properties',
            command=self.update_material)
        self.update_button.grid(row=11, column=1, columnspan=2)
        self.update_button = tk.Button(
            root, text='Update Material Properties',
            command=self.update_material)
        self.update_button.grid(row=11, column=1, columnspan=2)

        self.confirm_button = tk.Button(
            root, text='Confirm\nAll\nSelected\nMaterial\nProperties',
            command=self.on_close)
        self.confirm_button.grid(row=0, column=3, rowspan=12)

        # Initial Material Data
        if not preselected_materials:
            self.materials = [None]
            self.add_material(None)
        else:
            for material in preselected_materials:
                self.add_material(material)
        self.selected_material = 0

        # Overwrite the root window's close action to call `self.on_close`
        self.root.protocol('WM_DELETE_WINDOW', self.on_close)

    def plot_reference_data(self):
        if self.ref_data_y is None:
            return
        handle = self.ax.plot(
            self.ref_data_x, self.ref_data_y, label=self.ref_data_label)
        self.legend_handles = handle
        self.ax.legend(handles=self.legend_handles)
        self.ax.set_xlabel('Energy (keV)')
        self.ax.set_ylabel('Intensity (a.u)')
        self.canvas.draw()

    def add_material(self, new_material=None):
        # Local modules
        from CHAP.edd.models import MaterialConfig

        if new_material is None:
            new_material = MaterialConfig(
                material_name='Ti64',
                sgnum=194,
                lattice_parameters=[2.9217, 4.66027]
            )
        self.materials.append(new_material)
        self.material_listbox.insert(tk.END, new_material.material_name)
        self.material_listbox.select_set(tk.END)
        self.on_material_select(None)
        self.update_plot()

    def remove_material(self):
        if self.selected_material is not None:
            self.materials.pop(self.selected_material)
            self.material_listbox.delete(self.selected_material)
            self.selected_material = None
            self.clear_fields()
            self.update_plot()

    def update_material(self):
        # Local modules
        from CHAP.edd.utils import make_material
#        from CHAP.utils.material import Material

        if self.selected_material is None:
            return

        material = self.materials[self.selected_material]
        try:
            # Retrieve values from fields
            name = self.fields['Material Name'].get()
            sgnum = int(self.fields['Space Group'].get())
            lattice_parameters = [
                float(self.fields[param].get())
                for param in ('a', 'b', 'c', 'alpha', 'beta', 'gamma')
            ]
            # Make a hexrd material from those values so we can
            # propagate any other updates required by the material's
            # symmetries
            _material = make_material(name, sgnum, lattice_parameters)
#            _material = Material.make_material(
#                name, sgnum=sgnum,
#                lattice_parameters_angstroms=lattice_parameters,
#                pos=['4a', '8c'])
#                #pos=[(0,0,0), (1/4, 1/4, 1/4), (3/4, 3/4, 3/4)])
            lattice_parameters_angstroms = lattice_parameters
            material.material_name = name
            material.sgnum = _material.sgnum
            material.lattice_parameters = [
                _material.latticeParameters[i].value for i in range(6)
            ]
            material._material = _material
            # If the updated field forces other field(s) to get new
            # values (because of space group symmetries), propagate
            # those new values to the gui entries too.
            for key, entry in self.fields.items():
                if key == 'Material Name':
                    continue
                entry.delete(0, tk.END)
            self.fields['Space Group'].insert(0, str(material.sgnum))
            for i, key in enumerate(('a', 'b', 'c', 'alpha', 'beta', 'gamma')):
                self.fields[key].insert(
                    0, str(_material.latticeParameters[i].value))

            # Update the listbox name display
            self.material_listbox.delete(self.selected_material)
            self.material_listbox.insert(
                self.selected_material, material.material_name)
            self.update_plot()
        except ValueError:
            messagebox.showerror(
                'Invalid input',
                'Please enter valid numbers for lattice parameters.')

    def on_material_select(self, event):
        if len(self.material_listbox.curselection()) == 0:
            # Listbox item deselection event can be ignored
            return
        # Update the selected material index
        self.selected_material = self.material_listbox.curselection()[0]
        material = self.materials[self.selected_material]
        self.clear_fields()
        self.fields['Material Name'].insert(0, material.material_name)
        self.fields['Space Group'].insert(0, str(material.sgnum))
        for i, key in enumerate(('a', 'b', 'c', 'alpha', 'beta', 'gamma')):
            self.fields[key].insert(
                0, str(material._material.latticeParameters[i].value))

    def clear_fields(self):
        for entry in self.fields.values():
            entry.delete(0, tk.END)

    def update_plot(self):
        # Local modules
        from CHAP.edd.utils import (
            get_unique_hkls_ds,
            get_peak_locations,
        )

        self.ax.cla()
        self.legend_handles = []
        self.plot_reference_data()  # Re-plot reference data

        # Plot each material's hkl peak locations
        for i, material in enumerate(self.materials):
            hkls, ds = get_unique_hkls_ds([material])
            E0s = get_peak_locations(ds, self.tth)
            for hkl, E0 in zip(hkls, E0s):
                if E0 < min(self.ref_data_x) or E0 > max(self.ref_data_x):
                    continue
                line = self.ax.axvline(
                    E0, c=f'C{i+1}', ls='--', lw=1,
                    label=material.material_name)
                self.ax.text(E0, 1, str(hkl)[1:-1], c=f'C{i+1}',
                             ha='right', va='top', rotation=90,
                             transform=self.ax.get_xaxis_transform())
            self.legend_handles.append(line)
        self.ax.legend(handles=self.legend_handles)
        self.canvas.draw()

    def on_close(self):
        """Handle closing the GUI and triggering the on_complete
        callback."""
        if self.on_complete:
            self.on_complete(self.materials, self.figure)
        self.root.destroy()  # Close the tkinter root window
        plt.close()

def run_material_selector(
        x, y, tth, preselected_materials=None, label='Reference Data',
        on_complete=None, interactive=False):
    """Run the MaterialParamSelector tkinter application.

    :param x: MCA channel energies.
    :type x: np.ndarray
    :param y: MCA intensities.
    :type y: np.ndarray
    :param tth: The (calibrated) 2&theta; angle.
    :type tth: float
    :param preselected_materials: Materials to get HKLs and lattice
        spacings for.
    :type preselected_materials: list[hexrd.material.Material], optional
    :param label: Legend label for the 1D plot of reference MCA data
        from the parameters `x`, `y`, defaults to `"Reference Data"`.
    :type label: str, optional
    :param on_complete: Callback function to handle completion of the
        material selection, defaults to `None`.
    :type on_complete: Callable, optional
    :param interactive: Show the plot and allow user interactions with
        the matplotlib figure, defaults to `False`.
    :type interactive: bool, optional
    :return: The selected materials for the strain analyses.
    :rtype: list[CHAP.edd.models.MaterialConfig]
    """
    # Initialize the main application window
    root = tk.Tk()

    # Create the material parameter selection GUI within the main
    # window
    # This GUI allows the user to adjust and visualize lattice
    # parameters and space group
    app = MaterialParamSelector(
        root, x, y, tth, preselected_materials, label, on_complete)

    if interactive:
        # If interactive mode is enabled, start the GUI event loop to
        # allow user interaction
        root.mainloop()
    else:
        # If not in interactive mode, immediately close the
        # application
        app.on_close()

def select_material_params(
        x, y, tth, preselected_materials=None, label='Reference Data',
        interactive=False, return_buf=False):
    """Interactively adjust the lattice parameters and space group for
    a list of materials. It is possible to add / remove materials from
    the list.

    :param x: MCA channel energies.
    :type x: np.ndarray
    :param y: MCA intensities.
    :type y: np.ndarray
    :param tth: The (calibrated) 2&theta angle.
    :type tth: float
    :param preselected_materials: Materials to get HKLs and
        lattice spacings for.
    :type preselected_materials: list[hexrd.material.Material],
        optional
    :param label: Legend label for the 1D plot of reference MCA data
        from the parameters `x`, `y`, defaults to `"Reference Data"`.
    :type label: str, optional
    :param interactive: Show the plot and allow user interactions with
        the matplotlib figure, defaults to `False`.
    :type interactive: bool, optional
        :param return_buf: Return an in-memory object as a byte stream
        represention of the Matplotlib figure, defaults to `False`.
    :type return_buf: bool, optional
    :return: The selected materials for the strain analyses and a byte
        stream represention of the Matplotlib figure if return_buf is
        `True` (`None` otherwise).
    :rtype: list[CHAP.edd.models.MaterialConfig],
        Union[io.BytesIO, None]
    """
    # Local modules
    from CHAP.utils.general import fig_to_iobuf

    # Run the MaterialParamSelector with the callback function to
    # handle the materials data and, if requested, the output figure
    materials = None
    figure = None
    def on_complete(_materials, _figure):
        nonlocal materials, figure
        materials = _materials
        figure = _figure

    run_material_selector(
        x, y, tth, preselected_materials, label, on_complete, interactive)

    if return_buf:
        buf = fig_to_iobuf(figure)
    else:
        buf = None

    return materials, buf


if __name__ == '__main__':
    # Local modules
    from CHAP.edd.models import MaterialConfig

    xx = np.linspace(40, 100, 100)
    yy = np.sin(xx)
    ttth = 5
    ppreselected_materials = [
        MaterialConfig(
            material_name='Ti64_orig',
            sgnum=194,
            lattice_parameters=[2.9217, 4.66027]
        )
    ]
    mmaterials = select_material_params(
        xx, yy, ttth, preselected_materials=ppreselected_materials,
        interactive=True)
    print(f'Returned materials: {mmaterials}')
