import numpy as np
from scipy.interpolate import interp1d

class ControlSchedule:
    """
    航空发动机控制计划 (Control Schedule)
    输入: PLA (角度 0~110)
    输出: Target NH, Target NL (0.0~1.0), Limit T5
    """
    def __init__(self):
        
        # ==========================================
        # 1. 定义控制计划数据点
        # ==========================================
        
        # --- NH (高压转速) 计划 ---
        # 依据图片1：在 PLA=80 时达到 1.0
        # 0 -> 14.99: 0
        # 15.0      : 0.7
        # 80.0      : 1.0
        self._pla_nh = np.array([0.0, 14.99, 15.0, 80.0, 110.0])
        self._val_nh = np.array([0.0, 0.0,   0.70, 1.00, 1.00])
        
        # --- NL (低压转速) 计划 ---
        # 依据图片2：在 PLA=90 时达到 1.0
        # 0 -> 14.99: 0
        # 15.0      : 0.5
        # 90.0      : 1.0
        self._pla_nl = np.array([0.0, 14.99, 15.0, 90.0, 110.0])
        self._val_nl = np.array([0.0, 0.0,   0.50, 1.00, 1.00])
        
        # --- 限制值 (T5, P3) 计划 ---
        # 保持原有逻辑：在 PLA=100 时达到最大限制
        self._pla_lim = np.array([0.0, 14.99, 15.0, 100.0, 110.0])
        
        # Limit T5 (温度限制) [K]
        # 15度时1200K, 100度时1500K
        self._val_t5  = np.array([1200.0, 1200.0, 1200.0, 1500.0, 1500.0])

        # Limit P3 (压力限制) [kPa]
        # 15度时2000kPa, 100度时2500kPa
        self._val_p3  = np.array([800.0, 800.0, 2000.0, 2500.0, 2500.0])


        # ==========================================
        # 2. 构建插值函数
        # ==========================================
        self.f_nh = interp1d(self._pla_nh, self._val_nh, kind='linear', fill_value="extrapolate")
        self.f_nl = interp1d(self._pla_nl, self._val_nl, kind='linear', fill_value="extrapolate")
        self.f_t5 = interp1d(self._pla_lim, self._val_t5, kind='linear', fill_value="extrapolate")
        self.f_p3 = interp1d(self._pla_lim, self._val_p3, kind='linear', fill_value="extrapolate")

    def get_targets(self, pla):
        """
        输入: pla (角度, 0.0 ~ 110.0)
        """
        # 稍微做个限幅，防止超出定义域太多
        pla = np.clip(pla, 0.0, 110.0)
        
        t_nh = float(self.f_nh(pla))
        t_nl = float(self.f_nl(pla))
        l_t5 = float(self.f_t5(pla))
        l_p3 = float(self.f_p3(pla))
        
        return t_nh, t_nl, l_t5, l_p3

# ==========================================
# 自测代码
# ==========================================
if __name__ == "__main__":
    import matplotlib.pyplot as plt
    
    sch = ControlSchedule()
    
    # 测试关键点
    test_plas = [0, 10, 14.9, 15.0, 15.1, 57.5, 100, 105]
    print(f"{'PLA':<10} {'NH':<10} {'NL':<10}")
    print("-" * 30)
    for p in test_plas:
        nh, nl, _, _ = sch.get_targets(p)
        print(f"{p:<10} {nh:<10.4f} {nl:<10.4f}")

    # 画图确认阶跃形状
    x = np.linspace(0, 110, 500)
    y_nh = [sch.get_targets(i)[0] for i in x]
    
    plt.figure()
    plt.plot(x, y_nh, label='Target NH')
    plt.axvline(15, color='r', linestyle='--', alpha=0.5, label='Idle Point (15 deg)')
    plt.title("PLA to Engine Speed Schedule")
    plt.xlabel("PLA (Degree)")
    plt.ylabel("Normalized Speed")
    plt.legend()
    plt.grid(True)
    plt.show()