src/streamlit_app_backup.py

"""
CASCADE Hyperlattice - Real-time 3D Visualization

Quine agent swarms navigating a 3D decision lattice with null gates.
Uses cascade-lattice for provenance and Rerun for 4D visualization.

🔮 GLASS BOX: Everything is observable. No hidden state.
"""

import streamlit as st
import plotly.graph_objects as go
import plotly.express as px
import numpy as np
import time
import colorsys
import json

from lattice import Hyperlattice
from swarm import QuineSwarm
from ipfs_sync import CollectiveMemory, SyncManager
from cascade_bridge import CascadeBridge
from champion_loader import get_champion_info, load_champion_module, get_champion_diagnostics, get_replicated_agents_info, get_download_status

try:
    import torch
    TORCH_AVAILABLE = True
    GPU_AVAILABLE = torch.cuda.is_available()
    GPU_NAME = torch.cuda.get_device_name(0) if GPU_AVAILABLE else None
    GPU_MEMORY = torch.cuda.get_device_properties(0).total_memory // (1024**3) if GPU_AVAILABLE else 0
except ImportError:
    TORCH_AVAILABLE = False
    GPU_AVAILABLE = False
    GPU_NAME = None
    GPU_MEMORY = 0

# Try to import cascade for tracing
try:
    from cascade import Tracer, CausationGraph, sdk_observe
    CASCADE_AVAILABLE = True
except ImportError:
    CASCADE_AVAILABLE = False
    sdk_observe = None

# ═══════════════════════════════════════════════════════════════════════════════
# PAGE CONFIG
# ═══════════════════════════════════════════════════════════════════════════════
st.set_page_config(
    page_title="CASCADE Hyperlattice",
    page_icon="🌌",
    layout="wide"
)

# ═══════════════════════════════════════════════════════════════════════════════
# LINEAGE COLOR SYSTEM
# ═══════════════════════════════════════════════════════════════════════════════
def get_lineage_color(root_id: str, generation: int, fitness: float, max_gen: int = 10):
    """
    Generate a color based on lineage with hue/saturation degradation.
    
    - Hue: Based on root agent ID (each original agent gets unique hue)
    - Saturation: Decreases with generation (replicant degradation)
    - Lightness: Based on fitness (healthier = brighter)
    """
    # Hash root_id to get consistent hue
    hue = (hash(root_id) % 360) / 360.0
    
    # Saturation degrades with generation (1.0 -> 0.3)
    saturation = max(0.3, 1.0 - (generation / max(max_gen, 1)) * 0.7)
    
    # Lightness based on fitness (normalized, range 0.4-0.9)
    lightness = 0.4 + min(1.0, max(0.0, fitness / 10.0)) * 0.5
    
    # Convert HSL to RGB
    r, g, b = colorsys.hls_to_rgb(hue, lightness, saturation)
    return f'rgb({int(r*255)},{int(g*255)},{int(b*255)})'


def get_lineage_colors_batch(lineage_data: dict, max_gen: int = 10):
    """Get colors for all agents based on lineage."""
    colors = []
    for agent_id, data in lineage_data.items():
        color = get_lineage_color(
            data['root_lineage'],
            data['generation'],
            data['fitness'],
            max_gen
        )
        colors.append(color)
    return colors


# ═══════════════════════════════════════════════════════════════════════════════
# THEMES
# ═══════════════════════════════════════════════════════════════════════════════
THEMES = {
    "neon_vampire": {
        "name": "🧛 Neon Vampire",
        "bg": "rgba(10,5,15,1)",
        "paper": "rgba(15,8,20,1)",
        "edges": "rgba(120,0,180,0.4)",
        "null_gates": "rgba(255,0,80,1)",
        "nodes": "Plasma",
        "agents": "Inferno",
        "accent": "#00ffff"
    },
    "cyberpunk": {
        "name": "🌆 Cyberpunk",
        "bg": "rgba(5,5,20,1)",
        "paper": "rgba(10,10,30,1)",
        "edges": "rgba(0,255,255,0.3)",
        "null_gates": "rgba(255,0,128,1)",
        "nodes": "Viridis",
        "agents": "Turbo",
        "accent": "#ff00ff"
    },
    "matrix": {
        "name": "💚 Matrix",
        "bg": "rgba(0,5,0,1)",
        "paper": "rgba(0,10,0,1)",
        "edges": "rgba(0,255,0,0.2)",
        "null_gates": "rgba(0,255,100,1)",
        "nodes": "Greens",
        "agents": "YlGn",
        "accent": "#00ff00"
    },
    "solar_flare": {
        "name": "☀️ Solar Flare",
        "bg": "rgba(20,5,0,1)",
        "paper": "rgba(30,10,0,1)",
        "edges": "rgba(255,100,0,0.3)",
        "null_gates": "rgba(255,200,0,1)",
        "nodes": "Hot",
        "agents": "YlOrRd",
        "accent": "#ff5500"
    },
    "ice_palace": {
        "name": "❄️ Ice Palace",
        "bg": "rgba(5,10,20,1)",
        "paper": "rgba(10,15,30,1)",
        "edges": "rgba(100,200,255,0.3)",
        "null_gates": "rgba(200,230,255,1)",
        "nodes": "Blues",
        "agents": "ice",
        "accent": "#ffffff"
    },
}

