import numpy as np
import matplotlib.pyplot as plt
from engine_env.engine_env import AeroEngineGymEnv

# ==========================================
# 1. 增量式 PID 类 (保持不变)
# ==========================================
class IncrementalPIDController:
    def __init__(self, kp, ki, kd, output_min=-1.0, output_max=1.0):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        
        self.min_val = output_min
        self.max_val = output_max
        self.error_prev = 0.0
        self.error_prev2 = 0.0
        self.current_output = 0.0

    def set_current_output(self, value):
        self.current_output = np.clip(value, self.min_val, self.max_val)
        self.error_prev = 0.0
        self.error_prev2 = 0.0

    def update(self, error, dt):
        delta_p = self.kp * (error - self.error_prev)
        delta_i = self.ki * error * dt
        delta_d = self.kd * (error - 2*self.error_prev + self.error_prev2) / dt
        
        delta_u = delta_p + delta_i + delta_d
        self.current_output += delta_u
        self.current_output = np.clip(self.current_output, self.min_val, self.max_val)
        
        self.error_prev2 = self.error_prev
        self.error_prev = error
        return self.current_output

# ==========================================
# 2. 主测试逻辑
# ==========================================
def run_pid_test():
    env = AeroEngineGymEnv()
    
    # -----------------------------------------------------
    # 【关键修改】定义正确的 PID 配对
    # -----------------------------------------------------
    
    # Loop 1: 目标 NH -> 控制 Wf (燃油)
    # 逻辑: NH 低 -> 加油 (正反馈)
    pid_nh_wf = IncrementalPIDController(kp=10.0, ki=6.0, kd=0.1) 
    
    # Loop 2: 目标 NL -> 控制 A8 (喷管)
    pid_nl_a8 = IncrementalPIDController(kp=-4.0, ki=3.0, kd=0.05)
    
    total_steps = 3000  # 60秒
    obs, info = env.reset(options={'pla': 0.0})
    
    # 状态标志位
    pid_initialized = False
    
    history = {
        'time': [], 'PLA': [], 'NH': [], 'Target_NH': [], 'NL': [], 'Target_NL': [],
        'Wf_Action': [], 'A8_Action': [], 'Reward': []
    }
    
    print("开始修正后的 PID 仿真 (NH->Wf, NL->A8)...")
    
    for step in range(total_steps):
        t = step * env.dt
        
        # --- A. 任务剖面 ---
        if t < 2.0: current_pla = 0.0
        elif t < 20.0: current_pla = 15.0 
        else: current_pla = 60.0
            
        env.set_pla(current_pla)
        
        # --- B. 获取状态 ---
        real_nh, real_nl = obs[0], obs[1]
        target_nh, target_nl = obs[4], obs[5]
        
        # 计算误差
        error_nh = target_nh - real_nh
        error_nl = target_nl - real_nl
        
        action = np.zeros(5) 
        
        # =========================================================
        # C. 分阶段控制
        # =========================================================
        
        # --- 阶段 1: 起动 (Open Loop) ---
        if current_pla >= 15.0 and real_nh < 0.60:
            action[2] = 0.8  # 起动机
            action[3] = -0.6 # 点火油量
            action[4] = 1.0  # 喷管全开
            pid_initialized = False 
            
        # --- 阶段 2: 闭环控制 (Closed Loop) ---
        elif current_pla >= 15.0 and real_nh >= 0.60:
            
            # 无扰切换
            if not pid_initialized:
                print(f"[Switch] t={t:.2f}s, 切入闭环。PID1: NH->Fuel(-0.6), PID2: NL->A8(1.0)")
                # NH控制燃油，继承 -0.6
                pid_nh_wf.set_current_output(-0.6)
                # NL控制喷管，继承 1.0
                pid_nl_a8.set_current_output(1.0)
                pid_initialized = True
            
            action[2] = -1.0 
            
            # 【核心修正】 PID 计算
            # Loop 1: 用 NH 的误差算 Wf
            wf_out = pid_nh_wf.update(error_nh, env.dt)
            
            # Loop 2: 用 NL 的误差算 A8
            a8_out = pid_nl_a8.update(error_nl, env.dt)
            
            action[3] = wf_out
            action[4] = a8_out
            
        # --- 阶段 3: 停机 ---
        else:
            action = np.array([0,0,-1,-1,1])
            pid_initialized = False

        # --- D. 执行 ---
        obs, reward, terminated, truncated, info = env.step(action)
        
        # 记录
        history['time'].append(t)
        history['PLA'].append(current_pla)
        history['NH'].append(real_nh)
        history['Target_NH'].append(target_nh)
        history['NL'].append(real_nl)
        history['Target_NL'].append(target_nl)
        history['Wf_Action'].append(action[3])
        history['A8_Action'].append(action[4])
        history['Reward'].append(reward)
        
        if terminated:
            print(f"Terminated at {t:.2f}s")
            break
            
    env.close()
    plot_results(history)

def plot_results(history):
    t = history['time']
    plt.figure(figsize=(12, 8))
    
    # 1. Speed
    plt.subplot(2, 2, 1)
    plt.plot(t, history['Target_NH'], 'r--', alpha=0.6)
    plt.plot(t, history['NH'], 'r', label='NH (Controlled by Fuel)')
    plt.plot(t, history['Target_NL'], 'b--', alpha=0.6)
    plt.plot(t, history['NL'], 'b', label='NL (Controlled by A8)')
    plt.legend()
    plt.title('Rotor Speeds')
    plt.grid(True)
    
    # 2. Action
    plt.subplot(2, 2, 2)
    plt.plot(t, history['Wf_Action'], 'orange', label='Fuel (Wf)')
    plt.plot(t, history['A8_Action'], 'green', label='Nozzle (A8)')
    plt.legend()
    plt.title('Control Actions')
    plt.ylim(-1.1, 1.1)
    plt.grid(True)
    
    # 3. PLA
    plt.subplot(2, 2, 3)
    plt.plot(t, history['PLA'], 'k')
    plt.title('PLA')
    plt.grid(True)

    # 4. Reward
    plt.subplot(2, 2, 4)
    plt.plot(t, history['Reward'])
    plt.title('Reward')
    plt.grid(True)
    
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    run_pid_test()