src/grid_dashboard.py

"""
Grid Dashboard - Interactive Resizable Panel System

🔮 GLASS BOX AI - Brady Bunch style grid with:
- Equilateral resizable cells
- Drag-to-swap with 3D connection visualization
- Edge-click resizing affecting adjacent cells
- Pre-emptive swap preview on hover

Uses streamlit-elements for the interactive grid.
"""

import streamlit as st
from streamlit_elements import elements, mui, html, dashboard, lazy, sync

# ═══════════════════════════════════════════════════════════════════════════════
# GRID LAYOUT CONFIGURATION
# ═══════════════════════════════════════════════════════════════════════════════

# Default grid layout - 3x3 Brady Bunch style
# Each panel has: x, y position, width, height
DEFAULT_LAYOUT = [
    # Row 1
    dashboard.Item("3d_graph", 0, 0, 4, 3, isDraggable=True, isResizable=True),
    dashboard.Item("champion", 4, 0, 2, 2, isDraggable=True, isResizable=True),
    dashboard.Item("agent", 6, 0, 2, 2, isDraggable=True, isResizable=True),
    # Row 2  
    dashboard.Item("lineages", 4, 2, 2, 2, isDraggable=True, isResizable=True),
    dashboard.Item("provenance", 6, 2, 2, 2, isDraggable=True, isResizable=True),
    # Row 3
    dashboard.Item("events", 0, 3, 2, 2, isDraggable=True, isResizable=True),
    dashboard.Item("stats", 2, 3, 2, 2, isDraggable=True, isResizable=True),
    dashboard.Item("traits", 4, 3, 2, 2, isDraggable=True, isResizable=True),
    dashboard.Item("system", 6, 3, 2, 2, isDraggable=True, isResizable=True),
]


def init_grid_state():
    """Initialize grid layout in session state."""
    if 'grid_layout' not in st.session_state:
        st.session_state.grid_layout = DEFAULT_LAYOUT
    if 'grid_maximized' not in st.session_state:
        st.session_state.grid_maximized = None  # Which panel is maximized


def handle_layout_change(updated_layout):
    """Handle layout changes from drag/resize."""
    st.session_state.grid_layout = updated_layout


def render_panel_header(title: str, icon: str, panel_id: str):
    """Render a panel header with maximize button."""
    return mui.Box(
        sx={
            "display": "flex",
            "justifyContent": "space-between",
            "alignItems": "center",
            "borderBottom": "1px solid rgba(255,255,255,0.1)",
            "pb": 1,
            "mb": 1,
            "cursor": "grab",
        },
        children=[
            mui.Typography(
                f"{icon} {title}",
                sx={"fontWeight": "bold", "fontSize": "0.9rem"}
            ),
            mui.IconButton(
                mui.icon.Fullscreen,
                size="small",
                onClick=lambda: toggle_maximize(panel_id),
                sx={"color": "rgba(255,255,255,0.5)"}
            )
        ]
    )


def toggle_maximize(panel_id: str):
    """Toggle maximized state for a panel."""
    if st.session_state.grid_maximized == panel_id:
        st.session_state.grid_maximized = None
    else:
        st.session_state.grid_maximized = panel_id


# ═══════════════════════════════════════════════════════════════════════════════
# PANEL CONTENT RENDERERS
# ═══════════════════════════════════════════════════════════════════════════════

def render_3d_graph_panel():
    """Render the 3D lattice visualization panel."""
    with mui.Card(
        sx={
            "height": "100%",
            "display": "flex",
            "flexDirection": "column",
            "bgcolor": "rgba(15,8,20,0.9)",
            "border": "1px solid rgba(120,0,180,0.3)",
        }
    ):
        render_panel_header("3D Lattice", "🌐", "3d_graph")
        mui.CardContent(
            sx={"flex": 1, "p": 1},
            children=[
                # Plotly chart will be embedded here
                html.Div("3D Graph Placeholder - Click agents to select")
            ]
        )