# ═══════════════════════════════════════════════════════════════════════════════
# SESSION STATE
# ═══════════════════════════════════════════════════════════════════════════════
if 'initialized' not in st.session_state:
    st.session_state.initialized = False
    st.session_state.lattice = None
    st.session_state.swarm = None
    st.session_state.cascade = None
    st.session_state.sync = None
    st.session_state.memory = None
    st.session_state.step_count = 0
    st.session_state.theme = 'neon_vampire'
    st.session_state.events = []
    st.session_state.trails = {}
    st.session_state.lineage_trails = {}  # Track trails per lineage
    st.session_state.tracer = None  # Cascade tracer for temporal sequencing
    st.session_state.rerun_logger = None  # Rerun 4D visualization
    st.session_state.rerun_enabled = False  # Toggle for Rerun logging
    st.session_state.hold_count = 0
    st.session_state.override_count = 0


def init_simulation(size=6, agents=5, gates=4, enable_rerun=False):
    """Initialize the simulation."""
    st.session_state.lattice = Hyperlattice(dimensions=(size, size, size))
    if gates > 5:
        st.session_state.lattice._add_null_gates(gates - 5)
    st.session_state.swarm = QuineSwarm(st.session_state.lattice, initial_agents=agents)
    st.session_state.memory = CollectiveMemory(use_ipfs=False)
    st.session_state.sync = SyncManager(st.session_state.memory)
    st.session_state.cascade = CascadeBridge(session_id=f"stream_{int(time.time())}")
    st.session_state.step_count = 0
    st.session_state.events = ["🌌 Simulation initialized!"]
    st.session_state.trails = {}
    st.session_state.lineage_trails = {}
    st.session_state.initialized = True
    st.session_state.hold_count = 0
    st.session_state.override_count = 0
    
    # Initialize cascade tracer if available
    if CASCADE_AVAILABLE:
        try:
            # Tracer requires a CausationGraph
            graph = CausationGraph()
            st.session_state.tracer = Tracer(graph)
            st.session_state.events.append("🔗 Cascade-lattice tracing enabled")
            
            # Log initialization to CASCADE
            if sdk_observe:
                sdk_observe(
                    model_id="hyperlattice_app",
                    inputs={"size": size, "agents": agents, "gates": gates},
                    outputs={"initialized": True},
                    latency_ms=0.0
                )
        except Exception as e:
            st.session_state.events.append(f"⚠️ Cascade tracer: {e}")
    
    # Initialize Rerun if enabled and available
    st.session_state.rerun_enabled = enable_rerun
    if enable_rerun and RERUN_AVAILABLE:
        try:
            st.session_state.rerun_logger = init_rerun_logger(
                application_id="cascade-hyperlattice",
                spawn=False,  # Don't spawn desktop app, use web viewer
                serve=True,   # Serve web viewer
                web_port=9090
            )
            if st.session_state.rerun_logger:
                st.session_state.events.append("👁️ Rerun 4D visualization enabled (port 9090)")
                
                # Log initial lattice structure
                _log_lattice_to_rerun()
        except Exception as e:
            st.session_state.events.append(f"⚠️ Rerun init: {e}")


def _log_lattice_to_rerun():
    """Log lattice structure to Rerun."""
    logger = st.session_state.rerun_logger
    if not logger or not st.session_state.lattice:
        return
    
    lattice = st.session_state.lattice
    data = lattice.to_plotly_data()
    
    # Convert edges to line strips
    edges = []
    x, y, z = data['edges']['x'], data['edges']['y'], data['edges']['z']
    for i in range(0, len(x) - 2, 3):  # Plotly uses None separators
        if x[i] is not None and x[i+1] is not None:
            edges.append([(x[i], y[i], z[i]), (x[i+1], y[i+1], z[i+1])])
    
    logger.log_lattice_edges(edges)
    
    # Log null gates
    null_gates = []
    nx, ny, nz = data['null_edges']['x'], data['null_edges']['y'], data['null_edges']['z']
    for i in range(0, len(nx) - 2, 3):
        if nx[i] is not None and nx[i+1] is not None:
            null_gates.append([(nx[i], ny[i], nz[i]), (nx[i+1], ny[i+1], nz[i+1])])
    
    logger.log_null_gates(null_gates)


