# Path: astroNN/lamost/lamost_shared.py
# def lamost_default_dr(dr=None):
#     """
#     Check if dr argument is provided, if none then use default
# 
#     :param dr: data release
#     :type dr: Union(int, NoneType)
#     :return: data release
#     :rtype: int
#     :History: 2018-May-13 - Written - Henry Leung (University of Toronto)
#     """
#     # enforce dr5 restriction
#     if dr is None:
#         dr = 5
#         print(f'dr is not provided, using default dr={dr}')
#     elif dr == 5:
#         pass
#     else:
#         raise ValueError('Only LAMOST DR5 is supported')
# 
#     return dr


import numpy as np
from astroNN.lamost.lamost_shared import lamost_default_dr
    # Partial Credit: https://github.com/chanconrad/slomp/blob/master/lamost.py
    w = np.exp(-0.5 * (wavelength[:, None] - wavelength[None, :]) ** 2 / L ** 2)
    denominator = np.dot(ivar, w.T)
    numerator = np.dot(flux * ivar, w.T)
    bad_pixel = denominator == 0
    smoothed = np.zeros(numerator.shape)
    smoothed[~bad_pixel] = numerator[~bad_pixel] / denominator[~bad_pixel]
    return smoothed


def pseudo_continuum(flux, ivar, wavelength=None, L=50, dr=None):
    """
    Pseudo-Continuum normalise a spectrum by dividing by a Gaussian-weighted smoothed spectrum.

    :param flux: The observed flux array.
    :type flux: ndarray
    :param ivar: The inverse variances of the fluxes.
    :type ivar: ndarray
    :param wavelength: An array of the wavelengths.
    :type wavelength: ndarray
    :param L: [optional] The width of the Gaussian in pixels.
    :type L: int
    :param dr: [optional] dara release
    :type dr: int
    :returns: Continuum normalized flux and flux uncerteinty
    :rtype: ndarray
    """

    # Partial Credit: https://github.com/chanconrad/slomp/blob/master/lamost.py
    if dr is None:

# Path: astroNN/lamost/lamost_shared.py
# def lamost_default_dr(dr=None):
#     """
#     Check if dr argument is provided, if none then use default
# 
#     :param dr: data release
#     :type dr: Union(int, NoneType)
#     :return: data release
#     :rtype: int
#     :History: 2018-May-13 - Written - Henry Leung (University of Toronto)
#     """
#     # enforce dr5 restriction
#     if dr is None:
#         dr = 5
#         print(f'dr is not provided, using default dr={dr}')
#     elif dr == 5:
#         pass
#     else:
#         raise ValueError('Only LAMOST DR5 is supported')
# 
#     return dr


# Path: astroNN/lamost/chips.py
import numpy as np
from astroNN.lamost.lamost_shared import lamost_default_dr

    # Partial Credit: https://github.com/chanconrad/slomp/blob/master/lamost.py
    w = np.exp(-0.5 * (wavelength[:, None] - wavelength[None, :]) ** 2 / L ** 2)
    denominator = np.dot(ivar, w.T)
    numerator = np.dot(flux * ivar, w.T)
    bad_pixel = denominator == 0
    smoothed = np.zeros(numerator.shape)
    smoothed[~bad_pixel] = numerator[~bad_pixel] / denominator[~bad_pixel]
    return smoothed


def pseudo_continuum(flux, ivar, wavelength=None, L=50, dr=None):
    """
    Pseudo-Continuum normalise a spectrum by dividing by a Gaussian-weighted smoothed spectrum.

    :param flux: The observed flux array.
    :type flux: ndarray
    :param ivar: The inverse variances of the fluxes.
    :type ivar: ndarray
    :param wavelength: An array of the wavelengths.
    :type wavelength: ndarray
    :param L: [optional] The width of the Gaussian in pixels.
    :type L: int
    :param dr: [optional] dara release
    :type dr: int
    :returns: Continuum normalized flux and flux uncerteinty
    :rtype: ndarray
    """

    # Partial Credit: https://github.com/chanconrad/slomp/blob/master/lamost.py
    if dr is None:

        dr = lamost_default_dr(dr)