# -*- coding: utf-8 -*-
# @Time    : 03/12/2024 19.52
# @Author  : ljc
# @FileName: uly_spect_read_lms.py
# @Software: PyCharm


# 1. 简介
"""
 Python conversion of the IDL uly_spect_read.pro .
目的：
    读取 LAMOST 内部数据, 并建立一个光谱字典结构.
    1) 读取 LAMOST 内部数据.
    2) 将 log10 对数下的真空波长转为 ln 对数下的空气波长.
    3) 截取指定波长范围内的 LAMOST 数据.
    4) 设置好像素点.
    5) 返回一个数据结构:
                     SignalIn = {"title": title,               # LAMOST 内部数据文件名称.
                                 "hdr": h,                     # LAMOST 内部数据头文件.
                                 "data": spec,                 # 指定波长范围内的 LAMOST 流量.
                                 "err": err,                   # 流量误差, 空值, LASP 没有使用 (但需注意 LAMOST 提供了流量误差).
                                 "wavelen": wavelen,           # 波长数组, 当采样=2时使用, LASP 没有设置.
                                 "goodpix": goodpix,           # 根据无穷值、Nan 值设置的待 MASK 的像素点索引值.
                                 "start":start,                # ln 对数下的空气波长, 如 start=8.2158041.
                                 "step": step,                 # ln 对数下的空气波长步长, 如 step=0.00023025851.
                                 "sampling": sampling,         # 波长采样方法, ln 对数波长, 即 sampling=1.
                                 "dof_factor": dof_factor      # 自由度因子. 实际像素数与独立测量数之比.
                                                               # 当光谱重采样到更小的像素时, 此参数会增加. LASP-MPFit 没有开启这个选项, 即设置为 1.
                                 }
函数：
    1) uly_spect_read_lss
    2) uly_spect_read_lms
解释：
    1) uly_spect_read_lss 函数: 读取 LAMOST fits 文件, 并返回包含 LAMOST 光谱信息的字典结构.
    2) uly_spect_read_lms 函数: 使用 uly_spect_read_lss 返回的 LAMOST 结构, 做数据处理, 并返回更新后的 LAMOST 数据结构.
"""


# 2. 调库
import numpy as np
from astropy.io import fits
import warnings
warnings.filterwarnings("ignore")


