""" Functions to wavelength calibrate the spectra. Finds lines in arc lamp spectra (ThAr), matches them to a line list based on a preliminary input wavelength solution, and then iteratively fits a polynomial solution.

Created by DMK on 11/11/2023

Last updated by DMK on 11/13/2023
"""

##### Imports
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from astropy.io import fits
from datetime import datetime
from matplotlib.backends.backend_pdf import PdfPages
from scipy import optimize, signal, stats, interpolate

import os

from tull_coude_reduction.modules import reduction_utils

##### Functions

def order_offset_with_wave_sol_guess( prelim_wavelengths, obs_flux, ref_wavelength, ref_flux, order_use_range = [ 12, 40 ], offset_test_radius = 10 ):
    """ Function to determine if there is an order offset between the wavelength solution guess and the observed spectrum.
    It uses a reference arc lamp spectrum to calculate residuals with the observed lamp spectrum while shifting the wavelength solution guess to determine the best fit order offset.
    It primarily works if the wavelength solution guess is already pretty good, and is used mostly for the wavelength solution guess having more orders than extracted (padding the calibration).

    Parameters
    ----------
    prelim_wavelengths : array
        The preliminary wavelength solution guess, with shape (number of orders, number of pixels). There can be a different number of orders than the extracted spectrum.
    obs_flux : array
        The observed arc lamp flux spectrum, with shape (number of extracted orders, number of pixel).
    ref_wavelength : array
        The wavelength array for the arc lamp reference spectrum.
    ref_flux : array
        The flux array for the arc lamp reference spectrum.
    order_use_range : list of int, optional
        List of minimum and maximum orders to calculate the order offset for. The default is [ 12, 40 ].
    offset_test_radius : int, optional
        The order offset +/- each order to calculate residuals for. The default is 10.

    Returns
    -------
    order_offset : int
        The minimum residual order offset between the wavelength solution guess and observed spectrum.
    """

    # The array of orders to calculate the best fit offset for
    orders_to_use = np.arange( *order_use_range )

    # The offsets to test
    order_offset_test_arr = np.arange( -offset_test_radius, offset_test_radius + 1 )

    # Offset to hold the offset with the minimum residual across test orders
    offset_with_min_resid = np.full( orders_to_use.size, np.nan )

    # Go through each of the orders to test
    for i_order, order in enumerate( orders_to_use ):

        sum_abs_residuals = np.full( order_offset_test_arr.size, np.nan )

        # Go through each offset
        for i_offset, offset in enumerate( order_offset_test_arr ):

            # "shift" the wavelength solution guess by the order offset amount
            test_wave = prelim_wavelengths[order+offset]

            # Interpolate the reference arc spectrum onto the wavelength solution guess
            ref_flux_interp = np.interp( test_wave, ref_wavelength, ref_flux )

            # Calculate the residuals between the reference and data spectrum -- normalize each by the median
            residuals = obs_flux[order] / np.nanmedian( obs_flux[order] ) - ref_flux_interp / np.nanmedian( ref_flux_interp )

            # Calculate and output the summed absolute residuals
            sum_abs_residuals[i_offset] = np.nansum( np.abs( residuals ) )

        # Find the offset with the minimum absolute residuals for this order
        offset_with_min_resid[i_order] = order_offset_test_arr[np.nanargmin( sum_abs_residuals )]

    # Find the minimum residual offset that is most common across all test orders
    unique_order_offsets, order_offset_counts = np.unique( offset_with_min_resid, return_counts = True )
    order_offset = int( unique_order_offsets[np.argmax(order_offset_counts)] )

    return order_offset

def get_flux_mad_from_spectral_chunks( obs_flux, chunk_size = 50 ):
    """ Function to get a noise estimate for the arc lamp spectrum by looking at the MAD of chunks of the spectrum.
    Designed to reduce the influence of bleeding over of saturated Ar lines (and oxide bands to a lesser extent)

    Parameters
    ----------
    flux : array
        The observed spectrum flux array.
    chunk_size : int, optional
        The number of pixels in the spectral chunks to calculate the MAD within. The default is 50.

    Returns
    -------
    median_chunk_mad : float
        The median MAD calculated across all the spectral chunks.
    """

    # The number of chunks from the input chunk size in pixels
    number_of_chunks = obs_flux.size // chunk_size

    # The array to hold the MAD for each of the spectral chunks
    mad_arr = np.full( number_of_chunks, np.nan )

    # Go through each of the chunks
    for i in range( number_of_chunks ):

        # The indices of the spectrum within this chunk
        ind_low, ind_high = chunk_size * i, chunk_size * ( i + 1 )

        # Calculate the scaled MAD
        mad_arr[i] = stats.median_abs_deviation( obs_flux[ind_low:ind_high], nan_policy = 'omit', scale = 'normal' )

    # Adopt the median of the chunk MAD array
    median_chunk_mad = np.nanmedian( mad_arr )

    return median_chunk_mad

def find_arc_lamp_line_pixel_centers( flux, config ):
    """ Function to find lines in an input lamp spectrum using scipy.signal's find_peaks algorithm, with options set in the config file.

    Parameters
    ----------
    flux : array
        The flux spectrum array.
    config : dict
        The overall config file defined using YAML with all parameters for running the reduction and analysis pipeline.

    Returns
    -------
    peak_pixels_initial_fit : array
        The array of pixel centroids found as lines in the input flux spectrum.
    """

    ### First prepare from input flux spectrum: normalize it and get noise estimate

    # Normalize the flux by its median
    flux = flux / np.nanmedian( flux )

    # Now get an estimate of the noise for peak finding threshold. Two current options: one is just from the MAD of the flux spectrum, the other divides the spectrum into chunks and takes the median of each chunk's MAD.
    if config['wavecal']['peak_threshold_mad_method'] == 'full_spectrum':

        mad_peak_threshold = stats.median_abs_deviation( flux, nan_policy = 'omit', scale = 'normal' )

    elif config['wavecal']['peak_threshold_mad_method'] == 'chunk_spectrum':

        mad_peak_threshold = get_flux_mad_from_spectral_chunks( flux )

    ### Now find the peaks using scipy signal's find peak algorithm, with the prominence threshold based on the "noise" estimate of the spectrum

    peak_pixels_initial, peak_properties = signal.find_peaks( flux, distance = config['wavecal']['lamp_line_min_separation_pix'], width = config['wavecal']['lamp_line_pix_width_limits'], prominence = config['wavecal']['lamp_line_peak_threshold_sigma'] * mad_peak_threshold )

    # Now go through each of the lines found and fit with a Gaussian to get the decimal pixel centroid

    peak_pixels_initial_fit = np.full( peak_pixels_initial.size, np.nan )

    for i_peak, peak in enumerate( peak_pixels_initial ):

        # Get the fit location: within +/- the config 'lamp_line_min_separation_pix' of the peak centroid
        fit_range = np.arange( peak - config['wavecal']['lamp_line_min_separation_pix'], peak + config['wavecal']['lamp_line_min_separation_pix'] + 1 )
        # Get rid of any nans!
        fit_range = np.intersect1d( fit_range, np.where( np.isfinite( flux ) )[0] )

        # The initial guess for 1D gaussian line parameters: the flux at the found centroid, the found int centroid, the minimum of the config 'lamp_line_pix_width_limits', and 0
        p_guess = [ flux[peak], peak, min( config['wavecal']['lamp_line_pix_width_limits'] ), 0 ]

        # Fit!
        try:
            line_fit, _ = optimize.curve_fit( reduction_utils.gaussian_1d, fit_range, flux[fit_range], p0 = p_guess )
        except:
            continue

        # Output the new fit centroid
        peak_pixels_initial_fit[i_peak] = line_fit[1]

    # Get rid of any nans!
    peak_pixels_initial_fit = peak_pixels_initial_fit[np.isfinite( peak_pixels_initial_fit )]

    return peak_pixels_initial_fit

def fit_wavelength_solution( pixel_centroids, prelim_wavelengths, line_list_wavelengths, config ):
    """ Function to fit a polynomial wavelength solution, iterating based on sigma rejection with the fit.

    Parameters
    ----------
    pixel_centroids : array
        The array of pixel centroid locations (dependent variable for the wavelength solution).
    prelim_wavelengths : array
        A preliminary guess for the wavelength solution (shape same as the output spectrum).
    line_list_wavelengths : array
        The array of line list wavelengths for the arc lamp used to match to.
    config : dict
        The overall config file defined using YAML with all parameters for running the reduction and analysis pipeline.

    Returns
    -------
    wave_poly_fit : array
        The final adopted polynomial fit coefficients for the wavelength solution (wavelength as a function of pixel). Size is the polynomial degree + 1.
    line_centroid_record : dict
        A dictionary containing a record of the fitting iterations. Each entry is a list of arrays, with each entry being a subsequent fitting iteration.
        Contains keys:
        'pixel': the pixel centroids used in the fit iteration.
        'wavelength': the corresponding wavelength centroids use in the fit iteration.
        'vel_resid': the residuals between the input wavelength centroids and polynomial fit values in velocity (km/s)
        'poly_coeffs': that iteration's polynomial fit coefficients. The last entry is the same as the output 'wave_poly_fit'
    """

    # Turn the pixel centroids into wavelength centroids based on the initial wavelength solution guess
    wavelength_centroids_init = np.interp( pixel_centroids, np.arange( prelim_wavelengths.size ), prelim_wavelengths )

    ### Now get the line list wavelength for the found peaks

    # Get the absolute difference between the initial fit wavelength centroids and the line list
    wave_diff_with_line_list = np.abs( wavelength_centroids_init[:,np.newaxis] - line_list_wavelengths )

    # Get the line list wavelengths for the closest match between observation and line list
    closest_line_wavelengths = line_list_wavelengths[np.nanargmin( wave_diff_with_line_list, axis =1 )]

    # Only use lines that have a match in the line list within some threshold of wavelength
    line_use_mask = np.abs( wavelength_centroids_init - closest_line_wavelengths ) < config['wavecal']['max_wave_diff_with_list']

    # Dictionary to hold the sequence of iterative wavelength solution fit results: the input pixel-wavelength pairs, velocity residuals, and polynomial coefficients
    # Initialize the pixel and wavelength lists with the array of pixel-wavelength pairs that pass the line list closesness cut
    line_centroid_record = { 'pixel': [ pixel_centroids[line_use_mask] ], 'wavelength': [ closest_line_wavelengths[line_use_mask] ],
                            'vel_resid': [], 'poly_coeffs': [] }

    ### Now run the loop for iteratively fitting the wavelength solution -- set break if the number of points include is fewer than the polynomial degree + 1
    while line_centroid_record['pixel'][-1].size >= ( config['wavecal']['wave_cal_poly_order'] + 1 ):

        # Fit a polynomial to the lines!
        wave_poly_fit = np.polyfit( line_centroid_record['pixel'][-1], line_centroid_record['wavelength'][-1], config['wavecal']['wave_cal_poly_order'] )
        line_centroid_record['poly_coeffs'].append( wave_poly_fit )

        # Calculate the fit polynomial at the found line peaks
        wave_poly_vals = np.polyval( wave_poly_fit, line_centroid_record['pixel'][-1] )

        # Calculate residuals between the catalogue wavelength and the polynomial fit, in velocity
        velocity_residuals = ( wave_poly_vals - line_centroid_record['wavelength'][-1] ) / line_centroid_record['wavelength'][-1] * 3e5
        line_centroid_record['vel_resid'].append( velocity_residuals )

        # Get the MAD of the velocity residuals
        velocity_residual_mad = stats.median_abs_deviation( velocity_residuals, scale = 'normal', nan_policy = 'omit' )

        # See how many lines are outside the sigma level defined in the config
        mad_reject = np.abs( velocity_residuals - np.nanmedian( velocity_residuals ) ) > config['wavecal']['vel_resid_sigma_reject'] * velocity_residual_mad

        # If there are no more points to reject -- break out of the loop!
        if mad_reject.sum() == 0 or (~mad_reject).sum() < ( config['wavecal']['wave_cal_poly_order'] + 1 ):
            break
        else:
            line_centroid_record['pixel'].append( line_centroid_record['pixel'][-1][~mad_reject] )
            line_centroid_record['wavelength'].append( line_centroid_record['wavelength'][-1][~mad_reject] )

    # Return the final polynomial fit coefficients separate from the record dictionary for ease of access

    return wave_poly_fit, line_centroid_record