def step():
    """Run one simulation step."""
    if not st.session_state.initialized:
        return
    
    swarm = st.session_state.swarm
    cascade = st.session_state.cascade
    logger = st.session_state.rerun_logger
    
    experiences = swarm.step()
    st.session_state.step_count += 1
    
    # Update Rerun time
    if logger:
        logger.set_time(step=st.session_state.step_count)
    
    # Get lineage data for trail tracking
    lineage_data = swarm.get_agent_lineage_data()
    
    for exp in experiences:
        cascade.log_decision(
            agent_id=exp.agent_id,
            action=exp.action,
            observation={'from': exp.from_node, 'to': exp.to_node},
            value_estimate=exp.reward,
            action_probs={'move': 0.7, 'null_transit': 0.3}
        )
        
        # Track trails with lineage info
        if hasattr(exp, 'to_node'):
            pos = st.session_state.lattice.get_node_position(exp.to_node)
            if pos is not None:
                # Regular trails
                if exp.agent_id not in st.session_state.trails:
                    st.session_state.trails[exp.agent_id] = []
                st.session_state.trails[exp.agent_id].append(pos)
                if len(st.session_state.trails[exp.agent_id]) > 15:
                    st.session_state.trails[exp.agent_id] = st.session_state.trails[exp.agent_id][-15:]
                
                # Lineage-aware trails (grouped by root)
                agent_data = lineage_data.get(exp.agent_id, {})
                root = agent_data.get('root_lineage', exp.agent_id)
                if root not in st.session_state.lineage_trails:
                    st.session_state.lineage_trails[root] = []
                st.session_state.lineage_trails[root].append({
                    'pos': pos,
                    'agent_id': exp.agent_id,
                    'gen': agent_data.get('generation', 0),
                    'time': st.session_state.step_count
                })
                # Keep last 50 trail points per lineage
                if len(st.session_state.lineage_trails[root]) > 50:
                    st.session_state.lineage_trails[root] = st.session_state.lineage_trails[root][-50:]
        
        if exp.action == 'null_transit':
            st.session_state.events.append(f"🌀 {exp.agent_id} null jump!")
    
    # Log agents to Rerun
    if logger and lineage_data:
        agents_batch = [
            {
                'agent_id': aid,
                'position': data['position'],
                'generation': data['generation'],
                'fitness': data['fitness'],
                'root_lineage': data['root_lineage']
            }
            for aid, data in lineage_data.items()
        ]
        logger.log_agent_batch(agents_batch)
        
        # Log swarm stats
        stats = swarm.get_swarm_stats()
        avg_fitness = sum(d['fitness'] for d in lineage_data.values()) / max(len(lineage_data), 1)
        logger.log_swarm_stats(
            num_agents=stats['num_agents'],
            avg_fitness=avg_fitness,
            max_generation=stats.get('max_generation', 0),
            total_forks=stats.get('total_forks', 0),
            total_holds=st.session_state.hold_count,
            total_overrides=st.session_state.override_count
        )
        
        # Log trails to Rerun
        for agent_id, trail_positions in st.session_state.trails.items():
            if len(trail_positions) >= 2:
                logger.log_trail(agent_id, trail_positions)
    
    # Forks
    new_agents = swarm.fork_at_null_gates()
    for a in new_agents:
        st.session_state.events.append(f"🧬 {a.agent_id} spawned (gen {a.generation})!")
    
    # Prune
    pruned = swarm.prune_weak(keep_top=15)
    if pruned:
        st.session_state.events.append(f"💀 Pruned {pruned}")
    
    # Keep events manageable
    st.session_state.events = st.session_state.events[-30:]
    
    # Log causation tree to Rerun periodically
    if logger and st.session_state.step_count % 10 == 0:
        tree = cascade.get_causation_tree()
        logger.log_causation_tree(tree)