# 3. 读取 LAMOST fits 文件, 并返回包含 LAMOST 光谱信息的字典结构
def uly_spect_read_lss(file_in, public, flux_median=False) -> dict:

    """
       读取 LAMOST fits 文件, 并返回包含 LAMOST 光谱信息的字典结构.

       输入参数:
       -----------
       file_in:
               LAMOST fits 文件路径.
       public:
               待测 LAMOST 光谱是内部数据、还是外部数据.
               如果是内部数据, 设置 public=0.
               如果是外部数据, 设置 public=1.
       flux_median:
                   是否除以光谱流量中值, 可以提升优化效率、可能会避免数值优化出错.
                   flux_median=False, 表示不除以流量中值.
                   flux_median=True, 表示除以流量中值.
                   默认不除以流量中值.

       输出参数:
       -----------
       SignalIn:
                待测光谱头文件、流量、转为 ln 对数下的波长起点、波长步长等信息组成的光谱字典结构.
    """

    # 3.1 读取 LAMOST fits 文件, 并获取流量、头文件信息
    f = fits.open(file_in)

    # 3.2 如果待测 LAMOST 光谱是内部数据, 则获取流量、头文件信息
    if public is False:
        # 3.2.1 如果需要除以流量中值, 则除以流量中值
        # 注意: 当中值不稳定时, 可以考虑其他分位数!!!!!!
        if flux_median is True:
            spec = f[0].data[0] / np.median(f[0].data[0])
        # 3.2.2 如果不需要除以流量中值, 则不除以流量中值
        if flux_median is False:
            spec = f[0].data[0]
        # 3.2.3 获取 LAMOST fits 文件的头文件
        h = f[0].header

    # 3.3 如果待测 LAMOST 光谱是外部数据, 则获取流量、头文件信息
    if public is True:
        # 3.3.1 如果需要除以流量中值, 则除以流量中值
        # 注意: 当中值不稳定时, 可以考虑其他分位数!!!!!!
        if flux_median is True:
            spec = f[1].data["FLUX"][0] / np.median(f[1].data["FLUX"][0])
            # spec = f[1].data["NORMALIZATION"][0]
        # 3.3.2 如果不需要除以流量中值, 则不除以流量中值
        if flux_median is False:
            spec = f[1].data["FLUX"][0]
        # 3.3.3 获取 LAMOST fits 文件的头文件
        h = f[0].header

    # 3.4 关闭 LAMOST fits 文件
    f.close()

    # 3.5 获取 LAMOST 光谱字典结构
    # LASP 使用的 uly_spect_read_lss 函数处理的光谱, 并返回的流量字典结构. 如需使用其他读取数据方法, 请参考原始 IDL 代码
    # 1) naxis=2 表示 LASP 使用 uly_spect_read_lss 处理 LAMOST 内部数据
    # 2) disp_type=1 表示使用 log10 对数波长
    # 3) disp_axis=1, 强制选择第 1 个轴为波长轴 (没意义, 为了与 IDL 一致, 这里也设置 disp_axis=1)
    # 4) vacuum=0, 表示 LAMOST 的波长是真空波长
    # 5) ax=1, 没有意义, 只是为了获取 LAMOST 头文件中的波长信息
    # 6) dof_factor=1, 无用, 只是一个系数, 设置为 1
    # if public is False:
    #     naxis = h['NAXIS']
    # if public is True:
    #     naxis = 2
    disp_type, disp_axis, vacuum, ax, dof_factor, sampling = 1, 1, 0, 1, 1, 1
    # 3.5.1 获取 LAMOST 的起始波长点 (log10 对数下)
    crval = np.double(h['CRVAL' + str(ax)])
    # 3.5.2 LAMOST fits 文件中, 对数波长下的波长增量 CD1_1, log10 对数下为 0.0001
    cdelt = np.double(h['CD' + str(ax) + '_' + str(ax)])
    # 3.5.3 LAMOST 的 cdelt 是第几个波长点, CRPIX1=1
    crpix = np.double(h['CRPIX' + str(ax)])
    # 3.5.4 LAMOST 的第 1 个波长点, crval1=3.5682
    crval = crval - (crpix - 1) * cdelt

    # 3.5.5 LAMOST 的 log10 对数波长转为 ln 对数波长, 并将真空波长转为空气波长
    if (disp_type == 1) and (crval <= 5):
        # 3.5.5.1 使用 ln 对数波长起始点, crval=8.2160841
        crval *= np.log(10)
        # 3.5.5.2 使用 ln 对数波长增量, cdelt=0.00023025851
        cdelt *= np.log(10)
    # 3.5.6 把 LAMOST 的波长起点使用空气波长表示, 即 crval=8.2158041
    crval -= 0.00028

    # 3.5.7 sampling=1 表示采样 ln 对数下的波长
    sampling = disp_type
    # 3.5.8 ln 对数下的波长起点
    start = crval
    # 3.5.9 ln 对数下的波长步长
    step = cdelt

    # 3.5.10 设置 LAMOST fits 文件路径、流量误差、等信息, 并以字典形式存储
    title, err, wavelen, goodpix = file_in, [], [], []
    SignalIn = {"title": title,              # LAMOST 内部数据文件名称.
                 "hdr": h,                   # LAMOST 内部数据头文件.
                 "data": spec,               # 指定波长范围内的 LAMOST 流量.
                 "err": err,                 # 流量误差, 空值, LASP 没有使用 (但需注意 LAMOST 提供了流量误差).
                 "wavelen": wavelen,         # 波长数组, 当采样=2时使用, LASP 没有设置.
                 "goodpix": goodpix,         # 根据无穷值、Nan 值设置的待 MASK 的像素点索引值.
                 "start": start,             # ln 对数下的空气波长, 如 start=8.2158041.
                 "step": step,               # ln 对数下的空气波长步长, 如 step=0.00023025851.
                 "sampling": sampling,       # 波长采样方法, ln 对数波长, 即 sampling=1.
                 "dof_factor": dof_factor    # 自由度因子. 实际像素数与独立测量数之比.
                                             # 当光谱重采样到更小的像素时, 此参数会增加. LASP-MPFit 没有开启这个选项, 即设置为 1.
                }

    # 3.6 返回 LAMOST 光谱字典结构
    return SignalIn


