src/app_dash.py

"""
CASCADE-LATTICE Chess - Dash Version
====================================
Proper callback-based visualization that doesn't shit itself.
"""

import dash
from dash import dcc, html, callback, Input, Output, State
import plotly.graph_objects as go
import chess
import chess.engine
import numpy as np
import asyncio
import platform
from pathlib import Path
import shutil
from dataclasses import dataclass
from typing import List, Optional

# ═══════════════════════════════════════════════════════════════════════════════
# STOCKFISH SETUP
# ═══════════════════════════════════════════════════════════════════════════════

PROJECT_ROOT = Path(__file__).parent.parent.resolve()
LOCAL_STOCKFISH = PROJECT_ROOT / "stockfish" / "stockfish-windows-x86-64-avx2.exe"

if LOCAL_STOCKFISH.exists():
    STOCKFISH_PATH = str(LOCAL_STOCKFISH)
elif shutil.which("stockfish"):
    STOCKFISH_PATH = shutil.which("stockfish")
else:
    STOCKFISH_PATH = "/usr/games/stockfish"

print(f"[STOCKFISH] {STOCKFISH_PATH}")

# Fix Windows asyncio for chess engine
if platform.system() == 'Windows':
    asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())

# Global engine (initialized once)
ENGINE = None
try:
    ENGINE = chess.engine.SimpleEngine.popen_uci(STOCKFISH_PATH)
    print("[ENGINE] Loaded OK")
except Exception as e:
    print(f"[ENGINE] Failed: {e}")

# ═══════════════════════════════════════════════════════════════════════════════
# CASCADE-LATTICE
# ═══════════════════════════════════════════════════════════════════════════════

CASCADE_AVAILABLE = False
HOLD = None
CAUSATION = None
TRACER = None

try:
    from cascade import Hold
    HOLD = Hold()
    CASCADE_AVAILABLE = True
    print("[CASCADE] Hold ready")
except ImportError:
    print("[CASCADE] Hold not available")

try:
    from cascade import CausationGraph
    CAUSATION = CausationGraph()
    print("[CASCADE] CausationGraph ready")
except:
    pass

try:
    from cascade import Tracer
    if CAUSATION:
        TRACER = Tracer(CAUSATION)
        print("[CASCADE] Tracer ready")
except:
    pass

# ═══════════════════════════════════════════════════════════════════════════════
# VISUAL THEME
# ═══════════════════════════════════════════════════════════════════════════════

# Board - rich wood tones
BOARD_LIGHT = '#D4A574'   # Maple
BOARD_DARK = '#8B5A2B'    # Walnut
BOARD_EDGE = '#4A3520'    # Dark frame

# Pieces - polished look  
WHITE_PIECE = '#FFFEF0'   # Ivory
WHITE_SHADOW = '#C0B090'  # Ivory shadow
BLACK_PIECE = '#2A2A2A'   # Ebony
BLACK_SHADOW = '#151515'  # Ebony shadow

# Accents
GOLD = '#FFD700'
CRIMSON = '#DC143C'
CYAN = '#00FFD4'
MAGENTA = '#FF00AA'

BG_COLOR = '#0a0a0f'
BG_GRADIENT = '#12121a'

# Piece geometry - height and multiple size layers for 3D effect
PIECE_CONFIG = {
    chess.PAWN:   {'h': 0.5, 'base': 8,  'mid': 6,  'top': 4,  'symbol': '♟'},
    chess.KNIGHT: {'h': 0.8, 'base': 10, 'mid': 8,  'top': 6,  'symbol': '♞'},
    chess.BISHOP: {'h': 0.9, 'base': 10, 'mid': 7,  'top': 5,  'symbol': '♝'},
    chess.ROOK:   {'h': 0.7, 'base': 10, 'mid': 9,  'top': 8,  'symbol': '♜'},
    chess.QUEEN:  {'h': 1.1, 'base': 12, 'mid': 9,  'top': 6,  'symbol': '♛'},
    chess.KING:   {'h': 1.2, 'base': 12, 'mid': 8,  'top': 5,  'symbol': '♚'},
}

# ═══════════════════════════════════════════════════════════════════════════════
# DATA
# ═══════════════════════════════════════════════════════════════════════════════

@dataclass
class MoveCandidate:
    move: str
    prob: float
    value: float
    from_sq: int
    to_sq: int
    is_capture: bool = False
    is_check: bool = False
    move_type: str = 'quiet'

@dataclass
class CascadeTrace:
    """Represents a cascade-lattice inference trace."""
    step: int
    operation: str
    inputs: List[str]
    output: str
    duration_ms: float
    confidence: float
    
# Global trace storage
TRACE_LOG: List[dict] = []
DECISION_TREE: List[dict] = []

# ═══════════════════════════════════════════════════════════════════════════════
# 3D VISUALIZATION
# ═══════════════════════════════════════════════════════════════════════════════

def square_to_3d(square: int):
    file = square % 8
    rank = square // 8
    return (file - 3.5, rank - 3.5, 0)

def create_board_traces():
    """Create 3D board with mesh tiles."""
    traces = []
    tile_size = 0.48
    tile_height = 0.08
    
    for rank in range(8):
        for file in range(8):
            cx, cy = file - 3.5, rank - 3.5
            is_light = (rank + file) % 2 == 1
            color = BOARD_LIGHT if is_light else BOARD_DARK
            
            # Top surface of tile
            traces.append(go.Mesh3d(
                x=[cx-tile_size, cx+tile_size, cx+tile_size, cx-tile_size],
                y=[cy-tile_size, cy-tile_size, cy+tile_size, cy+tile_size],
                z=[tile_height, tile_height, tile_height, tile_height],
                i=[0, 0], j=[1, 2], k=[2, 3],
                color=color, opacity=0.95, flatshading=True,
                hoverinfo='skip', showlegend=False
            ))
    
    # Board frame/border
    border = 4.3
    frame_z = tile_height / 2
    traces.append(go.Scatter3d(
        x=[-border, border, border, -border, -border],
        y=[-border, -border, border, border, -border],
        z=[frame_z]*5, mode='lines',
        line=dict(color=GOLD, width=3),
        hoverinfo='skip', showlegend=False
    ))
    
    # File labels (a-h)
    for i, label in enumerate('abcdefgh'):
        traces.append(go.Scatter3d(
            x=[i-3.5], y=[-4.6], z=[0.1], mode='text', text=[label],
            textfont=dict(size=11, color='#666'), showlegend=False
        ))
    # Rank labels (1-8)
    for i in range(8):
        traces.append(go.Scatter3d(
            x=[-4.6], y=[i-3.5], z=[0.1], mode='text', text=[str(i+1)],
            textfont=dict(size=11, color='#666'), showlegend=False
        ))
    
    return traces