def create_plot():
    """Create the 3D visualization with lineage-aware coloring and clickable causal routes."""
    if not st.session_state.initialized:
        return go.Figure()
    
    lattice = st.session_state.lattice
    swarm = st.session_state.swarm
    cascade = st.session_state.cascade
    theme = THEMES[st.session_state.theme]
    data = lattice.to_plotly_data()
    
    fig = go.Figure()
    
    # Edges (regular lattice connections)
    fig.add_trace(go.Scatter3d(
        x=data['edges']['x'], y=data['edges']['y'], z=data['edges']['z'],
        mode='lines', line=dict(color=theme['edges'], width=1),
        hoverinfo='none', name='Edges'
    ))
    
    # Null gates (wormhole connections)
    fig.add_trace(go.Scatter3d(
        x=data['null_edges']['x'], y=data['null_edges']['y'], z=data['null_edges']['z'],
        mode='lines', line=dict(color=theme['null_gates'], width=5),
        hoverinfo='none', name='Null Gates'
    ))
    
    # Nodes
    fig.add_trace(go.Scatter3d(
        x=data['nodes']['x'], y=data['nodes']['y'], z=data['nodes']['z'],
        mode='markers',
        marker=dict(size=3, color=data['nodes']['colors'], colorscale=theme['nodes'], opacity=0.5),
        hoverinfo='none', name='Nodes'
    ))
    
    # ═══════════════════════════════════════════════════════════════════════════
    # CLICKABLE CAUSATION ROUTES - Draw decision chain connections
    # ═══════════════════════════════════════════════════════════════════════════
    if cascade:
        lineage_data = swarm.get_agent_lineage_data()
        decision_nodes = cascade.decision_nodes
        
        # Collect causation edges with midpoints for clicking
        causal_edge_x, causal_edge_y, causal_edge_z = [], [], []
        midpoint_x, midpoint_y, midpoint_z = [], [], []
        midpoint_texts = []
        midpoint_customdata = []
        
        for node_id, node in decision_nodes.items():
            if node.parent_id and node.parent_id in decision_nodes:
                parent = decision_nodes[node.parent_id]
                
                # Get positions from agent lineage data
                agent_pos = None
                parent_pos = None
                
                # Find agent positions at time of decision
                if node.agent_id in lineage_data:
                    agent_pos = lineage_data[node.agent_id]['position']
                if parent.agent_id in lineage_data:
                    parent_pos = lineage_data[parent.agent_id]['position']
                
                # If same agent, we can trace via trail
                if agent_pos is not None and parent_pos is not None:
                    # Draw the causation edge
                    causal_edge_x.extend([parent_pos[0], agent_pos[0], None])
                    causal_edge_y.extend([parent_pos[1], agent_pos[1], None])
                    causal_edge_z.extend([parent_pos[2], agent_pos[2], None])
                    
                    # Add clickable midpoint
                    mid_x = (parent_pos[0] + agent_pos[0]) / 2
                    mid_y = (parent_pos[1] + agent_pos[1]) / 2
                    mid_z = (parent_pos[2] + agent_pos[2]) / 2
                    midpoint_x.append(mid_x)
                    midpoint_y.append(mid_y)
                    midpoint_z.append(mid_z)
                    
                    # Build hover text with full causation info
                    hover_text = (
                        f"<b>🔗 CAUSAL ROUTE</b><br>"
                        f"<b>From:</b> {parent.node_id}<br>"
                        f"<b>To:</b> {node.node_id}<br>"
                        f"<b>Agent:</b> {node.agent_id}<br>"
                        f"<b>Action:</b> {node.action}<br>"
                        f"<b>Value:</b> {node.value_estimate:.3f}<br>"
                        f"<b>Merkle:</b> {node.merkle_hash}<br>"
                        f"<b>Parent Merkle:</b> {parent.merkle_hash}<br>"
                        f"<b>Probs:</b> {json.dumps(node.action_probs)[:50]}..."
                    )
                    midpoint_texts.append(hover_text)
                    midpoint_customdata.append({
                        'type': 'causal_route',
                        'from_node': parent.node_id,
                        'to_node': node.node_id,
                        'agent_id': node.agent_id,
                        'action': node.action,
                        'merkle': node.merkle_hash,
                        'parent_merkle': parent.merkle_hash,
                        'value': node.value_estimate,
                        'probs': node.action_probs
                    })
        
        # Draw causation edges
        if causal_edge_x:
            fig.add_trace(go.Scatter3d(
                x=causal_edge_x, y=causal_edge_y, z=causal_edge_z,
                mode='lines',
                line=dict(color='rgba(255,215,0,0.6)', width=3, dash='dot'),
                hoverinfo='none',
                name='Causal Routes'
            ))
        
        # Draw clickable midpoints (these are what you click to drill into)
        if midpoint_x:
            fig.add_trace(go.Scatter3d(
                x=midpoint_x, y=midpoint_y, z=midpoint_z,
                mode='markers',
                marker=dict(
                    size=6,
                    color='gold',
                    symbol='diamond',
                    opacity=0.9,
                    line=dict(width=1, color='white')
                ),
                hoverinfo='text',
                text=midpoint_texts,
                customdata=midpoint_customdata,
                name='🔗 Click to Drill'
            ))
    
    # Lineage-aware trails (colored by root lineage with gradient for time)
    for root_id, trail_data in st.session_state.lineage_trails.items():
        if len(trail_data) >= 2:
            # Sort by time for proper trail rendering
            sorted_trail = sorted(trail_data, key=lambda x: x['time'])
            
            # Create gradient opacity based on recency
            trail_x = [p['pos'][0] for p in sorted_trail]
            trail_y = [p['pos'][1] for p in sorted_trail]
            trail_z = [p['pos'][2] for p in sorted_trail]
            
            # Get base color for this lineage
            base_color = get_lineage_color(root_id, 0, 5.0, 10)
            
            fig.add_trace(go.Scatter3d(
                x=trail_x, y=trail_y, z=trail_z,
                mode='lines+markers',
                line=dict(color=base_color, width=2),
                marker=dict(
                    size=[2 + (i / len(sorted_trail)) * 4 for i in range(len(sorted_trail))],
                    color=base_color,
                    opacity=0.6
                ),
                hoverinfo='text',
                text=[f"{p['agent_id']} (gen {p['gen']}) t={p['time']}" for p in sorted_trail],
                showlegend=False,
                name=f'Trail: {root_id}'
            ))
    
    # Agents with lineage-based coloring
    lineage_data = swarm.get_agent_lineage_data()
    if lineage_data:
        # Find max generation for saturation scaling
        max_gen = max(d['generation'] for d in lineage_data.values()) if lineage_data else 1
        max_gen = max(max_gen, 1)  # Avoid division by zero
        
        agent_ids = list(lineage_data.keys())
        ax = [lineage_data[aid]['position'][0] for aid in agent_ids]
        ay = [lineage_data[aid]['position'][1] for aid in agent_ids]
        az = [lineage_data[aid]['position'][2] for aid in agent_ids]
        
        # Generate colors based on lineage
        colors = []
        sizes = []
        hover_texts = []
        symbols = []
        
        for aid in agent_ids:
            d = lineage_data[aid]
            color = get_lineage_color(d['root_lineage'], d['generation'], d['fitness'], max_gen)
            colors.append(color)
            
            # Size based on generation (originals bigger, replicants smaller)
            size = 14 - min(d['generation'], 8)  # 14 for gen 0, down to 6 for gen 8+
            sizes.append(size)
            
            # Symbol: diamond for originals, circle for replicants
            symbols.append('diamond' if d['generation'] == 0 else 'circle')
            
            # Hover info
            hover_texts.append(
                f"<b>{aid}</b><br>"
                f"Lineage: {d['root_lineage']}<br>"
                f"Gen: {d['generation']}<br>"
                f"Fitness: {d['fitness']:.2f}<br>"
                f"Hash: {d['quine_hash'][:8]}..."
            )
        
        # Plot originals (generation 0) separately with diamond markers
        orig_indices = [i for i, aid in enumerate(agent_ids) if lineage_data[aid]['generation'] == 0]
        if orig_indices:
            fig.add_trace(go.Scatter3d(
                x=[ax[i] for i in orig_indices],
                y=[ay[i] for i in orig_indices],
                z=[az[i] for i in orig_indices],
                mode='markers',
                marker=dict(
                    size=[sizes[i] for i in orig_indices],
                    color=[colors[i] for i in orig_indices],
                    symbol='diamond',
                    line=dict(width=2, color='white')
                ),
                hoverinfo='text',
                text=[hover_texts[i] for i in orig_indices],
                customdata=[agent_ids[i] for i in orig_indices],  # For click events
                name='Original Agents'
            ))
        
        # Plot replicants with circle markers (smaller, desaturated)
        rep_indices = [i for i, aid in enumerate(agent_ids) if lineage_data[aid]['generation'] > 0]
        if rep_indices:
            fig.add_trace(go.Scatter3d(
                x=[ax[i] for i in rep_indices],
                y=[ay[i] for i in rep_indices],
                z=[az[i] for i in rep_indices],
                mode='markers',
                marker=dict(
                    size=[sizes[i] for i in rep_indices],
                    color=[colors[i] for i in rep_indices],
                    symbol='circle',
                    opacity=0.8
                ),
                hoverinfo='text',
                text=[hover_texts[i] for i in rep_indices],
                customdata=[agent_ids[i] for i in rep_indices],  # For click events
                name='Replicants'
            ))
    
    # Layout
    fig.update_layout(
        scene=dict(
            xaxis=dict(visible=False, backgroundcolor=theme['bg']),
            yaxis=dict(visible=False, backgroundcolor=theme['bg']),
            zaxis=dict(visible=False, backgroundcolor=theme['bg']),
            bgcolor=theme['bg'],
            camera=dict(eye=dict(
                x=1.5 * np.cos(st.session_state.step_count * 0.03),
                y=1.5 * np.sin(st.session_state.step_count * 0.03),
                z=0.8
            ))
        ),
        paper_bgcolor=theme['paper'],
        margin=dict(l=0, r=0, t=0, b=0),
        height=600,
        showlegend=False
    )
    
    return fig


