#!/usr/bin/env python3
"""
JspaceAI —— 全部接入版（具身 Agent + 完整神经系统 + 自主心智）

接入全部 5 个感官输入 + 4 个输出执行器 + 神经系统（小脑/海马体/基底神经节/中枢神经）
+ 自主心智（好奇心驱动 + 状态持久化 + 自我模型 + 元学习）。

模式：
    --mode test:    子系统自检（权限 + I/O + 模型 + 海马体回忆）
    --mode live:    实时具身循环 + 自主心智（默认安全：不控制鼠标键盘）
    --mode safe:    同 live 但更保守（更高动作门控阈值）

运行：
    python main.py --mode test
    python main.py --mode live --steps 100
    python main.py --mode safe --steps 50
    python main.py --mode live --steps 200 --unsafe   # 允许鼠标键盘输出

再次运行会从上次状态继续（海洋不蒸发）。
"""
from __future__ import annotations
import argparse
import torch
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import time

from jspaceai import (
    MultimodalConfig, MultimodalJSpaceModel, EmbodiedAgent,
    WorkspaceRuntime, PLATFORM,
    get_screen_size, print_permission_guide,
    check_camera_permission, check_microphone_permission,
    check_input_monitoring_permission,
)


def get_config() -> MultimodalConfig:
    return MultimodalConfig(
        vocab_size=50, embed_dim=16, input_dim=8, workspace_dim=64,
        expert_dim=24, num_experts=12, num_wells=4, ode_steps=3,
        dt=0.1, tau_w=0.3, jacobian_sparsity=16, noise_std=0.01,
        img_size=32, audio_frame_size=1024, keyboard_vocab=128,
    )


def test_subsystems():
    """子系统自检"""
    print("=" * 60)
    print(f"具身 Agent 子系统测试 | 平台: {PLATFORM}")
    print("=" * 60)

    print("\n1. 权限检查:")
    checks = {
        '摄像头': check_camera_permission(),
        '麦克风': check_microphone_permission(),
        '键盘/鼠标监听': check_input_monitoring_permission(),
    }
    for name, ok in checks.items():
        print(f"   {name}: {'OK' if ok else '需要权限'}")
    if not all(checks.values()):
        print_permission_guide()

    sw, sh = get_screen_size()
    print(f"   屏幕尺寸: {sw}x{sh}")

    print("\n2. 模型 + Agent 初始化:")
    config = get_config()
    model = MultimodalJSpaceModel(config)
    print(f"   模型参数: {sum(p.numel() for p in model.parameters()):,}")
    print(f"   专家分工: {model.expert_modality}")

    agent = EmbodiedAgent(
        model, device='cpu',
        enable_mouse_output=False, enable_keyboard_output=False,
        enable_audio_output=True, enable_screen_output=False,
    )
    print(f"   小脑参数: {sum(p.numel() for p in agent.cerebellum.parameters()):,}")
    print(f"   海马体容量: {agent.hippocampus.capacity}")
    print(f"   反射弧数: {len(agent.cns.reflexes)}")

    print("\n3. 单步循环测试（2秒采集）:")
    agent.senses.start()
    time.sleep(2)
    for _ in range(5):
        info = agent.step_once()
        print(f"   step {info['step']:2d} | mod {info['modality']:8s} | "
              f"||w|| {info['w_norm']:.3f} | action {info['action']['action_name']} | "
              f"executed {info['action']['executed']} | mem {info['memories_count']}")
        time.sleep(0.5)
    agent.senses.stop()
    agent.audio_actuator.stop()

    print("\n4. 海马体回忆测试:")
    if agent.hippocampus and agent.hippocampus.size() > 0:
        w = agent.state['w'][0].cpu().numpy()
        for i, m in enumerate(agent.hippocampus.recall(w, top_k=3)):
            print(f"   记忆 {i}: sim={m['similarity']:.3f} ctx={m['context']}")

    print("\n所有子系统测试完成")