def wavelength_solution_post_process( wave_sol_coeffs ):

    # The output set of wavelength solution polynomial coefficients -- will have rejected orders replaced with the fit to the coefficient vs. order
    output_wave_sol_coeffs = wave_sol_coeffs.copy()

    # Go through each of the polynomial coefficients
    for i_par in range( wave_sol_coeffs.shape[1] ):

        # Sigma reject fit to the polynomial coefficient vs. order
        fit, x_use, y_use = reduction_utils.polynomial_fit_sigma_reject( np.arange( wave_sol_coeffs.shape[0] ), wave_sol_coeffs[:,i_par], 4, 5, 3, return_data = True )

        # The orders that were rejected in the fits
        reject_orders = np.setdiff1d( np.arange( wave_sol_coeffs.shape[0] ), x_use )

        # Replace those coeffs with the fit values
        output_wave_sol_coeffs[reject_orders,i_par] = np.polyval( fit, reject_orders )

    return output_wave_sol_coeffs

def interpolate_wavelength_solution( jd_to_interpolate, jd_reference, wavelength_solution_reference ):
    """ Function to interpolate the wavelength solutions measured with the reference sources (e.g. arc lamps) to the times of science observations.

    Parameters
    ----------
    jd_to_interpolate : float
        The JD of the observation to interpolate to.
    jd_reference : array
        The array of JDs of the reference source observations.
    wavelength_solution_reference : array
        The wavelength solution for the reference source observations.

    Returns
    -------
    wavelength_solution_interpolate : array
        The interpolated wavelength solution at the input observation time.
    wavelength_solution_flag : str
        A string denoting what wavelength solution was adopted. Either 'LINTERP' to denote that it was linearly interpolated, or 'CLOSEST" to denote that it could not be interpolated and the closest reference solution was adopted.
    """

    # First check if the observation is bounded by reference spectra -- if not just adopt the closest wavelength solution and return a flag
    if jd_reference.min() > jd_to_interpolate or jd_reference.max() < jd_to_interpolate:

        wavelength_solution_interpolate = wavelength_solution_reference[np.argmin( np.abs( jd_to_interpolate - jd_reference ) )]

        wavelength_solution_flag = 'CLOSEST'

    # If it is bounded by arcs -- linear interpolation
    else:
        interp_fn = interpolate.interp1d( jd_reference, wavelength_solution_reference, axis = 0, kind = 'linear' )

        wavelength_solution_interpolate = interp_fn( jd_to_interpolate ).astype( float )

        wavelength_solution_flag = 'LINTERP'

    return wavelength_solution_interpolate, wavelength_solution_flag

### Plotting functions

def plot_spectra_zoom_windows( obs_wavelength, obs_flux, ref_wavelength, ref_flux, lines_used_wavelength, file_name, window_size = 10, number_of_subplots = 6 ):
    """ Function to plot zoom-ins of the final adopted wavelength calibrated observed spectrum compared to a reference source spectrum. Each page has multiple panels with small cut-out of the spectrum.
    This is to highlight the quality of the wavelength solution, and see how it compares to the reference for individual lines across the spectrum.

    Parameters
    ----------
    obs_wavelength : array
        The final adopted wavelength solution for the observed spectrum.
    obs_flux : array
        The observed flux array.
    ref_wavelength : array
        The reference source wavelength array.
    ref_flux : array
        The reference source flux array.
    lines_used_wavelength : array
        The wavelengths of the reference lines from the line list used to calibrate the wavelength solution.
    file_name : str
        The plot file name to write to.
    window_size : float, optional
        The wavelength width of each individual spectrum window. The default is 10.
    number_of_subplots : int, optional
        The number of spectrum window panels on each page. The default is 6.

    Returns
    -------
    None.
    """

    number_of_pages = int( np.ceil( np.ceil( np.ptp( obs_wavelength ) / window_size ) / number_of_subplots ) )

    subplot_wave_start = obs_wavelength.min()

    with PdfPages( file_name ) as pdf:

        for i_page in range( number_of_pages ):

            # fig = plt.figure( figsize = ( 12, 8 ), num = 1, clear = True )
            fig = plt.figure( num = 1, clear = True )

            fig.set_size_inches( 12, 8 )

            i_subplot = 1

            while subplot_wave_start <= obs_wavelength.max() and i_subplot <= number_of_subplots:

                subplot_wave_end = subplot_wave_start + 12.0

                obs_plot_loc  = np.where( ( obs_wavelength >= subplot_wave_start ) & ( obs_wavelength <= subplot_wave_end ) )[0]
                ref_plot_loc  = np.where( ( ref_wavelength >= subplot_wave_start ) & ( ref_wavelength <= subplot_wave_end ) )[0]
                line_plot_loc = np.where( ( lines_used_wavelength >= subplot_wave_start ) & ( lines_used_wavelength <= subplot_wave_end ) )[0]

                fig.add_subplot( 230 + i_subplot )

                plt.plot( obs_wavelength[obs_plot_loc], ( obs_flux / np.nanmedian( obs_flux ) )[obs_plot_loc], '#323232', lw = 1.25 )

                plt.plot( ref_wavelength[ref_plot_loc], ( ref_flux / np.nanmedian( ref_flux ) )[ref_plot_loc] * 5, '#bf3465', lw = 1 )

                for line_wavelength in lines_used_wavelength[line_plot_loc]:
                    plt.axvline( x = line_wavelength, c = '#1c6ccc', lw = 0.75, ls = '--' )

                plt.yticks( [], [] )
                plt.gca().set_yscale( 'log' )

                subplot_wave_start += window_size
                i_subplot += 1

            pdf.savefig( bbox_inches = 'tight', pad_inches = 0.05 )
            # fig.clear()
            # plt.close( fig )

    return None

def plot_wavelength_fit_iteration_spectra( fit_record, obs_flux, ref_wavelength, ref_flux, file_name ):
    """ Function to plot the wavelength-calibrated observed spectra for each of the fitting iterations compared to a reference spectrum. This is to highlight the quality of the solution.
    It is a multi-page PDF plot where each page is a separate iteration (working backwards -- the first page is the adopted iteration).

    Parameters
    ----------
    fit_record : dict
        A dictionary containing a record of the lines used and fit coefficients for each of the fitting iterations. This is the output from fit_wavelength_solution.
    obs_flux : array
        The observed flux array.
    ref_wavelength : array
        The reference source wavelength array.
    ref_flux : array
        The reference source flux array.
    file_name : str
        The plot file name to write to.

    Returns
    -------
    None.
    """

    ### Wrap everything in one multi-page PDF -- one page for each iteration
    with PdfPages( file_name ) as pdf:

        ### Go through each iteration, but reverse them -- the last iteration is plotted on the first page, and so on
        for i_iter in reversed( range( len( fit_record['pixel'] ) ) ):

            # Make the wavelength solution for this iteration
            wavelength_solution = np.polyval( fit_record['poly_coeffs'][i_iter], np.arange( obs_flux.size ) )

            # Make the figure
            fig = plt.figure( num = 1, clear = True )
            fig.set_size_inches( 12, 6 )

            # Plot the spectra! Normalize by the median
            plt.plot( wavelength_solution, obs_flux / np.nanmedian( obs_flux ), '#323232', lw = 1.25, label = 'Data Flux' )

            # Plot the reference ThAr spectrum. Normalize by the median and then x10 so it is offset from the data flux
            ref_plot_loc = np.where( ( ref_wavelength >= wavelength_solution.min() - 5 ) & ( ref_wavelength <= wavelength_solution.max() + 5 ) )[0]
            plt.plot( ref_wavelength[ref_plot_loc], ref_flux[ref_plot_loc] / np.nanmedian( ref_flux[ref_plot_loc] ) * 10, '#bf3465', lw = 1.0, label = 'Reference Arc Spectrum' )

            # Plot the lines
            for line_wavelength in fit_record['wavelength'][i_iter]:
                plt.axvline( x = line_wavelength, c = '#1c6ccc', lw = 0.75, ls = '--' )

            ### Labels and such
            plt.xlabel( 'Wavelength (${\\rm\AA}$)' )
            plt.ylabel( 'Flux' )
            plt.title( 'Lines Used: {}'.format( fit_record['wavelength'][i_iter].size ) )

            plt.legend( fontsize = 'small' )

            plt.gca().set_yscale( 'log' )

            pdf.savefig( bbox_inches = 'tight', pad_inches = 0.05 )
            # plt.gcf().clear()
            # plt.close( plt.gcf() )

    plt.cla()

    return None