# ═══════════════════════════════════════════════════════════════════════════════
# UI
# ═══════════════════════════════════════════════════════════════════════════════
st.title("🌌 CASCADE Hyperlattice")
st.caption("🔮 GLASS BOX AI - Every decision is transparent and drillable")

# Import diagnostic functions
from champion_loader import get_champion_diagnostics, get_replicated_agents_info, get_download_status

# Session state for selection
if 'selected_agent' not in st.session_state:
    st.session_state.selected_agent = None
if 'selected_route' not in st.session_state:
    st.session_state.selected_route = None  # For causal route drilldown
if 'diagnostic_tab' not in st.session_state:
    st.session_state.diagnostic_tab = 'overview'

# Sidebar
with st.sidebar:
    st.header("🎨 Theme")
    theme = st.selectbox("Theme", list(THEMES.keys()), format_func=lambda x: THEMES[x]['name'])
    if theme != st.session_state.theme:
        st.session_state.theme = theme
    
    st.divider()
    st.header("⚙️ Setup")
    size = st.slider("Lattice Size", 3, 10, 6)
    agents = st.slider("Agents", 1, 10, 5)
    gates = st.slider("Null Gates", 2, 10, 4)
    
    if st.button("🔄 Initialize", type="primary", width='stretch'):
        init_simulation(size, agents, gates, enable_rerun=False)
    
    st.divider()
    st.header("▶️ Animation")
    speed = st.slider("Speed", 1, 20, 8)
    
    col1, col2 = st.columns(2)
    play = col1.button("▶️ Play", width='stretch')
    stop = col2.button("⏹️ Stop", width='stretch')
    
    if st.button("⚡ Step", width='stretch'):
        step()
    
    # Quick stats
    st.divider()
    if st.session_state.initialized:
        stats = st.session_state.swarm.get_swarm_stats()
        c1, c2, c3 = st.columns(3)
        c1.metric("Step", st.session_state.step_count)
        c2.metric("Agents", stats['num_agents'])
        c3.metric("MaxGen", stats.get('max_generation', 0))
        
        # Integration status
        st.divider()
        st.caption("🔗 Integrations")
        cols = st.columns(2)
        cols[0].write(f"CASCADE {'✅' if CASCADE_AVAILABLE else '❌'}")
        cols[1].write(f"HF {'✅' if get_download_status().get('downloaded') else '⏳'}")
        
        # GPU Status
        st.divider()
        if GPU_AVAILABLE:
            st.success(f"🚀 GPU: {GPU_NAME}")
            st.caption(f"VRAM: {GPU_MEMORY} GB")
        else:
            st.info("💻 CPU Mode")
            if TORCH_AVAILABLE:
                st.caption("PyTorch available, no GPU detected")

# ═══════════════════════════════════════════════════════════════════════════════
# MAIN LAYOUT - Graph + Diagnostic Panels
# ═══════════════════════════════════════════════════════════════════════════════
main_col, diag_col = st.columns([3, 2])

with main_col:
    # 3D Graph with click interaction
    st.subheader("🌐 Lattice Visualization")
    st.caption("👆 Click any agent directly in the graph to inspect")
    
    plot_placeholder = st.empty()
    
    # Show selected agent stats (no dropdown - click graph instead)
    if st.session_state.initialized:
        lineage_data = st.session_state.swarm.get_agent_lineage_data()
        
        # Clear selection button
        if st.session_state.selected_agent:
            col_info, col_clear = st.columns([4, 1])
            with col_info:
                if st.session_state.selected_agent in lineage_data:
                    data = lineage_data[st.session_state.selected_agent]
                    st.success(f"🎯 Selected: **{st.session_state.selected_agent}** (Gen {data.get('generation', 0)})")
            with col_clear:
                if st.button("❌ Clear", help="Clear selection"):
                    st.session_state.selected_agent = None
                    st.rerun()
        
        # Quick stats bar for selected agent
        if st.session_state.selected_agent and st.session_state.selected_agent in lineage_data:
            data = lineage_data[st.session_state.selected_agent]
            stat_cols = st.columns(4)
            stat_cols[0].metric("Gen", data.get('generation', 0))
            stat_cols[1].metric("Fitness", f"{data.get('fitness', 0):.2f}")
            pos = data.get('position', (0, 0, 0))
            stat_cols[2].metric("Pos", f"({pos[0]:.0f},{pos[1]:.0f},{pos[2]:.0f})")
            stat_cols[3].metric("Children", data.get('num_children', 0))

