# -*- coding: utf-8 -*-
# @Time    : 2024/12/31 11:30
# @Author  : ljc
# @FileName: test_LAMOST.py
# @Software: PyCharm


# 1. 简介
"""
 无用文件 .
目的:
    增加所需功能后, 通过测试 1 条光谱的大气参数与 RV, 调试代码.
函数:
    f
解释:
    f 函数:
          1) 返回该光谱所测参数
          2) 所测参数存为 csv 文件
          3) 可视化光谱残差
"""


# 2. 导库
import pandas as pd
from uly_read_lms.uly_spect_read_lms import uly_spect_read_lms
from uly_tgm_alone.ulyss import ulyss
from tqdm import *
from file_paths import TGM_MODEL_FILE, TEST_DATA_DIR
import math
import time
from filelock import FileLock


# 3. 获取待测光谱数据信息
LAMOST_data = pd.read_csv("D:/Spectral data/PhD/巡天/LAMOST/LAMOST DR10 low/LASP_分组优化/LAMOST_DR10_APOGEE_DR16_CFI.csv", dtype={"obsid_1": str})
LAMOST_data = LAMOST_data.query("obsid_1 == '378009124'")
LAMOST_data.index = [i for i in range(len(LAMOST_data))]
model_file = TGM_MODEL_FILE()
file_in = "D:/Spectral data/PhD/巡天/LAMOST/LAMOST DR10 low/LASP_分组优化/LAMOST_DR10_APOGEE_DR16/"

# 4. 推断参数存为 csv 文件
data = pd.DataFrame(columns=["obsid", "pred_Rv", "pred_Teff", "pred_logg", "pred_FeH",
                             "pred_Rv_err", "pred_Teff_err", "pred_logg_err", "pred_FeH_err",
                             "use_time"])
data.to_csv('test_data.csv')


# 5. 定义每条待测光谱推断参数时所调用的函数
time0 = time.time()
def f(i):
    try:
        # 5.1 获取待测 LAMOST 样本的光谱字典结构
        # 注意:
        # 1) 待拟合的光谱波长范围: lmin、lmax, 当前版本支持单一区间. 如果需要设置分段区间, 建议精心设置 goodpix_
        # 2) LAMOST 内部数据与官网发布的数据 fits 文件有差异, 若读取官网发布的数据 (设置 public=True), 若读取内部数据 (设置 public=False)
        # 3) 是否使用中值归一化光谱 (默认使用, 即 flux_median=False, 在后续代码中设置 specs = torch.tensor(spectra_data / np.median(spectra_data, axis=1).reshape(-1, 1)).to(device, dtype=dtype))
        inspectr = uly_spect_read_lms(lmin=[4200],                                                           # 待拟合的光谱波长起点
                                      lmax=[5700],                                                           # 待拟合的光谱波长终点
                                      file_in=file_in + str(LAMOST_data.loc[i]["obsid_1"]) + ".fits.gz",     # 待测光谱文件地址
                                      public=True,                                                           # LAMOST 内部还是官网发布的光谱
                                      flux_median=True                                                       # 是否使用中值归一化光谱
                                      )
        # 5.2 LAMOST 待测光谱字典结构以及待测光谱初始值等信息传入 ulyss 函数，获得最终推断参数.
        # 注意:
        # 1) inspectr 为 LAMOST 样本的光谱字典结构
        # 2) mdegree 为消除 TGM 模型光谱与 LAMOST 光谱形状差异的勒让德多项式阶数
        # 3) clean 是否迭代剔除 TGM 与 LAMOST 光谱流量差异较大的点, 默认不剔除 (clean=False)
        # 4) quiet 是否打印参数推断结果 (False: 不打印; True: 打印)
        # 5) full_output 是否显示 curve_fit 是否成功推断出了恒星大气物理参数, 默认不显示 (full_output=False)
        # 6) plot_fitting=True 可视化光谱拟合结果, 推荐检验单条样本时使用
        # 7) 默认使用 CFI 参数初始值作为迭代起始点, 也可自行设置初始值
        Rv, Teff, logg, FeH, Rv_err, Teff_err,\
        logg_err, FeH_err, used_time, loss = ulyss(inspectr=inspectr,                         # 待测光谱字典结构
                                                   mdegree=50,                                # 勒让德多项式阶数
                                                   allow_polynomial_reduction=False,          # 是否允许多项式阶数减少
                                                   sampling_function=None,                    # 插值方法
                                                   clean=False,                               # 是否迭代剔除 TGM 与 LAMOST 光谱流量差异较大的点
                                                   quiet=True,                                # 是否打印参数推断结果
                                                   IDL_fit_method=True,                       # 是否使用 IDL 方法
                                                   full_output=True,                          # 是否显示 curve_fit 是否成功推断出了恒星大气物理参数
                                                   plot_fitting=True,                         # 是否可视化光谱拟合结果
                                                   # t_guess=LAMOST_data.loc[i]["CFI_Teff"],  # Teff 使用 CFI 参数初始值
                                                   # l_guess=LAMOST_data.loc[i]["CFI_logg"],  # logg 使用 CFI 参数初始值
                                                   # z_guess=LAMOST_data.loc[i]["CFI_FeH"],   # [Fe/H] 使用 CFI 参数初始值
                                                   t_guess=5000,                              # Teff 使用固定点初始值
                                                   l_guess=3,                                 # logg 使用固定点初始值
                                                   z_guess=-0.5,                              # [Fe/H] 使用固定点初始值
                                                   model_file=model_file                      # 光谱模型文件
                                                   )
    except:
        Rv, Teff, logg, FeH, Rv_err, Teff_err, logg_err, FeH_err, used_time, loss = -9999,  -9999, -9999, -9999, -9999, -9999, -9999, -9999, -9999, -9999
    # 5.3 待测参数存入已保存的 csv 文件
    new_parameters = pd.DataFrame(data={"obsid": LAMOST_data.loc[i]["obsid_1"], "pred_Rv": Rv, "pred_Teff": Teff, "pred_logg": logg, "pred_FeH": FeH,
                                        "pred_Rv_err": Rv_err, "pred_Teff_err": Teff_err, "pred_logg_err": logg_err,
                                        "pred_FeH_err": FeH_err, "use_time": used_time},
                                  columns=["obsid", "pred_Rv", "pred_Teff", "pred_logg", "pred_FeH",
                                           "pred_Rv_err", "pred_Teff_err", "pred_logg_err", "pred_FeH_err",
                                           "use_time"],
                                  index=[i])
    with FileLock('test_data.csv.lock'):
        new_parameters.to_csv('test_data.csv', mode='a', header=0)


# 6. 测试 1 条
f(0)