# -*- coding: utf-8 -*-
# @Time    : 10/12/2024 14.52
# @Author  : ljc
# @FileName: uly_tgm.py
# @Software: PyCharm


# 1. 简介
"""
 Python conversion of the IDL uly_tgm.pro .
目的：
    定义一个 TGM (Teff - log g - [Fe/H]) 组件.
函数：
    1) uly_tgm
解释：
    1) uly_tgm 函数: 定义 TGM (Teff - log g - [Fe/H]) 组件, 包含这 3 个恒星参数初始值、恒星参数区间等初始信息.
"""


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


# 3. 定义 uly_tgm 函数
def uly_tgm(model_file=None, t_limits=(3100.0, 40000.0), l_limits=(-0.25, 5.9), z_limits=(-2.5, 1.0),
            t_guess=None, l_guess=None, z_guess=None, fixpar=None, lim_weight=(0, np.finfo(float).max)) -> dict:

    """
      定义一个 TGM (Teff - log g - [Fe/H]) 组件的函数.
      1) 单组件 cmp: 恒星参数初始猜测仅有 1 组.
      2) 多组件 cmp: 恒星参数初始猜测有多组(目前支持 1 组).
      3) ULySS 多组件情况时, 最终推断的恒星参数不是选择各组件卡方最小时的参数, 而是基于多组参数加权得到的.
      因此, 使用多组件时, 需要设置权重 weight.
      4) LASP 根据 CFI 初始化恒星参数, 即初始猜测大多使用 1 组, 因此该代码适用于单组件情况! 如需要多组件设置, 请参考 IDL 原始代码.

      输入参数：
      -----------
      model_file:
                 指定 TGM 模型的 FITS 文件.
      t_limits:
               Teff 的推断范围.
      l_limits:
                log g 的推断范围.
      z_limits:
               [Fe/H] 的推断范围.
      t_guess:
              Teff 的初始猜测.
      l_guess:
              log g 的初始猜测.
      z_guess:
              [Fe/H] 的初始猜测.
      fixpar:
             参数固定标志数组. 如果设置固定, 则相对应参数不参与推断.
      lim_weight:
                 组件权重范围.

      输出参数:
      -----------
      cmp:
          定义的 cmp 字典.
    """

    # 3.1 使用默认模型文件路径
    model_file = os.path.join(model_file)

    # 3.2 参数初始化
    # 对 Teff 取对数, 保持与 IDL 代码一致
    if t_guess is None:
        t_guess = np.log(7000.0)
    else:
        t_guess = np.log(t_guess)
    if l_guess is None:
        l_guess = 2.0
    else:
        l_guess = float(l_guess)
    if z_guess is None:
        z_guess = 0.0
    else:
        z_guess = float(z_guess)

    # 3.3 设置 TGM 组件的主要结构字段
    # 1) 单组件: 恒星参数初始猜测仅有 1 组.
    # 2) 多组件: 恒星参数初始猜测有多组.
    # 3) ULySS 多组件情况时, 最终推断的恒星参数不是选择各组件卡方最小时的参数, 而是基于多组参数加权得到的.
    # 因此, 使用多组件时, 需要设置权重 weight.
    # 4) LASP 根据 CFI 初始化恒星参数, 即初始猜测大多使用 1 组, 因此该代码适用于单组件情况! 如需要多组件设置, 请参考 IDL 原始代码.
    cmp = {
           'init_fun': 'uly_tgm_init',               # 初始化函数名称.
           "init_data": {"model": model_file,        # 传递给初始化函数的数据, 包含模型文件、LSF 文件及重采样选项.
                             "lsf_file": "no_lsf",   # 降低分辨率的文件, LASP 没有使用线宽函数.
                             "rebin_coef": 0},       # 重采样选项.
           "eval_fun": "",                           # 计算 TGM 光谱的评估函数, 调用 uly_tgm_eval 函数.
           "eval_data": "",                          # 传递给评估函数的数据, 存储 TGM 模型信息, 如系数、波长等.
           'para': "",                               # 待测恒星参数.
           'start': 0.0,                             # 波长起始值.
           'step': 0.0,                              # 波长步长.
           'npix': 0,                                # 流量维度.
           'sampling': -1,                           # 波长采样模式 (0: 线性, 1: 对数, 2: 非均匀波长).
           'mask': None,                             # 像素掩码, 标记好坏像素点.

                                                     # LASP 使用的是单组件, 不是多组件, 即权重设置没有发挥意义!!!!!!
           'weight': 0.0,                            # 组件权重.
           'e_weight': 0.0,                          # 权重误差.
           'l_weight': 0.0,                          # 流量比重.
           'lim_weig':  np.finfo(np.float64).max * 
           np.array([0, 1]),                         # 权重上下限.
          }

    # 3.4 定义 TGM 参数(Teff、log g、[Fe/H])的结构
    params = [
        {'name': 'Teff', 'unit': 'K', 'guess': t_guess, 'step': 0.005, 'limits': np.log(t_limits), 'limited': [1, 1],
         'fixed': 0, 'value': 0.0, 'error': 0.0, 'dispf': 'exp'},    # Teff 参数
      
        {'name': 'Logg', 'unit': 'cm/s2', 'guess': l_guess, 'step': 0.01, 'limits': l_limits, 'limited': [1, 1],
         'fixed': 0, 'value': 0.0, 'error': 0.0},                    # log g 参数
      
        {'name': '[Fe/H]', 'unit': 'dex', 'guess': z_guess, 'step': 0.01, 'limits': z_limits, 'limited': [1, 1],
         'fixed': 0, 'value': 0.0, 'error': 0.0}                     # [Fe/H] 参数
        ]

    # 3.5 如果提供了固定参数数组 fixpar, 将其应用到参数结构中
    if fixpar is not None:
        for i, fixed in enumerate(fixpar):
            params[i]['fixed'] = fixed
    
    # 3.6 设置组件权重范围.
    if lim_weight is not None:
        cmp["lim_weig"] = lim_weight

    # 3.7 将参数列表分配给组件结构中的 para 字段
    cmp['para'] = params

    # 3.8 返回组件结构
    return cmp