"""
Example: Integrating DCMN with Neuro-Symbolic Planner
Demonstrates how to use the Decentralized Causal Memory Networks
"""

import os
import sys
import json
import time
from pathlib import Path

# add parent directory so we can import our modules
sys.path.append(str(Path(__file__).parent))

from neuro_symbolic_planner import NeuroSymbolicPlanner
from dcmn_causal_memory import (
    integrate_dcmn_with_planner,
    CausalKnowledgeAsset,
    CausalTriple,
    CausalRelationType,
    DomainParanet
)

def demonstrate_dcmn_integration():
    """Demonstrate DCMN integration with the planner"""
    
    print("🚀 Initializing Neuro-Symbolic Planner with DCMN...")
    
    # set up the planner with your api key
    api_key = os.getenv("GROQ_API_KEY", "your-api-key-here")
    planner = NeuroSymbolicPlanner(api_key)
    
    # add all the dcmn learning components
    planner = integrate_dcmn_with_planner(planner)
    
    print("✅ DCMN integration complete!")
    print(f"   - Causal Graph Memory: {planner.causal_memory}")
    print(f"   - Multi-Agent Validator: {planner.validator}")
    print(f"   - Domain Paranets: {list(planner.paranets.keys())}")
    
    # example 1: save a successful planning experience
    print("\n📝 Example 1: Storing Causal Knowledge from Planning")
    
    task = "Move block A from table to block B"
    plan = ["pickup(block_a)", "stack(block_a, block_b)"]
    
    # pretend we executed this plan and see what happened
    execution_trace = [
        {
            'action': 'pickup(block_a)',
            'state': 'holding(block_a) ∧ clear(table)',
            'confidence': 0.95
        },
        {
            'action': 'stack(block_a, block_b)',
            'state': 'on(block_a, block_b) ∧ clear(block_a)',
            'confidence': 0.92
        }
    ]
    
    # save what we learned as causal knowledge
    planner.store_causal_experience(task, plan, 'success', execution_trace)
    
    print(f"   ✓ Stored causal knowledge for: {task}")
    
    # example 2: use multiple agents to validate plans
    print("\n🤝 Example 2: Multi-Agent Validation")
    
    # add some agents that understand the blocks domain
    blocks_paranet = planner.paranets['blocks']
    blocks_paranet.add_member("expert_agent_1", 0.9)
    blocks_paranet.add_member("expert_agent_2", 0.85)
    blocks_paranet.add_member("expert_agent_3", 0.8)
    
    print(f"   - Added {len(blocks_paranet.members)} expert agents to blocks domain")
    
    # example 3: look up causal knowledge we learned before
    print("\n🔍 Example 3: Retrieving Causal Knowledge")
    
    query = "How to stack blocks"
    results = planner.drag.federated_query(query, domain='blocks')
    
    if results:
        top_asset = results[0]
        print(f"   - Found relevant knowledge: {top_asset.task_description}")
        print(f"   - Consensus score: {top_asset.get_consensus_score():.2f}")
        print(f"   - Causal triples:")
        for triple in top_asset.causal_triples[:3]:
            print(f"     • {triple.subject} {triple.predicate.value} {triple.object}")
    
    # example 4: find causal relationships between actions
    print("\n🔗 Example 4: Discovering Causal Paths")
    
    # Add more causal knowledge for path finding
    causal_triples = [
        CausalTriple("pickup(X)", CausalRelationType.ENABLES, "holding(X)"),
        CausalTriple("holding(X)", CausalRelationType.ENABLES, "stack(X,Y)"),
        CausalTriple("stack(X,Y)", CausalRelationType.PRODUCES, "on(X,Y)"),
        CausalTriple("on(X,Y)", CausalRelationType.REQUIRES, "clear(X)")
    ]
    
    # Create and store asset with these relationships
    path_asset = CausalKnowledgeAsset(
        asset_id="",
        causal_triples=causal_triples,
        plan_trace=["pickup", "stack"],
        domain="blocks",
        task_description="Generic stacking knowledge",
        success_rate=0.95,
        creation_time=time.time(),
        creator_id="system"
    )
    
    planner.causal_memory.store_asset(path_asset)
    
    # Find causal paths
    paths = planner.retriever.find_causal_paths("pickup(X)", "on(X,Y)")
    
    if paths:
        print(f"   - Found {len(paths)} causal paths from pickup to stacking:")
        for i, path in enumerate(paths[:2]):
            print(f"   Path {i+1}:")
            for triple in path:
                print(f"     → {triple.subject} {triple.predicate.value} {triple.object}")
    
    # Example 5: Augment plan with causal insights
    print("\n💡 Example 5: Augmenting Plans with Causal Knowledge")
    
    original_plan = ["move_to(kitchen)", "open_fridge", "take(milk)", "close_fridge"]
    augmented_plan = planner.drag.augment_plan_with_causal_knowledge(original_plan, 'household')
    
    print("   Original plan:")
    for action in original_plan:
        print(f"     - {action}")
    
    print("\n   Augmented plan with causal insights:")
    for item in augmented_plan:
        if item.startswith("#"):
            print(f"     {item}")
        else:
            print(f"     - {item}")
    
    # Example 6: Domain expertise tracking
    print("\n📊 Example 6: Domain Expertise Tracking")
    
    # Simulate performance updates
    blocks_paranet.update_expertise("expert_agent_1", 0.95)
    blocks_paranet.update_expertise("expert_agent_1", 0.98)
    blocks_paranet.update_expertise("expert_agent_2", 0.75)
    
    top_experts = blocks_paranet.get_experts(2)
    print(f"   Top experts in blocks domain:")
    for expert_id in top_experts:
        print(f"     - {expert_id}: expertise score {blocks_paranet.members[expert_id]:.2f}")
    
    # Example 7: Causal graph statistics
    print("\n📈 Example 7: Causal Knowledge Graph Statistics")
    
    print(f"   - Total causal assets: {len(planner.causal_memory.assets)}")
    print(f"   - Total causal triples: {sum(len(asset.causal_triples) for asset in planner.causal_memory.assets.values())}")
    print(f"   - RDF graph size: {len(planner.causal_memory.causal_graph)} triples")
    
    # Save causal knowledge to file
    print("\n💾 Saving causal knowledge...")
    
    causal_data = {
        'assets': [
            {
                'id': asset.asset_id,
                'domain': asset.domain,
                'task': asset.task_description,
                'triples': [
                    {
                        'subject': t.subject,
                        'predicate': t.predicate.value,
                        'object': t.object,
                        'confidence': t.confidence
                    }
                    for t in asset.causal_triples
                ],
                'consensus_score': asset.get_consensus_score()
            }
            for asset in planner.causal_memory.assets.values()
        ]
    }
    
    with open('dcmn_causal_knowledge.json', 'w') as f:
        json.dump(causal_data, f, indent=2)
    
    print("   ✓ Saved causal knowledge to dcmn_causal_knowledge.json")
    
    return planner

