#! /usr/bin/env python
# _*_ coding: utf-8 _*_
"""
author:     Chen XuHui
date:       2024/7/1
file:       load_ckpt.py
version:
doc:        尝试load ckpt.目前已知:

路径问题:
1. ray自动保存在 pretrained_path:
    必须将 pretrained_path 细化成如下形式, 才可以用 Policy.from_checkpoint() API成功加载.
    pretrained_path = os.path.join(
        work_dir, 'results', 'kz', 'train', 'train_agent',
        'PPO_2024-07-01_10-50-33', 'PPO_KZ_env_af688_00000_0_2024-07-01_10-50-33', 'checkpoint_000000'
    )  # 实际使用中难以给出后3项
    返回值为 policy: dict[str, "Policy"].
2. Policy.export_checkpoint(p_ckpt) 保存在 p_ckpt:
    可以用 Policy.from_checkpoint() API成功加载, 返回值为 policy: "Policy"
3. Algorithm.save(a_folder) 保存在 a_folder:
    可以用 Policy.from_checkpoint() API成功加载, 返回值为 policy: dict[str, "Policy"].

"""
import os

from ray.rllib.algorithms.algorithm import Algorithm
from ray.rllib.algorithms.ppo.ppo import PPOConfig
from ray.rllib.policy.policy import Policy

from corekit.utils import DESKTOP, WORK_DIR
from corekit.utils import get_debug_logger

_logger = get_debug_logger()

p_ckpt = os.path.join(DESKTOP, 'temp', 'p-ckpt')  # 用在 Policy 的API
p_model = os.path.join(DESKTOP, 'temp', f'p-model')
a_folder = os.path.join(DESKTOP, 'temp', 'a-folder')  # 用在 Algorithm 的API

pretrained_path = os.path.join(
    WORK_DIR, 'results', 'demo', 'train_results',
    'PPO_2024-07-01_10-50-33', 'PPO_KZ_env_af688_00000_0_2024-07-01_10-50-33', 'checkpoint_000000'
)


def try_load(ckpt_path: str, path_name: str):
    path = path_name
    _logger.debug(f'尝试 load {path}')

    try:
        _logger.debug('尝试 Policy.from_checkpoint API')
        _policy: dict[str, "Policy"] = Policy.from_checkpoint(ckpt_path)
        _logger.info(f'Policy.from_checkpoint 加载 {path} 成功')
    except Exception as _e:
        # _logger.error(_e)
        _logger.warning(f'Policy.from_checkpoint 加载 {path} 失败')
        _policy = {}

    try:
        _logger.debug('尝试 Algorithm.from_checkpoint API')
        _algorithm0: Algorithm = Algorithm.from_checkpoint(ckpt_path)
        _logger.info(f'Algorithm.from_checkpoint 加载 {path} 成功')
    except Exception as _e:
        # _logger.error(_e)
        _logger.warning(f'Algorithm.from_checkpoint 加载 {path} 失败')
        _algorithm0 = None

    try:
        _logger.debug('尝试 Algorithm.load_checkpoint API')
        alg = PPOConfig.from_dict({}).build()
        _r1 = alg.load_checkpoint(ckpt_path)
        _logger.info(f'Algorithm.load_checkpoint 加载 {path} 成功')
    except Exception as _e:
        # _logger.error(_e)
        _logger.warning(f'Algorithm.load_checkpoint 加载 {path} 失败')
        _r1 = None

    try:
        _logger.debug('尝试 Algorithm.restore API')
        alg = PPOConfig.from_dict({}).build()
        _r2 = alg.restore(ckpt_path)
        _logger.info(f'Algorithm.restore 加载 {path} 成功')
    except Exception as _e:
        # _logger.error(_e)
        _logger.warning(f'Algorithm.restore 加载 {path} 失败')
        _r2 = None

    return _policy, _algorithm0, _r1, _r2


# policy, algorithm, r1, r2 = try_load(pretrained_path, 'pretrained_path')
policy, algorithm, r1, r2 = try_load(p_ckpt, 'p_ckpt')
# policy, algorithm, r1, r2 = try_load(a_folder, 'a_folder')