def create_piece_traces(board: chess.Board):
    """Create holographic crystal-style chess pieces using rotated point silhouettes."""
    traces = []
    
    # Define piece silhouettes as point patterns (normalized 0-1 range)
    # These trace the outline/shape of each piece type
    PIECE_SILHOUETTES = {
        chess.PAWN: [
            (0.5, 0.0), (0.3, 0.1), (0.25, 0.2), (0.3, 0.3), (0.35, 0.4),
            (0.4, 0.5), (0.35, 0.6), (0.3, 0.7), (0.4, 0.8), (0.5, 0.9),
            (0.6, 0.8), (0.7, 0.7), (0.65, 0.6), (0.6, 0.5), (0.65, 0.4),
            (0.7, 0.3), (0.75, 0.2), (0.7, 0.1), (0.5, 0.0)
        ],
        chess.ROOK: [
            (0.25, 0.0), (0.25, 0.15), (0.35, 0.15), (0.35, 0.25), (0.25, 0.25),
            (0.25, 0.85), (0.35, 0.9), (0.35, 1.0), (0.45, 1.0), (0.45, 0.9),
            (0.55, 0.9), (0.55, 1.0), (0.65, 1.0), (0.65, 0.9), (0.75, 0.85),
            (0.75, 0.25), (0.65, 0.25), (0.65, 0.15), (0.75, 0.15), (0.75, 0.0)
        ],
        chess.KNIGHT: [
            (0.25, 0.0), (0.25, 0.2), (0.3, 0.35), (0.25, 0.5), (0.3, 0.6),
            (0.35, 0.7), (0.3, 0.8), (0.4, 0.9), (0.5, 1.0), (0.6, 0.95),
            (0.7, 0.85), (0.75, 0.7), (0.7, 0.55), (0.65, 0.4), (0.7, 0.25),
            (0.75, 0.15), (0.75, 0.0)
        ],
        chess.BISHOP: [
            (0.3, 0.0), (0.25, 0.15), (0.3, 0.3), (0.35, 0.45), (0.4, 0.6),
            (0.35, 0.7), (0.4, 0.8), (0.5, 0.95), (0.6, 0.8), (0.65, 0.7),
            (0.6, 0.6), (0.65, 0.45), (0.7, 0.3), (0.75, 0.15), (0.7, 0.0)
        ],
        chess.QUEEN: [
            (0.25, 0.0), (0.2, 0.15), (0.25, 0.3), (0.3, 0.5), (0.25, 0.65),
            (0.3, 0.75), (0.2, 0.85), (0.35, 0.9), (0.5, 1.0), (0.65, 0.9),
            (0.8, 0.85), (0.7, 0.75), (0.75, 0.65), (0.7, 0.5), (0.75, 0.3),
            (0.8, 0.15), (0.75, 0.0)
        ],
        chess.KING: [
            (0.25, 0.0), (0.2, 0.15), (0.25, 0.35), (0.3, 0.55), (0.35, 0.7),
            (0.3, 0.8), (0.45, 0.85), (0.45, 0.92), (0.4, 0.92), (0.4, 0.97),
            (0.5, 1.0), (0.6, 0.97), (0.6, 0.92), (0.55, 0.92), (0.55, 0.85),
            (0.7, 0.8), (0.65, 0.7), (0.7, 0.55), (0.75, 0.35), (0.8, 0.15), (0.75, 0.0)
        ],
    }
    
    PIECE_HEIGHT = {
        chess.PAWN: 0.4,
        chess.KNIGHT: 0.55,
        chess.BISHOP: 0.55,
        chess.ROOK: 0.5,
        chess.QUEEN: 0.65,
        chess.KING: 0.7,
    }
    
    NUM_ROTATIONS = 6  # How many rotated planes per piece
    
    for sq in chess.SQUARES:
        piece = board.piece_at(sq)
        if not piece:
            continue
            
        cx, cy, _ = square_to_3d(sq)
        silhouette = PIECE_SILHOUETTES[piece.piece_type]
        height = PIECE_HEIGHT[piece.piece_type]
        color = '#E8E8E8' if piece.color == chess.WHITE else '#404040'
        
        # Create rotated copies of the silhouette
        for rot in range(NUM_ROTATIONS):
            angle = (rot / NUM_ROTATIONS) * np.pi  # 0 to 180 degrees
            
            xs, ys, zs = [], [], []
            for px, pz in silhouette:
                # Offset from center (-0.5 to +0.5 range)
                local_x = (px - 0.5) * 0.4  # Scale to fit in square
                local_z = pz * height + 0.1  # Height above board
                
                # Rotate around Y (vertical) axis
                rotated_x = local_x * np.cos(angle)
                rotated_y = local_x * np.sin(angle)
                
                xs.append(cx + rotated_x)
                ys.append(cy + rotated_y)
                zs.append(local_z)
            
            traces.append(go.Scatter3d(
                x=xs, y=ys, z=zs,
                mode='lines',
                line=dict(color=color, width=2),
                opacity=0.7,
                hoverinfo='skip',
                showlegend=False
            ))
    
    return traces
    
    return traces

def create_arc(from_sq, to_sq, is_white, num_points=40):
    """Create smooth arc trajectory with proper curve."""
    x1, y1, _ = square_to_3d(from_sq)
    x2, y2, _ = square_to_3d(to_sq)
    dist = np.sqrt((x2-x1)**2 + (y2-y1)**2)
    # Cap height to stay within z bounds - short moves get small arcs, long moves capped
    height = min(1.8, max(0.8, dist * 0.35))
    
    t = np.linspace(0, 1, num_points)
    x = x1 + (x2 - x1) * t
    y = y1 + (y2 - y1) * t
    # Arc always goes UP (positive z) - offset from board surface
    z = 0.3 + height * np.sin(np.pi * t)
    return x, y, z

def create_move_arcs(candidates: List[MoveCandidate], is_white: bool):
    """Create glowing move arc traces."""
    traces = []
    
    # Color gradient based on rank
    if is_white:
        colors = [CYAN, '#00CCAA', '#009988', '#006666', '#004444']
    else:
        colors = [MAGENTA, '#CC0088', '#990066', '#660044', '#440022']
    
    for i, cand in enumerate(candidates):
        x, y, z = create_arc(cand.from_sq, cand.to_sq, is_white)
        color = colors[min(i, len(colors)-1)]
        width = max(4, cand.prob * 15)
        opacity = max(0.5, cand.prob * 1.2)
        
        # Outer glow (wide, transparent)
        traces.append(go.Scatter3d(
            x=x, y=y, z=z, mode='lines',
            line=dict(color=color, width=width+6),
            opacity=opacity*0.2, hoverinfo='skip', showlegend=False
        ))
        # Mid glow
        traces.append(go.Scatter3d(
            x=x, y=y, z=z, mode='lines',
            line=dict(color=color, width=width+2),
            opacity=opacity*0.4, hoverinfo='skip', showlegend=False
        ))
        # Core line
        traces.append(go.Scatter3d(
            x=x, y=y, z=z, mode='lines',
            line=dict(color=color, width=width),
            opacity=opacity,
            name=f"{cand.move} ({cand.prob*100:.0f}%)",
            hovertemplate=f"<b>{cand.move}</b><br>Probability: {cand.prob*100:.1f}%<br>Eval: {cand.value:+.2f}<extra></extra>"
        ))
        
        # Start point (piece origin)
        traces.append(go.Scatter3d(
            x=[x[0]], y=[y[0]], z=[z[0]], mode='markers',
            marker=dict(size=5, color=color, symbol='circle', opacity=opacity),
            hoverinfo='skip', showlegend=False
        ))
        
        # End point (destination) - larger, prominent
        traces.append(go.Scatter3d(
            x=[x[-1]], y=[y[-1]], z=[z[-1]], mode='markers',
            marker=dict(size=8, color=color, symbol='diamond',
                       line=dict(color='white', width=1)),
            hoverinfo='skip', showlegend=False
        ))
        
        # Drop line to board (shows landing square)
        x2, y2, _ = square_to_3d(cand.to_sq)
        traces.append(go.Scatter3d(
            x=[x2, x2], y=[y2, y2], z=[z[-1], 0.15],
            mode='lines', line=dict(color=color, width=2, dash='dot'),
            opacity=opacity*0.5, hoverinfo='skip', showlegend=False
        ))
    
    return traces

# ═══════════════════════════════════════════════════════════════════════════════
# TACTICAL VISUALIZATION
# ═══════════════════════════════════════════════════════════════════════════════

def get_piece_attacks(board: chess.Board, square: int) -> List[int]:
    """Get squares attacked by piece at given square."""
    piece = board.piece_at(square)
    if not piece:
        return []
    return list(board.attacks(square))