with diag_col:
    # Diagnostic Panels - ALL DATA RESTORED
    st.subheader("📊 Diagnostics")
    
    # ═══════════════════════════════════════════════════════════════════════════
    # CHAMPION MODEL - FULL DATA
    # ═══════════════════════════════════════════════════════════════════════════
    with st.expander("🏆 Champion Model", expanded=True):
        diagnostics = get_champion_diagnostics()
        
        if diagnostics.get('status') == 'loaded':
            # Identity
            identity = diagnostics.get('identity', {})
            col_a, col_b = st.columns(2)
            with col_a:
                st.metric("Generation", identity.get('generation', '?'))
                st.metric("Brain Type", identity.get('brain_type', '?'))
            with col_b:
                fitness = identity.get('fitness', 0)
                st.metric("Fitness", f"{fitness:.4f}" if isinstance(fitness, (int, float)) else fitness)
                st.metric("Timestamp", identity.get('timestamp', 'N/A'))
            
            # Quine Hash
            qhash = identity.get('quine_hash', 'unknown')
            st.code(f"Quine: {qhash}", language=None)
            
            # Architecture - FULL
            st.markdown("**🏗️ Architecture**")
            arch = diagnostics.get('architecture', {})
            st.write(f"Brain Type: {arch.get('brain_type', 'N/A')}")
            st.write(f"DreamerV3 RSSM: {'✅' if arch.get('has_rssm') else '❌'}")
            st.write(f"Dreamer Brain: {'✅' if arch.get('has_dreamer') else '❌'}")
            st.write(f"Hidden Size: {arch.get('hidden_size', 'N/A')}")
            c1, c2 = st.columns(2)
            with c1:
                st.write(f"LoRA Rank: {arch.get('lora_rank', 'N/A')}")
                st.write(f"LoRA Alpha: {arch.get('lora_alpha', 'N/A')}")
            with c2:
                st.write(f"Latent Dim: {arch.get('latent_dim', 'N/A')}")
                st.write(f"Action Dim: {arch.get('action_dim', 'N/A')}")
            
            # Capabilities - FULL
            st.markdown("**⚙️ Capabilities**")
            caps = diagnostics.get('capabilities', {})
            cap_cols = st.columns(2)
            with cap_cols[0]:
                st.write(f"CapsuleAgent: {'✅' if caps.get('CapsuleAgent') else '❌'}")
                st.write(f"forward(): {'✅' if caps.get('forward') else '❌'}")
                st.write(f"forward_hold(): {'✅' if caps.get('forward_hold') else '❌'}")
                st.write(f"clone: {'✅' if caps.get('clone') else '❌'}")
            with cap_cols[1]:
                st.write(f"export_pt: {'✅' if caps.get('export_pt') else '❌'}")
                st.write(f"export_onnx: {'✅' if caps.get('export_onnx') else '❌'}")
                st.write(f"verify_quine: {'✅' if caps.get('verify_quine') else '❌'}")
                st.write(f"bridge: {'✅' if caps.get('bridge') else '❌'}")
                st.write(f"export_interface: {'✅' if caps.get('export_interface') else '❌'}")
            
            # Evolved Traits - FULL
            traits = diagnostics.get('traits', {})
            if traits:
                st.markdown("**🧬 Evolved Traits**")
                trait_cols = st.columns(3)
                trait_items = list(traits.items())
                for i, (trait, value) in enumerate(trait_items):
                    with trait_cols[i % 3]:
                        if isinstance(value, float):
                            st.write(f"{trait}: {value:.6f}")
                        else:
                            st.write(f"{trait}: {value}")
            
            # HOLD System - FULL
            hold = diagnostics.get('hold_system')
            if hold:
                st.markdown("**🛑 HOLD System**")
                st.write(f"Enabled: {'✅' if hold.get('enabled') else '❌'}")
                st.write(f"Timeout: {hold.get('timeout', 30)}s")
                h1, h2 = st.columns(2)
                h1.metric("Hold Count", hold.get('hold_count', 0))
                h2.metric("Override Count", hold.get('override_count', 0))
                if hold.get('last_merkle'):
                    st.code(f"Last Merkle: {hold['last_merkle']}", language=None)
            
            # Provenance - FULL
            prov = diagnostics.get('provenance', {})
            if prov:
                st.markdown("**🔗 Provenance**")
                st.write(f"Genesis Root: {prov.get('genesis_root', 'unknown')}")
                if prov.get('parent_hash'):
                    st.write(f"Parent Hash: {prov['parent_hash']}")
            
            # Download info
            dl = diagnostics.get('download', {})
            if dl.get('downloaded'):
                st.success(f"✅ Downloaded from {dl.get('repo', 'HF')}")
                if dl.get('filename'):
                    st.caption(f"File: {dl['filename']}")
        else:
            st.warning(f"⚠️ {diagnostics.get('error', 'Champion not loaded')}")
            if st.button("📥 Download Champion"):
                from champion_loader import download_champion
                download_champion()
                st.rerun()
    
    # ═══════════════════════════════════════════════════════════════════════════
    # SELECTED AGENT - FULL DATA
    # ═══════════════════════════════════════════════════════════════════════════
    if st.session_state.selected_agent and st.session_state.initialized:
        agent_id = st.session_state.selected_agent
        lineage_data = st.session_state.swarm.get_agent_lineage_data()
        
        if agent_id in lineage_data:
            data = lineage_data[agent_id]
            gen = data.get('generation', 0)
            
            with st.expander(f"👤 Agent: {agent_id}", expanded=True):
                # Header
                if gen == 0:
                    st.success(f"◆ **{agent_id}** (Original)")
                else:
                    st.info(f"● **{agent_id}** (Replicant Gen {gen})")
                
                # Core stats
                c1, c2 = st.columns(2)
                c1.metric("Generation", gen)
                c2.metric("Fitness", f"{data.get('fitness', 0):.3f}")
                
                # Position
                pos = data.get('position', (0, 0, 0))
                st.write(f"**Position**: ({pos[0]:.1f}, {pos[1]:.1f}, {pos[2]:.1f})")
                
                # Lineage - FULL
                st.markdown("#### 🧬 Lineage")
                st.write(f"**Root**: {data.get('root_lineage', 'self')}")
                st.write(f"**Parent Hash**: {data.get('parent_hash', 'genesis')[:16] if data.get('parent_hash') else 'genesis'}...")
                st.write(f"**Children**: {data.get('num_children', 0)}")
                
                # Quine Hash - FULL
                qhash = data.get('quine_hash', 'unknown')
                st.code(f"Hash: {qhash}", language=None)
                
                # Decision history from cascade bridge - FULL
                if st.session_state.cascade:
                    chain = st.session_state.cascade.get_agent_chain(agent_id)
                    if chain:
                        st.markdown(f"#### 📜 Decision History ({len(chain)} decisions)")
                        for decision in chain[-10:]:  # Last 10
                            st.write(f"• {decision.get('action')} → val:{decision.get('value', 0):.3f}")
        else:
            st.warning("Agent not found in swarm")
    else:
        with st.expander("👤 Agent Inspector", expanded=False):
            st.caption("Click an agent in the graph to select it")
    
    # ═══════════════════════════════════════════════════════════════════════════
    # LINEAGES - FULL DATA
    # ═══════════════════════════════════════════════════════════════════════════
    with st.expander("🧬 All Lineages", expanded=False):
        if st.session_state.initialized:
            agents_info = get_replicated_agents_info(st.session_state.swarm)
            
            if agents_info:
                # Summary stats
                st.write(f"**Total Agents**: {len(agents_info)}")
                originals = sum(1 for a in agents_info if a['is_original'])
                replicants = len(agents_info) - originals
                st.write(f"**Originals**: {originals} | **Replicants**: {replicants}")
                
                # Group by root lineage
                lineages = {}
                for agent in agents_info:
                    root = agent['root_lineage']
                    if root not in lineages:
                        lineages[root] = []
                    lineages[root].append(agent)
                
                st.markdown("---")
                for root_id, members in lineages.items():
                    st.markdown(f"**🌳 {root_id}** ({len(members)} agents)")
                    for m in members:
                        prefix = "◆" if m['is_original'] else "●"
                        col1, col2, col3 = st.columns([2, 1, 1])
                        col1.write(f"{prefix} {m['agent_id']}")
                        col2.write(f"Gen {m['generation']}")
                        col3.write(f"Fit: {m['fitness']:.3f}")
                        if st.button(f"Inspect", key=f"inspect_{m['agent_id']}"):
                            st.session_state.selected_agent = m['agent_id']
                            st.rerun()
            else:
                st.caption("No agents in swarm")
        else:
            st.caption("Initialize simulation first")
    
    # ═══════════════════════════════════════════════════════════════════════════
    # PROVENANCE - FULL DATA
    # ═══════════════════════════════════════════════════════════════════════════
    with st.expander("🔗 Merkle Provenance", expanded=False):
        if st.session_state.initialized and st.session_state.cascade:
            receipt = st.session_state.cascade.get_session_receipt()
            
            st.metric("Session", receipt.get('session_id', 'unknown')[:20] + "...")
            st.code(f"Merkle Root: {receipt.get('merkle_root', 'N/A')}", language=None)
            
            c1, c2 = st.columns(2)
            c1.metric("Decisions", receipt.get('total_decisions', 0))
            c2.metric("Agents", receipt.get('total_agents', 0))
            
            st.write(f"**Ghost Tape Events**: {receipt.get('ghost_tape_events', 0)}")
            st.write(f"**CASCADE Mode**: {'✅' if receipt.get('cascade_mode') else '❌'}")
            
            # Causation tree - FULL
            st.markdown("#### 🌲 Causation Tree")
            tree = st.session_state.cascade.get_causation_tree()
            if tree.get('roots'):
                for root in tree['roots'][:5]:  # First 5 roots
                    st.json(root)
            else:
                st.caption("No decisions logged yet")
        else:
            st.caption("Initialize simulation to view provenance")
    
    # ═══════════════════════════════════════════════════════════════════════════
    # EVENT LOG - FULL DATA
    # ═══════════════════════════════════════════════════════════════════════════
    with st.expander("📜 Event Log", expanded=False):
        if st.session_state.events:
            st.write(f"**Total Events**: {len(st.session_state.events)}")
            st.markdown("---")
            for event in reversed(st.session_state.events[-30:]):  # Last 30
                st.write(event)
        else:
            st.caption("No events yet")
    
    # ═══════════════════════════════════════════════════════════════════════════
    # SWARM STATS - NEW SECTION
    # ═══════════════════════════════════════════════════════════════════════════
    with st.expander("📈 Swarm Statistics", expanded=False):
        if st.session_state.initialized:
            stats = st.session_state.swarm.get_swarm_stats()
            
            c1, c2, c3 = st.columns(3)
            c1.metric("Total Agents", stats.get('num_agents', 0))
            c2.metric("Max Generation", stats.get('max_generation', 0))
            c3.metric("Avg Fitness", f"{stats.get('avg_fitness', 0):.3f}")
            
            st.write(f"**Step**: {st.session_state.step_count}")
            # Use dimensions tuple instead of size attribute
            dims = st.session_state.lattice.dimensions if st.session_state.lattice else (0,0,0)
            st.write(f"**Lattice Size**: {dims[0]}×{dims[1]}×{dims[2]}")
            st.write(f"**Null Gates**: {len(st.session_state.lattice.null_gates) if st.session_state.lattice else 0}")
            
            # Agent distribution by generation
            if stats.get('generation_distribution'):
                st.markdown("**Generation Distribution**")
                for gen, count in sorted(stats['generation_distribution'].items()):
                    st.write(f"Gen {gen}: {count} agents")
        else:
            st.caption("Initialize simulation first")
    
    # ═══════════════════════════════════════════════════════════════════════════
    # GPU/SYSTEM INFO
    # ═══════════════════════════════════════════════════════════════════════════
    with st.expander("🖥️ System Info", expanded=False):
        if GPU_AVAILABLE:
            st.success(f"🚀 GPU: {GPU_NAME}")
            st.write(f"**VRAM**: {GPU_MEMORY} GB")
        else:
            st.info("💻 CPU Mode")
        
        st.write(f"**PyTorch**: {'✅' if TORCH_AVAILABLE else '❌'}")
        st.write(f"**CASCADE**: {'✅' if CASCADE_AVAILABLE else '❌'}")
        
        dl = get_download_status()
        st.write(f"**Champion Downloaded**: {'✅' if dl.get('downloaded') else '❌'}")
        if dl.get('path'):
            st.caption(f"Path: {dl['path']}")