def render_champion_panel(diagnostics: dict):
    """Render champion model diagnostics panel."""
    with mui.Card(
        sx={
            "height": "100%",
            "display": "flex",
            "flexDirection": "column",
            "bgcolor": "rgba(15,8,20,0.9)",
            "border": "1px solid rgba(255,215,0,0.3)",
            "overflow": "auto",
        }
    ):
        render_panel_header("Champion", "🏆", "champion")
        
        if diagnostics.get('status') == 'loaded':
            identity = diagnostics.get('identity', {})
            arch = diagnostics.get('architecture', {})
            
            mui.CardContent(
                sx={"flex": 1, "p": 1, "fontSize": "0.75rem"},
                children=[
                    mui.Typography(f"Gen: {identity.get('generation', '?')}", variant="body2"),
                    mui.Typography(f"Brain: {identity.get('brain_type', '?')}", variant="body2"),
                    mui.Typography(f"Fitness: {identity.get('fitness', 0):.4f}", variant="body2"),
                    mui.Divider(sx={"my": 0.5}),
                    mui.Typography(f"Hidden: {arch.get('hidden_size', 'N/A')}", variant="caption"),
                    mui.Typography(f"LoRA: {arch.get('lora_rank', 'N/A')}", variant="caption"),
                    mui.Typography(f"RSSM: {'✅' if arch.get('has_rssm') else '❌'}", variant="caption"),
                    mui.Divider(sx={"my": 0.5}),
                    mui.Chip(
                        label=diagnostics.get('download', {}).get('repo', 'HF'),
                        size="small",
                        color="success",
                        sx={"fontSize": "0.65rem"}
                    )
                ]
            )
        else:
            mui.CardContent(children=[
                mui.Typography("Champion not loaded", color="warning.main")
            ])


def render_agent_panel(agent_data: dict, agent_id: str):
    """Render selected agent panel."""
    with mui.Card(
        sx={
            "height": "100%",
            "display": "flex",
            "flexDirection": "column",
            "bgcolor": "rgba(15,8,20,0.9)",
            "border": "1px solid rgba(0,200,255,0.3)",
            "overflow": "auto",
        }
    ):
        render_panel_header("Agent", "👤", "agent")
        
        if agent_data:
            gen = agent_data.get('generation', 0)
            prefix = "◆" if gen == 0 else "●"
            
            mui.CardContent(
                sx={"flex": 1, "p": 1, "fontSize": "0.75rem"},
                children=[
                    mui.Typography(
                        f"{prefix} {agent_id[:16]}...",
                        variant="body2",
                        sx={"fontWeight": "bold"}
                    ),
                    mui.Typography(f"Gen: {gen}", variant="body2"),
                    mui.Typography(f"Fitness: {agent_data.get('fitness', 0):.3f}", variant="body2"),
                    mui.Typography(f"Children: {agent_data.get('num_children', 0)}", variant="body2"),
                    mui.Divider(sx={"my": 0.5}),
                    mui.Typography(f"Root: {agent_data.get('root_lineage', 'self')[:12]}...", variant="caption"),
                    mui.Typography(f"Hash: {agent_data.get('quine_hash', '?')[:16]}...", variant="caption"),
                ]
            )
        else:
            mui.CardContent(children=[
                mui.Typography("Click agent in graph", variant="caption", color="text.secondary")
            ])


def render_lineages_panel(lineages: dict):
    """Render lineages overview panel."""
    with mui.Card(
        sx={
            "height": "100%",
            "display": "flex",
            "flexDirection": "column",
            "bgcolor": "rgba(15,8,20,0.9)",
            "border": "1px solid rgba(0,255,100,0.3)",
            "overflow": "auto",
        }
    ):
        render_panel_header("Lineages", "🧬", "lineages")
        
        mui.CardContent(
            sx={"flex": 1, "p": 1, "fontSize": "0.7rem", "overflow": "auto"},
            children=[
                *[
                    mui.Box(
                        children=[
                            mui.Typography(
                                f"🌳 {root_id[:10]}... ({len(members)})",
                                variant="caption",
                                sx={"fontWeight": "bold"}
                            ),
                            *[
                                mui.Typography(
                                    f"{'◆' if m.get('is_original') else '●'} Gen{m.get('generation', 0)}",
                                    variant="caption",
                                    sx={"ml": 1}
                                )
                                for m in members[:3]
                            ]
                        ]
                    )
                    for root_id, members in list(lineages.items())[:5]
                ] if lineages else [mui.Typography("No lineages", variant="caption")]
            ]
        )