def plot_wavelength_fit_iteration_residuals( fit_record, vel_resid_sigma_reject, file_name ):
    """ Function to plot the velocity residuals for the iterative wavelength solution fitting. It is a multi-page PDF plot where each page is a separate iteration (working backwards -- the first page is the adopted iteration).

    Parameters
    ----------
    fit_record : dict
        A dictionary containing a record of the lines used and fit coefficients for each of the fitting iterations. This is the output from fit_wavelength_solution.
    file_name : str
        The plot file name to write to.
    vel_resid_sigma_reject : float
        The sigma-level used for rejection, to plot horizontal lines and highlight points that are rejected in each iteration.

    Returns
    -------
    None.
    """

    ### Wrap everything in one multi-page PDF -- one page for each iteration
    with PdfPages( file_name ) as pdf:

        ### Go through each iteration, but reverse them -- the last iteration is plotted on the first page, and so on
        for i_iter in reversed( range( len( fit_record['pixel'] ) ) ):

            # Make the figure
            fig = plt.figure( num = 1, clear = True )
            fig.set_size_inches( 12, 6 )

            # Plot the velocity residuals
            plt.plot( fit_record['wavelength'][i_iter], fit_record['vel_resid'][i_iter], 'o', c = '#dfa5e5', mec = '#323232', mew = 0.5 )

            # Get the velocity residual MAD
            velocity_residual_mad = stats.median_abs_deviation( fit_record['vel_resid'][i_iter], scale = 'normal', nan_policy = 'omit' )

            # Plot x's over the points that are rejected!
            lines_rejected = np.where( np.abs( fit_record['vel_resid'][i_iter] - np.nanmedian( fit_record['vel_resid'][i_iter] ) ) > vel_resid_sigma_reject * velocity_residual_mad )[0]

            if lines_rejected.size > 0:
                plt.plot( fit_record['wavelength'][i_iter][lines_rejected], fit_record['vel_resid'][i_iter][lines_rejected], 'x', c = '#bf3465', mew = 1.25, ms = 8 )

            # Plot horizontal lines at the +/- sigma rejection level
            for i in [ -1, 1 ]:
                plt.axhline( y = np.nanmedian( fit_record['vel_resid'][i_iter] ) + i * vel_resid_sigma_reject * velocity_residual_mad, ls = '--', lw = 1.25, c = '#323232' )

            ### Labels and such
            plt.xlabel( 'Wavelength (${\\rm\AA}$)' )
            plt.ylabel( 'Fit Velocity Residual (km/s)' )

            plt.title( 'Lines Used: {}, Lines Rejected: {}, Residuals Sigma: {:.2f} km/s, Cut-off Sigma: {}'.format( fit_record['wavelength'][i_iter].size, lines_rejected.size, velocity_residual_mad, vel_resid_sigma_reject ) )

            pdf.savefig( bbox_inches = 'tight', pad_inches = 0.05 )
            # plt.gcf().clear()
            # plt.close( plt.gcf() )

    return None

##### Main wrapper script for wavelength calibration

def wavelength_solution( file_indices, header_df, config ):
    """ Main function to run for fitting the wavelength solution of a set of reference source spectra (e.g. arc lamps).
    Each reference source observation is read in, lines are found and matched to a line list, and the wavelength solution is iteratively fit with a polynomial. The wavelength solution is appended to the spectrum FITS file.

    Parameters
    ----------
    file_indices : array
        The file indices (in the header information file) of the reference source observations for wavelength solution fitting.
    header_df : pandas DataFrame
        The compiled information from the file headers.
    config : dict
        The overall config file defined using YAML with all parameters for running the reduction and analysis pipeline.

    Returns
    -------
    None.
    """

    ##### First make the wavelength solution!

    ### Read things in that we need for fitting the wavelength solution

    # The preliminary wavelength solution guess
    wavelength_solution_guess = np.load( os.path.join( config['general']['code_data_dir'], config['wavecal']['wave_sol_guess'] ) )

    # The arc lamp line list
    lamp_line_list_file = pd.read_csv( os.path.join( config['general']['code_data_dir'], config['wavecal']['line_list'] ) )
    lamp_line_list = lamp_line_list_file['wavelength'].values

    # Read in the wavelength vs flux CSV of the ThAr photron reference spectrum
    arc_ref_spectrum = pd.read_csv( os.path.join( config['general']['code_data_dir'], config['wavecal']['arc_ref_file'] ) )

    ### Go through each of the input file indices for the arc lamps to use
    for i_file in file_indices:

        file_in = fits.open( os.path.join( config['general']['reduction_dir'], 'spectrum_files', 'tullcoude_{}_spectrum.fits'.format( header_df['file_token'].values[i_file] ) ) )

        # Make directory for this frame's wavelength solution
        frame_dir_path = os.path.join( config['general']['reduction_dir'], 'wavecal', header_df['file_token'].values[i_file] )

        # Make the sub-directories for each of the types of plots
        for plot_dir_name in [ 'fit_residuals', 'fit_iter_spectra', 'adopted_sol_spectra_zoom' ]:
            os.makedirs( os.path.join( frame_dir_path, plot_dir_name ), exist_ok = True )

        ### Get the order offset to be applied to the preliminary solution if config is flagged to
        if config['wavecal']['use_prelim_sol_order_offset']:
            order_offset = order_offset_with_wave_sol_guess( wavelength_solution_guess, file_in['extracted flux'].data, arc_ref_spectrum['wavelength'].values, arc_ref_spectrum['flux'].values )

            # Don't calibrate with this file if the offset is bad! If it would be accessing an unallowed index of the wavelength solution guess (negative or larger than shape)
            if order_offset < 0 or ( order_offset + file_in[1].data.shape[0] ) > wavelength_solution_guess.shape[0]:
                continue
        else:
            order_offset = 0

        all_orders_wave_sol_poly_coeffs = np.full( ( config['trace']['number_of_orders'], config['wavecal']['wave_cal_poly_order'] + 1 ), np.nan )
        all_orders_wave_sol = np.full( file_in[1].data.shape, np.nan )

        # Go order by order and get wavelength solution
        for order in range( config['trace']['number_of_orders'] ):

            ### Fit the wavelength solution

            # Get the peak centroids
            lamp_line_pixel_centroids = find_arc_lamp_line_pixel_centers( file_in[1].data[order], config )

            # Fit the wavelength solution
            wavelength_solution_poly_coeffs, wavelength_solution_fit_record = fit_wavelength_solution( lamp_line_pixel_centroids, wavelength_solution_guess[order+order_offset], lamp_line_list, config )

            all_orders_wave_sol_poly_coeffs[order] = wavelength_solution_poly_coeffs
            all_orders_wave_sol[order] = np.polyval( wavelength_solution_poly_coeffs, np.arange( all_orders_wave_sol.shape[1] ) )

            ### Make any plots!

            # The fit velocity residuals
            plot_file_name = os.path.join( frame_dir_path, 'fit_residuals', 'fit_residuals_order_{}.pdf'.format( order ) )
            plot_wavelength_fit_iteration_residuals( wavelength_solution_fit_record, config['wavecal']['vel_resid_sigma_reject'], plot_file_name )

            # The spectra for each of the fits
            plot_file_name = os.path.join( frame_dir_path, 'fit_iter_spectra', 'fit_iter_spectra_order_{}.pdf'.format( order ) )
            plot_wavelength_fit_iteration_spectra( wavelength_solution_fit_record, file_in[1].data[order], arc_ref_spectrum['wavelength'].values, arc_ref_spectrum['flux'].values, plot_file_name )

            # Plot order spectrum with zoom-in windows to highlight the solution
            plot_file_name = os.path.join( frame_dir_path, 'adopted_sol_spectra_zoom', 'spectra_zoom_order_{}.pdf'.format( order ) )
            plot_spectra_zoom_windows( all_orders_wave_sol[order], file_in[1].data[order], arc_ref_spectrum['wavelength'].values, arc_ref_spectrum['flux'].values, wavelength_solution_fit_record['wavelength'][-1], plot_file_name )

        ### Output this frame's wavelength solution!

        # Re-build the output file rather than append to the input, in case of multiple runs
        output_file = fits.HDUList( [ file_in[0], file_in['extracted flux'], file_in['extracted flux error'], fits.ImageHDU( all_orders_wave_sol, name = 'wavelength' ) ] )

        # Add a header keyword to the primary HDU for the polynomial degree of the wavelength solution
        output_file[0].header['WAVPOLYD'] = ( config['wavecal']['wave_cal_poly_order'], 'Polynomial degree of wavelength solution' )

        # Add a history entry to the primary HDU to mark that it is wavelength calibrated
        output_file[0].header['HISTORY'] = 'Spectrum wavelength calibrated on {}'.format( datetime.strftime( datetime.now(), '%Y/%m/%d' ) )

        # Write out the file with the wavelength solution -- overwrite the previous file
        output_file.writeto( os.path.join( config['general']['reduction_dir'], 'spectrum_files', 'tullcoude_{}_spectrum.fits'.format( header_df['file_token'].values[i_file] ) ), overwrite = True )

    return None

def wavelength_calibrate( obj_file_indices, ref_file_indices, header_df, config ):
    """ Main function to run to wavelength calibrate the observed science spectra.
    The wavelength solutions fit for each of the reference source observations are interpolated onto the observation times of the other spectra.

    Parameters
    ----------
    obj_file_indices : array
        The file indices (in the header information file) of all observations with extracted spectra (including science frames and wavelength reference sources).
    ref_file_indices : array
        The file indices (in the header information file) of the reference source observations for wavelength solution fitting.
    header_df : pandas DataFrame
        The compiled information from the file headers.
    config : dict
        The overall config file defined using YAML with all parameters for running the reduction and analysis pipeline.

    Returns
    -------
    None.

    """
    ### Read in the arc wavelength solutions to use for interpolation

    arc_wavelength_solutions = []

    for i_file in ref_file_indices:

        file_name = os.path.join( config['general']['reduction_dir'], 'spectrum_files', 'tullcoude_{}_spectrum.fits'.format( header_df['file_token'].values[i_file] ) )

        file_in = fits.getdata( file_name, extname = 'wavelength' )

        arc_wavelength_solutions.append( file_in )

    arc_wavelength_solutions = np.array( arc_wavelength_solutions )

    ### Go through each of the input file indices for the arc lamps to use
    for i_file in obj_file_indices:

        # Interpolate the wavelength solution
        obj_wavelength_solution, obj_wavelength_solution_flag = interpolate_wavelength_solution( header_df['obs_jd'].values[i_file], header_df['obs_jd'].values[ref_file_indices], arc_wavelength_solutions )

        ### Output the wavelength calibrated file

        # Read in the spectrum file to append wavelength solution to
        file_name = os.path.join( config['general']['reduction_dir'], 'spectrum_files', 'tullcoude_{}_spectrum.fits'.format( header_df['file_token'].values[i_file] ) )
        file_in   = fits.open( file_name )

        # Build the output file
        output_file = fits.HDUList( [ file_in[0], file_in['extracted flux'], file_in['extracted flux error'], fits.ImageHDU( obj_wavelength_solution, name = 'wavelength' ) ] )

        # Add a header keyword to the primary HDU for the polynomial degree of the wavelength solution
        output_file[0].header['WAVPOLYD'] = ( config['wavecal']['wave_cal_poly_order'], 'Polynomial degree of wavelength solution' )

        # Add a header keyword for the wavelength solution interpolation: whether it is linearly interpolated or the closest reference wavelength solution was adopted
        output_file[0].header['WAVTYPE']  = ( obj_wavelength_solution_flag, 'Wavelength solution type, interpolation or closest adoption')

        # Add a history entry to the primary HDU to mark that it is wavelength calibrated
        output_file[0].header['HISTORY'] = 'Spectrum wavelength calibrated on {}'.format( datetime.strftime( datetime.now(), '%Y/%m/%d' ) )

        # Write out the file with the wavelength solution -- overwrite the previous file
        output_file.writeto( os.path.join( config['general']['reduction_dir'], 'spectrum_files', 'tullcoude_{}_spectrum.fits'.format( header_df['file_token'].values[i_file] ) ), overwrite = True )

    return None


