from PyQt6.QtGui import *
from PyQt6.QtWidgets import *
from PyQt6.QtCore import *


import os,sys
import numpy as np
import pyqtgraph as pg
import datetime,time

from design_files.Raw_data_design import Ui_MainWindow


def smooth_transition(x, a, b, c, x0):
    """
    A function that is constant before t0 and transitions smoothly into a linear function after t0.
    
    Parameters:
    t  : array-like, time values
    t0 : float, transition point
    a  : float, constant value before t0
    b  : float, slope of the linear function after t0
    c  : float, controls the smoothness of the transition
    
    Returns:
    array-like, function values
    """
    return a + (b * (x - x0)) / (1 + np.exp(-c * np.clip(x - x0, -100, 100)))

# Define linear fit function:
def linear_fit(x, a, b):
    '''
    Linear function for fitting of CRAFT data
    '''
    return a+b*x

class RawdataWindow(QMainWindow, Ui_MainWindow):
    def __init__(self, *args, **kwargs):
        # Get the current script's directory
        self.current_dir = os.path.dirname(os.path.abspath(__file__))
        # Get the parent directory by going one level up
        self.parent_dir = os.path.dirname(self.current_dir)

        
        #import Gui from QT designer file
        super(RawdataWindow, self).__init__(*args, **kwargs)
        self.setupUi(self)

        #setup plots
        self.plot_widgets = [self.graphWidget_AMR_mag,self.graphWidget_AMR_x,self.graphWidget_AMR_y,self.graphWidget_AMR_z,self.graphWidget_FG,self.graphWidget_T] #list of all plots for easier handling
        self.plot_colors = ['#a6cee3', '#1f78b4','#b2df8a','#33a02c','#fb9a99','#e31a1c','#fdbf6f','#ff7f00','#cab2d6','#6a3d9a','#e7298a','#b15928','#016c59','#d9d9d9','#000000'] #list of colors for plot
        self.plots_T = [i for i in range(8)] #will contain all plots in temperature grapgh widget. Is already the correct lengths, so the right amount of plots are created for the widget. Contains numbers from 0 to the number of plots so the elements can also be used as index below.
        self.plots_fluxgate = [i for i in range(3)] #will contain all plots in Fluxgate grapgh widget. Is already the correct lengths, so the right amount of plots are created for the widget. Contains numbers from 0 to the number of plots so the elements can also be used as index below.
        self.plots_AMR_x = [i for i in range(15)] #will contain all plots in AMR_x grapgh widget. Is already the correct lengths, so the right amount of plots are created for the widget. Contains numbers from 0 to the number of plots so the elements can also be used as index below.
        self.plots_AMR_y = [i for i in range(15)] #will contain all plots in AMR_y grapgh widget. Is already the correct lengths, so the right amount of plots are created for the widget. Contains numbers from 0 to the number of plots so the elements can also be used as index below.
        self.plots_AMR_z = [i for i in range(15)] #will contain all plots in AMR_z grapgh widget. Is already the correct lengths, so the right amount of plots are created for the widget. Contains numbers from 0 to the number of plots so the elements can also be used as index below.
        self.plots_AMR_mag = [i for i in range(15)] #will contain all plots in AMR_mag grapgh widget. Is already the correct lengths, so the right amount of plots are created for the widget. Contains numbers from 0 to the number of plots so the elements can also be used as index below.
        self.plots_grad = [i for i in range(1)] #will contain all plots in grad_loc grapgh widget. Is already the correct lengths, so the right amount of plots are created for the widget. Contains numbers from 0 to the number of plots so the elements can also be used as index below.
        self.plots_ind_T = [i for i in range(8)]
        self.plot_sets = [self.plots_AMR_mag, self.plots_AMR_x, self.plots_AMR_y, self.plots_AMR_z, self.plots_fluxgate, self.plots_T]#list of all plot-sets. Each plot-set contains all plots of one graph widget
        self.pens = [pg.mkPen(color = c, width = 2) for c in self.plot_colors]
        self.titles = ['AMR Magnitude', 'AMR x', 'AMR y', 'AMR z', 'Fluxgates', 'Temperature'] #Plot titles
        self.ylabel = ['B [µT]', 'B [µT]', 'B [µT]', 'B [µT]', 'B [µT]', 'Temperature [K]'] #y axis labels for plots
        self.axis = [] #list of axis. They are needed to set x axis to time. List is filled below

        #all plots with time on x axis
        for i in range(len(self.plot_widgets)): #fill list with axis
         self.axis.append(pg.DateAxisItem())

        for widget, title, label, axis, plot_set in zip(self.plot_widgets, self.titles, self.ylabel, self.axis, self.plot_sets):
            #setup background, title, labels, grid, and time axis
            widget.setBackground('w')
            widget.setTitle(title)
            widget.setLabel('left', label)
            widget.setLabel('bottom', 'Temperature')
            widget.showGrid(x= True, y= True, alpha = 0.5)
            widget.setAxisItems({'bottom':axis})

            temp = [time.time(),time.time()-1]

            for i in plot_set: #since plot_set so far just contains numbers from 0 to the number of desired plots, the elements can also be used as index
                plot_set[i] = (widget.plot(temp,[1,0],pen = self.pens[i]))
                # plot_set[i].clear()

        #local gradient plot 
        self.graphWidget_grad_loc.setBackground('w')
        self.graphWidget_grad_loc.setTitle("local gradient")
        self.graphWidget_grad_loc.setLabel('bottom', 'Sensor (1 at top edge)')
        self.graphWidget_grad_loc.setLabel('left', 'dT/dz [K/cm]')
        self.plots_grad[0] = (self.graphWidget_grad_loc.plot([1,0],[1,0],color = self.plot_colors[0],symbol = 'x',symbolsize = 20))   

        #define signals and slots
        
        #define constants
        self.mcol = ["#0072BD","#D95319","#EDB120","#7E2F8E","#77AC30","#4DBEEE","#A2142F","#0072BD","#D95319","#EDB120","#7E2F8E","#77AC30","#4DBEEE","#A2142F"] #define matlab colors
        self.marker = ['o','s','t','d','+','p','x','arrow_up','t1','h','crosshair','t3','star','arrow_down']
        self.pen_lines = pg.mkPen(color = 'k', width = 1)
        self.pen_B_lines = pg.mkPen(color = 'r', width = 1)
        self.lines = [] #list will store vertical lines for later delation


    def update_plots(self,times,data,grad_loc,indices, Wave_mag, popt_T1,popt_T2, act_sens, T_c):
        #clear temperature plot
        self.graphWidget_T.clear()
        
        #calculate AMR magnitude               
        AMR_mag = np.zeros(shape = (len(data[:,0]),15))
        for j in range(15):
            for i in range(len(data[:,0])):
                AMR_mag[i,j] = np.sqrt(data[i,j+8]**2 + data[i,j+23]**2 + data[i,j+38]**2)
        timestamps = [t.timestamp() for t in times]

        #plot horizontal line in temperature plot as T_c
        if type(T_c) == float:
            self.graphWidget_T.addLine(y = T_c, pen = self.pen_lines)
        else:
            self.graphWidget_T.addLine(y = T_c[0], pen = self.pen_lines)
            self.graphWidget_T.addLine(y = T_c[1], pen = self.pen_lines)

        #plot vertical lines for temperature at indices in temperature, AMR_x, AMR_y, AMR_z and fluxgate plots
        #delete lines that are stored in self.xlines
        for line in self.lines:
            for widget in self.plot_widgets:
                widget.removeItem(line)
        for widget in self.plot_widgets:#create new lines
            for i in indices[:-3]: #ignore the last 3 indices which correspond to the magnetic field
                self.lines.append(widget.addLine(x = timestamps[i], pen = self.pen_lines))

        #plot vertical lines for magnetic field at last three indices in temperature, AMR_x, AMR_y, AMR_z and fluxgate plots
        for widget in self.plot_widgets:#create new lines
            # for i in indices[-3]: #To plot all just uncomment
            self.lines.append(widget.addLine(x = timestamps[indices[-2]], pen = self.pen_B_lines)) #only plot where expelled flux point is taken, since the other two and at the beginning and end anyways.

        # #plot horizontal line in AMR plots at wave magnitude
        # self.lines.append(self.graphWidget_AMR_x.addLine(y = Wave_mag[0], pen = self.pen_B_lines))
        # self.lines.append(self.graphWidget_AMR_y.addLine(y = Wave_mag[1], pen = self.pen_B_lines))
        # self.lines.append(self.graphWidget_AMR_z.addLine(y = Wave_mag[2], pen = self.pen_B_lines))

        #set data in AMR plots
        for set,ydata in zip(self.plot_sets,[AMR_mag, data[:,8:23], data[:,23:38], data[:,38:53], data[:,53:56]]):
            for i,plot in enumerate(set):
                    plot.setData(timestamps,ydata[:,i])
        
        #set data in temperature plot
        for i in act_sens:
            self.graphWidget_T.plot(timestamps,data[:,i-1], pen = self.pens[i-1])
        
        #set data in gradient plot
        if type(popt_T1) == int: #if CRAFT data is analysed, local gradients are not plotted
            x = [i+1 for i in range(8)]
            self.plots_grad[0].setData(x,grad_loc)

        #plot fits in temperature plot from popt_T1 and popt_T2 values, if they are given (not zero)
        if  type(popt_T1) == np.ndarray:
            x = np.linspace(0, popt_T1[-1] - popt_T1[-2], 100)
            x_plot = np.linspace(popt_T1[-2], popt_T1[-1], 100)
            self.graphWidget_T.plot(x = x_plot, 
                                    y = smooth_transition(x, *popt_T1[:4]), pen = self.pens[-1])
            self.graphWidget_T.plot(x = x_plot, 
                                    y = smooth_transition(x, *popt_T2[:4]), pen = self.pens[-1])