def render_provenance_panel(receipt: dict):
    """Render Merkle provenance panel."""
    with mui.Card(
        sx={
            "height": "100%",
            "display": "flex",
            "flexDirection": "column",
            "bgcolor": "rgba(15,8,20,0.9)",
            "border": "1px solid rgba(255,100,0,0.3)",
            "overflow": "auto",
        }
    ):
        render_panel_header("Provenance", "🔗", "provenance")
        
        mui.CardContent(
            sx={"flex": 1, "p": 1, "fontSize": "0.75rem"},
            children=[
                mui.Typography(f"Decisions: {receipt.get('total_decisions', 0)}", variant="body2"),
                mui.Typography(f"Agents: {receipt.get('total_agents', 0)}", variant="body2"),
                mui.Typography(f"Ghost Tape: {receipt.get('ghost_tape_events', 0)}", variant="body2"),
                mui.Divider(sx={"my": 0.5}),
                mui.Typography(
                    f"Merkle: {receipt.get('merkle_root', 'N/A')[:20]}...",
                    variant="caption",
                    sx={"fontFamily": "monospace"}
                ),
            ]
        )


def render_events_panel(events: list):
    """Render event log panel."""
    with mui.Card(
        sx={
            "height": "100%",
            "display": "flex",
            "flexDirection": "column",
            "bgcolor": "rgba(15,8,20,0.9)",
            "border": "1px solid rgba(200,200,200,0.3)",
            "overflow": "auto",
        }
    ):
        render_panel_header("Events", "📜", "events")
        
        mui.CardContent(
            sx={"flex": 1, "p": 1, "fontSize": "0.65rem", "overflow": "auto"},
            children=[
                *[
                    mui.Typography(event, variant="caption", sx={"display": "block"})
                    for event in reversed(events[-10:])
                ] if events else [mui.Typography("No events", variant="caption")]
            ]
        )


def render_stats_panel(stats: dict, step_count: int):
    """Render swarm statistics panel."""
    with mui.Card(
        sx={
            "height": "100%",
            "display": "flex",
            "flexDirection": "column",
            "bgcolor": "rgba(15,8,20,0.9)",
            "border": "1px solid rgba(100,100,255,0.3)",
        }
    ):
        render_panel_header("Statistics", "📈", "stats")
        
        mui.CardContent(
            sx={"flex": 1, "p": 1, "fontSize": "0.75rem"},
            children=[
                mui.Typography(f"Step: {step_count}", variant="body2"),
                mui.Typography(f"Agents: {stats.get('num_agents', 0)}", variant="body2"),
                mui.Typography(f"Max Gen: {stats.get('max_generation', 0)}", variant="body2"),
                mui.Typography(f"Avg Fit: {stats.get('avg_fitness', 0):.3f}", variant="body2"),
            ]
        )


def render_traits_panel(traits: dict):
    """Render evolved traits panel."""
    with mui.Card(
        sx={
            "height": "100%",
            "display": "flex",
            "flexDirection": "column",
            "bgcolor": "rgba(15,8,20,0.9)",
            "border": "1px solid rgba(255,0,200,0.3)",
            "overflow": "auto",
        }
    ):
        render_panel_header("Traits", "🧬", "traits")
        
        mui.CardContent(
            sx={"flex": 1, "p": 1, "fontSize": "0.65rem", "overflow": "auto"},
            children=[
                *[
                    mui.Typography(
                        f"{k}: {v:.4f}" if isinstance(v, float) else f"{k}: {v}",
                        variant="caption",
                        sx={"display": "block"}
                    )
                    for k, v in list(traits.items())[:8]
                ] if traits else [mui.Typography("No traits", variant="caption")]
            ]
        )


