#!/usr/bin/env python3
"""
AETHER Demo Script.
Demonstrates the full neuro-symbolic self-evolving architecture.
"""

import torch
import logging
import json
import time

from aether.core import AetherCore, AetherConfig
from aether.knowledge import KnowledgeGraphEngine

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("AETHER.Demo")


def demo_knowledge_graph():
    logger.info("\n" + "=" * 60)
    logger.info("DEMO: Knowledge Graph Engine")
    logger.info("=" * 60)
    
    kg = KnowledgeGraphEngine(embedding_dim=128, num_relations=20)
    
    facts = [
        ("AETHER", "is_a", "Agent"),
        ("AETHER", "has_capability", "Reasoning"),
        ("AETHER", "has_capability", "Learning"),
        ("AETHER", "has_capability", "Evolution"),
        ("Agent", "requires", "Memory"),
        ("Reasoning", "uses", "Knowledge_Graph"),
        ("Learning", "uses", "Neural_Networks"),
        ("Evolution", "optimizes", "Architecture"),
        ("Neural_Networks", "implement", "Learning"),
        ("Knowledge_Graph", "implements", "Reasoning"),
    ]
    
    for h, r, t in facts:
        kg.add_fact(h, r, t, confidence=1.0)
        logger.info(f"  Added: ({h}, {r}, {t})")
    
    logger.info("\n  Query: 'AETHER capabilities'")
    result = kg.query("AETHER has_capability")
    for r in result["results"]:
        logger.info(f"    -> {r['tail']} (confidence={r['confidence']:.2f}, source={r['source']})")
    
    logger.info("\n  Symbolic query: 'AETHER is_a ?'")
    symbolic = kg.reason_symbolic("AETHER", "is_a")
    for r in symbolic:
        logger.info(f"    -> {r['tail']} (path={r['path']})")
    
    stats = kg.stats()
    logger.info(f"\n  KG Stats: {json.dumps(stats, indent=2)}")
    return kg


def demo_memory_system():
    logger.info("\n" + "=" * 60)
    logger.info("DEMO: CoALA Memory System")
    logger.info("=" * 60)
    
    from aether.memory import CoALAMemory
    
    memory = CoALAMemory(capacity=16)
    
    memory.store({"task": "Solve math problem", "priority": "high"}, "working")
    memory.store({"task": "Update knowledge graph", "priority": "medium"}, "working")
    memory.store({"task": "Evolve architecture", "priority": "high"}, "working")
    
    for i in range(5):
        memory.store({
            "episode": i,
            "action": f"action_{i}",
            "reward": 0.5 + i * 0.1,
        }, "episodic")
    
    memory.store({
        "gravity": "9.8 m/s^2",
        "pi": "3.14159",
        "euler": "2.71828",
    }, "semantic")
    
    logger.info("\n  Retrieve 'task' from working memory:")
    results = memory.retrieve("task", memory_type="working", top_k=3)
    for r in results:
        logger.info(f"    -> {r}")
    
    logger.info("\n  Retrieve 'episode' from episodic memory:")
    results = memory.retrieve("episode", memory_type="episodic", top_k=3)
    for r in results:
        logger.info(f"    -> {r}")
    
    logger.info(f"\n  Memory stats: working={len(memory.working)}, episodic={len(memory.episodic)}")
    return memory


def demo_agents():
    logger.info("\n" + "=" * 60)
    logger.info("DEMO: Agent Orchestration")
    logger.info("=" * 60)
    
    from aether.agents import AetherAgentOrchestrator
    
    config = AetherConfig(
        macro_policy_dim=128,
        micro_policy_dim=64,
        num_agents=4,
    )
    
    orchestrator = AetherAgentOrchestrator(config)
    
    logger.info("\n  Executing task: 'Build a reasoning system'")
    result = orchestrator.execute("Build a reasoning system", None, {})
    
    logger.info(f"\n  Output: {result['output'][:100]}...")
    logger.info(f"  Blueprint: {result['blueprint']}")
    logger.info(f"  Agents used: {list(result['agent_outputs'].keys())}")
    logger.info(f"  Routing weights: {[f'{w:.3f}' for w in result['routing_weights']]}")
    
    logger.info("\n  Running BabyAGI loop for objective: 'Research self-evolution'")
    babyagi_result = orchestrator.run_babyagi("Research self-evolution", max_iterations=3)
    logger.info(f"  Completed {babyagi_result['iterations']} iterations")
    logger.info(f"  Tasks completed: {len(babyagi_result['completed_tasks'])}")
    
    return orchestrator


def demo_evolution():
    logger.info("\n" + "=" * 60)
    logger.info("DEMO: Evolutionary Engine")
    logger.info("=" * 60)
    
    from aether.evolution import AetherEvolutionEngine
    
    config = AetherConfig(
        population_size=8,
        mutation_rate=0.15,
        learning_rate=2e-5,
        macro_policy_dim=256,
        micro_policy_dim=128,
        num_agents=4,
    )
    
    engine = AetherEvolutionEngine(config)
    
    logger.info("\n  Generating candidate architectures...")
    candidates = engine.generate_candidates(config, population_size=8)
    logger.info(f"  Generated {len(candidates)} candidates")
    
    fitness_scores = [0.3 + 0.5 * torch.rand(1).item() for _ in candidates]
    logger.info(f"  Fitness scores: {[f'{s:.3f}' for s in fitness_scores]}")
    
    logger.info("\n  Selecting candidates (Performance-Novelty)...")
    selected = engine.select(candidates, fitness_scores, alpha_exploration=0.3)
    logger.info(f"  Selected {len(selected)} candidates")
    
    logger.info("\n  Applying constrained mutations...")
    mutated = engine.mutate(selected, mutation_rate=0.2)
    logger.info(f"  Produced {len(mutated)} mutants")
    
    engine.update_archive(candidates, fitness_scores)
    logger.info(f"\n  Archive stats: {json.dumps(engine.get_diversity_stats(), indent=2)}")
    
    return engine