class Vehicle():
    # Class Methods/ Attributes

    # Here self is passed as an argument because instance is passed as first argument
    def type_of_vehicle(self):  # Without self it throws an error
        print('I have a scooty')
        print("Vidit has a car")

    # print(2*4) # uncomment, you'll see this runs first

motor = Vehicle()
# print(motor)  # Print Memory location of "motor"
motor.type_of_vehicle()


import logging
import json
import os
import random
from typing import Dict
from typing import List
from typing import Iterable
from typing import Optional
from typing import Tuple
from pathlib import Path

from .analyses import AnalysisResultsStore
from .advisers import AdvisersResultsStore
from .dependency_monkey_reports import DependencyMonkeyReportsStore
from .inspections import InspectionStore
from .provenance import ProvenanceResultsStore
from .revsolvers import RevSolverResultsStore
from .security_indicators import SIAggregatedStore, SecurityIndicatorsResultsStore
from .solvers import SolverResultsStore

from .graph import GraphDatabase

_LOGGER = logging.getLogger(__name__)
_RANDOMIZE_LISTING = bool(int(os.getenv("THOTH_STORAGES_RANDOMIZE_LISTING", 0)))

def sync_adviser_documents(
    document_ids: Optional[Iterable[str]] = None,
    force: bool = False,
    graceful: bool = False,
    graph: Optional[GraphDatabase] = None,
    is_local: bool = False,
) -> Tuple[int, int, int, int]:
    """Sync adviser documents into graph."""
    if is_local and not document_ids:
        raise ValueError(
            "Cannot sync documents from local directory without explicitly specifying a list of documents to be synced"
        )

    if not graph:
        graph = GraphDatabase()
        graph.connect()

    if not is_local:
        adviser_store = AdvisersResultsStore()
        adviser_store.connect()

    listing = document_ids or adviser_store.get_document_listing()

    if _RANDOMIZE_LISTING:
        random.shuffle(list(listing))

    processed, synced, skipped, failed = 0, 0, 0, 0
    for document_id in listing:
        processed += 1

        if force or not graph.adviser_document_id_exist(os.path.basename(document_id)):
            try:
                if is_local:
                    _LOGGER.debug("Loading document from a local file: %r", document_id)
                    with open(document_id, "r") as document_file:
                        document = json.loads(document_file.read())
                else:
                    _LOGGER.info(
                        "Syncing adviser document from %r with id %r to graph", adviser_store.ceph.host, document_id
                    )
                    document = adviser_store.retrieve_document(document_id)

                graph.sync_adviser_result(document)
                synced += 1
            except Exception:
                if not graceful:
                    raise

                _LOGGER.exception("Failed to sync adviser result with document id %r", document_id)
                failed += 1
        else:
            _LOGGER.info(f"Sync of adviser document with id {document_id!r} skipped - already synced")
            skipped += 1

    return processed, synced, skipped, failed

def sync_solver_documents(
    document_ids: Optional[Iterable[str]] = None,
    force: bool = False,
    graceful: bool = False,
    graph: Optional[GraphDatabase] = None,
    is_local: bool = False,
) -> Tuple[int, int, int, int]:
    """Sync solver documents into graph."""
    if is_local and not document_ids:
        raise ValueError(
            "Cannot sync documents from local directory without explicitly specifying a list of documents to be synced"
        )

    if not graph:
        graph = GraphDatabase()
        graph.connect()

    if not is_local:
        solver_store = SolverResultsStore()
        solver_store.connect()

    listing = document_ids or solver_store.get_document_listing()

    if _RANDOMIZE_LISTING:
        random.shuffle(list(listing))

    processed, synced, skipped, failed = 0, 0, 0, 0
    for document_id in listing:
        processed += 1
        if force or not graph.solver_document_id_exists(os.path.basename(document_id)):
            try:
                if is_local:
                    _LOGGER.debug("Loading document from a local file: %r", document_id)
                    with open(document_id, "r") as document_file:
                        document = json.loads(document_file.read())
                else:
                    _LOGGER.info(
                        "Syncing solver document from %r with id %r to graph", solver_store.ceph.host, document_id
                    )
                    document = solver_store.retrieve_document(document_id)

                graph.sync_solver_result(document, force=force)
                synced += 1
            except Exception:
                if not graceful:
                    raise

                _LOGGER.exception("Failed to sync solver result with document id %r", document_id)
                failed += 1
        else:
            _LOGGER.info(f"Sync of solver document with id {document_id!r} skipped - already synced")
            skipped += 1

    return processed, synced, skipped, failed

def sync_revsolver_documents(
    document_ids: Optional[Iterable[str]] = None,
    force: bool = False,
    graceful: bool = False,
    graph: Optional[GraphDatabase] = None,
    is_local: bool = False,
) -> Tuple[int, int, int, int]:
    """Sync reverse solver documents into graph."""
    if is_local and not document_ids:
        raise ValueError(
            "Cannot sync documents from local directory without explicitly specifying a list of documents to be synced"
        )

    if not graph:
        graph = GraphDatabase()
        graph.connect()

    if not is_local:
        revsolver_store = RevSolverResultsStore()
        revsolver_store.connect()

    listing = document_ids or revsolver_store.get_document_listing()

    if _RANDOMIZE_LISTING:
        random.shuffle(list(listing))

    processed, synced, skipped, failed = 0, 0, 0, 0
    for document_id in listing:
        processed += 1
        try:
            if is_local:
                _LOGGER.debug("Loading document from a local file: %r", document_id)
                with open(document_id, "r") as document_file:
                    document = json.loads(document_file.read())
            else:
                _LOGGER.info(
                    "Syncing reverse solver document from %r with id %r to graph",
                    revsolver_store.ceph.host,
                    document_id,
                )
                document = revsolver_store.retrieve_document(document_id)

            graph.sync_revsolver_result(document)
            synced += 1
        except Exception:
            if not graceful:
                raise

            _LOGGER.exception("Failed to sync reverse solver result with document id %r", document_id)
            failed += 1

    return processed, synced, skipped, failed

def sync_analysis_documents(
    document_ids: Optional[Iterable[str]] = None,
    force: bool = False,
    graceful: bool = False,
    graph: Optional[GraphDatabase] = None,
    is_local: bool = False,
) -> Tuple[int, int, int, int]:
    """Sync image analysis documents into graph."""
    if is_local and not document_ids:
        raise ValueError(
            "Cannot sync documents from local directory without explicitly specifying a list of documents to be synced"
        )

    if not graph:
        graph = GraphDatabase()
        graph.connect()

    if not is_local:
        analysis_store = AnalysisResultsStore()
        analysis_store.connect()

    listing = document_ids or analysis_store.get_document_listing()

    if _RANDOMIZE_LISTING:
        random.shuffle(list(listing))

    processed, synced, skipped, failed = 0, 0, 0, 0
    for document_id in listing:
        processed += 1

        if force or not graph.analysis_document_id_exist(os.path.basename(document_id)):
            try:
                if is_local:
                    _LOGGER.debug("Loading document from a local file: %r", document_id)
                    with open(document_id, "r") as document_file:
                        document = json.loads(document_file.read())
                else:
                    _LOGGER.info(
                        "Syncing analysis document from %r with id %r to graph", analysis_store.ceph.host, document_id
                    )
                    document = analysis_store.retrieve_document(document_id)

                graph.sync_analysis_result(document)
                synced += 1
            except Exception:
                if not graceful:
                    raise

                _LOGGER.exception("Failed to sync analysis result with document id %r", document_id)
                failed += 1
        else:
            _LOGGER.info(f"Sync of analysis document with id {document_id!r} skipped - already synced")
            skipped += 1

    return processed, synced, skipped, failed

def sync_provenance_checker_documents(
    document_ids: Optional[Iterable[str]] = None,
    force: bool = False,
    graceful: bool = False,
    graph: Optional[GraphDatabase] = None,
    is_local: bool = False,
) -> Tuple[int, int, int, int]:
    """Sync provenance check documents into graph."""
    if is_local and not document_ids:
        raise ValueError(
            "Cannot sync documents from local directory without explicitly specifying a list of documents to be synced"
        )

    if not graph:
        graph = GraphDatabase()
        graph.connect()

    if not is_local:
        provenance_check_store = ProvenanceResultsStore()
        provenance_check_store.connect()

    listing = document_ids or provenance_check_store.get_document_listing()

    if _RANDOMIZE_LISTING:
        random.shuffle(list(listing))

    processed, synced, skipped, failed = 0, 0, 0, 0
    for document_id in listing:
        processed += 1

        if force or not graph.provenance_checker_document_id_exist(os.path.basename(document_id)):
            try:
                if is_local:
                    _LOGGER.debug("Loading document from a local file: %r", document_id)
                    with open(document_id, "r") as document_file:
                        document = json.loads(document_file.read())
                else:
                    _LOGGER.info(
                        "Syncing provenance-checker document from %r with id %r to graph",
                        provenance_check_store.ceph.host,
                        document_id,
                    )
                    document = provenance_check_store.retrieve_document(document_id)

                graph.sync_provenance_checker_result(document)
                synced += 1
            except Exception:
                if not graceful:
                    raise

                _LOGGER.exception("Failed to sync provenance-checker result with document id %r", document_id)
                failed += 1
        else:
            _LOGGER.info(f"Sync of provenance-checker document with id {document_id!r} skipped - already synced")
            skipped += 1

    return processed, synced, skipped, failed