def demonstrate_advanced_features(planner):
    """Demonstrate advanced DCMN features"""
    
    print("\n🔬 Advanced DCMN Features")
    
    # Feature 1: Cross-domain causal transfer
    print("\n🔄 Feature 1: Cross-Domain Causal Transfer")
    
    # Create cooking domain knowledge that can transfer to general tasks
    cooking_triples = [
        CausalTriple("heat(water)", CausalRelationType.PRODUCES, "hot(water)"),
        CausalTriple("hot(water)", CausalRelationType.ENABLES, "brew(coffee)"),
        CausalTriple("brew(coffee)", CausalRelationType.PRODUCES, "ready(coffee)")
    ]
    
    cooking_asset = CausalKnowledgeAsset(
        asset_id="",
        causal_triples=cooking_triples,
        plan_trace=["heat", "brew"],
        domain="cooking",
        task_description="Make coffee by heating water",
        success_rate=0.9,
        creation_time=time.time(),
        creator_id="cooking_expert"
    )
    
    planner.causal_memory.store_asset(cooking_asset)
    
    # Query with a general heating task
    general_results = planner.drag.federated_query("heat liquid", domain=None)
    print(f"   - Found {len(general_results)} cross-domain results for heating")
    
    # Feature 2: Causal link prediction
    print("\n🔮 Feature 2: Causal Link Prediction")
    
    # Predict causal relationship between new actions
    action1 = "open(door)"
    action2 = "enter(room)"
    
    relation, confidence = planner.causal_memory.discover_causal_links(action1, action2)
    print(f"   - Predicted: {action1} {relation.value} {action2}")
    print(f"   - Confidence: {confidence:.2f}")
    
    # Feature 3: Consensus evolution
    print("\n📊 Feature 3: Consensus Evolution Tracking")
    
    # Show how consensus changes with more validators
    sample_asset = list(planner.causal_memory.assets.values())[0]
    initial_consensus = sample_asset.get_consensus_score()
    
    # Add more validations
    new_validators = [
        ("validator_1", 0.85),
        ("validator_2", 0.90),
        ("validator_3", 0.88)
    ]
    
    is_valid, new_consensus = planner.validator.validate_asset(sample_asset, new_validators)
    
    print(f"   - Initial consensus: {initial_consensus:.2f}")
    print(f"   - New consensus after {len(new_validators)} validators: {new_consensus:.2f}")
    print(f"   - Validation status: {'Valid' if is_valid else 'Invalid'}")
    
    # Feature 4: Export to RDF
    print("\n🌐 Feature 4: RDF Export for Interoperability")
    
    # Export causal knowledge to RDF format
    rdf_graph = planner.causal_memory.causal_graph
    rdf_output = rdf_graph.serialize(format='turtle')
    
    with open('dcmn_knowledge.ttl', 'w') as f:
        f.write(rdf_output)
    
    print(f"   - Exported {len(rdf_graph)} RDF triples to dcmn_knowledge.ttl")
    print(f"   - Compatible with OriginTrail and other RDF-based systems")

if __name__ == "__main__":
    # Run demonstration
    print("=" * 60)
    print("DCMN Integration Demonstration")
    print("Decentralized Causal Memory Networks for Planning")
    print("=" * 60)
    
    try:
        planner = demonstrate_dcmn_integration()
        demonstrate_advanced_features(planner)
        
        print("\n✨ Demonstration complete!")
        print("\nKey DCMN Benefits:")
        print("  • Verifiable causal knowledge with cryptographic hashes")
        print("  • Multi-agent consensus validation")
        print("  • Domain-specific expertise networks (Paranets)")
        print("  • Graph neural network-based retrieval")
        print("  • Decentralized knowledge sharing")
        print("  • RDF interoperability with OriginTrail")
        
    except Exception as e:
        print(f"\n❌ Error during demonstration: {e}")
        print("Make sure you have set GROQ_API_KEY environment variable")