sample_data.py

"""
Sample Data Generation

This module generates reproducible sample data for the ASI Chain simulation,
including agents, blocks, transactions, governance proposals, and cross-chain transfers.
"""

import random
import numpy as np
from typing import Dict, List, Tuple

from agents import generate_agents, Agent
from blockdag import BlockDAG
from shards import create_default_shards, Economy
from governance import Governance, create_sample_proposals
from memory import HypergraphMemory
from bridge import CrossChainBridge


def set_random_seeds(seed: int = 42):
    """
    Set random seeds for reproducibility.
    
    Args:
        seed: Random seed value
    """
    random.seed(seed)
    np.random.seed(seed)


def generate_simulation_data(num_agents: int = 10,
                            num_blocks: int = 20,
                            num_proposals: int = 5,
                            num_bridge_transfers: int = 15,
                            seed: int = 42) -> Dict:
    """
    Generate complete simulation data for all components.
    
    Args:
        num_agents: Number of agents to generate
        num_blocks: Number of blocks in BlockDAG
        num_proposals: Number of governance proposals
        num_bridge_transfers: Number of cross-chain transfers
        seed: Random seed for reproducibility
        
    Returns:
        Dictionary containing all simulation components
    """
    # Set seeds
    set_random_seeds(seed)
    
    print("Generating simulation data...")
    
    # ========== AGENTS ==========
    print(f"  Creating {num_agents} agents...")
    agents = generate_agents(num_agents, seed=seed)
    
    # ========== BLOCK DAG ==========
    print(f"  Building BlockDAG with {num_blocks} blocks...")
    dag = BlockDAG()
    
    for i in range(num_blocks):
        # Random transactions for this block
        num_txs = random.randint(2, 8)
        transactions = []
        
        for _ in range(num_txs):
            sender = random.choice(agents)
            receiver = random.choice([a for a in agents if a.agent_id != sender.agent_id])
            
            transactions.append({
                'from': sender.agent_id,
                'from_name': sender.name,
                'to': receiver.agent_id,
                'to_name': receiver.name,
                'amount': random.uniform(1.0, 100.0),
                'signature': sender.sign_action(f"tx_{i}")
            })
        
        dag.add_block(transactions=transactions, data={'block_index': i})
    
    # ========== SHARD ECONOMY ==========
    print(f"  Simulating shard economy...")
    shards = create_default_shards()
    economy = Economy(shards)
    
    # Simulate activity
    economy.simulate_all_shards(agents)
    
    # ========== GOVERNANCE ==========
    print(f"  Creating {num_proposals} governance proposals...")
    governance = Governance()
    sample_proposal_data = create_sample_proposals()
    
    proposals = []
    for prop_data in sample_proposal_data[:num_proposals]:
        proposal = governance.create_proposal(
            title=prop_data['title'],
            description=prop_data['description'],
            options=prop_data['options']
        )
        proposals.append(proposal)
        
        # Random agents vote
        num_voters = random.randint(max(3, num_agents // 2), num_agents)
        voters = random.sample(agents, k=num_voters)
        
        for agent in voters:
            choice = random.choice(proposal.options)
            governance.cast_vote(agent, proposal.proposal_id, choice)
        
        # Close voting
        governance.close_proposal(proposal.proposal_id)
    
    # ========== HYPERGRAPH MEMORY ==========
    print(f"  Building hypergraph memory...")
    memory = HypergraphMemory()
    
    # Add agent nodes
    for agent in agents:
        memory.add_node(agent.agent_id, 'agent', 
                       name=agent.name, 
                       reputation=agent.reputation)
    
    # Add shard nodes
    for shard in shards:
        memory.add_node(shard.name, 'shard',
                       name=shard.name,
                       type=shard.shard_type)
    
    # Add asset nodes
    num_assets = 5
    assets = []
    for i in range(num_assets):
        asset_id = f"asset_{i}"
        asset_name = f"Asset-{chr(65 + i)}"  # A, B, C, etc.
        memory.add_node(asset_id, 'asset',
                       name=asset_name,
                       value=random.randint(100, 1000))
        assets.append(asset_id)
    
    # Add relationships
    
    # Trust relations (agents trust other agents)
    for _ in range(num_agents + 5):
        a1, a2 = random.sample(agents, 2)
        memory.add_relation(
            a1.agent_id, a2.agent_id, 'trusts',
            strength=random.uniform(0.3, 1.0)
        )
    
    # Ownership relations (agents own assets)
    for _ in range(num_assets * 2):
        agent = random.choice(agents)
        asset = random.choice(assets)
        memory.add_relation(
            agent.agent_id, asset, 'owns',
            amount=random.randint(1, 20)
        )
    
    # Collaboration relations (agents collaborate)
    for _ in range(num_agents // 2):
        a1, a2 = random.sample(agents, 2)
        memory.add_relation(
            a1.agent_id, a2.agent_id, 'collaborates',
            project=f"project_{random.randint(1, 10)}"
        )
    
    # Usage relations (agents use shards)
    for agent in agents:
        # Each agent uses 1-2 shards
        num_shards_used = random.randint(1, min(2, len(shards)))
        used_shards = random.sample(shards, num_shards_used)
        
        for shard in used_shards:
            memory.add_relation(
                agent.agent_id, shard.name, 'uses',
                frequency=random.choice(['low', 'medium', 'high'])
            )
    
    # Trade relations (agents trade with each other)
    for _ in range(num_agents):
        a1, a2 = random.sample(agents, 2)
        memory.add_relation(
            a1.agent_id, a2.agent_id, 'trades',
            volume=random.uniform(10.0, 500.0)
        )
    
    # ========== CROSS-CHAIN BRIDGE ==========
    print(f"  Simulating cross-chain bridge with {num_bridge_transfers} transfers...")
    bridge = CrossChainBridge(chains=['Ethereum', 'BSC', 'Polygon'])
    
    # Set initial balances
    for agent in agents:
        for chain in bridge.chains:
            initial_balance = random.uniform(100.0, 1000.0)
            bridge.set_balance(agent.agent_id, chain, initial_balance)
    
    # Perform transfers
    for _ in range(num_bridge_transfers):
        agent = random.choice(agents)
        from_chain = random.choice(bridge.chains)
        to_chain = random.choice([c for c in bridge.chains if c != from_chain])
        
        # Transfer a portion of balance
        current_balance = bridge.get_balance(agent.agent_id, from_chain)
        if current_balance > 10:
            amount = random.uniform(5.0, min(100.0, current_balance * 0.4))
            bridge.transfer(agent.agent_id, from_chain, to_chain, amount, agent.name)
    
    print("Simulation data generation complete!\n")
    
    # ========== RETURN ALL COMPONENTS ==========
    return {
        'agents': agents,
        'dag': dag,
        'economy': economy,
        'shards': shards,
        'governance': governance,
        'proposals': proposals,
        'memory': memory,
        'bridge': bridge
    }


def get_simulation_summary(sim_data: Dict) -> Dict:
    """
    Get a summary of simulation data.
    
    Args:
        sim_data: Simulation data dictionary from generate_simulation_data
        
    Returns:
        Dictionary with summary statistics
    """
    agents = sim_data['agents']
    dag = sim_data['dag']
    economy = sim_data['economy']
    governance = sim_data['governance']
    memory = sim_data['memory']
    bridge = sim_data['bridge']
    
    return {
        'agents': {
            'count': len(agents),
            'avg_reputation': sum(a.reputation for a in agents) / len(agents),
            'total_transactions': sum(len(a.transaction_history) for a in agents)
        },
        'blockdag': dag.get_statistics(),
        'economy': economy.get_total_statistics(),
        'governance': governance.get_statistics(),
        'memory': memory.get_statistics(),
        'bridge': bridge.get_statistics()
    }


def simulate_reputation_decay(agents: List[Agent], 
                             time_steps: int = 10,
                             decay_rate: float = 0.02) -> List[Dict]:
    """
    Simulate reputation decay over time for visualization.
    
    Args:
        agents: List of agents
        time_steps: Number of time steps to simulate
        decay_rate: Decay rate per time step
        
    Returns:
        List of dictionaries with reputation history
    """
    history = []
    
    for step in range(time_steps + 1):
        for agent in agents:
            history.append({
                'time_step': step,
                'agent_id': agent.agent_id,
                'agent_name': agent.name,
                'reputation': agent.reputation
            })
        
        # Apply decay for next step (except on last step)
        if step < time_steps:
            for agent in agents:
                agent.apply_decay(decay_rate=decay_rate, time_elapsed=1.0)
    
    return history


def create_sankey_data(bridge: CrossChainBridge) -> Dict:
    """
    Create Sankey diagram data from bridge transactions.
    
    Args:
        bridge: CrossChainBridge instance
        
    Returns:
        Dictionary with Sankey diagram data
    """
    completed_txs = bridge.get_transaction_log(status='completed')
    
    # Build source-target-value triplets
    flows = {}
    
    for tx in completed_txs:
        key = (tx['from_chain'], tx['to_chain'])
        flows[key] = flows.get(key, 0) + tx['net_amount']
    
    # Create node labels
    chains = sorted(set(bridge.chains))
    node_labels = chains
    node_indices = {chain: i for i, chain in enumerate(chains)}
    
    # Create links
    sources = []
    targets = []
    values = []
    
    for (from_chain, to_chain), amount in flows.items():
        sources.append(node_indices[from_chain])
        targets.append(node_indices[to_chain])
        values.append(amount)
    
    return {
        'labels': node_labels,
        'sources': sources,
        'targets': targets,
        'values': values
    }


if __name__ == "__main__":
    # Demo usage
    print("=== Sample Data Generation Demo ===\n")
    
    # Generate full simulation
    sim_data = generate_simulation_data(
        num_agents=10,
        num_blocks=20,
        num_proposals=5,
        num_bridge_transfers=15,
        seed=42
    )
    
    print("\n" + "="*60 + "\n")
    
    # Print summary
    summary = get_simulation_summary(sim_data)
    
    print("Simulation Summary:\n")
    
    print("Agents:")
    print(f"  Count: {summary['agents']['count']}")
    print(f"  Avg Reputation: {summary['agents']['avg_reputation']:.2f}")
    print(f"  Total Transactions: {summary['agents']['total_transactions']}")
    
    print("\nBlockDAG:")
    for key, value in summary['blockdag'].items():
        print(f"  {key}: {value}")
    
    print("\nEconomy:")
    for key, value in summary['economy'].items():
        print(f"  {key}: {value}")
    
    print("\nGovernance:")
    for key, value in summary['governance'].items():
        print(f"  {key}: {value}")
    
    print("\nMemory:")
    for key, value in summary['memory'].items():
        print(f"  {key}: {value}")
    
    print("\nBridge:")
    for key, value in summary['bridge'].items():
        if isinstance(value, dict):
            print(f"  {key}:")
            for k, v in value.items():
                print(f"    {k}: {v}")
        else:
            print(f"  {key}: {value}")
    
    print("\n" + "="*60 + "\n")
    
    # Test reputation decay simulation
    print("Testing reputation decay simulation...")
    agents = sim_data['agents'][:3]  # Use first 3 agents
    
    # Reset reputation for demo
    for agent in agents:
        agent.reputation = 75.0
    
    decay_history = simulate_reputation_decay(agents, time_steps=5, decay_rate=0.05)
    
    print(f"Generated {len(decay_history)} reputation data points")
    print("\nSample decay data (first 6 rows):")
    for row in decay_history[:6]:
        print(f"  Step {row['time_step']}: {row['agent_name']} = {row['reputation']:.2f}")