import numpy as np
import scipy.signal as sig
import matplotlib.backends.backend_qtagg as qt
import matplotlib.figure as fig
import utils

rect = [0.125,0.11,0.775,0.77]

calib_top  = np.load( 'asset/calib_top.npy' )
calib_down = np.load( 'asset/calib_down.npy' )
reference  = np.load( 'asset/reference.npy' )

mean_calib_top  = np.mean( calib_top  , axis = 0 )
mean_calib_down = np.mean( calib_down , axis = 0 )

peaks_calib_top = utils.retrieve_peaks( mean_calib_top )
peaks_reference = utils.retrieve_peaks( reference[1] , window_size = 1 , max_window_size = 1 )

polyval = np.polyfit( peaks_calib_top , peaks_reference[ 2 : ] , 1 )

peaks_values = [ reference[ 1 , i ] for i in np.array( peaks_reference ).astype( int ) ]
sorting = np.argsort( peaks_values )[ :: -1 ]
wavelength = [ reference[ 0 , i ] for i in np.array( [ peaks_reference[ j ] for j in sorting ] ).astype( int ) ]

polyval_wavelength = np.polyfit( peaks_reference , np.sort( wavelength ) , 1 )

wavelength = np.polyval( polyval_wavelength , np.polyval( polyval , np.arange( 0 , len( mean_calib_top ) , 1 ) ) )

"""
manager = qt.FigureManager(
    qt.FigureCanvas(
        fig.Figure(
            figsize = ( 10 , 5 )
        ),
    ),
    0,
)

manager.canvas.figure.add_subplot( 2 , 1 , 1 , xlim = ( 3800 , 4000 ) , xmargin = 0 ).plot(
    wavelength    ,
    mean_calib_top,
)
manager.canvas.figure.axes[0].set_title( 'raw data' )

manager.canvas.figure.add_subplot( 2 , 1 , 2 , xlim = ( 3800 , 4000 ) , xmargin = 0 ).plot(
    reference[0],
    reference[1],
)
manager.canvas.figure.axes[1].set_title( 'reference' )

manager.show()
manager.start_main_loop()
"""
for peak in np.array( peaks_calib_top ).astype( int ):
    first = peak
    peak , old = first - 2 , first - 1
    while mean_calib_top[ peak ] <= mean_calib_top[ old ]:
        old   = peak
        peak -= 1
    mean_calib_top[ peak : first ] = mean_calib_top[ peak ] * np.ones( first - peak )

    peak , old = first + 2 , first + 1
    while mean_calib_top[ peak ] <= mean_calib_top[ old ]:
        old   = peak
        peak += 1
    mean_calib_top[ first : peak ] = mean_calib_top[ peak ] * np.ones( peak - first )

manager = qt.FigureManager(
    qt.FigureCanvas(
        fig.Figure(
            figsize = ( 10 , 5 )
        ),
    ),
    0,
)
manager.canvas.figure.gca().plot(
    mean_calib_top
)
manager.canvas.figure.gca().plot(
    sig.medfilt( mean_calib_top , 71 )
)
manager.show()
manager.start_main_loop()

manager = qt.FigureManager(
    qt.FigureCanvas(
        fig.Figure(),
    ),
    1,
)
signal_calib = sig.convolve(
    sig.medfilt( mean_calib_top , 71 ),
    np.ones( 50 ),
    'same'
)
manager.canvas.figure.gca().plot(
    signal
)
manager.canvas.figure.gca().vlines(
    [
        np.argmin( signal[ 25 : -25 ] )
    ]                           ,
    np.min( signal[ 25 : -25 ] ),
    np.max( signal[ 25 : -25 ] ),
    color = 'red'               ,
)
manager.show()
manager.start_main_loop()