# Author: Alex DelFranco
# Advisor: Rafa Martin Domenech
# Institution: Center for Astrophysics | Harvard & Smithsonian
# Date: 22 August 2022

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import ipywidgets as widgets
import copy as cp
import pickle as pkl

import importlib as imp
from base import widget_base
from baseline import baseline
from fit import fit
from int_peaks import int_peaks
from labels import labels
from display_tools import display_tools
from export import export

from scipy import integrate as scpyintegrate
import functools

class main_menu(widget_base):
    def __init__(self,data,load_file=None,fig=None,ax=None,exp_path=''):
        '''
        Initializes a widget instance
        '''
        # Close any earlier plots
        plt.close()

        # Initialize widget parameters
        self.widget_on = True
        self.active_widget = self
        self.prev_widget = self
        self.menu,self.sub_menu,self.info_menu = widgets.HBox(),widgets.HBox(),widgets.HBox()

        # Check for an load dictionary
        if load_file is not None:
            plt.ioff()
            load = cp.deepcopy(pkl.load(open(load_file,'rb')))
            fig = load['fig']; ax = fig.axes
            plt.ion()
        else: load = None

        # Initialize figure
        self.fig_init(fig,ax)

        if load is None:
            # Initialize artist list
            self.artists = dict.fromkeys(self.fig.axes,{})
            # Import the data from the initialized instance
            self.artists[self.ax]['Input Data'],self.input_data = self.setup_data(data)
            self.artists['Primary Artists'] = {self.ax:'Input Data'}
            self.artists['Interactive Axes'] = [self.ax]

            # Initialize data list
            self.data = dict.fromkeys(self.fig.axes,{})

        else:
            self.artists = load['artists_global']
            self.data = load['data_global']

        # Initialize all sub_menus
        self.baseline = baseline(self.fig,menu=self.sub_menu,artists_global=self.artists,data_global=self.data,load=load)
        self.fit = fit(self.fig,menu=self.sub_menu,param_menu=self.info_menu,artists_global=self.artists,data_global=self.data,load=load)
        self.int_peaks = int_peaks(self.fig,menu=self.sub_menu,artists_global=self.artists,data_global=self.data,load=load)
        self.labels = labels(self.fig,menu=self.sub_menu,artists_global=self.artists,data_global=self.data,load=load)
        self.display_tools = display_tools(self.fig,menu=self.sub_menu,artists_global=self.artists)
        self.export = export(main_menu=self,menu=self.sub_menu,exp_path=exp_path)

        # Widget list
        self.widgets = {
        'baseline':self.baseline,
        'fit':self.fit,
        'int_peaks':self.int_peaks,
        'labels':self.labels,
        'display_tools':self.display_tools,
        'export':self.export
        }

        # Setup the main menu and a blank sub_menu
        self.button_list = self.setup_buttons()
        self.place_menu(self.menu,self.button_list)
        self.place_menu(self.sub_menu,[])
        self.place_menu(self.info_menu,[])
        self.fig_show(self.fig)
        plt.show()

    # Define all main menu buttons and their functionality
    def setup_buttons(self):

        ##########################################
        ## ON OFF BUTTON
        ##########################################

        # Add an on/off toggle to activate or deactive the widget
        self.on_off_button = widgets.ToggleButton(description='On',value=True,layout = widgets.Layout(width='45px'))

        def on_on_off_button_clicked(b, self = self):

            # If the widget is on, turn it off
            if self.widget_on:
                # Set the description
                self.on_off_button.description = 'Off'
                # Turn off the main widget
                self.deactivate(main=True)
                # Turn off all buttons in the main menu
                for tb in self.button_list:
                    # Turn it off if except for the selected menu
                    if tb != self.on_off_button: tb.value = False
                # Deactivate the current active widget
                if self.active_widget != self: self.active_widget.deactivate()
            # If the widget is off, turn it on
            else:
                self.on_off_button.description = 'On'
                self.activate()

        self.on_off_button.observe(on_on_off_button_clicked)

        ##########################################
        ## BUTTON TEMPLATE
        ##########################################

        def menu_button(self,menuclass,button):
            '''
            Generic button function for all new buttons
            '''
            def func(b,self=self):
                if not self.widget_on:
                    for tb in self.button_list:
                        if tb != self.on_off_button: tb.value = False
                    return
                # If the widget is on, turn it off
                if menuclass.widget_on:
                    # Turn off the widget
                    menuclass.deactivate()
                    # Set the previous widget
                    self.prev_widget = menuclass
                    # Set the active widget to the main menu
                    self.active_widget = self
                # If the widget is off, turn it on
                else:
                    # For each button in the main menu
                    for tb in self.button_list:
                        # Turn it off if except for the selected menu
                        if tb != button and tb != self.on_off_button: tb.value = False
                    # Set the active widget to the selected menu
                    self.active_widget = menuclass
                    # If moving to the export menu
                    if self.active_widget == self.display_tools: self.display_tools.update(self.artists,None)
                    # Update the menu
                    self.prev_widget.update(self.artists,self.data)
                    # Turn off all the menu's buttons
                    for tb in menuclass.toggle_buttons: tb.value = False
                    # Turn on the selected menu
                    menuclass.activate()

                    # Listen for a click
                    self.listen(menuclass)
            return func

        ##########################################
        ## ADD WIDGETS
        ##########################################

        # Baseline toggle
        self.baseline_button = widgets.ToggleButton(description='Baseline',value=False)
        self.baseline_button.observe(menu_button(self,self.baseline,self.baseline_button), names=['value'])

        ##########################################

        # Peaks toggle
        self.int_peaks_button = widgets.ToggleButton(description='Integrate/Peaks',value=False)
        self.int_peaks_button.observe(menu_button(self,self.int_peaks,self.int_peaks_button), names=['value'])

        # ##########################################

        # Integrate toggle
        self.fit_button = widgets.ToggleButton(description='Fit',value=False)
        self.fit_button.observe(menu_button(self,self.fit,self.fit_button), names=['value'])

        # ##########################################

        # Labels toggle
        self.labels_button = widgets.ToggleButton(description='Labels',value=False)
        self.labels_button.observe(menu_button(self,self.labels,self.labels_button), names=['value'])

        # ##########################################
        #
        # Display toggle
        self.display_tools_button = widgets.ToggleButton(description='Display',value=False)
        self.display_tools_button.observe(menu_button(self,self.display_tools,self.display_tools_button), names=['value'])

        # ##########################################
        #
        # Export toggle
        self.export_button = widgets.ToggleButton(description='Export',value=False)
        self.export_button.observe(menu_button(self,self.export,self.export_button), names=['value'])

        ##########################################
        ## BUTTON LIST
        ##########################################

        return [self.on_off_button,self.baseline_button,self.int_peaks_button,self.fit_button,self.labels_button,self.display_tools_button,self.export_button]

    def setup_data(self,data):
        # Setup data
        data_dict = {}
        if isinstance(data,mpl.lines.Line2D):
            data_dict['xpoints'],data_dict['ypoints'] = data.get_data()
            artist = data
        elif isinstance(data,list) or isinstance(data,tuple) or isinstance(data,np.ndarray):
            if isinstance(data[0],mpl.lines.Line2D):
                data_dict['xpoints'],data_dict['ypoints'] = data[0].get_data()
                artist = data[0]
            elif len(data) == 2:
                data_dict['xpoints'],data_dict['ypoints'] = data
                artist = self.ax.plot(data_dict['xpoints'],data_dict['ypoints'])[0]
        return artist,data_dict

    # Initialize all the default colors and color targets
    def color_init(self):
        self.color_target = 'Marker - Primary'
        self.primary_color = 'deepskyblue'
        self.new_color = 'red'
        self.secondary_color = 'green'
        self.fit_color = 'darkviolet'
        self.bscolor = 'blue'
        self.fill_color = 'skyblue'
        self.marker = '*'

    def data_tree(self,*kwargs):
        self.active_widget.update(self.artists,self.data)
        ddict = self.get_data()
        for arg in kwargs:
            if arg in ddict.keys():
                ddict = ddict[arg]
        self.printdict(ddict)

    def artist_tree(self,*kwargs):
        self.active_widget.update(self.artists,self.data)
        adict = self.get_artists()
        for arg in kwargs:
            if arg in adict.keys():
                adict = adict[arg]
        self.printdict(adict)

    def get_data(self):
        # Update the data
        self.active_widget.update(self.artists,self.data)
        # Get the list of axes
        axes = self.get_axes()
        # Create a new data dictionary for data export
        exp_data = {}
        # Fill the new dictionary
        for ind,ax in enumerate(axes):
            exp_data[f'Axis {ind+1}'] = self.data[ax]
        # Return the export dictionary
        return exp_data

    def get_artists(self):
        # Update the artists
        self.active_widget.update(self.artists,self.data)
        # Get the list of axes
        axes = self.get_axes()
        # Create a new data dictionary for data export
        exp_artists = {}
        # Fill the new dictionary
        for ind,ax in enumerate(axes):
            exp_artists[f'Axis {ind+1}'] = self.artists[ax]
        # Return the export dictionary
        return exp_artists

    def export_dict(self):
        self.active_widget.update(self.artists,self.data)
        export = {
        'artists_global': self.artists,
        'data_global': self.data,
        }
        for key in self.widgets.keys():
            if key == 'export' or key == 'display_tools':continue
            w_export = {
            'data': self.widgets[key].get_data(),
            'artists': self.widgets[key].get_artists(),
            'style': self.widgets[key].style,
            'info': self.widgets[key].info
            }
            export[key] = w_export

        export['fig'] = self.fig

        return export

'''

'''