def sync_dependency_monkey_documents(
    document_ids: Optional[Iterable[str]] = None,
    force: bool = False,
    graceful: bool = False,
    graph: Optional[GraphDatabase] = None,
    is_local: bool = False,
) -> Tuple[int, int, int, int]:
    """Sync dependency monkey reports into graph database."""
    if is_local and not document_ids:
        raise ValueError(
            "Cannot sync documents from local directory without explicitly specifying a list of documents to be synced"
        )

    if not graph:
        graph = GraphDatabase()
        graph.connect()

    if not is_local:
        dependency_monkey_reports_store = DependencyMonkeyReportsStore()
        dependency_monkey_reports_store.connect()

    listing = document_ids or dependency_monkey_reports_store.get_document_listing()

    if _RANDOMIZE_LISTING:
        random.shuffle(list(listing))

    processed, synced, skipped, failed = 0, 0, 0, 0
    for document_id in listing:
        processed += 1

        if force or not graph.dependency_monkey_document_id_exists(os.path.basename(document_id)):
            try:
                if is_local:
                    _LOGGER.debug("Loading document from a local file: %r", document_id)
                    with open(document_id, "r") as document_file:
                        document = json.loads(document_file.read())
                else:
                    _LOGGER.info(
                        f"Syncing dependency monkey report document from "
                        f"{dependency_monkey_reports_store.ceph.host} with id "
                        f"{document_id}, to graph"
                    )
                    document = dependency_monkey_reports_store.retrieve_document(document_id)

                graph.sync_dependency_monkey_result(document)
                synced += 1
            except Exception:
                if not graceful:
                    raise

                _LOGGER.exception("Failed to sync dependency-monkey result with document id %r", document_id)
                failed += 1
        else:
            _LOGGER.info(f"Sync of dependency-monkey document with id {document_id!r} skipped - already synced")
            skipped += 1

    return processed, synced, skipped, failed

def sync_inspection_documents(
    document_ids: Optional[Iterable[str]] = None,
    force: bool = False,
    graceful: bool = False,
    graph: Optional[GraphDatabase] = None,
    is_local: bool = False,
) -> Tuple[int, int, int, int]:
    """Sync observations made on Amun into graph database."""
    if is_local and not document_ids:
        raise ValueError(
            "Cannot sync documents from local directory without explicitly specifying a list of documents to be synced"
        )

    if not graph:
        graph = GraphDatabase()
        graph.connect()

    listing = document_ids or InspectionStore.iter_inspections()

    if _RANDOMIZE_LISTING:
        random.shuffle(list(listing))

    processed, synced, skipped, failed = 0, 0, 0, 0
    for inspection_document_id in listing:

        if not is_local:
            results = []
            inspection_store = InspectionStore(inspection_id=inspection_document_id)
            inspection_store.connect()

            number_results = inspection_store.results.get_results_count()
        else:
            main_repo = Path(f"{inspection_document_id}/results")
            results = [int(repo.name) for repo in main_repo.iterdir()]
            # dir entry should be numbers

        if number_results > 0:

            for inspection_result_number in results or range(number_results):

                processed += 1
                try:
                    if force or not graph.inspection_document_id_result_number_exists(
                        inspection_document_id=inspection_document_id, inspection_result_number=inspection_result_number
                    ):
                        if is_local:

                            inspection_specification_path = f"{inspection_document_id}/build/specification"
                            _LOGGER.debug(
                                "Loading specification document from a local file: %r", inspection_specification_path
                            )

                            inspection_result_path = (
                                f"{inspection_document_id}/results/{inspection_result_number}/result"
                            )
                            _LOGGER.debug("Loading result document from a local file: %r", inspection_result_path)

                            with open(inspection_specification_path, "r") as document_file:
                                inspection_specification_document = json.loads(document_file.read())

                            with open(inspection_result_path, "r") as document_file:
                                inspection_result_document = json.loads(document_file.read())
                        else:
                            _LOGGER.info(
                                "Syncing analysis document from %r with id %r and number %r to graph",
                                inspection_store.results.ceph.host,
                                inspection_document_id,
                                inspection_result_number,
                            )

                            inspection_specification_document = inspection_store.retrieve_specification()
                            inspection_result_document = inspection_store.results.retrieve_result(
                                inspection_result_number
                            )

                        inspection_document = {
                            "document_id": inspection_document_id,
                            "result_number": inspection_result_number,
                            "specification": inspection_specification_document,
                            "result": inspection_result_document,
                        }

                        graph.sync_inspection_result(inspection_document)
                        synced += 1

                    else:
                        _LOGGER.info(
                            f"Sync of results n.{inspection_result_number!r}"
                            f" from inspection id {inspection_document_id!r} skipped - already synced"
                        )
                        skipped += 1

                except Exception:
                    if not graceful:
                        raise

                    _LOGGER.exception(
                        f"Failed to sync results n.{inspection_result_number!r}"
                        f" from inspection id {inspection_document_id!r}"
                    )
                    failed += 1

        else:
            _LOGGER.info(f"inspection_document_id: {inspection_document_id!r} - does not have any results.")

    return processed, synced, skipped, failed

def sync_security_indicators_documents(
    document_ids: Optional[Iterable[str]] = None,
    force: bool = False,
    graceful: bool = False,
    graph: Optional[GraphDatabase] = None,
    is_local: bool = False,
) -> Tuple[int, int, int, int]:
    """Sync security indicators results into graph."""
    if is_local and not document_ids:
        raise ValueError(
            "Cannot sync documents from local directory without explicitly specifying a list of documents to be synced"
        )

    if not graph:
        graph = GraphDatabase()
        graph.connect()

    listing = document_ids or SecurityIndicatorsResultsStore.iter_security_indicators()

    if _RANDOMIZE_LISTING:
        random.shuffle(list(listing))

    processed, synced, skipped, failed = 0, 0, 0, 0
    for security_indicator_id in listing:

        processed += 1

        if force or not graph.si_aggregated_document_id_exists(security_indicator_id):
            try:
                if is_local:
                    _LOGGER.debug("Loading document from a local file: %r", security_indicator_id)
                    with open(security_indicator_id, "r") as document_file:
                        aggregated_document = json.loads(document_file.read())
                else:
                    si_aggregated_store = SIAggregatedStore(security_indicator_id=security_indicator_id)
                    si_aggregated_store.connect()
                    _LOGGER.info(
                        "Syncing analysis document from %r with id %r to graph",
                        si_aggregated_store.ceph.host,
                        security_indicator_id,
                    )

                    aggregated_document = si_aggregated_store.retrieve_document()

                graph.sync_security_indicator_aggregated_result(aggregated_document)
                synced += 1
            except Exception:
                if not graceful:
                    raise

                _LOGGER.exception("Failed to sync analysis result with document id %r", security_indicator_id)
                failed += 1
        else:
            _LOGGER.info(f"Sync of analysis document with id {security_indicator_id!r} skipped - already synced")
            skipped += 1

    return processed, synced, skipped, failed

# Corresponding mapping of document prefix (before the actual unique document identifier part) to
# functions which can handle the given document sync.

HANDLERS_MAPPING = {
    "adviser": sync_adviser_documents,
    "dependency": sync_dependency_monkey_documents,
    "inspection": sync_inspection_documents,
    "package": sync_analysis_documents,
    "provenance": sync_provenance_checker_documents,
    "solver": sync_solver_documents,
    "revsolver": sync_revsolver_documents,
    "security": sync_security_indicators_documents,
}

def sync_documents(
    document_ids: Optional[Iterable[str]] = None,
    force: bool = False,
    graceful: bool = False,
    graph: Optional[GraphDatabase] = None,
    is_local: bool = False,
) -> Dict[str, Tuple[int, int, int, int]]:
    """Sync documents based on document type.

    If no list of document ids is provided, all documents will be synced.
    >>> from thoth.storages.sync import sync_documents
    >>> sync_documents(["adviser-efa7213babd12911", "package-extract-f8e354d9597a1203"])
    """
    handlers: Dict[str, List[str]] = dict()
    if document_ids:
        handlers = {key: [] for key in HANDLERS_MAPPING.keys()}
        assert handlers is not None
        for doc in document_ids:
            try:
                handlers[doc[doc.rfind("/") + 1 : doc.find("-")]].append(doc)
                # Basename for local syncs, document_id should not have slash otherwise.
            except KeyError:
                error_msg = f"No handler defined for document identifier {doc}"
                if not graceful:
                    raise ValueError(error_msg)
                _LOGGER.error(error_msg)

    else:
        static_handlers = dict.fromkeys(HANDLERS_MAPPING, None)

    stats = dict.fromkeys(HANDLERS_MAPPING, (0, 0, 0, 0))
    for handler, documents in (handlers or static_handlers).items():
        if documents != []:
            stats[handler] = HANDLERS_MAPPING[handler](
                documents, force=force, graceful=graceful, graph=graph, is_local=is_local
            )

    return stats


# Given an array of integers, return a new array such that each element at index i of the new array is the product of all the numbers in the original array except the one at i.

# For example, if our input was [1, 2, 3, 4, 5], the expected output would be [120, 60, 40, 30, 24]. If our input was [3, 2, 1], the expected output would be [2, 3, 6].

# Follow-up: what if you can't use division?

from typing import List

def product_pos(numbers: List) -> List:
    prefix: List[int] = [1]*len(numbers)
    postfix: List[int] = [1] * len(numbers)
    result: List[int] = [0] * len(numbers)
    # create prefix and postfix products
    for i in range(1, len(numbers)):
        prefix[i] = prefix[i-1] * numbers[i-1]
    for i in range(len(numbers)-2, -1, -1):
        postfix[i] = postfix[i+1] * numbers[i+1]

    for i in range(0, len(numbers)):
        result[i] = prefix[i]*postfix[i]
    return result

def main():
    test1: List[int] = [1, 2, 3, 4, 5]
    test2: List[int] = [3, 2, 1]

    assert product_pos(test1) == [120, 60, 40, 30, 24]
    assert product_pos(test2) == [2, 3, 6]

if __name__ == "__main__":
    main()