# 4. 使用 uly_spect_read_lss 返回的 LAMOST 结构, 做数据处理, 并返回更新后的 LAMOST 数据结构
def uly_spect_read_lms(lmin, lmax, file_in, public, flux_median, overwrite=True) -> dict:

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

       输入参数:
       -----------
       lmin:
            输入待拟合的光谱波长范围最小值, 目前支持单个值 (如 lmin=[4200]), 不支持分段列表 (分段列表实现不难, 但需要修改代码或修改 MASK 等).
       lmax:
            输入待拟合的光谱波长范围最大值, 目前支持单个值 (如 lmax=[5700]), 不支持分段列表 (分段列表实现不难, 但需要修改代码或修改 MASK 等).
       file_in:
               LAMOST fits 文件路径.
       public:
              待测 LAMOST 光谱是内部数据、还是外部数据.
              如果是内部数据, 设置 public=0.
              如果是外部数据, 设置 public=1.
       flux_median:
                   是否除以光谱流量中值, 可以提升优化效率、可能会避免数值优化出错.
                   flux_median=False, 表示不除以流量中值.
                   flux_median=True, 表示除以流量中值.
                   默认不除以流量中值.
       overwrite:
                 表示该函数将在输入的 LAMOST 光谱结构基础上, 更新相关信息.

       输出参数:
       -----------
       SignalOut:
                 更新后的头文件、流量、转为 ln 对数下的波长起点、波长步长等信息组成的光谱结构.
    """

    # 4.1 使用 uly_spect_read_lss 函数获取并定义 LAMOST 光谱字典结构
    SignalIn = uly_spect_read_lss(file_in=file_in,         # LAMOST fits 文件路径
                                  public=public,           # 待测 LAMOST 光谱是内部数据、还是外部数据
                                  flux_median=flux_median  # 是否除以光谱流量中值
                                  )
    
    # 4.2 启动数据更新
    if overwrite == True:
        SignalOut = SignalIn

    # 4.3 获取 wr, 并指定 ln 对数波长
    if SignalIn["sampling"] == 1:
        # 4.3.1 LAMOST 的流量维度
        ntot = len(SignalIn["data"])
        # 4.3.2 LASP 读取的 LAMOST 内部数据 sampling=1
        # 注意: 这里将 ln 对数波长转为线性波长
        wr = np.exp(SignalIn["start"] + [0, (ntot - 1) * SignalIn["step"]])

    # 4.4 LASP 设置拟合的波长范围
    wr[0], wr[1] = min(lmin), max(lmax)
   
    # 4.5 如果 wr 不为空
    if len(wr) != 0:
        # 4.5.1 如果为 ln 对数波长
        if SignalOut["sampling"] == 1:
            # 4.5.1.1 将 wr 转为 ln 对数下的波长
            wr_ = np.log(wr)
            # 4.5.1.2  TGM 光谱左侧 MASK 多少点, 根据 ln 对数波长计算
            # 注意: 
            # 1) step 为 LAMOST 的 log10 对数波长间隔转为了 ln 对数波长间隔, 即 ln_step = ln(10) * log10_step
            # 2) 为了避免多删除流量点, 使用 np.floor 向下取整, 而不是 np.ceil 向上取整
            # 3) np.floor 向下取整, 为了避免计算误差, 形如 1.999999 被取为 1 的情况, 加上 0.01 作为补偿
            nummin = np.floor((wr_[0] - SignalOut["start"]) / SignalOut["step"] + 0.01)
            # 4.5.1.3 波长点数量
            npix = len(SignalOut["data"])
            # wr_ 只设置了一个波长区间, 即两个值
            if len(wr_) == 2:
                # 4.5.1.4 计算 TGM 光谱右侧 MASK 多少点
                # 注意: 
                # 1) 为了避免多删除流量点, 使用 np.ceil 向上取整, 而不是 np.floor 向下取整
                # 2) np.ceil 向上取整, 为了避免计算误差, 形如 2.0000001 被取为 3 的情况, 减去 0.01 作为补偿
                nummax = np.ceil((wr_[1] - SignalOut["start"]) / SignalOut["step"] - 0.01)
                # 4.5.1.5 右侧 MASK 点超过 LAMOST 波长右边界
                if nummax >= npix:
                    # 最大不超过右边界
                    nummax = npix-1
            else:
                # 4.5.1.6 如果 wr_ 不是两个值
                nummax = npix-1

            # 4.5.1.7 分析完输入的波长范围与 LAMOST 内部数据的波长范围后, 确定最终的波长起点, 并更新 start
            SignalOut["start"] += nummin * SignalOut["step"]
            
            # 4.5.1.8 计算 LAMOST 内部 fits 文件的波长点数量
            # s = len(SignalOut["data"])
            # lmin=[4200], lmax=[5700] 之间的波长点数量
            # s = nummax - nummin + 1

            # 4.5.1.9 分析完输入的波长范围与 LAMOST 内部数据的波长范围后, 确定最终的波长点对应的流量值
            # 请注意 IDL 索引与 Python 不一样, 这里需要 +1, 并更新 data
            SignalOut["data"] = SignalOut["data"][int(nummin): int(nummax) + 1]
    
    # 4.6 MASK 波长范围内的总元素数量
    ntot = len(SignalOut["data"])

    # 4.7 MASK 无穷大或 Nan 像素点
    def where_finite(spect_data):
        # 4.7.1 获取有效数据的索引
        good = np.where(np.isfinite(spect_data))[0].tolist()
        # 4.7.2 获取有效数据的数量
        cnt = len(good)
        # 4.7.3 获取无效数据的索引
        nans = np.where(~np.isfinite(spect_data))[0].tolist()
        # 4.7.4 获取无效数据的数量
        nnans = len(nans)
        return good, cnt, nans, nnans
    good, cnt, nans, nnans = where_finite(SignalOut["data"])
    
    # 4.8 如果存在无效数据
    if nnans > 0:
        # 4.8.1 如果没有设置好像素点, 则非异常值点即为好像素点
        if len(SignalOut["goodpix"]) == []:
            SignalOut["goodpix"] = good
        """
        # LASP 没有设置 goodpix 字段. 如有需要, 请参考原始 IDL 代码
        # else:
        #     maskI = np.zeros(ntot, dtype=np.uint8)
        #     maskI[spect["goodpix"]] = 1
        #     maskI[nans] = 0
        #     spect["goodpix"] = np.where(maskI=1)
        """

        # 4.8.2 获取异常值后面的数据
        next_ = nans + 1
        # 4.8.2.1 如果最后 1 个流量点为异常值
        if next_[len(next_) - 1] == ntot:
            # 4.8.2.1.1 next_ 最后 1 个索引减 1
            next_[len(next_) - 1] = nans[len(next_) - 1]
        # 4.8.2.1.2 使用异常值后面的数据填充异常值
        SignalOut["data"][nans] = SignalOut["data"][next_]

        # 4.8.2.2 再次检查是否还存在无穷或 NaN 值
        nans = np.where(~np.isfinite(SignalOut["data"]))[0]
        nnans = len(nans)
        if nnans > 0:
            # 4.8.2.2.1 使用异常值前面的数据填充异常值
            prev = nans - 1
            if prev[0] < 0:
                prev[0] = nans[1]
            # 4.8.2.2.2 使用异常值前面的数据填充异常值
            SignalOut["data"][nans] = SignalOut["data"][prev]

        # 4.8.2.3 再次检查是否还存在无穷或 NaN 值
        nans = np.where(~np.isfinite(SignalOut["data"]))[0]
        nnans = len(nans)
        if nnans > 0:
            # 4.8.2.3.1 后、前填充都不行了, 就使用 0 填充
            SignalOut["data"][nans] = 0

    # 4.9 MASK 波长范围、MASK 异常值后的波长点总数量
    ntot = len(SignalOut["data"])
    # 4.9.1 ln 对数下的所有波长点
    Pix_gal = SignalIn["start"] + np.linspace(start=0, stop=ntot-1, num=ntot) * SignalIn["step"]
    # 4.9.2 波长起点与终点转为 ln 对数波长
    # 4.9.2.1 如果 lmin 为空, 则 lmn 为波长起点
    if len(lmin) == 0:
        lmn = [SignalOut["start"]]
    else:
        # 4.9.2.2 如果 lmin 不为空, 则 lmn 为 lmin
        lmn = lmin
        # 4.9.2.2.1 如果 sampling=1, 则 lmn 为 ln 对数波长
        if SignalOut["sampling"] == 1:
            lmn = np.log(lmn)
    # 4.9.3 如果 lmax 为空, 则 lmx 为波长终点
    if len(lmax) == 0:
        lmx =[SignalOut["start"] + SignalOut["step"] * (ntot-1)]
    else:
        # 4.9.3.1 如果 lmax 不为空, 则 lmx 为 lmax
        lmx = lmax
        # 4.9.3.2 如果 sampling=1, 则 lmx 为 ln 对数波长
        if SignalOut["sampling"] == 1:
            lmx = np.log(lmx)

    # 4.10 满足指定波长范围内的索引
    # 第 1 个位置是 0, 即不是好像素点, 因为重采样插值可能在端点处不稳定
    good = [0]
    for i in range(len(lmn)):
        # 4.10.1 获取指定波长范围内的索引
        # 注意: 
        # 1) 第 1 个位置是 0, 即不是好像素点, 因为重采样插值可能在端点处不稳定
        # 2) 对于 LAMOST 光谱波长特点, 最后 1 个位置小于 lmx, 所以最后 1 个位置不是好像素点
        # 3) 就算最后 1 个位置小于 lmx, 考虑到重采样插值的稳定性, 我建议最后 1 个位置也要 MASK
        good = good + np.where((Pix_gal > lmn[i]) & (Pix_gal < lmx[i]))[0].tolist()
    # 4.10.2 判断好像素点与指定波长范围内的索引是否存在
    if (len(SignalOut["goodpix"]) == 0) & (len(good) > 1):
        # 4.10.2.1 创建一个全 0 的数组, 用于存储 MASK 信息
        maskI = np.zeros(ntot, dtype=np.uint8)
        # 4.10.2.2 将 good 索引对应的值设置为 1
        # 注意: 对于 LAMOST 光谱波长特点, 第 1 个与最后 1 个流量点被 MASK
        # 1) 第 1 个位置是 0, 所以从第 2 个位置开始设置
        # 2) 最后 1 个位置也是 0, 这个与 good 索引设置对应
        maskI[good[1:]] = 1
        # 4.10.2.3 获取 MASK 的索引
        SignalOut["goodpix"] = np.where(maskI == 1)[0].tolist()
    # 4.10.3 不满足 (len(SignalOut["goodpix"]) == 0) & (len(good) > 1), 即无穷或 Nan 存在, 此时需 goodpix 索引与 good 索引均是好点的保留
    else:
        # 4.10.3.1 创建一个全 0 的数组, 用于存储 MASK 信息
        maskI = np.zeros(ntot, dtype=np.uint8)
        # 4.10.3.2 将 goodpix 索引对应的值设置为 1
        maskI[SignalOut["goodpix"]] = 1
        # 4.10.3.3 将 good 索引对应的值设置为 1
        # 注意: 对于 LAMOST 光谱波长特点, 第 1 个与最后 1 个流量点被 MASK
        # 1) 第 1 个位置是 0, 所以从第 2 个位置开始设置
        # 2) 最后 1 个位置也是 0, 这个与 good 索引设置对应
        maskI[good[1:]] += 1
        # 4.10.3.4 获取 MASK 的索引
        SignalOut["goodpix"] = np.where(maskI == 2)[0].tolist()
    # 4.10.4 将波长转为线性波长
    lamrange = np.exp([SignalOut["start"] + np.arange(SignalOut["data"].size) * SignalOut["step"]])[0]
    
    # 4.11 返回更新后的 LAMOST 光谱字典结构
    return SignalOut