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


# 1. 简介
"""
 Python conversion of the IDL uly_spect_alloc.pro .
目的：
    初始化 LAMOST 光谱数据结构. 该函数在 Python 语言中不重要, 但为了与 LASP-MPFit 一致, 这里保留.
函数：
    1) uly_spect_alloc
解释：
    1) uly_spect_alloc 函数: 初始化 LAMOST 光谱的字典结构.
"""


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


# 3. 更新 LAMOST 光谱字典结构中的数据
def uly_spect_alloc(TITLE=None, DATA=None, START=None, STEP=None, SAMPLING=None, ERR=None,
                    WAVELEN=None, GOODPIX=None, HEADER=None, DOF_FACTOR=None, SPECTRUM=None) -> dict:

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

       输入参数:
       -----------
       TITLE:
             LAMOST 内部数据文件名称.
       DATA:
            指定波长范围内的 LAMOST 流量.
       START:
             ln 对数下的空气波长, 如 start=8.2158041.
       STEP:
            ln 对数下的空气波长步长, 如 step=0.00023025851.
       SAMPLING:
                波长采样方法, ln 对数波长, 即 sampling=1.
       ERR:
           空值, LASP 没有使用(但需注意 LAMOST 提供了流量误差).
       WAVELEN:
               波长数组, 当采样=2 时使用, LASP 没有设置.
       GOODPIX:
               根据无穷值、Nan 值设置的待 MASK 的像素点索引值.
       HEADER:
              LAMOST 内部数据头文件.
       DOF_FACTOR:
                  自由度因子. 实际像素数与独立测量数(即探测器上的像素数)之比.
                  当光谱重采样到更小的像素时, 此参数会增加. LASP 没有开启这个选项, 即设置为 1.
       SPECTRUM:
                光谱结构, 用此输入变量的内容初始化新结构.

       输出参数:
       -----------
       spectrum:
                初始化的 LAMOST 光谱结构.
    """

    # 3.1 初始化 LAMOST 光谱字典结构
    spectrum = {"title": None,        # LAMOST 内部数据文件名称
                "hdr": None,          # LAMOST 内部数据头文件
                "data": None,         # 指定波长范围内的 LAMOST 流量
                "err": None,          # 空值, LASP 没有使用 (但需注意 LAMOST 提供了流量误差)
                "wavelen": None,      # 波长数组, 当采样=2 时使用, LASP 没有设置
                "goodpix": None,      # 根据无穷值、Nan 值设置的 MASK 的像素点索引值
                "start": 1.0,         # ln 对数下的空气波长, 会根据采样方式更新
                "step": 1.0,          # ln 对数下的空气波长步长, 会根据采样方式更新
                "sampling": 1,        # 波长采样方法, ln 对数波长, 即 sampling=1, 会根据采样方式更新
                "dof_factor": 1       # 自由度因子. 实际像素数与独立测量数之比
                }

    # 3.2 如果提供了 SPECTRUM 信息, 即 LAMOST 光谱结构, 更新
    if SPECTRUM is not None:
        if SPECTRUM["hdr"] is not None:
            spectrum["hdr"] = SPECTRUM["hdr"]
        if SPECTRUM["data"] is not None:
            spectrum["data"] = SPECTRUM["data"]
        if SPECTRUM["err"] is not None:
            spectrum["err"] = SPECTRUM["err"]

        spectrum["title"] = SPECTRUM["title"]
        spectrum["start"] = SPECTRUM["start"]
        spectrum["step"] = SPECTRUM["step"]
        spectrum["sampling"] = SPECTRUM["sampling"]
        spectrum["dof_factor"] = SPECTRUM["dof_factor"]

        if 'wavelen' in SPECTRUM:
            if SPECTRUM["wavelen"] is not None:
                spectrum["wavelen"] = SPECTRUM["wavelen"]
        if 'goodpix' in SPECTRUM:
            if SPECTRUM["goodpix"] is not None:
                spectrum["goodpix"] = SPECTRUM["goodpix"]
        if 'mask' in SPECTRUM:
            if SPECTRUM["mask"] is not None:
                spectrum["mask"] = np.where(SPECTRUM["mask"] == 1)[0].tolist()

    # 3.3 如果提供了 TITLE 信息, 则更新
    if TITLE is not None:
        spectrum["title"] = TITLE

    # 3.4 如果提供了 HEADER 信息, 则更新
    if HEADER is not None:
        spectrum["header"] = HEADER
    
    # 3.5 如果提供了 DATA 信息, 则更新
    if DATA is not None:
        spectrum["data"] = DATA
    
    # 3.6 如果提供了 ERR 信息, 则更新
    if ERR is not None:
        spectrum["err"] = ERR
    
    # 3.7 如果提供了 WAVELEN 信息, 则更新
    if WAVELEN is not None:
        if SAMPLING is not None:
            if SAMPLING != 2:
                raise ValueError("inconsistency, if WAVELEN is passed, SAMPLING must be 2 or ommitted!")
            SAMPLING = 2
            spectrum["wavelen"] = WAVELEN
    
    # 3.8 如果提供了 GOODPIX 信息, 则更新
    if GOODPIX is not None:
        spectrum["goodpix"] = GOODPIX
    
    # 3.9 如果提供了 START 信息, 则更新
    if START is not None:
        spectrum["start"] = START
    
    # 3.10 如果提供了 STEP 信息, 则更新
    if STEP is not None:
        spectrum["step"] = STEP
    
    # 3.11 如果提供了 SAMPLING 信息, 则更新
    if SAMPLING is not None:
        spectrum["sampling"] = SAMPLING
    
    # 3.12 如果提供了 DOF_FACTOR 信息, 则更新
    if DOF_FACTOR is not None:
        spectrum["dof_factor"] = DOF_FACTOR

    # 3.13 返回初始化、更新后的 LAMOST 光谱字典结构
    return spectrum