def extractYumeineijiworksWordpressCom(item):
	'''
	Parser for 'yumeineijiworks.wordpress.com'
	'''

	vol, chp, frag, postfix = extractVolChapterFragmentPostfix(item['title'])
	if not (chp or vol) or "preview" in item['title'].lower():
		return None

	tagmap = [
		('hagure seirei ino shinsatsu kiroku',                                    'hagure seirei ino shinsatsu kiroku',                                                   'translated'),
		('homeless tensei: isekai de jiyuu sugiru majutsu jisoku seikatsu',       'homeless tensei: isekai de jiyuu sugiru majutsu jisoku seikatsu',                      'translated'),
		('chocolate bliss',                                                       'chocolate bliss',                                                                      'translated'),
		('kajiya de hajimeru isekai slow life',                                   'kajiya de hajimeru isekai slow life',                                                  'translated'),
		('glitzheria',                                                            'glitzheria',                                                                           'translated'),
		('PRC',       'PRC',                      'translated'),
		('Loiterous', 'Loiterous',                'oel'),
	]

	for tagname, name, tl_type in tagmap:
		if tagname in item['tags']:
			return buildReleaseMessageWithType(item, name, vol, chp, frag=frag, postfix=postfix, tl_type=tl_type)

	return False


import requests
import json
import sys
import boto3
from decimal import Decimal
from boto3.dynamodb.conditions import Key
import datetime
from selenium import webdriver
from webdriver_manager.chrome import ChromeDriverManager
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.by import By
import time
import pandas as pd
import random
from bs4 import BeautifulSoup as bs

# connect to dynamoDB
try:
    table = boto3.resource('dynamodb').Table('product-data')
except Exception as e:
    print(e)
    sys.exit(-1)

def lambda_handler(event=None, context=None):
    # Ideally we would get a query/store search parameter from event through api gateway
        # Ex: {'query' : 'apple', 'store' : 'vendor1'}
    query = event.get('queryStringParameters') and event['queryStringParameters'].get('query')
    store = event.get('queryStringParameters') and event['queryStringParameters'].get('store')
    category = event.get('queryStringParameters') and event['queryStringParameters'].get('category')
    #Default
    if query is None:
        query = "milk"

    if store is None:
        store = "vendor1"

    if category is None:
        category = 'milk_cheese_eggs'

    response = query_api(query, store, category)
    return response

def query_vendor1(query):
    headers = {'User-Agent': 'Mozilla/5.0', }
    response = None
    store_id = '5659'  # Arroyo Grande vendor 1 Store ID
    proxies = { 'http': 'http://10.10.1.10:3128' }
    # the below url only grabs 60 items for now...
    vendor = ""
    url = f'https://www.{vendor}.com/grocery/v4/api/products/search?count=60&offset=0&page=1&storeId={store_id}&query={query}'
    response = requests.get(url, headers=headers, proxies=proxies)
    json_response = json.loads(response.text, parse_float=Decimal)
    products = json_response['products']
    return products

def query_vendor2(query):
    headers = {'User-Agent': 'Mozilla/5.0', }
    response = None
    store_id = '2759'

    provider = "" # NOTE: vendor has a provider
    vendor = "" # NOTE: provide own vendor
    key = "" # NOTE: requires an api key
    visitorId = "" # NOTE: requires a visitor id
    url = f"https://{provider}.{vendor}.com/{provider}_aggregations/v1/web/plp_search_v1?key={key}&channel=WEB&count=24&keyword={query}" \
    f"&page={query}&pricing_store_id={store_id}&store_ids={store_id}&visitor_id={visitorId}"

    response = requests.get(url, headers=headers)
    json_response = json.loads(response.text, parse_float=Decimal)
    products = json_response['data']['search']['products']
    return products

def query_vendor3(query):
    '''
    Queries vendor 3, an Online marketplace based in the _________.
    Offers prices for ___, ________-, and ____ for their respective currencies

    Currently:
        Api only retreives offers from ___. Possible change in cookie values could retrieve other countries
        Limits high-rate requests, so the requests getting specific product specifications must be spread out
        If less specific data is ok, just the search result can be used. It contains the product name and price
    '''
    def vendor3_await_response():
        time.sleep(random.randint(27, 47)) # Scatters requests to not get 429 code. These numbers work consistently and were derived experimentally

    headers = {'User-Agent': 'Mozilla/5.0', }
    response = None
    #country_codes = ["___", "______", "______"]
    limit = 50  # default is 50, max is 150
    page = 1  # default is 1
    search_query = query.replace("-", " ")   # for better results
    vendor3 = ""
    search_url = f"https://www.{vendor3}.com/_svc/catalog/api/v2/u/search?limit={limit}&page={page}&q={search_query}"
    search_response = requests.get(search_url, headers=headers)
    if search_response.status_code != 200:
        raise Exception(f"Vendor 3 API Request Failed. Status Code {search_response.status_code}")
    vendor3_await_response()
    json_search_response = json.loads(search_response.text, parse_float=Decimal)
    results = json_search_response['hits']
    total_pages = json_search_response["nbPages"]  # not used, but can be for more results
    products = []
    for i, result in enumerate(results):
        product_url = f"https://www.{vendor3}.com/_svc/catalog/api/v2/u/{result['url']}/{result['sku']}/p?o={result['offer_code']}"
        product_response = requests.get(product_url, headers=headers)
        # print(f'{product_response.status_code} for vendor 3 product #{i+1}')
        if product_response.status_code != 200:
            print(f"Vendor 3 product #{i+1} for query, \"{query}\", returned reponse code {product_response.status_code}")
            if product_response.status_code == 429:
                for i in range(2):  # wait 2x the normal amount
                    vendor3_await_response()
            else:
                vendor3_await_response()
            break
        print(f'Vendor 3 product #{i+1} successfully retrieved')
        json_product_response = json.loads(product_response.text, parse_float=Decimal)
        products.append(json_product_response['product'])
        vendor3_await_response()
    return products

def query_vendor4(query):
    options = webdriver.ChromeOptions()
    options.add_argument('headless')
    options.add_argument("--log-level=3")

    vendor4 = ""

    # load page
    driver = webdriver.Chrome(ChromeDriverManager().install(), options=options)
    driver.get(f'https://www.{vendor4}.com.mx/en/search/?ref_p=ser_tp&q={query}&sort_by=relevance')
    # time.sleep(3)

    wait = WebDriverWait(driver, 10)
    # elements=wait.until(EC.element_to_be_selected(driver.find_elements_by_class_name("item")))

    wait.until(EC.presence_of_element_located((By.CLASS_NAME, "item")))
    elements = driver.find_elements_by_class_name("item")
    product_list = []
    for item in elements:
        if len(product_list) >= 5:
            break
        # initial page of multiple items

        time.sleep(10)
        min_detail = item.find_elements_by_class_name("product-min-detail")[0]
        description = min_detail.find_element_by_tag_name("a").get_attribute('title')
        price = min_detail.find_element_by_class_name("product-price").text
        image = min_detail.parent.find_element_by_class_name('product-image') \
         .find_element_by_tag_name("a").find_element_by_tag_name("img").get_attribute('src')
        description = min_detail.find_element_by_tag_name("a").get_attribute('title')

        # specific item data
        url = min_detail.find_element_by_tag_name("a").get_attribute('href')
        # need new driver otherwise encounter stale element reference
        driver2 = webdriver.Chrome(ChromeDriverManager().install(), options=options)
        driver2.get(url)
        time.sleep(10)

        # time.sleep(3)
        wait = WebDriverWait(driver, 10)
        # wait.until(EC.element_to_be_selected(driver2.find_element_by_id('product-attribute-specs-table'))).get_attribute('outerHTML')
        wait.until(EC.presence_of_element_located((By.ID, 'product-attribute-specs-table')))
        table = driver2.find_element_by_id('product-attribute-specs-table').get_attribute('outerHTML')

        driver2.quit()

        product_details = pd.read_html(table)[0]
        print(product_details)
        json = product_details.to_json(orient='index')
        record = {'description' : description, 'price' : price, 'image' : image, 'specifications' : json}
        product_list.append(record)

    raise Exception("goodnight")
    driver.quit()
    return product_list
# For each store we want to get a few basic data points including
#   * product id / sku
#   * product image
#   * general info: name, sales quantity, dimensions/volume etc
#   * PRICE
#   * Date

def insert_vendor1_products(query, products, category):
    if len(products) == 0:
        print('No query results found')
        return None

    try:
        product_list = []
        stores = []
        store_info = {'name': 'vendor1'}
        product_metadata = {'product': query}

        for product in products:
            if 'sku' in product or 'USItemId':
                product_info = {
                    'sku': product['sku'] if 'sku' in product.keys() else product['USItemId'],
                    'info': product['basic'],  # vendor 1 specific
                    'pricing': product['store']['price']
                }
                product_list.append(product_info)

        existing_record = table.query(KeyConditionExpression=Key('query').eq(query))

        stores = len(existing_record['Items']) > 0 and existing_record['Items'][0]['stores']
        if stores:
            stores['vendor1']  = product_list
            product_metadata['stores'] = stores
        else:
            product_metadata['stores'] = {'vendor1' : product_list}

        # currencies = len(existing_record['Items']) > 0 and existing_record['Items'][0]['currencies']
        # if currencies:
        #     product_metadata['currencies']['vendor1'] = 'USD'
        # else:
        #     product_metadata['currencies'] = {'vendor1' : 'USD'}

        product_metadata['category'] = category
        product_metadata['query'] = query # dynamo partition key
        product_metadata['date-collected'] = datetime.datetime.now().strftime("%Y-%m-%d")

        response = table.put_item(Item=product_metadata)

        return response

    except Exception as e:
        print("Could not insert into db", e)
        return None

def insert_vendor2_products(query, products, category):
    if len(products) == 0:
        print('No query results found')
        return None

    try:
        product_list = []
        stores = []
        store_info = {'name': 'vendor2'}
        product_metadata = {'product': query,}

        for product in products:
            if 'tcin' in product:
                product_info = {
                    'tcin' : product['tcin'],
                    'info' : product['item'],
                    'pricing' : product['price']
                }
            product_list.append(product_info)

        existing_record = table.query(KeyConditionExpression=Key('query').eq(query))

        stores = len(existing_record['Items']) > 0 and existing_record['Items'][0]['stores']
        if stores:
            stores['vendor2']  = product_list
            product_metadata['stores'] = stores
        else:
            product_metadata['stores'] = {'vendor2' : product_list, }

        # currencies = len(existing_record['Items']) > 0 and existing_record['Items'][0]['currencies']
        # if currencies:
        #     product_metadata['currencies']['vendor2'] = 'USD'
        # else:
        #     product_metadata['currencies'] = {'vendor2' : 'USD'}

        product_metadata['category'] = category
        product_metadata['query'] = query # dynamo partition key
        product_metadata['date-collected'] = datetime.datetime.now().strftime("%Y-%m-%d")
        response = table.put_item(Item=product_metadata)

        return response

    except Exception as e:
        print("Could not insert into db", e)
        return None

