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

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

peaks_reference_sorted  = np.loadtxt( 'calibration/OHP_Hg.calib' )
peaks_reference         = np.sort( peaks_reference_sorted )

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

peaks_calib_top  = np.array(
    utils.retrieve_peaks( mean_calib_top )
).astype( int )
peaks_calib_down = np.array(
    utils.retrieve_peaks( mean_calib_down )
).astype( int )

"""
Signal without peaks
"""

def remove_peaks( signal , peaks ):
    peakless_signal = signal.copy()
    for peak in peaks:
        first = peak
        peak , old = first - 2 , first - 1
        while peakless_signal[ peak ] <= peakless_signal[ old ]:
            old   = peak
            peak -= 1
        peakless_signal[ peak : first ] = peakless_signal[ peak ] * np.ones( first - peak )

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

    return sig.medfilt( peakless_signal , 111 )
peakless_calib_top = remove_peaks( mean_calib_top , peaks_calib_top )

def get_extremities( signal ):
    """
    It's possible to have an idea of the part the spectrum begins with the
    small large peak at the start of it, the peak at 3810 A should be inside
    it. The same goes for the area at the end.
    ONLY TRUE FOR THE CURRENT CALIBRATION (Hg)
    """
    argmin = np.argmin(
        signal[ : len( signal ) // 2 ]
    )
    argmax = np.argmax( peakless_calib_top[ : argmin ] )
    peaks_inside = np.where(
        np.logical_and(
            argmax < peaks_calib_top,
            argmin > peaks_calib_top,
        )
    )[0]
    if len( peaks_inside ) < 1:
        raise Exception( 'unknown plage, cannot autocalibrate' )
    first_peak = peaks_inside[0]

    """
    The next peak after the minimum at the end of the spectrum should be
    5079 A.
    ONLY TRUE FOR THE CURRENT CALIBRATION (Hg)
    """
    argmin_1 = np.argmin(
        peakless_calib_top[ len( peakless_calib_top ) // 2 : ]
    ) + len( peakless_calib_top ) // 2

    peaks_inside = np.where(
        argmin_1 < peaks_calib_top
    )[0]
    if len( peaks_inside ) < 1:
        raise Exception( 'unknown plage, cannot autocalibrate' )
    return ( first_peak , peaks_inside[0] )

first_peak_ref = 3 # hard-coded
last_peak_ref = 20 # hard-coded
first_peak_cal , last_peak_cal = get_extremities( peakless_calib_top )

"""
delete if too much peak in calib
"""

peaks_calib_top = peaks_calib_top[ first_peak_cal : last_peak_cal ]
peaks_reference = peaks_reference[ first_peak_ref : last_peak_ref ]

diff_ref  = ( peaks_reference[ 1 : ] - peaks_reference[ : -1 ] ).astype( float )
diff_ref -= np.min( diff_ref )
diff_ref /= np.max( diff_ref )
diff_cal  = ( peaks_calib_top[ 1 : ] - peaks_calib_top[ : -1 ] ).astype( float )
diff_cal -= np.min( diff_cal )
diff_cal /= np.max( diff_cal )

peaks = [ -1 ]
for i in range( len( diff_ref ) ):
    good , sum_ = -1 , 0
    for j in range( peaks[ - 1 ] + 1 , len( diff_cal ) ):
        sum_ += diff_cal[ j ]
        if sum_ - diff_ref[ i ] >  0.002:
            print( sum_ - diff_ref[ i ] )
            raise Exception( 'reference peak not found' )
        if sum_ - diff_ref[ i ] > - 0.002:
            good = j
            break
    if good == -1:
        raise Exception( 'reference peak not found and not exceeded' )
    peaks.append( good )
peaks.append( peaks[-1] + 1 )

peaks_calib_top = np.array(
    [ peaks_calib_top[ i ] for i in peaks[ 1 : ] ]
).astype( int )

"""
Calibration
"""

polyval_wavelength = np.polyfit(
    peaks_calib_top,
    peaks_reference,
    1              ,
)

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

manager = qt.FigureManager(
    qt.FigureCanvas(
        fig.Figure(
            figsize = ( 10 , 5 )
        ),
    ),
    0,
)
manager.canvas.figure.gca().plot(
    peaks_reference,
    peaks_reference - np.polyval( polyval_wavelength , peaks_calib_top ),
    'o-'           ,
)
manager.show()
manager.start_main_loop()