# -*- coding: utf-8 -*-
# @Time    : 07/12/2024 23.21
# @Author  : ljc
# @FileName: uly_spect_get.py
# @Software: PyCharm


# 1. 简介
"""
 Python conversion of the IDL uly_spect_get.pro .
目的：
    从光谱结构中检索信息.
函数：
    1) uly_spect_get
解释：
    1) uly_spect_get 函数: 返回从光谱字典结构中检索到的信息.
"""


# 2. 导库
import numpy as np
import warnings
warnings.filterwarnings("ignore")


# 3. 更新光谱字典结构
def uly_spect_get(SignalIn=None, WAVERANGE=None, GOODPIX=None, HDR=None, MASK=None) -> tuple[np.ndarray, np.ndarray, str, np.ndarray]:

    """
       使用 uly_spect_read_lss 返回的 LAMOST 结构, 做数据处理, 并返回更新后的 LAMOST 数据结构.

       输入参数:
       -----------
       SignalIn:
                输入 LAMOST 字典结构.
       WAVERANGE:
                 给出波长范围, 返回包含 2 个元素的波长数组.
       GOODPIX:
               从 LAMOST 字典结构中返回好像素列表. 如果未设置, 则假定所有像素都是好的.
       HDR:
           返回 LAMOST 字典结构的 fits 头文件. 如果不存在则返回空字符串.
       MASK:
            返回包含像素掩码的数组 (1 表示好像素, 0 表示坏像素), 掩码长度与光谱向量相同.
            如果输入光谱结构不包含好像素列表, 则假定所有像素都是好的.

       输出参数:
       -----------
       WAVERANGE, GOODPIX, HDR, MASK:
                                     依据调用 uly_spect_get 函数时设置的变量, 使用相对应的返回值.
    """

    # 3.1 如果指定使用波长范围, 则计算光谱结构中的波长信息, 并返回波长范围
    if WAVERANGE is not None:
        # 3.1.1 线性采样
        if SignalIn["sampling"] == 0:
            WAVERANGE = [SignalIn["start"], SignalIn["start"] + (SignalIn["data"].shape[0] - 1) * SignalIn["step"]]
        # 3.1.2 ln 对数采样
        if SignalIn["sampling"] == 1:
            WAVERANGE = np.exp([SignalIn["start"], SignalIn["start"] + (SignalIn["data"].shape[0] - 1) * SignalIn["step"]])
        # 3.1.3 非均匀采样
        if SignalIn["sampling"] == 2:
            WAVERANGE = [SignalIn["wavelen"][0], SignalIn["wavelen"][-1]]

    # 3.2 如果指定使用好坏像素列表, 则设置后并返回
    if GOODPIX is not None:
        # 3.2.1 如果已设置好像素列表, 则直接返回
        if len(GOODPIX) != 0:
            GOODPIX = SignalIn["goodpix"]
        # 3.2.2 如果未设置好像素列表, 则假定所有像素都是好的
        else:
            npix = SignalIn["data"].shape[0]
            GOODPIX = np.arange(npix)

    # 3.3 如果指定使用 fits 头文件, 则设置后并返回
    # 注意: 光谱的 hdr 字段不必要非设置不可, 应根据需求决定是否设置
    if HDR is not None:
        # 3.3.1 如果已设置 fits 头文件, 则直接返回
        if SignalIn["hdr"] != "":
            HDR = SignalIn["hdr"]
        # 3.3.2 如果未设置 fits 头文件, 则返回空字符串
        else:
            HDR = ""

    # 3.4 如果指定使用 MASK 列表, 则设置后并返回
    if MASK is not None:
        MASK = np.zeros(SignalIn["data"].shape[0], dtype=np.uint8)
        # 3.4.1 如果已设置好像素列表, 则根据好像素列表索引设置 MASK
        if len(SignalIn["goodpix"]) > 0:
            MASK[SignalIn["goodpix"]] = 1
        # 3.4.2 如果未设置好像素列表, 则假定所有像素都是好的, 即 MASK 全为 1
        else:
            MASK = MASK + 1

    # 3.5 返回波长范围, 好坏像素列表, fits 头文件, MASK
    return WAVERANGE, GOODPIX, HDR, MASK