def insert_vendor3_products(query, products, category):
    # if len(products) == 0:
    #     print('No query results found')
    #     return None

    # try:
        product_list = []
        stores = []
        store_info = {'name': 'vendor3'}
        product_metadata = {'product': query}

        for product in products:
            if 'sku' in product:
                try:
                    product_info = {
                        'sku' : product['sku'],
                        'brand' : product['brand'],
                        'name' : product['product_title'],
                        'specifications' : product['specifications'],
                        'pricing' : {'price' : product['variants'][0]['offers'][0]['price'],
                                    'sale_price' : product['variants'][0]['offers'][0]['sale_price']},
                        'images' : {'image_keys' : product['image_keys']},
                        'url' : product['variants'][0]['offers'][0]['url'] + "/" + product["sku"]
                    }
                except Exception as e:
                    print(f"vendor3: product {product['sku']} raised: {e}")
                    continue

            product_list.append(product_info)

        existing_record = table.query(KeyConditionExpression=Key('query').eq(query))

        stores = len(existing_record['Items']) > 0 and existing_record['Items'][0]['stores']

        if stores:
            stores['vendor3']  = product_list
            product_metadata['stores'] = stores
        else:
            product_metadata['stores'] = {'vendor3' : product_list,}

        # currencies = len(existing_record['Items']) > 0 and existing_record['Items'][0]['currencies']
        # if currencies:
        #     product_metadata['currencies']['vendor3'] = 'USD'
        # else:
        #     product_metadata['currencies'] = {'vendor3' : 'USD'}

        product_metadata['category'] = category
        product_metadata['query'] = query # dynamo partition key
        product_metadata['date-collected'] = datetime.datetime.now().strftime("%Y-%m-%d")
        response = table.put_item(Item=product_metadata)

        return response

    # except Exception as e:
    #     print("Could not insert into db", e)
    #     return None

def insert_vendor4_products(query, products, category):
    if len(products) == 0:
        print('No query results found')
        return None

    try:
        product_list = []
        stores = []
        store_info = {'name': 'vendor4'}
        product_metadata = {'product': query}

        for product in products:
            # record = {'description' : description, 'price' : price, 'image' : image, 'specifications' : json}
            product_info = {
                'description' : product['description'],
                'price' : product['price'],
                'image' : product['image'],
                'specifications' : product['specifications'],
            }
            product_list.append(product_info)

        existing_record = table.query(KeyConditionExpression=Key('query').eq(query))

        stores = len(existing_record['Items']) > 0 and existing_record['Items'][0]['stores']
        if stores:
            stores['vendor4']  = product_list
            product_metadata['stores'] = stores
        else:
            product_metadata['stores'] = {'vendor4' : product_list,}

        # currencies = len(existing_record['Items']) > 0 and existing_record['Items'][0]['currencies']
        # if currencies:
        #     product_metadata['currencies']['vendor4'] = 'USD'
        # else:
        #     product_metadata['currencies'] = {'vendor4' : 'USD'}

        product_metadata['category'] = category
        product_metadata['query'] = query # dynamo partition key
        product_metadata['date-collected'] = datetime.datetime.now().strftime("%Y-%m-%d")
        response = table.put_item(Item=product_metadata)

        return response

    except Exception as e:
            print("Could not insert into db", e)
            return None

def query_api(query, store, category):
    products = None
    response = None

    if store == 'vendor1':
        products = query_vendor1(query)
        response = insert_vendor1_products(query, products, category)
    elif store == 'vendor2':
        products = query_vendor2(query)
        response = insert_vendor2_products(query, products, category)
    elif store == 'vendor3':
        products = query_vendor3(query)
        response = insert_vendor3_products(query, products, category)
    elif store == 'vendor4':
        products = query_vendor4(query)
        response = insert_vendor4_products(query, products, category)
    else:
        raise ValueError('unsupported store')

    return response and response['ResponseMetadata']['HTTPStatusCode']

# print(lambda_handler({"queryStringParameters" : {"query" : "toaster", "store" : "vendor2",
# 'category' : 'small_electric_household_appliances'}}, None))

# print(lambda_handler({"queryStringParameters" : {"query" : "milk", "store" : "vendor3",
# 'category' : 'milk_cheese_eggs'}}, None))
# # print(lambda_handler({'queryStringParameters' : {'query' : 'mutton', 'store' : 'vendor1'}}, None))

# just took out 'rice' as the first element of the list from 'bread_cereal'
'''
foods = {
    'bread_cereal' : ['rice', 'cereal', 'bread', 'bakery', 'pasta', 'couscous'],
    'meat' : ['beef', 'veal', 'lamb', 'goat', 'chicken', 'poultry'],
    'fish_seafood' : ['fish', 'seafood', 'frozen-seafood', 'frozen-fish'],
    'milk_cheese_eggs' : ['milk', 'cheese', 'curd', 'eggs'],
    'oils_fats' : ['butter', 'margarine', 'oil'],
    'vegetables' : ['vegetables', 'potatoes', 'tuber-vegetables']
}

household_appliances = {
    'small_electric_household_appliances' : ['coffee maker', 'food processor',
    'mixer', 'electric grill', 'microwave', 'fryer', 'rice cooker']
}
'''

household_appliances = {
    'small_electric_household_appliances' : ['microwave', 'fryer', 'rice cooker']
}

store = "vendor3"

# for category in foods:
#     for query in foods[category]:
#         print("Beginning -> " + store + ": " + category + ": " + query)
#         print(lambda_handler({"queryStringParameters" : {"query" : query, "store" : store,
#                                                          'category' : category}}, None))
#         print("Completed -> " + store + ": " + category + ": " + query)

for category in household_appliances:
    for query in household_appliances[category]:
        print("Beginning -> " + store + ": " + category + ": " + query)
        print(lambda_handler({"queryStringParameters" : {"query" : query, "store" : store,
                                                         'category' : category}}, None))
        print("Completed -> " + store + ": " + category + ": " + query)

# stores = ['vendor2', 'vendor3']

# for store in stores:
#     for category, item_list in household_appliances.items():
#         for item in item_list:

#             if store != 'vendor4':
#                 time.sleep(random.randint(5,10))
#             try:
#                 response_code = lambda_handler({'queryStringParameters' :
#                     {'query' : item, 'store' : store, 'category' : category}}, None)
#                 print("{} {}".format(item, store))

#             except Exception as e:
#                 print('FAILED: {}       {}\n Error: {}'.format(store, item, e))

#         if response_code != 200:
#             print('FAILED: {}       {}'.format(store, item))


from Bot import Bot
from SuperBot import SuperBot
from FlyingBot import FlyingBot

print("Please enter the bot's name")
bots_name = input()
print("Please enter the bot's age")
bots_age = int(input())
print("Please enter the bot's energy level")
bots_energy_level = int(input())
print("Please enter the bots shield level")
bots_shield_level = int(input())

New_Bot = Bot(bots_name,bots_age,bots_energy_level,bots_shield_level)
New_Super_Bot = SuperBot()
New_Flying_Bot = FlyingBot()

#Displaying methods
New_Bot.display_name()
New_Bot.display_age()
New_Bot.dispaly_energy()
New_Bot.display_shield()
New_Bot.display_summary()

#Returning methods
New_Bot.get_age()
New_Bot.get_name()
New_Bot.get_energy()
New_Bot.get_shield()
New_Flying_Bot.get_hover_distane()
New_Super_Bot.get_super_power_level()

#Incrementing and Decrementing Methods
New_Bot.decrement_energy(2)
New_Bot.decrement_shield(2)
New_Bot.increment_energy(3)
New_Bot.increment_shield(3)
New_Bot.increment_age()

#Setting values methods

New_Bot.set_name("Boop")
New_Super_Bot.set_super_power_level(3)
New_Flying_Bot.set_hover_distance(4)


import numpy as np
from scipy import signal
import random

class Game:
    """Implementation of Conway's Game of Life using a list of ints.

    Args:
        x (int): X-Axis size.
        y (int): Y-Axis size.
        seed (int, optional): Grid generator seed.
    """

    def __init__(self, x, y):
        self.grid = np.random.randint(2, size=(y, x))

    def print_grid(self):
        """Print out the grid into the console."""

        for row in range(len(self.grid)):
            for col in range(len(self.grid[0])):
                print(self.grid[row][col], end=" ")
            print("\n", end="")
        print("\n", end="")

    def cycle(self):
        """Cycle through all of the positions in the grid."""

        kernel = np.array([[1,1,1],[1,0,1],[1,1,1]])
        neighbors = signal.convolve2d(self.grid, kernel, "same")

        for row in range(self.grid.shape[0]):
            for col in range(self.grid.shape[1]):
                if self.grid[row,col] == 1 and (neighbors[row,col] == 2 or neighbors[row,col] == 3):
                    self.grid[row,col] = 1
                elif self.grid[row,col] == 0 and neighbors[row,col] == 3:
                    self.grid[row,col] = 1
                else:
                    self.grid[row,col] = 0

    def clear(self):
        self.grid = np.zeros((self.grid.shape[1], self.grid.shape[0]))


with open('text.txt', 'r', encoding='utf-8') as file_descriptor:
    for line in file_descriptor:
        print(line)

class Hello:
    def __enter__(self):
        print('вход в блок')

    def __exit__(self, exp_type, exp_value, traceback):
        print('выход из блока')

    def __del__(self):
        print('деструктор')

with Hello():
     print('мой код')

'''
вход в блок
мой код
выход из блока
деструктор
'''

import contextlib
@contextlib.contextmanager
def context():
    print('вход в блок')
    try:
        yield {}
    except RuntimeError as err:
        print('error: ', err)
    finally:
        print('выход из блока')
'''
вход в блок
блок
выход из блока
'''

# nested

first, second = context(), context()
with first as first:
    with second as second:
        print('внутри блока')


import re
from src.helpers.inputs_and_validations import get_string_input, get_float_input
from src.controllers.courses import Courses, list_course_in_tabular_form
from src.controllers.earning import Earning
from src.controllers.faq import Faq
from src.controllers.feedback import Feedback
from src.controllers.mentor import Mentor
from src.controllers.student import Student
from src.models.database import DatabaseConnection
from src.models.fetch_json_data import JsonData
from src.utils.menu_prompt_functions import admin_menu, student_menu, mentor_menu, visitor_menu

DatabaseConnection = DatabaseConnection()

get_query = JsonData.load_data()

student = Student()
mentor = Mentor()
feedback = Feedback()
course = Courses()
faq = Faq()
earning = Earning()

