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 = np.array(
    utils.retrieve_peaks( mean_calib_top )
).astype( int )
peaks_reference = np.array(
    utils.retrieve_peaks(
        reference[1]       ,
        window_size = 1    ,
        max_window_size = 1,
    )
).astype( int )

"""
Signal without peaks
"""

peakless_calib_top = mean_calib_top.copy()
for peak in peaks_calib_top:
    first = peak
    peak , old = first - 2 , first - 1
    while peakless_calib_top[ peak ] <= peakless_calib_top[ old ]:
        old   = peak
        peak -= 1
    peakless_calib_top[ peak : first ] = peakless_calib_top[ peak ] * np.ones( first - peak )

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

peakless_calib_top = sig.medfilt( peakless_calib_top , 111 )

"""
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(
    peakless_calib_top[ : len( peakless_calib_top ) // 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 Error( 'unknown plage, cannot autocalibrate' )
first_peak_cal = peaks_inside[0]
first_peak_ref = 3 # hard-coded

"""
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 Error( 'unknown plage, cannot autocalibrate' )
last_peak_cal = peaks_inside[0]
last_peak_ref = 20 # hard-coded

print( first_peak_cal , last_peak_cal )
polyval = np.polyfit(
    peaks_calib_top[ first_peak_cal : last_peak_cal ],
    peaks_reference[ first_peak_ref : last_peak_ref ],
    1                                                ,
) # We suppose there is as much calib peak than ref peak for now

peaks_values = [ reference[ 1 , i ] for i in peaks_reference ]
sorting = np.argsort( peaks_values )[ :: -1 ]
wavelength = np.array( [
    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.gca().plot(
    wavelength    ,
    mean_calib_top,
)
manager.canvas.figure.gca().plot(
    wavelength  ,
    peakless_calib_top,
)
manager.canvas.figure.gca().vlines(
    [
        wavelength[ argmin ]  ,
        wavelength[ argmax ]  ,
        wavelength[ argmin_1 ],
    ]                           ,
    np.min( peakless_calib_top ),
    np.max( peakless_calib_top ),
    color = 'red'               ,
)
manager.show()
manager.start_main_loop()