# check_initial_state.py
"""
检查初始状态是否为“必死状态”（conservatively dead）：
1) 实例化 base_params（优先使用 rl_dqn_env 中提供的 base_params 或 UFParams）
2) 实例化环境类 UFSuperCycleEnv(base_params)
3) 调用 env.generate_initial_state() 生成 env.current_params（不调用 reset()）
4) 用最保守策略 (L_s=3600s, t_bw_s=60s) 连续模拟 max_steps 次，
   若任意一次 is_dead_cycle(info) 返回 False 则判定为必死（返回 True），否则返回 False。
"""

from typing import Any
import copy
import traceback

# 从 rl_dqn_env 导入必需项
try:
    from DQN_env import (
        simulate_one_supercycle,
        is_dead_cycle,
        UFSuperCycleEnv,
        UFParams,       # 如果模块里有 UFParams 类就导入
        base_params     # 如果模块直接提供 base_params 实例也尝试导入
    )
except Exception:
    # 有可能某些名字不存在 —— 我们会稍后用回退方案处理
    # 先导入模块并再尝试访问属性，确保错误信息更友好
    import importlib
    rl = importlib.import_module("rl_dqn_env")
    simulate_one_supercycle = getattr(rl, "simulate_one_supercycle", None)
    is_dead_cycle = getattr(rl, "is_dead_cycle", None)
    UFSuperCycleEnv = getattr(rl, "UFSuperCycleEnv", None)
    UFParams = getattr(rl, "UFParams", None)
    base_params = getattr(rl, "base_params", None)

# 检查导入完整性
_missing = []
if simulate_one_supercycle is None:
    _missing.append("simulate_one_supercycle")
if is_dead_cycle is None:
    _missing.append("is_dead_cycle")
if UFSuperCycleEnv is None:
    _missing.append("UFSuperCycleEnv")
if _missing:
    raise ImportError(f"无法从 rl_dqn_env 导入以下必要项: {', '.join(_missing)}")

def is_dead_initial_state_env(env: UFSuperCycleEnv, max_steps: int = 15,
                              L_s: int = 4200, t_bw_s: int = 50,
                              verbose: bool = True) -> bool:
    """
    使用 env.current_params 作为初始状态判断是否为必死状态（保守策略）。

    参数:
        env: 已实例化的 UFSuperCycleEnv（必须包含 generate_initial_state() 与 current_params）
        max_steps: 模拟步数（默认 15）
        L_s: 过滤时长（s），保守值 3600
        t_bw_s: 物理反洗时长（s），保守值 60
        verbose: 是否打印每步结果

    返回:
        True 表示必死（conservatively dead）
        False 表示可行
    """
    # 1) 确保 env 有 current_params，并且 generate_initial_state 可用
    if not hasattr(env, "generate_initial_state"):
        raise AttributeError("env 缺少 generate_initial_state() 方法。")
    # 生成初始状态（不会调用 reset）
    env.generate_initial_state()

    if not hasattr(env, "current_params"):
        raise AttributeError("env.generate_initial_state() 未设置 env.current_params。")

    curr_p = copy.deepcopy(env.current_params)

    for step in range(1, max_steps + 1):
        try:
            info, next_params = simulate_one_supercycle(curr_p, L_s, t_bw_s)
        except Exception as e:
            # 如果 simulate 出错，把异常视为“失败”（保守处理）
            if verbose:
                print(f"[Step {step}] simulate_one_supercycle 抛出异常，视为失败。异常信息：{e}")
                traceback.print_exc()
            return True

        success = is_dead_cycle(info)  # True 表示成功循环

        if verbose:
            print(f"[Step {step}] 循环结果：{'成功' if success else '失败'}")
            # 如果 info 中有关键诊断字段，打印简要信息
            try:
                print(f"     TMP0: {info.get('TMP0')},max_TMP: {info.get('max_TMP_during_filtration')}, recovery: {info.get('recovery')}, "
                      f"R0: {info.get('R0')}, R_after_ceb: {info.get('R_after_ceb')}")
            except Exception:
                pass

        if not success:
            if verbose:
                print(f"在第 {step} 步检测到失败，判定为必死初始状态（conservatively dead）。")
            return True

        # 否则继续，用 next_params 作为下一步起始参数
        curr_p = next_params

    if verbose:
        print(f"{max_steps} 步均成功，初始状态判定为可行（non-dead）。")
    return False


if __name__ == "__main__":
    print("=== check_initial_state.py: 使用 env.generate_initial_state() 检查初始状态是否为必死 ===")

    try:
        # 1) 构造 base_params
        if base_params is not None:
            bp = base_params
            print("使用 rl_dqn_env 中提供的 base_params。")
        elif UFParams is not None:
            bp = UFParams()  # 使用默认构造
            print("使用 UFParams() 构造 base_params 的实例。")
        else:
            raise ImportError("无法构造 base_params：rl_dqn_env 中既无 base_params 也无 UFParams。")

        # 2) 实例化环境类（将 base_params 传入构造器）
        env = UFSuperCycleEnv(bp)
        print("已实例化 UFSuperCycleEnv 环境。")

        # 3) 调用 env.generate_initial_state() 并检查 env.current_params 是否为必死
        dead = is_dead_initial_state_env(env, max_steps=getattr(env, "max_episode_steps", 15),
                                        L_s=6000, t_bw_s=40, verbose=True)

        print("\n=== 判定结果 ===")
        if dead:
            print("当前生成的初始状态为【必死状态】（conservatively dead）。")
        else:
            print("当前生成的初始状态为【可行状态】（non-dead）。")

    except Exception as e:
        print("脚本执行出现错误：", e)
        traceback.print_exc()