class EntryMenu:
    # def __init__(self, role, user_id):
    #     self.role = role
    #     self.user_id = user_id
    #
    #     if role == 1:
    #         self.prompt_admin_menu(role, user_id)
    #     elif role == 2:
    #         # student
    #         self.prompt_student_menu(role, user_id)
    #     elif role == 3:
    #         # mentor
    #         self.prompt_mentor_menu(role, user_id)
    #     elif role == 4:
    #         # visitor
    #         self.prompt_visitor_menu(role, user_id)

    def prompt_admin_menu(self, role, user_id):
        pending_course_count = self.check_pending_courses()
        print(pending_course_count)
        if pending_course_count > 0:
            self.list_pending_course(pending_course_count)
        print(admin_menu)
        try:
            user_input = self.input_choice()
            while user_input != 7:
                if user_input == 1:
                    student.view_student_details()
                elif user_input == 2:
                    self.view_course_feedback(role)
                elif user_input == 3:
                    mentor.add_mentor()
                elif user_input == 4:
                    course.list_course(role, user_id)
                elif user_input == 5:
                    self.input_delete_course_name(user_id)
                elif user_input == 6:
                    earning.calculate_all_mentor_earning()
                else:
                    print("You entered wrong choice, please try again.. ")
                print(admin_menu)
                user_input = self.input_choice()
        except:
            print("You entered a wrong choice, please try again ...")
            self.prompt_admin_menu(role, user_id)

    def prompt_student_menu(self, role, user_id):
        print(student_menu)
        try:
            user_input = self.input_choice()
            while user_input != 6:
                if user_input == 1:
                    course.list_course(role, user_id)
                elif user_input == 2:
                    self.list_purchasable_course(user_id)
                elif user_input == 3:
                    self.input_course_name_to_study_from(user_id)
                elif user_input == 4:
                    self.view_course_feedback(role)
                elif user_input == 5:
                    self.add_course_feedback(user_id)
                else:
                    print("You entered wrong choice, please try again.. ")
                print(student_menu)
                user_input = self.input_choice()

        except:
            print("You entered a wrong choice, please try again ...")
            self.prompt_student_menu(role, user_id)

    def prompt_mentor_menu(self, role, user_id):
        print(mentor_menu)
        try:
            user_input = self.input_choice()

            while user_input != 11:
                if user_input == 1:
                    self.input_course_details(user_id)
                elif user_input == 2:
                    earning.calculate_mentor_earning(user_id)
                elif user_input == 3:
                    mentor.list_course(2, user_id)
                elif user_input == 4:
                    mentor.list_course(3, user_id)
                elif user_input == 5:
                    self.view_course_feedback(role)
                elif user_input == 6:
                    self.input_course_faq(user_id)
                elif user_input == 7:
                    self.view_course_faq(user_id)
                elif user_input == 8:
                    self.list_purchasable_course(user_id)
                elif user_input == 9:
                    self.add_course_feedback(user_id)
                elif user_input == 10:
                    self.input_course_name_to_study_from(user_id)
                else:
                    print("You entered wrong choice, please try again.. ")
                print(mentor_menu)
                user_input = self.input_choice()

        except:
            print("You entered a wrong choice, please try again ...")
            self.prompt_mentor_menu(role, user_id)

    def prompt_visitor_menu(self, role, user_id):
        print(visitor_menu)
        try:
            user_input = self.input_choice()
            while user_input != 5:
                if user_input == 1:
                    course.list_course(role, user_id)
                elif user_input == 2:
                    self.view_course_faq(user_id)
                elif user_input == 3:
                    self.list_purchasable_course(user_id)
                elif user_input == 4:
                    self.view_course_feedback(role)
                else:
                    print("You entered wrong choice, please try again.. ")
                print(visitor_menu)
                user_input = self.input_choice()

        except:
            print("You entered a wrong choice, please try again ...")
            self.prompt_visitor_menu(role, user_id)

    def input_choice(self):
        user_input = int(input("Please enter your choice : "))
        if user_input <= 0:
            print("input cannot be less than 0.. please try again. ")
            return self.input_choice()
        elif user_input > 0:
            return user_input
        else:
            print("Please enter valid number...")
            return self.input_choice()

    def input_course_details(self, user_id):
        self.course_name = self.input_course_name()
        is_valid_course_name = DatabaseConnection.get_from_db(get_query.get("GET_DETAILS_COURSES"), (self.course_name,))
        if is_valid_course_name:
            print("Same name course already exists. Please enter a different name.")
            return
        self.content = self.input_course_content()
        self.duration = self.input_course_duration()
        self.price = self.input_course_price()

        course.add_course(user_id, self.course_name, self.content, self.duration, self.price)

    def input_course_duration(self):
        user_input = int(input("Enter duration of course (in hours) : "))
        if user_input <= 0:
            print("input cannot be less than 0.. please try again. ")
            return self.input_course_duration()
        elif user_input > 0:
            return user_input
        else:
            print("Enter duration of course (in hours) : ")
            return self.input_course_duration()

    def input_course_price(self):
        user_input = int(input("Enter price of course (in Rs.) : "))
        if user_input <= 0:
            print("input cannot be less than 0.. please try again. ")
            return self.input_course_price()
        elif user_input > 0:
            return user_input
        else:
            print("Please enter price of course (in Rs.) : ")
            return self.input_course_price()

    def input_course_name(self):
        user_input = input("Enter name of course : ")
        user_input = user_input.strip()
        regex = "^[A-Za-z0-9- ]+$"
        if not re.fullmatch(regex, user_input) or user_input == '':
            print("Please enter valid characters.")
            user_input = self.input_course_name()

        return user_input

    def input_course_content(self):
        try:
            user_input = input("Enter content of course : ")
            if user_input.strip() == '':
                print("Input cannot be empty. Please try again.")
                user_input = self.input_course_content()
            return user_input
        except:
            print("Wrong input made... please try again. ")
            return self.input_course_content()

    def input_delete_course_name(self, user_id):
        content = course.list_course(2, user_id)
        if content is None:
            return
        course_to_delete = self.input_course_name()
        is_valid_course_name, course_id = self.check_valid_course(course_to_delete, content)
        if is_valid_course_name:
            course_to_delete = is_valid_course_name
            course.delete_course(course_to_delete)

    def list_purchasable_course(self, user_id):
        content = course.list_course(4, user_id)
        if content is None:
            return
        purchase_course_name = self.input_purchase_course_name()
        is_valid_course_name, course_id = self.check_valid_course(purchase_course_name, content)
        if is_valid_course_name:
            course.purchase_course(user_id, course_id)

    def check_valid_course(self, course_name, content):
        is_valid_course = False
        course_id = None
        for row in content:
            if row[1].lower() == course_name.lower():
                is_valid_course = True
                course_name = row[1]
                course_id = row[0]

        if not is_valid_course:
            print("No such course exists")
            return False
        else:
            return [course_name, course_id]

    def view_course_feedback(self, user_id):
        content = course.list_course(2, user_id)
        if content is None:
            return
        course_name = get_string_input("Enter the name of course you wish to view feedback of : ")
        is_valid_course_name, course_id = self.check_valid_course(course_name, content)
        if is_valid_course_name:
            feedback.view_course_feedback(course_id)

    def add_course_feedback(self, user_id):
        content = course.view_purchased_course(user_id)
        if content is None:
            return
        course_name = get_string_input("Enter the name of course you wish to add feedback of : ")
        is_valid_course_name, course_id = self.check_valid_course(course_name, content)
        print("Inside add course feedback")
        #
        # print(course_id)
        print(get_query.get("CHECK_IF_FEEDBACK_PRESENT"))
        content = DatabaseConnection.get_from_db(get_query.get("CHECK_IF_FEEDBACK_PRESENT"),(course_id,user_id))

        # print(content)
        # if content is None:
        #     print("You've already added feedback for this course")
        #     return
        if is_valid_course_name:
            self.input_course_feedback(course_id, user_id)

    def input_course_feedback(self, course_id, user_id):
        # content = DatabaseConnection.get_from_db(get_query.get("CHECK_IS_FEEDBACK_ALREADY_ADDED"), (course_id, user_id))

        print("**** Add feedback ****")
        rating = float(input("Enter rating out of 5 : "))
        while rating > 5 or rating <= 0:
            rating = get_float_input("Enter rating out of 5 : ")
            if rating > 0 and rating < 5:
                break
        comments = input("Enter any comment : ")
        if comments == "":
            comments = "No comments."
        feedback.add_course_feedback(course_id, rating, comments, user_id)

    def view_course_faq(self, user_id):
        content = course.list_course(4, user_id)
        if content is None:
            return
        course_name = get_string_input("Enter the name of course you wish to view FAQ of : ")
        is_valid_course_name, course_id = self.check_valid_course(course_name, content)
        if is_valid_course_name:
            faq.view_faq(course_name)

    def input_course_name_to_study_from(self, user_id):
        content = course.view_purchased_course(user_id)
        if content is None:
            return
        course_name = get_string_input("Enter the name of course you wish to study from : ")
        is_valid_course_name, course_id = self.check_valid_course(course_name, content)
        if is_valid_course_name:
            course_name = is_valid_course_name
            course.view_course_content(course_name)

    def check_pending_courses(self):
        query = get_query["PENDING_STATUS"]
        result = DatabaseConnection.get_from_db(query=query, val=("pending",))
        pending_course_count = 0
        if result:
            pending_course_count = result[0][0]

        return pending_course_count

    def list_pending_course(self, pending_course_count):
        print("**************************")
        print("Pending Notification : ")
        query = get_query.get("PENDING_STATUS")
        result = DatabaseConnection.get_from_db(query, ("pending",))

        print("Course details : \n")
        course_name = result[0][1]
        headers = ["Name", "Duration (in months)", "Price"]
        included_columns = [1, 3, 4]
        list_course_in_tabular_form(query, headers, "grid", included_columns, ("pending",))
        course.approve_course(course_name, pending_course_count)

    def input_course_faq(self, user_id):
        content = DatabaseConnection.get_from_db(get_query.get("GET_FAQ_DETAILS"), (user_id,))

        if len(content) == 0 or content is None:
            print("No course exists")
            return
        query = get_query.get("GET_FAQ_DETAILS")
        headers = ["Name", "Duration (in hrs )", "Price (in Rs.)", "Rating"]
        included_columns = [4, 6, 7, 8]
        list_course_in_tabular_form(query, headers, "grid", included_columns, user_id)
        course_name = get_string_input("Enter the name of the course in which you want to add FAQ : ")
        faq.add_faq(content, course_name)

    def input_purchase_course_name(self):
        try:
            user_input = input("Enter the name of course you wish to purchase : ")
            if user_input.strip() == '':
                print("Input cannot be empty. Please try again.")
                user_input = self.input_purchase_course_name()
            return user_input
        except:
            print("Wrong input made... please try again. ")
            return self.input_purchase_course_name()