#!/usr/bin/env python3
"""
DiscoRL ↔ Gym 集成 - 快速参考

使用方式:
  python scripts/QUICK_START.py
"""

quick_ref = """

╔════════════════════════════════════════════════════════════════════════════╗
║                    DiscoRL × Gym 集成 - 快速参考卡                         ║
╚════════════════════════════════════════════════════════════════════════════╝

┌──────────────────────────────────────────────────────────────────────────┐
│ 1️⃣  验证安装 (5 分钟)                                                    │
└──────────────────────────────────────────────────────────────────────────┘

$ python scripts/test_disco_setup.py

预期输出:
  ✓ All core components working!
  ✓ 所有 5 个测试通过


┌──────────────────────────────────────────────────────────────────────────┐
│ 2️⃣  在 CartPole 上训练 (1 分钟)                                          │
└──────────────────────────────────────────────────────────────────────────┘

$ python scripts/train_disco_cartpole.py

配置:
  - batch_size=4
  - trajectory_length=32
  - num_iterations=50
  
预期输出:
  ✓ Training Complete
  ✓ Final avg reward: ~0.97


┌──────────────────────────────────────────────────────────────────────────┐
│ 3️⃣  验证集成 (30 秒)                                                     │
└──────────────────────────────────────────────────────────────────────────┘

$ python scripts/poc_integration.py

预期输出:
  ✓ Success! DiscoRL ↔ Gym integration works!


╔════════════════════════════════════════════════════════════════════════════╗
║                            核心代码段
╚════════════════════════════════════════════════════════════════════════════╝

┌──────────────────────────────────────────────────────────────────────────┐
│ A) 环境设置                                                              │
└──────────────────────────────────────────────────────────────────────────┘

from disco_cartpole_env import DiscoCartPoleEnv

env = DiscoCartPoleEnv(batch_size=4, max_steps=500)
obs_spec = env.single_observation_spec()
act_spec = env.single_action_spec()


┌──────────────────────────────────────────────────────────────────────────┐
│ B) 代理创建                                                              │
└──────────────────────────────────────────────────────────────────────────┘

from disco_rl import agent as disco_agent

agent_settings = disco_agent.get_settings_disco()
agent = disco_agent.Agent(
    single_observation_spec=obs_spec,
    single_action_spec=act_spec,
    agent_settings=agent_settings,
    batch_axis_name=None,
)

learner_state = agent.initial_learner_state(rng_key)
actor_state = agent.initial_actor_state(rng_key)


┌──────────────────────────────────────────────────────────────────────────┐
│ C) 数据收集                                                              │
└──────────────────────────────────────────────────────────────────────────┘

state, timestep = env.reset(rng_key=subkey)

for t in range(trajectory_length):
    # 代理推理
    actor_timestep, actor_state = agent.actor_step(
        learner_state.params,
        rng,
        timestep,
        actor_state,
    )
    
    # 环境步进
    state, timestep = env.step(state, actor_timestep.actions)
    
    # 记录数据
    observations.append(timestep.observation['observation'])
    actions.append(actor_timestep.actions)
    rewards.append(timestep.reward)
    # ... 等


┌──────────────────────────────────────────────────────────────────────────┐
│ D) 参数更新                                                              │
└──────────────────────────────────────────────────────────────────────────┘

from disco_rl import types

rollout = types.ActorRollout(
    observations=jnp.stack(observations),
    actions=jnp.stack(actions),
    rewards=jnp.stack(rewards),
    discounts=jnp.stack(discounts),
    agent_outs=agent_outs_stacked,
    logits=jnp.stack(logits),
    states=actor_state,
)

new_learner_state, new_actor_state, logs = agent.learner_step(
    rng=rng,
    rollout=rollout,
    learner_state=learner_state,
    agent_net_state=actor_state,
    update_rule_params=update_rule_params,
    is_meta_training=False,
)


╔════════════════════════════════════════════════════════════════════════════╗
║                        常见问题与答案
╚════════════════════════════════════════════════════════════════════════════╝

Q: 如何用于自定义环境?
A: 1. cp disco_cartpole_env.py disco_custom_env.py
   2. 修改 __init__ 中的环境创建逻辑
   3. 调整 action_spec 和 observation_spec

Q: 如何加载预训练权重?
A: 预训练权重加载目前在开发中
   临时解决: 使用随机初始化的参数

Q: 如何扩展到多个 GPU?
A: 1. 移除 os.environ['JAX_PLATFORMS'] = 'cpu'
   2. 使用 jax.device_count() 获取设备数
   3. 在 agent.learner_step 中设置 batch_axis_name='devices'

Q: 性能太慢怎么办?
A: • 增加 batch_size (更多并行环境)
   • 减少 trajectory_length
   • 使用 GPU (移除 CPU-only 设置)
   • 减少网络大小 (调整 agent_settings)

Q: CartPole 不难吗?
A: CartPole 是验证集成的好工具
   一旦工作，应用到实际环境:
   • gym_env_250326_erase.py (自定义任务)
   • 或任何其他 Gym 兼容环境


╔════════════════════════════════════════════════════════════════════════════╗
║                        文件参考
╚════════════════════════════════════════════════════════════════════════════╝

核心文件 (必需):
  disco_cartpole_env.py      环境适配器 ← 为自定义环境修改这个
  
工具文件:
  disco_weights.py           权重加载
  train_disco_cartpole.py    训练循环
  test_disco_setup.py        测试套件
  
文档:
  INTEGRATION_GUIDE.py       详细指南
  COMPLETION_SUMMARY.py      完成报告
  QUICK_START.py             本文件

配置文件 (如需):
  config_*.json              在 configs/ 中


╔════════════════════════════════════════════════════════════════════════════╗
║                        关键数据类型
╚════════════════════════════════════════════════════════════════════════════╝

types.EnvironmentTimestep:
  observation: dict {'observation': Array([B, ...], float32)}
  step_type: Array([B], int32)          0=MID, 1=LAST
  reward: Array([B], float32)

types.ActorTimestep:
  observations: dict
  actions: Array([B], int32)             动作索引
  agent_outs: dict                       策略网络输出
  logits: Array([B, num_actions])
  ...

types.ActorRollout:
  observations, actions, rewards, discounts, agent_outs, logits, states
  (所有都在时间维度堆叠: [T, B, ...])


╔════════════════════════════════════════════════════════════════════════════╗
║                        环境规格 (CartPole)
╚════════════════════════════════════════════════════════════════════════════╝

观测空间: Box(4,)     ← 杆角度、角速度、车位置、车速度
动作空间: Discrete(2) ← 0=向左推, 1=向右推
奖励:     +1.0        ← 每一步 (最多 500 步)
完成:     当角度 > 24° 或位置超出界限


╔════════════════════════════════════════════════════════════════════════════╗
║                        常用命令
╚════════════════════════════════════════════════════════════════════════════╝

# 快速验证
python scripts/test_disco_setup.py

# 完整训练 (50 iter, 4 batch)
python scripts/train_disco_cartpole.py

# 端到端演示
python scripts/poc_integration.py

# 查看详细文档
python scripts/INTEGRATION_GUIDE.py | less

# 查看完成报告  
python scripts/COMPLETION_SUMMARY.py | less

# 运行此快速参考
python scripts/QUICK_START.py


╔════════════════════════════════════════════════════════════════════════════╗
║                            总结
╚════════════════════════════════════════════════════════════════════════════╝

✓ DiscoRL (JAX) ↔ Gym (任何环境)

✓ 完整的训练循环

✓ 预验证的代码

✓ 可立即复用的模板

准备开始? 运行:
  python scripts/test_disco_setup.py

有问题? 查看:
  python scripts/INTEGRATION_GUIDE.py

"""

print(quick_ref)

# 如果用户想保存
import sys
if len(sys.argv) > 1 and sys.argv[1] == '--save':
    with open('/home/frank14f/Frank_LBM/scripts/QUICK_START.txt', 'w') as f:
        f.write(quick_ref)
    print("✓ Saved to QUICK_START.txt")