import colorsys
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import zscore
from PyQt5.QtCore import Qt
import matplotlib.patches as patches



class FluorescencePlotWidget(FigureCanvas):
    """this class is used to display the fluorescence of the selected cell across days. It also allows to select a region of interest (ROI) on the fluorescence plot and zoom in on the selected ROI"""
    def __init__(self, all_fluorescence=None, all_ops=None, colors=None, all_stat_t2p=None):
        self.fig, self.ax_fluorescence = plt.subplots(1, 1)
        super().__init__(self.fig)
        self.all_fluorescence = all_fluorescence
        self.all_ops=all_ops
        self.fig.set_facecolor('black')
        self.colors = colors
        self.all_stat_t2p= all_stat_t2p
        
        self.rect = patches.Rectangle((0,0), 1, 1, color='white', linewidth=2) 
        
        self.fig.canvas.mpl_connect('button_press_event', self.on_press)
        self.fig.canvas.mpl_connect('button_release_event', self.on_release)
        self.fig.canvas.mpl_connect('key_press_event', self.on_key_press)
        self.fig.canvas.mpl_connect('key_press_event', self.on_enter_pressed)
        self.setFocusPolicy(Qt.StrongFocus)
        self.setFocus()
        
        self.initial_xlim = None
        self.initial_ylim = None 
        self.cmd_pressed = False
        
    def draw_rectangle(self):
        """it draws a rectangle on the fluorescence plot"""
        #if the rectangle is already on the plot, it is removed
        if self.rect in self.ax_fluorescence.patches:
            self.ax_fluorescence.patches.remove(self.rect)
        self.rect = patches.Rectangle((self.x0, self.y0), self.x1 - self.x0, self.y1 - self.y0, fill=None, linewidth=2, edgecolor='white') #create a rectangle, fill is set to None to make the rectangle transparent
        self.rect.set_zorder(10)
        self.ax_fluorescence.add_patch(self.rect)
        self.draw()
        
    def on_key_press(self, event):
        """it allows to reset the zoom of the fluorescence plot. It is called when the Command key and - key are pressed."""
        if event.key == 'cmd+-':
            self.ax_fluorescence.set_xlim(self.initial_xlim)
            self.ax_fluorescence.set_ylim(self.initial_ylim)
            self.fig.canvas.draw()
    
    def draw_point(self):
        if hasattr(self, 'point') and self.point in self.ax_fluorescence.collections:
            print(self.ax_fluorescence.collections)
            self.ax_fluorescence.collections.remove(self.point)
        self.point = self.ax_fluorescence.scatter([self.x0], [self.y0], s=5, color='w')  # Create a new point
        self.draw()
  
    def on_press(self, event):
        """it allows to draw a point on the fluorescence plot. It is called when the mouse button is pressed."""
        self.x0 = event.xdata #x coordinate of the mouse cursor
        self.y0 = event.ydata  #y coordinate of the mouse cursor
        self.draw_point()

    def on_release(self, event):
        """It allows to draw a rectangle on the fluorescence plot. It is called when the mouse button is released."""
        self.x1 = event.xdata 
        self.y1 = event.ydata
        self.rect.set_width(self.x1 - self.x0) 
        self.rect.set_height(self.y1 - self.y0)
        self.rect.set_xy((self.x0, self.y0)) #set the position of the rectangle
        self.draw_rectangle()
        
        
    def on_enter_pressed(self, event):
        """it allows to zoom in of the fluorescence plot. It is called when the Command key and + key are pressed."""
        if event.key == 'cmd+=':
            self.ax_fluorescence.set_xlim(self.x0, self.x1)
            self.ax_fluorescence.set_ylim(self.y0, self.y1)
            self.rect.set_visible(False) 
            self.point.set_visible(False)
            self.fig.canvas.draw()


    def display_all_fluorescence(self, selected_cell_index):
        """it plots the fluroescence of the selected cell across days where each curve being a different day (the curve at the top of the plot is the first day)"""
        
        if self.all_fluorescence is not None and selected_cell_index is not None:
            self.ax_fluorescence.clear()
            self.ax_fluorescence.set_facecolor('black')
            self.ax_fluorescence.tick_params(axis='x', colors='white') 
            self.ax_fluorescence.tick_params(axis='y', colors='white')  
            self.ax_fluorescence.xaxis.label.set_color('white') 
            self.ax_fluorescence.yaxis.label.set_color('white')
            self.ax_fluorescence.spines['bottom'].set_color('#666') 
        
            
            for i, fluorescence_data in list(enumerate(reversed(self.all_fluorescence))):
                fluorescence_zscore = zscore(fluorescence_data, axis=1, ddof=1) #zscore is used to normalize the fluorescence data
                offset = i * 12 #
                y_values = fluorescence_zscore[selected_cell_index, :] + offset 
                color = self.colors[selected_cell_index] 
                #create a gradient of colors for the curves (the darkest shade is for the last day and the lightest shade is for the first day)
                if i == 0:
                    color = color 
                else:
                    h, l, s = colorsys.rgb_to_hls(*color)  
                    l_range = 1 - (l + 0.05)  
                    l_add = l_range/len(self.all_fluorescence) 
                    
                    adjusted_luminosity = l + (l_add *i) 
                    color = colorsys.hls_to_rgb(h, adjusted_luminosity, s)
                ops=self.all_ops[i]
                fs = ops['fs']
                tstamps = np.arange(0,self.all_fluorescence[i].shape[1])/fs
                if len(tstamps) != len(y_values):
                    raise ValueError(f"tstamps and y_values must have the same length, but have lengths {len(tstamps)} and {len(y_values)}")
                self.ax_fluorescence.plot(tstamps,y_values, label=f'Curve {i + 1}', color= color)
            
            self.ax_fluorescence.set_yticklabels([])
            self.ax_fluorescence.get_yaxis().set_visible(False)
            self.ax_fluorescence.set_xlabel('Seconds')
            self.initial_xlim=self.ax_fluorescence.get_xlim()
            self.initial_ylim=self.ax_fluorescence.get_ylim()
            self.fig.tight_layout() 
            self.draw()
            
         
        pass