def render_system_panel(gpu_available: bool, gpu_name: str, gpu_memory: int):
    """Render system info panel."""
    with mui.Card(
        sx={
            "height": "100%",
            "display": "flex",
            "flexDirection": "column",
            "bgcolor": "rgba(15,8,20,0.9)",
            "border": "1px solid rgba(50,200,50,0.3)",
        }
    ):
        render_panel_header("System", "🖥️", "system")
        
        mui.CardContent(
            sx={"flex": 1, "p": 1, "fontSize": "0.75rem"},
            children=[
                mui.Chip(
                    label=f"🚀 {gpu_name[:20]}" if gpu_available else "💻 CPU",
                    size="small",
                    color="success" if gpu_available else "default",
                    sx={"fontSize": "0.6rem", "mb": 1}
                ),
                mui.Typography(f"VRAM: {gpu_memory}GB" if gpu_available else "No GPU", variant="caption"),
            ]
        )


# ═══════════════════════════════════════════════════════════════════════════════
# MAIN GRID RENDERER
# ═══════════════════════════════════════════════════════════════════════════════

def render_grid_dashboard(
    diagnostics: dict,
    agent_data: dict,
    agent_id: str,
    lineages: dict,
    receipt: dict,
    events: list,
    stats: dict,
    step_count: int,
    traits: dict,
    gpu_available: bool,
    gpu_name: str,
    gpu_memory: int,
    plot_figure=None
):
    """
    Render the full interactive grid dashboard.
    
    All panels are equilateral, resizable, and swappable.
    """
    init_grid_state()
    
    # Custom CSS for drag visualization
    custom_css = """
    <style>
    .react-grid-item.react-draggable-dragging {
        box-shadow: 0 0 30px rgba(120, 0, 180, 0.8);
        z-index: 1000;
        opacity: 0.9;
    }
    .react-grid-item.react-grid-placeholder {
        background: rgba(120, 0, 180, 0.3) !important;
        border: 2px dashed rgba(255, 255, 255, 0.5);
        border-radius: 8px;
    }
    .react-grid-item > .react-resizable-handle {
        background: rgba(120, 0, 180, 0.5);
    }
    .react-grid-item:hover {
        box-shadow: 0 0 15px rgba(120, 0, 180, 0.4);
    }
    </style>
    """
    st.markdown(custom_css, unsafe_allow_html=True)
    
    with elements("grid_dashboard"):
        # Dashboard grid with drag-drop-resize
        with dashboard.Grid(
            st.session_state.grid_layout,
            onLayoutChange=handle_layout_change,
            cols=8,
            rowHeight=80,
            draggableHandle=".drag-handle",
            compactType="vertical",
            preventCollision=False,
            margin=[8, 8],
        ):
            # 3D Graph Panel
            with mui.Box(key="3d_graph", sx={"height": "100%"}):
                render_3d_graph_panel()
            
            # Champion Panel
            with mui.Box(key="champion", sx={"height": "100%"}):
                render_champion_panel(diagnostics)
            
            # Agent Panel
            with mui.Box(key="agent", sx={"height": "100%"}):
                render_agent_panel(agent_data, agent_id)
            
            # Lineages Panel
            with mui.Box(key="lineages", sx={"height": "100%"}):
                render_lineages_panel(lineages)
            
            # Provenance Panel
            with mui.Box(key="provenance", sx={"height": "100%"}):
                render_provenance_panel(receipt)
            
            # Events Panel
            with mui.Box(key="events", sx={"height": "100%"}):
                render_events_panel(events)
            
            # Stats Panel
            with mui.Box(key="stats", sx={"height": "100%"}):
                render_stats_panel(stats, step_count)
            
            # Traits Panel
            with mui.Box(key="traits", sx={"height": "100%"}):
                render_traits_panel(traits)
            
            # System Panel
            with mui.Box(key="system", sx={"height": "100%"}):
                render_system_panel(gpu_available, gpu_name, gpu_memory)