def get_attackers(board: chess.Board, square: int, color: bool) -> List[int]:
    """Get pieces of given color attacking a square."""
    return list(board.attackers(color, square))

def create_threat_beams(board: chess.Board):
    """Create downward beams showing attack vectors - underneath the board."""
    traces = []
    
    for sq in chess.SQUARES:
        piece = board.piece_at(sq)
        if not piece:
            continue
        
        attacks = get_piece_attacks(board, sq)
        if not attacks:
            continue
        
        px, py, _ = square_to_3d(sq)
        color = 'rgba(255,215,0,0.15)' if piece.color == chess.WHITE else 'rgba(220,20,60,0.15)'
        
        for target_sq in attacks:
            tx, ty, _ = square_to_3d(target_sq)
            # Beam goes from piece position DOWN through board to show threat zone
            traces.append(go.Scatter3d(
                x=[px, tx], y=[py, ty], z=[-0.05, -0.2],
                mode='lines', line=dict(color=color, width=1),
                hoverinfo='skip', showlegend=False
            ))
    
    return traces

def create_tension_lines(board: chess.Board):
    """Create lines between pieces that threaten each other (mutual tension)."""
    traces = []
    tension_pairs = set()
    
    for sq in chess.SQUARES:
        piece = board.piece_at(sq)
        if not piece:
            continue
        
        # Check if this piece attacks any enemy pieces
        for target_sq in board.attacks(sq):
            target_piece = board.piece_at(target_sq)
            if target_piece and target_piece.color != piece.color:
                # Create unique pair key
                pair = tuple(sorted([sq, target_sq]))
                if pair not in tension_pairs:
                    tension_pairs.add(pair)
                    
                    x1, y1, _ = square_to_3d(sq)
                    x2, y2, _ = square_to_3d(target_sq)
                    
                    # Check if it's mutual (both threaten each other)
                    mutual = target_sq in board.attacks(sq) and sq in board.attacks(target_sq)
                    
                    if mutual:
                        color = MAGENTA
                        width = 3
                        opacity = 0.6
                    else:
                        color = CRIMSON if piece.color == chess.BLACK else GOLD
                        width = 2
                        opacity = 0.4
                    
                    # Tension line slightly above board
                    traces.append(go.Scatter3d(
                        x=[x1, x2], y=[y1, y2], z=[0.8, 0.8],
                        mode='lines', line=dict(color=color, width=width),
                        opacity=opacity, hoverinfo='skip', showlegend=False
                    ))
    
    return traces

def create_hanging_indicators(board: chess.Board):
    """Mark pieces that are hanging (attacked but not defended)."""
    traces = []
    
    for sq in chess.SQUARES:
        piece = board.piece_at(sq)
        if not piece:
            continue
        
        # Count attackers and defenders
        enemy_color = not piece.color
        attackers = len(board.attackers(enemy_color, sq))
        defenders = len(board.attackers(piece.color, sq))
        
        if attackers > 0 and defenders == 0:
            # Hanging! Draw danger indicator
            x, y, _ = square_to_3d(sq)
            
            # Pulsing ring under the piece
            theta = np.linspace(0, 2*np.pi, 20)
            ring_x = x + 0.3 * np.cos(theta)
            ring_y = y + 0.3 * np.sin(theta)
            ring_z = [0.05] * len(theta)
            
            traces.append(go.Scatter3d(
                x=ring_x, y=ring_y, z=ring_z,
                mode='lines', line=dict(color=CRIMSON, width=4),
                opacity=0.8, hoverinfo='skip', showlegend=False
            ))
            
            # Exclamation point above
            traces.append(go.Scatter3d(
                x=[x], y=[y], z=[1.0],
                mode='text', text=['⚠'],
                textfont=dict(size=16, color=CRIMSON),
                hoverinfo='skip', showlegend=False
            ))
    
    return traces

def create_king_safety(board: chess.Board):
    """Create safety dome around kings showing their safety status."""
    traces = []
    
    for color in [chess.WHITE, chess.BLACK]:
        king_sq = board.king(color)
        if king_sq is None:
            continue
        
        kx, ky, _ = square_to_3d(king_sq)
        
        # Count attackers around king zone
        danger_level = 0
        for sq in chess.SQUARES:
            file_dist = abs(chess.square_file(sq) - chess.square_file(king_sq))
            rank_dist = abs(chess.square_rank(sq) - chess.square_rank(king_sq))
            if file_dist <= 1 and rank_dist <= 1:  # King zone
                enemy_attackers = len(board.attackers(not color, sq))
                danger_level += enemy_attackers
        
        # Dome color based on safety
        if board.is_check() and board.turn == color:
            dome_color = CRIMSON
            opacity = 0.5
        elif danger_level > 5:
            dome_color = '#FF6600'  # Orange - danger
            opacity = 0.3
        elif danger_level > 2:
            dome_color = GOLD  # Yellow - caution
            opacity = 0.2
        else:
            dome_color = '#00FF00'  # Green - safe
            opacity = 0.15
        
        # Draw dome as circles at different heights
        for h in [0.3, 0.5, 0.7]:
            radius = 0.6 * (1 - h/1.5)  # Smaller at top
            theta = np.linspace(0, 2*np.pi, 24)
            dome_x = kx + radius * np.cos(theta)
            dome_y = ky + radius * np.sin(theta)
            dome_z = [h] * len(theta)
            
            traces.append(go.Scatter3d(
                x=dome_x, y=dome_y, z=dome_z,
                mode='lines', line=dict(color=dome_color, width=2),
                opacity=opacity, hoverinfo='skip', showlegend=False
            ))
    
    return traces

def create_candidate_origins(board: chess.Board, candidates: List[MoveCandidate]):
    """Highlight pieces that are generating the top candidate moves."""
    traces = []
    if not candidates:
        return traces
    
    # Collect unique origin squares from candidates
    origins = {}
    for i, cand in enumerate(candidates):
        sq = cand.from_sq
        if sq not in origins:
            origins[sq] = i  # Store best rank
    
    for sq, rank in origins.items():
        x, y, _ = square_to_3d(sq)
        
        # Intensity based on rank (top candidate = brightest)
        intensity = 1.0 - (rank * 0.15)
        color = f'rgba(0,255,212,{intensity * 0.6})'  # Cyan glow
        
        # Glow ring around the piece
        theta = np.linspace(0, 2*np.pi, 20)
        ring_x = x + 0.35 * np.cos(theta)
        ring_y = y + 0.35 * np.sin(theta)
        ring_z = [0.08] * len(theta)
        
        traces.append(go.Scatter3d(
            x=ring_x, y=ring_y, z=ring_z,
            mode='lines', line=dict(color=color, width=3 + (5-rank)),
            hoverinfo='skip', showlegend=False
        ))
    
    return traces

# ═══════════════════════════════════════════════════════════════════════════════
# STATE-SPACE LATTICE GRAPH
# ═══════════════════════════════════════════════════════════════════════════════

