app.py

"""
Nexus-Nano Inference API (Fixed)
Ultra-lightweight single-file engine with proper error handling
"""

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import onnxruntime as ort
import numpy as np
import chess
import time
import logging
import os
from pathlib import Path
from typing import Optional, Tuple

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# ==================== NANO ENGINE (Single File) ====================

class NexusNanoEngine:
    """Ultra-lightweight chess engine"""
    
    PIECE_VALUES = {
        chess.PAWN: 1, chess.KNIGHT: 3, chess.BISHOP: 3,
        chess.ROOK: 5, chess.QUEEN: 9, chess.KING: 0
    }
    
    def __init__(self, model_path: str):
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"Model not found: {model_path}")
        
        logger.info(f"Loading model from {model_path}...")
        logger.info(f"Model size: {os.path.getsize(model_path)/(1024*1024):.2f} MB")
        
        sess_options = ort.SessionOptions()
        sess_options.intra_op_num_threads = 2
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        
        self.session = ort.InferenceSession(
            model_path,
            sess_options=sess_options,
            providers=['CPUExecutionProvider']
        )
        
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
        self.nodes = 0
        
        logger.info("✅ Nexus-Nano engine loaded")
    
    def fen_to_tensor(self, fen: str) -> np.ndarray:
        board = chess.Board(fen)
        tensor = np.zeros((1, 12, 8, 8), dtype=np.float32)
        
        piece_map = {
            chess.PAWN: 0, chess.KNIGHT: 1, chess.BISHOP: 2,
            chess.ROOK: 3, chess.QUEEN: 4, chess.KING: 5
        }
        
        for sq, piece in board.piece_map().items():
            r, f = divmod(sq, 8)
            ch = piece_map[piece.piece_type] + (6 if piece.color == chess.BLACK else 0)
            tensor[0, ch, r, f] = 1.0
        
        return tensor
    
    def evaluate(self, board: chess.Board) -> float:
        self.nodes += 1
        tensor = self.fen_to_tensor(board.fen())
        output = self.session.run([self.output_name], {self.input_name: tensor})
        score = float(output[0][0][0]) * 400.0
        return -score if board.turn == chess.BLACK else score
    
    def order_moves(self, board: chess.Board, moves):
        """Simple MVV-LVA ordering"""
        scored = []
        for m in moves:
            s = 0
            if board.is_capture(m):
                v = board.piece_at(m.to_square)
                a = board.piece_at(m.from_square)
                if v and a:
                    s = self.PIECE_VALUES.get(v.piece_type, 0) * 10
                    s -= self.PIECE_VALUES.get(a.piece_type, 0)
            if m.promotion == chess.QUEEN:
                s += 90
            scored.append((s, m))
        scored.sort(key=lambda x: x[0], reverse=True)
        return [m for _, m in scored]
    
    def alpha_beta(
        self, 
        board: chess.Board, 
        depth: int, 
        alpha: float, 
        beta: float
    ) -> Tuple[float, Optional[chess.Move]]:
        
        if board.is_game_over():
            return (-10000 if board.is_checkmate() else 0), None
        
        if depth == 0:
            return self.evaluate(board), None
        
        moves = list(board.legal_moves)
        if not moves:
            return 0, None
        
        moves = self.order_moves(board, moves)
        
        best_move = moves[0]
        best_score = float('-inf')
        
        for move in moves:
            board.push(move)
            score, _ = self.alpha_beta(board, depth - 1, -beta, -alpha)
            score = -score
            board.pop()
            
            if score > best_score:
                best_score = score
                best_move = move
            
            alpha = max(alpha, score)
            if alpha >= beta:
                break
        
        return best_score, best_move
    
    def search(self, fen: str, depth: int = 3):
        board = chess.Board(fen)
        self.nodes = 0
        
        moves = list(board.legal_moves)
        if len(moves) == 0:
            return {'best_move': '0000', 'evaluation': 0.0, 'nodes': 0}
        if len(moves) == 1:
            return {
                'best_move': moves[0].uci(),
                'evaluation': round(self.evaluate(board) / 100.0, 2),
                'nodes': 1
            }
        
        best_move = moves[0]
        best_score = float('-inf')
        
        for d in range(1, depth + 1):
            try:
                score, move = self.alpha_beta(board, d, float('-inf'), float('inf'))
                if move:
                    best_move = move
                    best_score = score
            except:
                break
        
        return {
            'best_move': best_move.uci(),
            'evaluation': round(best_score / 100.0, 2),
            'depth': d,
            'nodes': self.nodes
        }


# ==================== FASTAPI APP ====================

app = FastAPI(
    title="Nexus-Nano API",
    description="Ultra-lightweight chess engine",
    version="1.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

engine = None


class MoveRequest(BaseModel):
    fen: str
    depth: Optional[int] = Field(3, ge=1, le=5)


class MoveResponse(BaseModel):
    best_move: str
    evaluation: float
    depth_searched: int
    nodes_evaluated: int
    time_taken: int


@app.on_event("startup")
async def startup():
    global engine
    logger.info("🚀 Starting Nexus-Nano API...")
    
    model_path = "/app/models/nexus_nano.onnx"
    
    # Debug: Check models directory
    if os.path.exists("/app/models"):
        logger.info(f"📂 Files in /app/models/:")
        for f in os.listdir("/app/models"):
            full_path = os.path.join("/app/models", f)
            size = os.path.getsize(full_path) / (1024*1024)
            logger.info(f"   - {f} ({size:.2f} MB)")
    else:
        logger.error("❌ /app/models/ directory not found!")
        raise FileNotFoundError("/app/models/ not found")
    
    # Load engine
    try:
        engine = NexusNanoEngine(model_path)
        logger.info("✅ Engine ready")
    except Exception as e:
        logger.error(f"❌ Failed to load: {e}", exc_info=True)
        raise


@app.get("/health")
async def health():
    return {
        "status": "healthy" if engine else "unhealthy",
        "model": "nexus-nano",
        "version": "1.0.0",
        "model_loaded": engine is not None
    }


@app.post("/get-move", response_model=MoveResponse)
async def get_move(req: MoveRequest):
    if not engine:
        raise HTTPException(503, "Engine not loaded")
    
    try:
        chess.Board(req.fen)
    except:
        raise HTTPException(400, "Invalid FEN")
    
    start = time.time()
    
    try:
        result = engine.search(req.fen, req.depth)
        elapsed = int((time.time() - start) * 1000)
        
        logger.info(
            f"Move: {result['best_move']} | "
            f"Eval: {result['evaluation']:+.2f} | "
            f"Nodes: {result['nodes']} | "
            f"Time: {elapsed}ms"
        )
        
        return MoveResponse(
            best_move=result['best_move'],
            evaluation=result['evaluation'],
            depth_searched=result['depth'],
            nodes_evaluated=result['nodes'],
            time_taken=elapsed
        )
    except Exception as e:
        logger.error(f"Search error: {e}", exc_info=True)
        raise HTTPException(500, str(e))


@app.get("/")
async def root():
    return {
        "name": "Nexus-Nano API",
        "version": "1.0.0",
        "model": "2.8M parameters",
        "speed": "Lightning-fast (0.2-0.5s)",
        "status": "healthy" if engine else "loading"
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")