# => makes the plot interactive
#%matplotlib notebook 
# inline makes the plots static
#%matplotlib inline 

import numpy as np
from astropy.io import fits
import matplotlib.pyplot as plt
from scipy.signal import find_peaks
import warnings
# filter astropy warning on fits headers
warnings.filterwarnings('ignore', category=UserWarning, append=True)

def plot_spectrum_panels(xaxis,data,num_panels=2):
    # not that useful, interactive plotting is better
    spectrum_data=np.copy(data)

    # Divide the spectrum into four panels
    panel_height = len(spectrum_data) // num_panels

    # Create subplots
    fig, axs = plt.subplots(num_panels, 1, figsize=(15, 6*num_panels))

    # Plot each panel
    for i in range(num_panels):
        start_idx = i * panel_height
        end_idx = (i + 1) * panel_height
        xaxis= list(range(len(spectrum_data)))
        axs[i].plot(xaxis[start_idx:end_idx],spectrum_data[start_idx:end_idx])
#        axs[i].set_xlabel('Wavelength (Angstroms)')
        axs[i].set_ylabel('counts(ADU)')
        axs[i].set_title(f'Panel {i + 1}')
        axs[i].grid(True)


def read_raw_spectrum(fits_file,get_header=0):
#    print("yoyoyo")
    # returns an index and the fluxes
    # Open the FITS file
    hdulist = fits.open(fits_file)
    header = hdulist[0].header
    # Assuming the spectrum data is in the first extension (HDU index 0)
    data = hdulist[0].data[0][0]
    xaxis= list(range(len(data)))
    if(get_header==0):
        return xaxis,data
    if(get_header==1):
        return xaxis,data,header

def write_fits_spectrum(target_file,data,header=[]):
    if(header==[]):
        hdu=fits.hdu.hdulist.PrimaryHDU([[data]])
    else:
        hdu=fits.hdu.hdulist.PrimaryHDU([[data]],header)
    hdul = fits.hdu.HDUList([hdu])
    hdul.writeto(target_file,overwrite=True)
    print('wrote ',target_file)
        
        
def read_raw_spectrum_old(fits_file,get_header=0):
    # returns an index and the fluxes
    # Open the FITS file
    hdulist = fits.open(fits_file)
    header = hdulist[0].header
    # Assuming the spectrum data is in the first extension (HDU index 0)
    data = hdulist[0].data[0]
    xaxis= list(range(len(data)))
    if(get_header==0):
        return xaxis,data
    if(get_header==1):
        return xaxis,data,header

def write_fits_spectrum_old(target_file,data,header=[]):
    if(header==[]):
        hdu=fits.hdu.hdulist.PrimaryHDU([data])
    else:
        hdu=fits.hdu.hdulist.PrimaryHDU([data],header)
    hdul = fits.hdu.HDUList([hdu])
    hdul.writeto(target_file,overwrite=True)
    print('wrote ',target_file)

def read_calibrated_spectrum(fits_file):
    # Open the FITS file
    with fits.open(fits_file) as hdul:
        # Extract the data and header from the primary HDU
        data = hdul[0].data[0]
        header = hdul[0].header

        # Extract necessary header keywords for calibration
        crval1 = header['CRVAL1']  # Reference value of the first pixel (wavelength)
        cdelt1 = header['CDELT1']  # Pixel scale (wavelength per pixel)
        naxis1 = header['NAXIS1']  # Number of pixels in the spectral axis

        # Create an array of calibrated wavelengths
        wavelengths = crval1 + np.arange(naxis1) * cdelt1
    return wavelengths,data


# Define the Gaussian function
def gaussian(x, *params):
    result= np.full(len(x),0.)
    for i in range(0, len(params), 3):
        amp, cen, wid = params[i:i+3]
        result += amp * np.exp(-(x - cen) ** 2 / (2 * wid ** 2))
    return result

# Define the Gaussian function
def gaussian_plus_bg_2(x, *params):
    # just like the gaussian but with a constant background under ti
    # should help fitting emission lines in objects
    # this is bad because there are multiple backgrounds => degenerate
    result= np.full(len(x),0.)
    for i in range(0, len(params), 4):
        amp, cen, wid , bg = params[i:i+4]
        result += amp * np.exp(-(x - cen) ** 2 / (2 * wid ** 2)) + bg
    return result

def gaussian_plus_bg(x, *params):
    # just like the gaussian but with a constant background under ti
    # should help fitting emission lines in objects
    result= np.full(len(x),0.)
    bg=params[0]
    result=result+bg
    for i in range(1, len(params), 3):
        amp, cen, wid = params[i:i+3]
        result += amp * np.exp(-(x - cen) ** 2 / (2 * wid ** 2))
    return result

def generate_first_guess(peaks):
#    npeaks=len(peaks[:3])
    npeaks=len(peaks)
    initial_cen=peaks
    initial_amp=np.full(npeaks,20000.)
    initial_wid=np.full(npeaks,1.)
#    print(initial_cen)
#    print(initial_amp)
#    print(initial_wid)
    first_guess=np.full((npeaks,3),0.)
    for i in range(npeaks):
        first_guess[i,0]=initial_amp[i]
        first_guess[i,1]=initial_cen[i]
        first_guess[i,2]=initial_wid[i]
    first_guess=np.reshape(first_guess,3*npeaks)
#    print(first_guess)
    return first_guess