def create_state_lattice(move_history: List[str], candidates: List[MoveCandidate] = None):
    """
    Create a state-space graph visualization above/below the board.
    Each move node is positioned ABOVE its destination square.
    White moves above board, black moves below.
    Edges trace the flow of the game through space.
    """
    traces = []
    if not move_history or len(move_history) < 1:
        return traces
    
    # Z heights - white moves float above, black sink below
    BASE_Z_WHITE = 2.2
    BASE_Z_BLACK = -0.3
    Z_LAYER_STEP = 0.25  # Each subsequent move goes higher/lower
    
    # Track all nodes for edge drawing
    all_nodes = []  # List of (x, y, z, move_uci, is_white, move_idx)
    
    for i, move_uci in enumerate(move_history):
        is_white = (i % 2 == 0)
        move_num = i // 2 + 1
        same_color_idx = i // 2  # How many moves of this color so far
        
        # Parse the UCI move to get destination square
        if len(move_uci) >= 4:
            to_sq_name = move_uci[2:4]  # e.g., "e4" from "e2e4"
            try:
                to_sq = chess.parse_square(to_sq_name)
                # Get 3D position of destination square
                x, y, _ = square_to_3d(to_sq)
            except:
                x, y = 0, 0
        else:
            x, y = 0, 0
        
        # Z position - stacks up/down with each move
        if is_white:
            z = BASE_Z_WHITE + (same_color_idx * Z_LAYER_STEP)
            node_color = '#FFFEF0'
            edge_color = 'rgba(255,254,240,0.6)'
            glow_color = 'rgba(0,255,212,0.4)'  # Cyan glow
        else:
            z = BASE_Z_BLACK - (same_color_idx * Z_LAYER_STEP)
            node_color = '#AAAAAA'
            edge_color = 'rgba(150,150,150,0.6)'
            glow_color = 'rgba(255,100,150,0.4)'  # Pink glow
        
        all_nodes.append((x, y, z, move_uci, is_white, i))
        
        # Draw vertical "drop line" from node to board
        board_z = 0.1
        traces.append(go.Scatter3d(
            x=[x, x], y=[y, y], z=[z, board_z if is_white else board_z],
            mode='lines',
            line=dict(color=glow_color, width=1),
            hoverinfo='skip', showlegend=False
        ))
        
        # Node marker
        traces.append(go.Scatter3d(
            x=[x], y=[y], z=[z],
            mode='markers+text',
            marker=dict(size=10, color=node_color, 
                       line=dict(width=2, color=glow_color),
                       symbol='diamond'),
            text=[f"{move_num}.{move_uci}" if is_white else move_uci],
            textposition='top center',
            textfont=dict(size=9, color=node_color, family='monospace'),
            hoverinfo='text',
            hovertext=f"{'White' if is_white else 'Black'} {move_num}: {move_uci}",
            showlegend=False
        ))
    
    # Draw edges connecting sequential moves (alternating colors)
    for i in range(len(all_nodes) - 1):
        x0, y0, z0, _, _, _ = all_nodes[i]
        x1, y1, z1, _, is_white_next, _ = all_nodes[i + 1]
        
        # Color based on who just moved (the edge leads TO the next move)
        edge_color = 'rgba(0,255,212,0.4)' if is_white_next else 'rgba(255,100,150,0.4)'
        
        traces.append(go.Scatter3d(
            x=[x0, x1], y=[y0, y1], z=[z0, z1],
            mode='lines',
            line=dict(color=edge_color, width=2),
            hoverinfo='skip', showlegend=False
        ))
    
    # Show candidate moves as potential branches from last position
    if candidates and all_nodes:
        last_x, last_y, last_z, _, last_white, _ = all_nodes[-1]
        is_white_turn = len(move_history) % 2 == 0  # Who moves next
        
        branch_z = (BASE_Z_WHITE + (len(move_history)//2) * Z_LAYER_STEP) if is_white_turn \
                   else (BASE_Z_BLACK - (len(move_history)//2) * Z_LAYER_STEP)
        
        for j, cand in enumerate(candidates[:5]):
            # Position at candidate's destination
            if len(cand.move) >= 4:
                try:
                    cand_to_sq = chess.parse_square(cand.move[2:4])
                    cx, cy, _ = square_to_3d(cand_to_sq)
                except:
                    continue
            else:
                continue
            
            cz = branch_z
            
            # Branch edge from last move to candidate
            traces.append(go.Scatter3d(
                x=[last_x, cx], y=[last_y, cy], z=[last_z, cz],
                mode='lines',
                line=dict(color='rgba(255,215,0,0.3)', width=1, dash='dot'),
                hoverinfo='skip', showlegend=False
            ))
            
            # Candidate node
            traces.append(go.Scatter3d(
                x=[cx], y=[cy], z=[cz],
                mode='markers+text',
                marker=dict(size=6, color=GOLD, opacity=0.5,
                           symbol='diamond'),
                text=[cand.move],
                textposition='top center',
                textfont=dict(size=7, color='rgba(255,215,0,0.6)', family='monospace'),
                hoverinfo='text',
                hovertext=f"Candidate: {cand.move} ({cand.cp}cp)",
                showlegend=False
            ))
    
    return traces


# Default camera position
DEFAULT_CAMERA = dict(
    eye=dict(x=0.8, y=-1.4, z=0.9),
    up=dict(x=0, y=0, z=1),
    center=dict(x=0, y=0, z=0)
)

def create_figure(board: chess.Board, candidates: List[MoveCandidate] = None, camera: dict = None):
    """Create the full 3D figure with polished visuals."""
    fig = go.Figure()
    
    # Use provided camera or default
    cam = camera if camera else DEFAULT_CAMERA
    
    # Layer 1: Board
    for trace in create_board_traces():
        fig.add_trace(trace)
    
    # Layer 2: Pieces (crystal style)
    for trace in create_piece_traces(board):
        fig.add_trace(trace)
    
    # Layer 3: Move arcs (when HOLD active)
    if candidates:
        # Highlight candidate origin pieces
        for trace in create_candidate_origins(board, candidates):
            fig.add_trace(trace)
        
        is_white = board.turn == chess.WHITE
        for trace in create_move_arcs(candidates, is_white):
            fig.add_trace(trace)
    
    # Turn indicator
    turn_text = "⚪ WHITE" if board.turn else "⚫ BLACK"
    turn_color = '#EEE' if board.turn else '#888'
    fig.add_trace(go.Scatter3d(
        x=[0], y=[5.2], z=[0.3], mode='text', text=[turn_text],
        textfont=dict(size=14, color=turn_color, family='monospace'),
        showlegend=False, hoverinfo='skip'
    ))
    
    # Move number
    fig.add_trace(go.Scatter3d(
        x=[0], y=[-5.2], z=[0.3], mode='text', 
        text=[f"Move {board.fullmove_number}"],
        textfont=dict(size=11, color='#666', family='monospace'),
        showlegend=False, hoverinfo='skip'
    ))
    
    fig.update_layout(
        scene=dict(
            xaxis=dict(range=[-5.5, 5.5], showgrid=False, showbackground=False, 
                      showticklabels=False, showline=False, zeroline=False, title=''),
            yaxis=dict(range=[-5.5, 5.5], showgrid=False, showbackground=False,
                      showticklabels=False, showline=False, zeroline=False, title=''),
            zaxis=dict(range=[-0.5, 3.0], showgrid=False, showbackground=False,
                      showticklabels=False, showline=False, zeroline=False, title=''),
            aspectmode='manual', 
            aspectratio=dict(x=1, y=1, z=0.35),
            camera=cam,
            bgcolor=BG_COLOR
        ),
        paper_bgcolor=BG_COLOR,
        plot_bgcolor=BG_COLOR,
        margin=dict(l=0, r=0, t=0, b=0),
        showlegend=False,
        height=650
    )
    return fig

# ═══════════════════════════════════════════════════════════════════════════════
# ENGINE + CASCADE HELPERS  
# ═══════════════════════════════════════════════════════════════════════════════

import time

def get_candidates_with_trace(board: chess.Board, num=5) -> tuple:
    """Get move candidates from Stockfish WITH cascade-lattice tracing."""
    global TRACE_LOG, DECISION_TREE
    
    trace_data = []
    decision_data = []
    start_time = time.perf_counter()
    
    if not ENGINE:
        return [], trace_data, decision_data
    
    try:
        # TRACE: Board state encoding
        t0 = time.perf_counter()
        fen = board.fen()
        trace_data.append({
            'step': 1, 'op': 'ENCODE', 'detail': f'FEN → tensor',
            'input': fen[:20] + '...', 'output': 'state_vec[768]',
            'duration': round((time.perf_counter() - t0) * 1000, 2),
            'confidence': 1.0
        })
        
        # TRACE: Engine analysis
        t1 = time.perf_counter()
        info = ENGINE.analyse(board, chess.engine.Limit(depth=10), multipv=num)
        trace_data.append({
            'step': 2, 'op': 'ANALYZE', 'detail': f'Stockfish depth=10',
            'input': 'state_vec', 'output': f'{len(info)} candidates',
            'duration': round((time.perf_counter() - t1) * 1000, 2),
            'confidence': 0.95
        })
        
        candidates = []
        total = 0
        
        # TRACE: Candidate scoring
        t2 = time.perf_counter()
        for i, pv in enumerate(info):
            move = pv['pv'][0]
            score = pv.get('score', chess.engine.Cp(0))
            
            if score.is_mate():
                value = 1.0 if score.mate() > 0 else -1.0
                eval_str = f"M{score.mate()}"
            else:
                cp = score.relative.score(mate_score=10000)
                value = max(-1, min(1, cp / 1000))
                eval_str = f"{cp:+d}cp"
            
            prob = 1.0 / (i + 1)
            total += prob
            
            cand = MoveCandidate(
                move=move.uci(), prob=prob, value=value,
                from_sq=move.from_square, to_sq=move.to_square,
                is_capture=board.is_capture(move),
                is_check=board.gives_check(move)
            )
            candidates.append(cand)
            
            # Decision tree entry
            decision_data.append({
                'move': move.uci(),
                'eval': eval_str,
                'prob': prob,
                'rank': i + 1,
                'capture': board.is_capture(move),
                'check': board.gives_check(move),
                'selected': i == 0
            })
        
        trace_data.append({
            'step': 3, 'op': 'SCORE', 'detail': f'Evaluate {len(candidates)} moves',
            'input': 'raw_candidates', 'output': 'scored_candidates',
            'duration': round((time.perf_counter() - t2) * 1000, 2),
            'confidence': 0.88
        })
        
        # Normalize probabilities
        for c in candidates:
            c.prob /= total
        for d in decision_data:
            d['prob'] /= total
        
        # TRACE: Hold decision
        t3 = time.perf_counter()
        if HOLD and candidates:
            # Use cascade-lattice Hold to potentially override
            hold_result = None
            try:
                # Hold.evaluate expects candidates, returns selected
                hold_result = HOLD.evaluate([c.move for c in candidates], 
                                           weights=[c.prob for c in candidates])
            except:
                pass
        
        trace_data.append({
            'step': 4, 'op': 'HOLD', 'detail': f'cascade.Hold decision gate',
            'input': 'scored_candidates', 'output': candidates[0].move if candidates else 'none',
            'duration': round((time.perf_counter() - t3) * 1000, 2),
            'confidence': candidates[0].prob if candidates else 0
        })
        
        # TRACE: Final selection
        total_time = (time.perf_counter() - start_time) * 1000
        trace_data.append({
            'step': 5, 'op': 'SELECT', 'detail': f'Final output',
            'input': candidates[0].move if candidates else '-', 
            'output': '✓ COMMITTED',
            'duration': round(total_time, 2),
            'confidence': 1.0
        })
        
        TRACE_LOG = trace_data
        DECISION_TREE = decision_data
        
        return candidates, trace_data, decision_data
        
    except Exception as e:
        print(f"[ENGINE] Error: {e}")
        return [], [], []

def get_candidates(board: chess.Board, num=5) -> List[MoveCandidate]:
    """Simple wrapper for backward compat."""
    candidates, _, _ = get_candidates_with_trace(board, num)
    return candidates

# ═══════════════════════════════════════════════════════════════════════════════
# DASH APP - TWO PANEL LAYOUT
# ═══════════════════════════════════════════════════════════════════════════════

app = dash.Dash(__name__, suppress_callback_exceptions=True)

# Styles
PANEL_STYLE = {
    'backgroundColor': '#0d0d12', 
    'borderRadius': '8px', 
    'padding': '15px',
    'border': '1px solid #1a1a2e'
}
TRACE_ROW_STYLE = {
    'display': 'flex', 'alignItems': 'center', 'padding': '8px 10px',
    'borderBottom': '1px solid #1a1a2e', 'fontFamily': 'monospace', 'fontSize': '12px'
}
BUTTON_BASE = {
    'margin': '5px', 'padding': '12px 24px', 'fontSize': '14px',
    'backgroundColor': '#1a1a2e', 'borderRadius': '4px', 
    'cursor': 'pointer', 'fontFamily': 'monospace'
}

app.layout = html.Div([
    # Header
    html.Div([
        html.H1("CASCADE // LATTICE", 
               style={'color': CYAN, 'margin': 0, 'fontFamily': 'Consolas, monospace',
                     'fontSize': '2.2em', 'letterSpacing': '0.1em', 'display': 'inline-block'}),
        html.Span(" × ", style={'color': '#333', 'fontSize': '1.5em', 'margin': '0 10px'}),
        html.Span("INFERENCE VISUALIZATION", 
                 style={'color': '#444', 'fontFamily': 'monospace', 'fontSize': '1.1em'})
    ], style={'textAlign': 'center', 'padding': '20px 0', 'borderBottom': '1px solid #1a1a2e'}),
    
    # Controls with loading indicator
    html.Div([
        html.Button("⏭ STEP", id='btn-step', n_clicks=0, 
                   style={**BUTTON_BASE, 'color': '#888', 'border': '1px solid #333'}),
        html.Button("⏸ HOLD", id='btn-hold', n_clicks=0,
                   style={**BUTTON_BASE, 'color': GOLD, 'border': f'1px solid {GOLD}'}),
        html.Button("▶▶ AUTO", id='btn-auto', n_clicks=0,
                   style={**BUTTON_BASE, 'color': CYAN, 'border': f'1px solid {CYAN}'}),
        html.Button("↺ RESET", id='btn-reset', n_clicks=0,
                   style={**BUTTON_BASE, 'color': CRIMSON, 'border': f'1px solid {CRIMSON}'}),
        # Loading spinner
        dcc.Loading(
            id='loading-indicator',
            type='circle',
            color=GOLD,
            children=html.Div(id='loading-output', style={'display': 'inline-block', 'marginLeft': '15px'})
        ),
    ], style={'textAlign': 'center', 'padding': '15px', 'display': 'flex', 
              'justifyContent': 'center', 'alignItems': 'center', 'gap': '5px'}),
    
    # Status bar
    html.Div(id='status', 
            style={'textAlign': 'center', 'color': '#666', 'padding': '10px',
                  'fontFamily': 'monospace', 'fontSize': '13px', 'backgroundColor': '#0a0a0f',
                  'borderTop': '1px solid #1a1a2e', 'borderBottom': '1px solid #1a1a2e'}),
    
    # ═══════════════════════════════════════════════════════════════════════════
    # MAIN THREE-COLUMN LAYOUT
    # ═══════════════════════════════════════════════════════════════════════════
    html.Div([
        # LEFT COLUMN - Engine & Game Info
        html.Div([
            # Engine Info Panel
            html.Div([
                html.Div([
                    html.Span("◈ ", style={'color': '#FF6B35'}),
                    html.Span("ENGINE", style={'color': '#888'})
                ], style={'fontFamily': 'monospace', 'fontSize': '14px', 'marginBottom': '12px',
                         'borderBottom': '1px solid #FF6B3533', 'paddingBottom': '8px'}),
                
                html.Div([
                    html.Div([
                        html.Span("Model", style={'color': '#555', 'width': '70px', 'display': 'inline-block'}),
                        html.Span("Stockfish 17", style={'color': '#FF6B35', 'fontWeight': 'bold'})
                    ], style={'marginBottom': '8px'}),
                    html.Div([
                        html.Span("Type", style={'color': '#555', 'width': '70px', 'display': 'inline-block'}),
                        html.Span("NNUE + Classical", style={'color': '#888'})
                    ], style={'marginBottom': '8px'}),
                    html.Div([
                        html.Span("Depth", style={'color': '#555', 'width': '70px', 'display': 'inline-block'}),
                        html.Span("10 ply", style={'color': CYAN})
                    ], style={'marginBottom': '8px'}),
                    html.Div([
                        html.Span("MultiPV", style={'color': '#555', 'width': '70px', 'display': 'inline-block'}),
                        html.Span("5 lines", style={'color': GOLD})
                    ], style={'marginBottom': '8px'}),
                    html.Div([
                        html.Span("Status", style={'color': '#555', 'width': '70px', 'display': 'inline-block'}),
                        html.Span("● READY" if ENGINE else "○ OFFLINE", 
                                 style={'color': '#0F0' if ENGINE else CRIMSON})
                    ]),
                ], style={'fontFamily': 'monospace', 'fontSize': '12px'})
            ], style={**PANEL_STYLE, 'marginBottom': '15px'}),
            
            # Cascade-Lattice Info Panel
            html.Div([
                html.Div([
                    html.Span("◈ ", style={'color': CYAN}),
                    html.Span("CASCADE-LATTICE", style={'color': '#888'})
                ], style={'fontFamily': 'monospace', 'fontSize': '14px', 'marginBottom': '12px',
                         'borderBottom': f'1px solid {CYAN}33', 'paddingBottom': '8px'}),
                
                html.Div([
                    html.Div([
                        html.Span("Package", style={'color': '#555', 'width': '70px', 'display': 'inline-block'}),
                        html.Span("cascade-lattice", style={'color': CYAN})
                    ], style={'marginBottom': '8px'}),
                    html.Div([
                        html.Span("Version", style={'color': '#555', 'width': '70px', 'display': 'inline-block'}),
                        html.Span("0.5.6", style={'color': '#888'})
                    ], style={'marginBottom': '8px'}),
                    html.Div([
                        html.Span("Hold", style={'color': '#555', 'width': '70px', 'display': 'inline-block'}),
                        html.Span("● ACTIVE" if HOLD else "○ OFF", 
                                 style={'color': '#0F0' if HOLD else '#555'})
                    ], style={'marginBottom': '8px'}),
                    html.Div([
                        html.Span("Causation", style={'color': '#555', 'width': '70px', 'display': 'inline-block'}),
                        html.Span("● TRACING" if CAUSATION else "○ OFF", 
                                 style={'color': MAGENTA if CAUSATION else '#555'})
                    ]),
                ], style={'fontFamily': 'monospace', 'fontSize': '12px'})
            ], style={**PANEL_STYLE, 'marginBottom': '15px'}),
            
            # Game State Panel
            html.Div([
                html.Div([
                    html.Span("◈ ", style={'color': GOLD}),
                    html.Span("GAME STATE", style={'color': '#888'})
                ], style={'fontFamily': 'monospace', 'fontSize': '14px', 'marginBottom': '12px',
                         'borderBottom': f'1px solid {GOLD}33', 'paddingBottom': '8px'}),
                
                html.Div(id='game-state-panel', children=[
                    html.Div([
                        html.Span("Turn", style={'color': '#555', 'width': '70px', 'display': 'inline-block'}),
                        html.Span("White", id='gs-turn', style={'color': '#FFF'})
                    ], style={'marginBottom': '8px'}),
                    html.Div([
                        html.Span("Move #", style={'color': '#555', 'width': '70px', 'display': 'inline-block'}),
                        html.Span("1", id='gs-movenum', style={'color': GOLD})
                    ], style={'marginBottom': '8px'}),
                    html.Div([
                        html.Span("Material", style={'color': '#555', 'width': '70px', 'display': 'inline-block'}),
                        html.Span("0", id='gs-material', style={'color': '#888'})
                    ], style={'marginBottom': '8px'}),
                    html.Div([
                        html.Span("Phase", style={'color': '#555', 'width': '70px', 'display': 'inline-block'}),
                        html.Span("Opening", id='gs-phase', style={'color': '#888'})
                    ]),
                ], style={'fontFamily': 'monospace', 'fontSize': '12px'})
            ], style=PANEL_STYLE),
            
        ], style={'flex': '0 0 220px', 'padding': '0 15px 0 0'}),
        
        # MIDDLE COLUMN - 3D Chess Board
        html.Div([
            dcc.Graph(id='chess-3d', figure=create_figure(chess.Board()), 
                      config={'displayModeBar': True, 'scrollZoom': True, 
                             'modeBarButtonsToRemove': ['toImage', 'sendDataToCloud']},
                      style={'height': '580px'}),
            # Move buttons (when HOLD)
            html.Div(id='move-buttons', style={'textAlign': 'center', 'padding': '10px'})
        ], style={'flex': '1', 'minWidth': '450px'}),
        
        # RIGHT COLUMN - Cascade Panel
        html.Div([
            # Cascade Trace Panel
            html.Div([
                html.Div([
                    html.Span("◈ ", style={'color': CYAN}),
                    html.Span("CAUSATION TRACE", style={'color': '#888'})
                ], style={'fontFamily': 'monospace', 'fontSize': '14px', 'marginBottom': '12px',
                         'borderBottom': f'1px solid {CYAN}33', 'paddingBottom': '8px'}),
                
                html.Div(id='cascade-trace', children=[
                    html.Div("Waiting for move...", style={'color': '#444', 'fontStyle': 'italic', 
                                                          'padding': '20px', 'textAlign': 'center'})
                ])
            ], style={**PANEL_STYLE, 'marginBottom': '15px'}),
            
            # Decision Tree Panel
            html.Div([
                html.Div([
                    html.Span("◈ ", style={'color': GOLD}),
                    html.Span("DECISION TREE", style={'color': '#888'})
                ], style={'fontFamily': 'monospace', 'fontSize': '14px', 'marginBottom': '12px',
                         'borderBottom': f'1px solid {GOLD}33', 'paddingBottom': '8px'}),
                
                html.Div(id='decision-tree', children=[
                    html.Div("No candidates yet", style={'color': '#444', 'fontStyle': 'italic',
                                                         'padding': '20px', 'textAlign': 'center'})
                ])
            ], style={**PANEL_STYLE, 'marginBottom': '15px'}),
            
            # Metrics Panel
            html.Div([
                html.Div([
                    html.Span("◈ ", style={'color': MAGENTA}),
                    html.Span("METRICS", style={'color': '#888'})
                ], style={'fontFamily': 'monospace', 'fontSize': '14px', 'marginBottom': '12px',
                         'borderBottom': f'1px solid {MAGENTA}33', 'paddingBottom': '8px'}),
                
                html.Div(id='metrics-panel', children=[
                    html.Div([
                        html.Span("Total Latency: ", style={'color': '#555'}),
                        html.Span("--ms", id='metric-latency', style={'color': CYAN})
                    ], style={'marginBottom': '5px'}),
                    html.Div([
                        html.Span("Candidates: ", style={'color': '#555'}),
                        html.Span("0", id='metric-candidates', style={'color': GOLD})
                    ], style={'marginBottom': '5px'}),
                    html.Div([
                        html.Span("Confidence: ", style={'color': '#555'}),
                        html.Span("--%", id='metric-confidence', style={'color': MAGENTA})
                    ]),
                ], style={'fontFamily': 'monospace', 'fontSize': '13px'})
            ], style=PANEL_STYLE),
            
        ], style={'flex': '1', 'minWidth': '350px', 'maxWidth': '450px', 'padding': '0 15px'}),
        
    ], style={'display': 'flex', 'padding': '20px', 'gap': '10px', 'alignItems': 'flex-start'}),
    
    # Hidden stores
    dcc.Store(id='board-fen', data=chess.STARTING_FEN),
    dcc.Store(id='candidates-store', data=[]),
    dcc.Store(id='trace-store', data=[]),
    dcc.Store(id='decision-store', data=[]),
    dcc.Store(id='is-held', data=False),
    dcc.Store(id='auto-play', data=False),
    dcc.Store(id='move-history', data=[]),
    dcc.Store(id='camera-store', data=None),  # Stores user's camera position
    
    # Auto-play interval
    dcc.Interval(id='auto-interval', interval=1200, disabled=True),
    
], style={'backgroundColor': BG_COLOR, 'minHeight': '100vh', 'padding': '0'})

@callback(
    Output('board-fen', 'data'),
    Output('candidates-store', 'data'),
    Output('trace-store', 'data'),
    Output('decision-store', 'data'),
    Output('is-held', 'data'),
    Output('move-history', 'data'),
    Output('auto-play', 'data'),
    Input('btn-step', 'n_clicks'),
    Input('btn-hold', 'n_clicks'),
    Input('btn-reset', 'n_clicks'),
    Input('btn-auto', 'n_clicks'),
    Input('auto-interval', 'n_intervals'),
    State('board-fen', 'data'),
    State('candidates-store', 'data'),
    State('is-held', 'data'),
    State('move-history', 'data'),
    State('auto-play', 'data'),
    prevent_initial_call=True
)
def handle_controls(step, hold, reset, auto, interval, fen, candidates, is_held, history, auto_play):
    ctx = dash.callback_context
    if not ctx.triggered:
        return fen, candidates, [], [], is_held, history, auto_play
    
    trigger = ctx.triggered[0]['prop_id'].split('.')[0]
    board = chess.Board(fen)
    
    if trigger == 'btn-reset':
        return chess.STARTING_FEN, [], [], [], False, [], False
    
    if trigger == 'btn-auto':
        return fen, [], [], [], False, history, not auto_play  # Toggle auto
    
    if trigger == 'btn-hold':
        if not is_held:
            cands, trace, decision = get_candidates_with_trace(board)
            return fen, [c.__dict__ for c in cands], trace, decision, True, history, False
        else:
            return fen, [], [], [], False, history, auto_play
    
    if trigger in ['btn-step', 'auto-interval']:
        if board.is_game_over():
            return fen, [], [], [], False, history, False
        
        cands, trace, decision = get_candidates_with_trace(board)
        if cands:
            move = chess.Move.from_uci(cands[0].move)
            board.push(move)
            history = history + [cands[0].move]
        return board.fen(), [], trace, decision, False, history, auto_play
    
    return fen, candidates, [], [], is_held, history, auto_play

@callback(
    Output('chess-3d', 'figure'),
    Output('loading-output', 'children'),
    Input('board-fen', 'data'),
    Input('candidates-store', 'data'),
    State('chess-3d', 'figure')  # Read current figure to get its camera
)
def update_figure(fen, candidates_data, current_fig):
    board = chess.Board(fen)
    candidates = [MoveCandidate(**c) for c in candidates_data] if candidates_data else None
    
    # Extract camera from current figure if it exists
    camera = None
    if current_fig and 'layout' in current_fig:
        scene = current_fig['layout'].get('scene', {})
        if 'camera' in scene:
            camera = scene['camera']
    
    return create_figure(board, candidates, camera), ""

@callback(
    Output('btn-hold', 'children'),
    Output('btn-hold', 'style'),
    Input('is-held', 'data')
)
def update_hold_button(is_held):
    if is_held:
        return "◉ HOLDING", {**BUTTON_BASE, 'color': '#000', 'backgroundColor': GOLD, 
                             'border': f'2px solid {GOLD}', 'fontWeight': 'bold'}
    else:
        return "⏸ HOLD", {**BUTTON_BASE, 'color': GOLD, 'border': f'1px solid {GOLD}'}

@callback(
    Output('status', 'children'),
    Input('board-fen', 'data'),
    Input('is-held', 'data'),
    Input('auto-play', 'data'),
    Input('move-history', 'data')
)
def update_status(fen, is_held, auto_play, history):
    board = chess.Board(fen)
    
    if board.is_game_over():
        result = board.result()
        return f"GAME OVER: {result}"
    
    turn = "WHITE" if board.turn else "BLACK"
    mode = "◉ HOLD ACTIVE - Select a move" if is_held else ("▶▶ AUTO" if auto_play else "MANUAL")
    return f"Move {board.fullmove_number} | {turn} | {mode}"

@callback(
    Output('auto-interval', 'disabled'),
    Input('auto-play', 'data')
)
def toggle_auto(auto_play):
    return not auto_play

@callback(
    Output('gs-turn', 'children'),
    Output('gs-turn', 'style'),
    Output('gs-movenum', 'children'),
    Output('gs-material', 'children'),
    Output('gs-material', 'style'),
    Output('gs-phase', 'children'),
    Input('board-fen', 'data')
)
def update_game_state(fen):
    board = chess.Board(fen)
    
    # Turn
    turn_text = "White" if board.turn else "Black"
    turn_style = {'color': '#FFF'} if board.turn else {'color': '#888'}
    
    # Move number
    move_num = str(board.fullmove_number)
    
    # Material count (simple piece values)
    piece_values = {chess.PAWN: 1, chess.KNIGHT: 3, chess.BISHOP: 3, 
                    chess.ROOK: 5, chess.QUEEN: 9, chess.KING: 0}
    white_material = sum(piece_values.get(p.piece_type, 0) 
                        for p in board.piece_map().values() if p.color == chess.WHITE)
    black_material = sum(piece_values.get(p.piece_type, 0) 
                        for p in board.piece_map().values() if p.color == chess.BLACK)
    diff = white_material - black_material
    if diff > 0:
        mat_text = f"+{diff} ⚪"
        mat_style = {'color': '#0F0'}
    elif diff < 0:
        mat_text = f"{diff} ⚫"
        mat_style = {'color': CRIMSON}
    else:
        mat_text = "Equal"
        mat_style = {'color': '#888'}
    
    # Game phase (rough estimate)
    total_pieces = len(board.piece_map())
    if total_pieces >= 28:
        phase = "Opening"
    elif total_pieces >= 14:
        phase = "Middlegame"
    else:
        phase = "Endgame"
    
    if board.is_check():
        phase = "⚠ CHECK"
    if board.is_game_over():
        phase = "Game Over"
    
    return turn_text, turn_style, move_num, mat_text, mat_style, phase

# ═══════════════════════════════════════════════════════════════════════════════
# CASCADE PANEL CALLBACKS
# ═══════════════════════════════════════════════════════════════════════════════

@callback(
    Output('cascade-trace', 'children'),
    Input('trace-store', 'data')
)
def render_trace(trace_data):
    if not trace_data:
        return html.Div("Waiting for move...", 
                       style={'color': '#444', 'fontStyle': 'italic', 'padding': '20px', 'textAlign': 'center'})
    
    rows = []
    for t in trace_data:
        # Color code by operation type
        op_colors = {'ENCODE': CYAN, 'ANALYZE': '#888', 'SCORE': GOLD, 'HOLD': MAGENTA, 'SELECT': '#0F0'}
        op_color = op_colors.get(t['op'], '#666')
        
        # Confidence bar
        conf_pct = t['confidence'] * 100
        
        row = html.Div([
            # Step number
            html.Span(f"{t['step']}", style={'color': '#444', 'width': '20px', 'marginRight': '10px'}),
            # Operation badge
            html.Span(t['op'], style={
                'backgroundColor': f'{op_color}22', 'color': op_color,
                'padding': '2px 8px', 'borderRadius': '3px', 'fontSize': '10px',
                'fontWeight': 'bold', 'width': '60px', 'textAlign': 'center', 'marginRight': '10px'
            }),
            # Detail
            html.Span(t['detail'], style={'color': '#888', 'flex': '1', 'fontSize': '11px'}),
            # Duration
            html.Span(f"{t['duration']}ms", style={'color': '#555', 'width': '60px', 'textAlign': 'right'}),
        ], style={**TRACE_ROW_STYLE})
        rows.append(row)
    
    # Total latency
    total = sum(t['duration'] for t in trace_data)
    rows.append(html.Div([
        html.Span("", style={'width': '90px'}),
        html.Span("TOTAL", style={'color': CYAN, 'fontWeight': 'bold', 'flex': '1'}),
        html.Span(f"{total:.1f}ms", style={'color': CYAN, 'fontWeight': 'bold', 'width': '60px', 'textAlign': 'right'})
    ], style={**TRACE_ROW_STYLE, 'borderBottom': 'none', 'backgroundColor': '#0a0a0f'}))
    
    return rows

@callback(
    Output('decision-tree', 'children'),
    Input('decision-store', 'data')
)
def render_decision_tree(decision_data):
    if not decision_data:
        return html.Div("No candidates yet", 
                       style={'color': '#444', 'fontStyle': 'italic', 'padding': '20px', 'textAlign': 'center'})
    
    rows = []
    for d in decision_data:
        is_selected = d.get('selected', False)
        bg_color = f'{CYAN}15' if is_selected else 'transparent'
        border_left = f'3px solid {CYAN}' if is_selected else '3px solid transparent'
        
        # Probability bar
        prob_pct = d['prob'] * 100
        
        row = html.Div([
            # Rank
            html.Span(f"#{d['rank']}", style={
                'color': CYAN if is_selected else '#555', 
                'width': '30px', 'fontWeight': 'bold' if is_selected else 'normal'
            }),
            # Move
            html.Span(d['move'], style={
                'color': '#FFF' if is_selected else '#888',
                'fontWeight': 'bold', 'width': '55px', 'fontFamily': 'monospace'
            }),
            # Eval
            html.Span(d['eval'], style={
                'color': GOLD if 'M' in str(d['eval']) else ('#0F0' if d['eval'][0] == '+' else CRIMSON),
                'width': '55px', 'textAlign': 'right'
            }),
            # Probability bar
            html.Div([
                html.Div(style={
                    'width': f'{prob_pct}%', 'height': '8px',
                    'backgroundColor': CYAN if is_selected else '#333',
                    'borderRadius': '2px'
                })
            ], style={'flex': '1', 'backgroundColor': '#1a1a2e', 'borderRadius': '2px', 'marginLeft': '10px'}),
            # Percentage
            html.Span(f"{prob_pct:.0f}%", style={'color': '#666', 'width': '40px', 'textAlign': 'right', 'marginLeft': '8px'}),
            # Flags
            html.Span(
                ("⚔" if d.get('capture') else "") + ("♚" if d.get('check') else ""),
                style={'color': CRIMSON, 'width': '25px', 'textAlign': 'right'}
            )
        ], style={
            'display': 'flex', 'alignItems': 'center', 'padding': '8px 10px',
            'backgroundColor': bg_color, 'borderLeft': border_left,
            'marginBottom': '4px', 'borderRadius': '3px',
            'fontFamily': 'monospace', 'fontSize': '12px'
        })
        rows.append(row)
    
    return rows

@callback(
    Output('metric-latency', 'children'),
    Output('metric-candidates', 'children'),
    Output('metric-confidence', 'children'),
    Input('trace-store', 'data'),
    Input('decision-store', 'data')
)
def update_metrics(trace_data, decision_data):
    if not trace_data:
        return "--ms", "0", "--%"
    
    total_latency = sum(t['duration'] for t in trace_data)
    num_candidates = len(decision_data) if decision_data else 0
    confidence = decision_data[0]['prob'] * 100 if decision_data else 0
    
    return f"{total_latency:.1f}ms", str(num_candidates), f"{confidence:.0f}%"

@callback(
    Output('move-buttons', 'children'),
    Input('candidates-store', 'data'),
    Input('is-held', 'data')
)
def show_move_buttons(candidates_data, is_held):
    if not is_held or not candidates_data:
        return []
    
    # Get whose turn it is from the current board state (we'll need this for styling)
    buttons = []
    for i, c in enumerate(candidates_data):
        prob_pct = c['prob'] * 100
        is_top = i == 0
        
        btn_style = {
            'margin': '5px', 'padding': '10px 20px', 'fontSize': '13px',
            'fontFamily': 'monospace', 'borderRadius': '4px', 'cursor': 'pointer',
            'backgroundColor': '#1a1a2e' if not is_top else f'{CYAN}22',
            'color': CYAN if is_top else '#888',
            'border': f'1px solid {CYAN}' if is_top else '1px solid #333'
        }
        
        btn = html.Button(
            f"{c['move']} ({prob_pct:.0f}%)",
            id={'type': 'move-btn', 'index': i},
            style=btn_style
        )
        buttons.append(btn)
    return buttons

# Move selection callback
@callback(
    Output('board-fen', 'data', allow_duplicate=True),
    Output('candidates-store', 'data', allow_duplicate=True),
    Output('trace-store', 'data', allow_duplicate=True),
    Output('decision-store', 'data', allow_duplicate=True),
    Output('is-held', 'data', allow_duplicate=True),
    Output('move-history', 'data', allow_duplicate=True),
    Input({'type': 'move-btn', 'index': dash.ALL}, 'n_clicks'),
    State('board-fen', 'data'),
    State('candidates-store', 'data'),
    State('trace-store', 'data'),
    State('decision-store', 'data'),
    State('move-history', 'data'),
    prevent_initial_call=True
)
def select_move(clicks, fen, candidates_data, trace_data, decision_data, history):
    ctx = dash.callback_context
    
    # Only proceed if an actual button was clicked
    if not ctx.triggered:
        raise dash.exceptions.PreventUpdate
    
    # Check if any click actually happened (not just initialization)
    if not clicks or not any(c for c in clicks if c):
        raise dash.exceptions.PreventUpdate
    
    # Find which button was clicked
    triggered_id = ctx.triggered[0]['prop_id']
    if triggered_id == '.':
        raise dash.exceptions.PreventUpdate
        
    import json
    try:
        idx = json.loads(triggered_id.split('.')[0])['index']
    except:
        raise dash.exceptions.PreventUpdate
    
    board = chess.Board(fen)
    if candidates_data and idx < len(candidates_data):
        move_uci = candidates_data[idx]['move']
        move = chess.Move.from_uci(move_uci)
        board.push(move)
        history = history + [move_uci]
    
    # Return with trace/decision preserved for display
    return board.fen(), [], trace_data, decision_data, False, history

if __name__ == '__main__':
    print("\n" + "="*50)
    print("CASCADE-LATTICE Chess")
    print("="*50)
    print("Open: http://127.0.0.1:8050")
    print("="*50 + "\n")
    # Note: debug=False to avoid Python 3.13 socket issues on Windows
    app.run(debug=False, port=8050)