# Auto-init if needed
if not st.session_state.initialized:
    init_simulation()

# Render plot with click selection enabled
selection = plot_placeholder.plotly_chart(
    create_plot(), 
    width='stretch', 
    key="main_plot",
    on_select="rerun",  # Enable selection events
    selection_mode="points"  # Select individual points
)

# Handle click selection - supports both agents and causal routes
if selection and hasattr(selection, 'selection') and selection.selection:
    points = selection.selection.get('points', [])
    if points:
        clicked_point = points[0]
        if 'customdata' in clicked_point and clicked_point['customdata']:
            custom = clicked_point['customdata']
            
            # Check if it's a causal route (dict with 'type' key)
            if isinstance(custom, dict) and custom.get('type') == 'causal_route':
                # Clicked on a causal route midpoint
                if custom != st.session_state.selected_route:
                    st.session_state.selected_route = custom
                    st.session_state.selected_agent = None  # Clear agent selection
                    st.rerun()
            else:
                # Clicked on an agent
                clicked_agent = custom[0] if isinstance(custom, list) else custom
                if clicked_agent != st.session_state.selected_agent:
                    st.session_state.selected_agent = clicked_agent
                    st.session_state.selected_route = None  # Clear route selection
                    st.rerun()

# ═══════════════════════════════════════════════════════════════════════════════
# CAUSAL ROUTE DRILLDOWN PANEL
# ═══════════════════════════════════════════════════════════════════════════════
if st.session_state.selected_route:
    route = st.session_state.selected_route
    st.markdown("---")
    st.subheader("🔗 Causal Route Details")
    
    col1, col2, col3 = st.columns([1, 2, 1])
    
    with col1:
        st.markdown("**FROM**")
        st.code(route.get('parent_merkle', 'N/A'), language=None)
        st.caption(route.get('from_node', ''))
    
    with col2:
        st.markdown(f"**→ Action: `{route.get('action', 'N/A')}` →**")
        st.metric("Value Estimate", f"{route.get('value', 0):.4f}")
        
        # Show action probabilities
        probs = route.get('probs', {})
        if probs:
            st.markdown("**Action Probabilities:**")
            for action, prob in probs.items():
                bar_len = int(prob * 30)
                bar = "█" * bar_len + "░" * (30 - bar_len)
                st.text(f"{action}: {bar} {prob:.2%}")
    
    with col3:
        st.markdown("**TO**")
        st.code(route.get('merkle', 'N/A'), language=None)
        st.caption(route.get('to_node', ''))
    
    # Full chain trace
    if st.session_state.cascade:
        st.markdown("**Full Decision Chain:**")
        chain = st.session_state.cascade.get_agent_chain(route.get('agent_id', ''))
        if chain:
            chain_str = " → ".join([f"`{n.get('action', '?')}`" for n in chain[-10:]])
            st.markdown(chain_str)
            st.caption(f"Showing last {min(len(chain), 10)} of {len(chain)} decisions")
    
    if st.button("❌ Close Route Details"):
        st.session_state.selected_route = None
        st.rerun()

# Animation loop
if play:
    for _ in range(200):  # Max 200 frames then stops
        step()
        plot_placeholder.plotly_chart(create_plot(), width='stretch', key=f"anim_{st.session_state.step_count}")
        time.sleep(1.0 / speed)