def demo_safety():
    logger.info("\n" + "=" * 60)
    logger.info("DEMO: Safety Sandbox")
    logger.info("=" * 60)
    
    from aether.safety import SafetySandbox
    
    sandbox = SafetySandbox(timeout=30.0)
    
    config = AetherConfig(population_size=8, mutation_rate=0.15, learning_rate=2e-5)
    
    logger.info("\n  Validating safe config...")
    valid = sandbox.validate_architecture(config)
    logger.info(f"  Result: {'PASS' if valid else 'FAIL'}")
    
    bad_config = AetherConfig(population_size=200, mutation_rate=0.8, learning_rate=0.1)
    
    logger.info("\n  Validating unsafe config...")
    valid = sandbox.validate_architecture(bad_config)
    logger.info(f"  Result: {'PASS' if valid else 'FAIL'}")
    
    logger.info("\n  Testing sandbox execution...")
    with sandbox.sandbox() as ctx:
        result = sum(range(100))
        ctx["modifications_attempted"].append("safe_operation")
        logger.info(f"    Safe operation result: {result}")
    
    logger.info(f"\n  Audit summary: {json.dumps(sandbox.get_audit_summary(), indent=2)}")
    
    return sandbox


def demo_full_aether():
    logger.info("\n" + "=" * 60)
    logger.info("DEMO: Full AETHER Integration")
    logger.info("=" * 60)
    
    config = AetherConfig(
        population_size=4,
        generations=3,
        mutation_rate=0.1,
        learning_rate=2e-5,
        macro_policy_dim=128,
        micro_policy_dim=64,
        num_agents=2,
        enable_self_modification=True,
    )
    
    aether = AetherCore(config=config, model_name="demo-model")
    
    facts = [
        ("AI", "is_a", "Technology"),
        ("AGI", "is_a", "AI"),
        ("AETHER", "is_a", "AGI"),
        ("AETHER", "integrates", "Neural_Networks"),
        ("AETHER", "integrates", "Symbolic_Reasoning"),
        ("Neural_Networks", "enables", "Learning"),
        ("Symbolic_Reasoning", "enables", "Explainability"),
    ]
    
    for h, r, t in facts:
        aether.knowledge.add_fact(h, r, t)
    
    tasks = [
        "What is AETHER?",
        "How does AETHER learn?",
        "Explain AETHER reasoning",
    ]
    
    for task in tasks:
        logger.info(f"\n  Task: {task}")
        result = aether.forward(task)
        logger.info(f"  Symbolic weight: {result['symbolic_weight']:.3f}")
        logger.info(f"  Neural weight: {result['neural_weight']:.3f}")
        logger.info(f"  KG context: {len(result['kg_context']['results'])} results")
    
    logger.info("\n  Self-reflection...")
    reflection = aether.self_reflect()
    logger.info(f"  Generation: {reflection['generation']}")
    logger.info(f"  Memory stats: {json.dumps(reflection['memory_stats'], indent=2)}")
    logger.info(f"  Neuro-symbolic balance: {json.dumps(reflection['neuro_symbolic_balance'], indent=2)}")
    
    logger.info("\n  Mock evolution...")
    def mock_eval(cfg):
        return 0.5 + 0.3 * torch.rand(1).item()
    
    evolve_result = aether.evolve(mock_eval, num_generations=2)
    logger.info(f"  Best fitness: {evolve_result['best_fitness']:.4f}")
    logger.info(f"  Generations evolved: {evolve_result['generations_evolved']}")
    logger.info(f"  Architecture changes: {len(evolve_result['architecture_history'])}")
    
    logger.info("\n  Exporting state...")
    state = aether.export_state()
    logger.info(f"  State keys: {list(state.keys())}")
    logger.info(f"  Metadata: {json.dumps(state['metadata'], indent=2)}")
    
    return aether


def main():
    logger.info("\n" + "=" * 70)
    logger.info("  AETHER: A Self-Evolving Neuro-Symbolic Architecture")
    logger.info("  Integrating AlphaEvolve + BabyAGI + HiMAC + GEA + Yunjue + ASI-Evolve")
    logger.info("  + CoALA + MLPO + Agentic Neural Networks + CoMAS")
    logger.info("=" * 70)
    
    start_time = time.time()
    
    demo_knowledge_graph()
    demo_memory_system()
    demo_agents()
    demo_evolution()
    demo_safety()
    demo_full_aether()
    
    elapsed = time.time() - start_time
    
    logger.info("\n" + "=" * 70)
    logger.info(f"  All demos completed in {elapsed:.2f}s")
    logger.info("=" * 70)
    logger.info("\n  Components verified:")
    logger.info("    [x] Knowledge Graph Engine (RGCN + ComplEx)")
    logger.info("    [x] CoALA Memory (Working + Episodic + Semantic + Procedural)")
    logger.info("    [x] Agent Orchestration (Hierarchical + BabyAGI + MLPO)")
    logger.info("    [x] Evolution Engine (MAP-Elites + Performance-Novelty)")
    logger.info("    [x] Safety Sandbox (Validation + Audit + Constraints)")
    logger.info("    [x] Full AETHER Integration")
    logger.info("\n  Ready for training with TRL GRPO!")


if __name__ == "__main__":
    main()