def live(n_steps: int, device: str, safe_mode: bool = False, unsafe: bool = False):
    """实时具身循环 + 自主心智"""
    print(f"平台: {PLATFORM}")
    sw, sh = get_screen_size()
    print(f"屏幕: {sw}x{sh}")

    allow_output = unsafe and not safe_mode
    if safe_mode:
        print("\n安全模式：不执行鼠标/键盘动作")
        print_permission_guide()
    elif not allow_output:
        print("\n默认模式：仅感知 + 音频/屏幕输出（不控制鼠标键盘）")
        print("  如需鼠标键盘输出，加 --unsafe")

    config = get_config()
    model = MultimodalJSpaceModel(config).to(device)
    mp = Path('outputs/multimodal_model.pt')
    if mp.exists():
        try:
            ckpt = torch.load(mp, map_location=device, weights_only=False)
            model.load_state_dict(ckpt['model'], strict=False)
            print(f"已加载模型: {mp}")
        except Exception:
            print("模型加载失败，随机初始化")
    else:
        print("未找到训练模型，随机初始化")
    model.eval()

    agent = EmbodiedAgent(
        model, device=device,
        enable_mouse_output=allow_output,
        enable_keyboard_output=allow_output,
        enable_audio_output=True,
        enable_screen_output=not safe_mode,
        risk_threshold=0.5 if safe_mode else 0.3,
    )
    runtime = WorkspaceRuntime(agent, save_dir='outputs/mind', device=device)

    print("\n" + "=" * 60)
    print("Workspace Runtime - 全感官具身循环")
    print("=" * 60)
    print("好奇心驱动 + 状态持久化 + 自我模型 + 元学习")
    print(f"运行 {n_steps} 步（Ctrl+C 中断，状态自动保存）\n")

    log = []

    def on_step(info):
        log.append(info)
        if info['step'] % 10 == 0:
            print(f"  step {info['step']:4d} | mod {info['modality']:8s} | "
                  f"||w|| {info['w_norm']:.3f} | curio {info['curiosity']:.3f} | "
                  f"success {info['success']:.2f} | focus={info['consensus_focus']} | "
                  f"weak={info['weakness']} | "
                  f"mem {info['memory_count']}")

    runtime.run(n_steps=n_steps, interval=0.2, save_every=30, on_step=on_step)

    print("\n" + runtime.introspect())

    # 可视化
    if log:
        fig, axes = plt.subplots(2, 3, figsize=(18, 10))
        steps = [s['step'] for s in log]

        # ||w|| 按模态着色
        ax = axes[0, 0]
        wn = [s['w_norm'] for s in log]
        mods = [s['modality'] for s in log]
        cm = {'image': 'green', 'screen': 'purple', 'audio': 'blue',
              'keyboard': 'orange', 'mouse': 'red', 'idle': 'gray'}
        colors = [cm.get(m, 'gray') for m in mods]
        ax.scatter(steps, wn, c=colors, alpha=0.7, s=30)
        ax.set_xlabel('Step'); ax.set_ylabel('||w||')
        ax.set_title('Workspace Norm (by modality)'); ax.grid(True, alpha=0.3)
        from matplotlib.patches import Patch
        ax.legend(handles=[Patch(facecolor=c, label=m) for m, c in cm.items()],
                  fontsize=7)

        # 好奇心 + 成功度
        ax = axes[0, 1]
        ax.plot(steps, [s['curiosity'] for s in log], 'r-', label='curiosity')
        ax.plot(steps, [s['success'] for s in log], 'g-', label='success')
        ax.legend(); ax.set_title('Curiosity & Success')
        ax.grid(True, alpha=0.3)

        # 模态分布
        ax = axes[0, 2]
        mc = {}
        for m in mods:
            mc[m] = mc.get(m, 0) + 1
        ax.bar(mc.keys(), mc.values(),
               color=[cm.get(m, 'gray') for m in mc.keys()])
        ax.set_title('Modality Distribution')

        # 自我模型
        ax = axes[1, 0]
        fc = log[-1]['self_confidence']
        ax.barh(list(fc.keys()), list(fc.values()),
                color=plt.cm.RdYlGn(list(fc.values())))
        ax.set_xlim(0, 1); ax.set_title('Self Model (confidence)')

        # 世界模型 loss
        ax = axes[1, 1]
        ax.plot(steps, [s['world_loss'] for s in log], 'orange')
        ax.set_title('World Model Loss'); ax.grid(True, alpha=0.3)

        # 记忆数
        ax = axes[1, 2]
        ax.plot(steps, [s['memory_count'] for s in log], 'teal')
        ax.set_title('Memory Count'); ax.grid(True, alpha=0.3)

        plt.tight_layout()
        Path('outputs').mkdir(exist_ok=True)
        plt.savefig('outputs/embodied_live.png', dpi=150, bbox_inches='tight')
        print(f"\n可视化: outputs/embodied_live.png")


def main():
    p = argparse.ArgumentParser(description='JspaceAI 全部接入版（具身 + 自主心智）')
    p.add_argument('--mode', default='test', choices=['test', 'live', 'safe'],
                   help='运行模式: test=自检, live=实时循环, safe=安全模式')
    p.add_argument('--steps', type=int, default=100, help='live/safe 模式步数')
    p.add_argument('--device', default='cpu', help='设备 (cpu/cuda/mps/auto)')
    p.add_argument('--unsafe', action='store_true',
                   help='允许鼠标键盘输出（默认禁用）')
    args = p.parse_args()

    dev = args.device
    if dev == 'auto':
        dev = 'cuda' if torch.cuda.is_available() else (
            'mps' if torch.backends.mps.is_available() else 'cpu')

    if args.mode == 'test':
        test_subsystems()
    elif args.mode == 'live':
        live(args.steps, dev, safe_mode=False, unsafe=args.unsafe)
    elif args.mode == 'safe':
        live(args.steps, dev, safe_mode=True, unsafe=False)


if __name__ == '__main__':
    main()
