"""
Brain v10.0 - VL-JEPA Paper-Faithful Implementation
====================================================
L14: Qwen3-VL-Embedding-2B (2048D, MRL 64-2048)
L15: Matrioshka MRL (Escala Adaptativa)
L16: Qwen3-VL-Reranker-2B (Multimodal)
L18.5: VL-JEPA Thought Predictor (Paper arXiv:2512.10942)
      - TransformerEncoder 4 layers
      - nomic-embed Y-Encoder
      - JEPA Loss (MSE + Cosine)

Arquitectura fiel al paper sin licencias Meta.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from fastapi import FastAPI, Form, HTTPException
from transformers import AutoProcessor, AutoTokenizer
import json
import gc
import numpy as np
from typing import Optional, List, Dict, Any, Union
from dataclasses import dataclass
from utils.bunker_client import BunkerClient

import warnings
import math
warnings.filterwarnings('ignore')

app = FastAPI(
    title="Brain - VL-JEPA Paper-Faithful",
    description="L14-L16 Qwen3-VL | L18.5 VL-JEPA (Transformer + nomic-embed)",
    version="10.0.0"
)

# ============================================================================
# CONFIGURACIÓN
# ============================================================================
print("=== BRAIN v10.0: VL-JEPA Paper-Faithful ===")

# L14: Qwen3-VL-Embedding-2B
EMBEDDING_MODEL_ID = "Qwen/Qwen3-VL-Embedding-2B"

# L16: Qwen3-VL-Reranker-2B
RERANKER_MODEL_ID = "Qwen/Qwen3-VL-Reranker-2B"

# L18.5: VL-JEPA Configuration (Paper-Faithful)
VLJEPA_CONFIG = {
    # Sequence-Based Configuration
    "node_dim": 256,          # Input Hypergraph Node dim
    "sem_dim": 1024,          # Input Semantic dim
    "hidden_dim": 512,        # Transformer d_model
    "num_heads": 8,           # Multi-head attention
    "num_layers": 8,          # Increased to 8 layers (Paper complexity)
    "dropout": 0.1,
    "action_dim": 16,         # Output action tokens
    "y_encoder_dim": 768,     # nomic-embed output dimension
    "seq_len": 10             # 8 nodes + 1 stats + 1 semantic = 10 tokens
}

# Fallbacks
FALLBACK_RERANKER = "cross-encoder/ms-marco-MiniLM-L-6-v2"
Y_ENCODER_MODEL = "nomic-ai/nomic-embed-text-v1.5"  # Open alternative to EmbeddingGemma

# Estado global
embedding_model = None
embedding_processor = None
reranker_model = None
reranker_processor = None
vljepa_model = None
y_encoder = None
optimizer = None  # Added for active training
device = None
bunker = BunkerClient(buffer_dir="_brain_buffer")


# MRL dimensions
MAX_DIM = 2048
DEFAULT_DIM = 1024

# ============================================================================
# L18.5: VL-JEPA PAPER-FAITHFUL IMPLEMENTATION
# ============================================================================

class PositionalEncoding(nn.Module):
    """Positional encoding for Transformer (paper-style)."""
    def __init__(self, d_model, max_len=512, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.unsqueeze(0))
    
    def forward(self, x):
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)

class VLJEPAPredictor(nn.Module):
    """
    VL-JEPA Predictor v10.1 (Paper-Faithful Sequence)
    
    Architecture:
    - Sequence-Based Input: [HG_Node_1...8, HG_Stats_9, Semantic_10]
    - 3-Way Projections -> 512D Tokens
    - 8-Layer Transformer Encoder with Self-Attention
    - Jointly attends to visual (hypergraph) and textual (Qwen) tokens.
    """
    
    def __init__(self, config):
        super().__init__()
        
        # 1. Hypergraph Projection (256D -> 512D)
        # Shared projection for all hypergraph tokens
        self.hg_proj = nn.Linear(config["node_dim"], config["hidden_dim"])
        
        # 2. Semantic Projection (1024D -> 512D)
        self.sem_proj = nn.Linear(config["sem_dim"], config["hidden_dim"])
        
        # 3. Learnable Query Token (optional, here we use Semantic as query)
        
        # Positional encoding
        self.pos_encoder = PositionalEncoding(
            config["hidden_dim"], 
            dropout=config["dropout"]
        )
        
        # Transformer Encoder (8 Layers)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=config["hidden_dim"],
            nhead=config["num_heads"],
            dim_feedforward=config["hidden_dim"] * 4,
            dropout=config["dropout"],
            activation='gelu',
            batch_first=True,
            norm_first=True
        )
        self.transformer = nn.TransformerEncoder(
            encoder_layer,
            num_layers=config["num_layers"],
            norm=nn.LayerNorm(config["hidden_dim"])
        )
        
        # Output projections
        self.output_proj = nn.Linear(config["hidden_dim"], config["y_encoder_dim"])
        self.action_head = nn.Linear(config["hidden_dim"], config["action_dim"])
        
        self.config = config
    
    def forward(self, hypergraph_input, semantic_context, return_hidden=False):
        """
        Args:
            hypergraph_input: [B, 2304] - Flattened sequence (9 tokens * 256D)
                 - [0:2048] Nodes (8 * 256D)
                 - [2048:2304] Stats (1 * 256D)
            semantic_context: [B, 1024] - Semantic embedding
            
        Returns:
            predicted_embedding: [B, 768] - Based on Semantic Token output
        """
        batch_size = hypergraph_input.size(0)
        
        # 1. Reshape Hypergraph Input -> 9 Tokens
        # Nodes: 8 tokens, Stats: 1 token
        hg_seq = hypergraph_input.view(batch_size, 9, 256)
        
        # 2. Project to Model Dim
        hg_tokens = self.hg_proj(hg_seq)  # [B, 9, 512]
        
        sem_token = self.sem_proj(semantic_context).unsqueeze(1)  # [B, 1, 512]
        
        # 3. Concatenate Sequence: [HG (9), SEM (1)] -> Length 10
        x = torch.cat([hg_tokens, sem_token], dim=1) # [B, 10, 512]
        
        # 4. Temporal/Positional Encoding
        x = self.pos_encoder(x)
        
        # 5. Transformer Pass (Self-Attention over SEQUENCE)
        hidden_seq = self.transformer(x)  # [B, 10, 512]
        
        # 6. Pooling / Prediction Strategy
        # Paper predicts target embedding based on query. 
        # Here "Semantic" is our query. We take the last token (Semantic position).
        query_output = hidden_seq[:, -1, :]  # [B, 512]
        
        # 7. Predictions
        predicted_embedding = self.output_proj(query_output) # [B, 768]
        predicted_embedding = F.normalize(predicted_embedding, p=2, dim=-1)
        
        action_logits = self.action_head(query_output)  # [B, 16]
        
        if return_hidden:
            return predicted_embedding, action_logits, hidden_seq
        return predicted_embedding, action_logits


class JEPALoss(nn.Module):
    """
    JEPA Training Loss (Paper-faithful)
    
    L = MSE(predicted, target) + λ * (1 - cosine_similarity)
    
    The model learns to predict the target embedding from Y-Encoder.
    """
    def __init__(self, lambda_cosine=0.1):
        super().__init__()
        self.lambda_cosine = lambda_cosine
    
    def forward(self, predicted, target):
        # MSE loss
        mse = F.mse_loss(predicted, target)
        
        # Cosine similarity loss (we want similarity = 1)
        cosine = F.cosine_similarity(predicted, target, dim=-1).mean()
        cosine_loss = 1 - cosine
        
        total = mse + self.lambda_cosine * cosine_loss
        
        return {
            "total": total,
            "mse": mse.item(),
            "cosine": cosine.item(),
            "cosine_loss": cosine_loss.item()
        }


# ============================================================================
# MODEL LOADING FUNCTIONS
# ============================================================================

def load_vljepa():
    """L18.5: Load VL-JEPA with Transformer + nomic-embed."""
    global vljepa_model, y_encoder, device
    
    if vljepa_model is not None:
        return True
    
    print("[L18.5] Loading VL-JEPA Paper-Faithful...")
    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Load VL-JEPA Predictor
        vljepa_model = VLJEPAPredictor(VLJEPA_CONFIG).to(device)
        vljepa_model.eval()
        
        # Load Y-Encoder (nomic-embed)
        from sentence_transformers import SentenceTransformer
        y_encoder = SentenceTransformer(Y_ENCODER_MODEL, trust_remote_code=True)
        
        # Count parameters
        params = sum(p.numel() for p in vljepa_model.parameters())
        print(f"[L18.5] ✅ VL-JEPA loaded: {params:,} params")
        print(f"[L18.5] ✅ Y-Encoder: {Y_ENCODER_MODEL}")
        return True
        
    except Exception as e:
        global last_error
        last_error = f"{type(e).__name__}: {str(e)}"
        print(f"[L18.5] ❌ Failed: {e}")
        return False

def unload_vljepa():
    global vljepa_model, y_encoder
    vljepa_model = None
    y_encoder = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    print("[L18.5] VL-JEPA unloaded")

# ============================================================================
# HYPERGRAPH ANALYSIS - For VL-JEPA Compatibility
# ============================================================================

def analyze_hypergraph_for_vljepa(hypergraph_state: np.ndarray) -> Dict:
    """
    Analiza si el estado del hipergrafo es óptimo para VL-JEPA v10.1.
    
    VL-JEPA v10.1 espera:
    - Sequence (2304D): [8 Tokens (256D) + 1 Stats (256D)]
    """
    input_dim = len(hypergraph_state)
    is_sequence = input_dim == 2304
    
    analysis = {
        "input_dimension": input_dim,
        "is_v10_1_sequence": is_sequence,
        "non_zero_ratio": np.count_nonzero(hypergraph_state) / input_dim,
        "norm": float(np.linalg.norm(hypergraph_state)),
        "issues": [],
        "recommendations": []
    }
    
    if not is_sequence:
        analysis["issues"].append(f"Expected 2304D Sequence, got {input_dim}D")
        analysis["recommendations"].append("Use hypergraph.get_context_for_vljepa() (v10.1)")
        analysis["vljepa_compatibility_score"] = 0.0
        return analysis
        
    # Sequence analysis
    # Reshape to (9, 256) conceptually
    state_matrix = hypergraph_state.reshape(9, 256)
    
    # Check 1: Sparsity of nodes (first 8 tokens)
    nodes_energy = np.linalg.norm(state_matrix[:8], axis=1)
    active_nodes = np.count_nonzero(nodes_energy > 0.01)
    
    if active_nodes < 3:
         analysis["issues"].append(f"Low history: Only {active_nodes}/8 active node steps")
         analysis["recommendations"].append("Accumulate more steps in hypergraph")

    # Check 2: Stats token (last one)
    stats_token = state_matrix[8]
    if np.linalg.norm(stats_token) < 0.001:
        analysis["issues"].append("Missing structural statistics in last token")
    
    # Check 3: Normalization
    if analysis["norm"] < 0.1 or analysis["norm"] > 100:
        analysis["issues"].append(f"Norm out of range: {analysis['norm']:.4f}")
    
    analysis["vljepa_compatibility_score"] = 1.0 - (len(analysis["issues"]) * 0.2)
    analysis["vljepa_compatibility_score"] = max(0, analysis["vljepa_compatibility_score"])
    
    return analysis

# ============================================================================
# ENDPOINTS - L14, L15, L16 (same as before, abbreviated)
# ============================================================================

# [Previous L14-L16 code would go here - keeping for brevity]
# Including: load_embedding_model, load_reranker_model, etc.

embedding_model = None
embedding_processor = None
reranker_model = None
reranker_processor = None

def load_embedding_model():
    global embedding_model, embedding_processor, device
    if embedding_model is not None:
        return True
    try:
        from sentence_transformers import SentenceTransformer
        embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
        embedding_processor = "fallback"
        device = torch.device("cpu")
        print("[L14] Using MiniLM fallback for demo")
        return True
    except:
        return False

def unload_embedding():
    global embedding_model, embedding_processor
    embedding_model = None
    embedding_processor = None
    gc.collect()

def load_reranker_model():
    global reranker_model, reranker_processor
    if reranker_model is not None:
        return True
    try:
        from sentence_transformers import CrossEncoder
        reranker_model = CrossEncoder(FALLBACK_RERANKER)
        reranker_processor = "fallback"
        print("[L16] Using MiniLM Reranker")
        return True
    except:
        return False

def unload_reranker():
    global reranker_model, reranker_processor
    reranker_model = None
    reranker_processor = None
    gc.collect()

def mrl_scale(embedding: np.ndarray, target_dim: int) -> np.ndarray:
    current = len(embedding)
    if current >= target_dim:
        scaled = embedding[:target_dim]
    else:
        scaled = np.zeros(target_dim, dtype=np.float32)
        scaled[:current] = embedding
    norm = np.linalg.norm(scaled)
    if norm > 0:
        scaled = scaled / norm
    return scaled

# ============================================================================
# API ENDPOINTS
# ============================================================================

@app.get("/")
def home():
    return {
        "status": "Brain v10.0: VL-JEPA Paper-Faithful",
        "architecture": {
            "L14": "Qwen3-VL-Embedding-2B (fallback: MiniLM)",
            "L15": "Matrioshka MRL",
            "L16": "Qwen3-VL-Reranker-2B (fallback: MiniLM)",
            "L18_5": "VL-JEPA (Transformer + nomic-embed)"
        },
        "paper": "arXiv:2512.10942",
        "vljepa_config": VLJEPA_CONFIG
    }

@app.post("/embed")
async def embed_text(
    text: str = Form(...),
    dim: int = Form(DEFAULT_DIM)
):
    if vljepa_model is not None:
        unload_vljepa()
    if reranker_model is not None:
        unload_reranker()
    
    if not load_embedding_model():
        raise HTTPException(503, "Embedding not available")
    
    embedding = embedding_model.encode(text, normalize_embeddings=True)
    scaled = mrl_scale(embedding, min(dim, MAX_DIM))
    
    return {
        "embedding": scaled.tolist(),
        "dimension": len(scaled),
        "layer": "L14-L15"
    }

@app.post("/predict_thought")
async def predict_thought(
    hypergraph_state: str = Form(..., description="JSON array 2304D sequence"),
    semantic_context: str = Form(..., description="JSON array 1024D semantic"),
):
    """
    L18.5: VL-JEPA Paper-Faithful Thought Predictor v10.1
    
    Input:
    - hypergraph_state (2304D): [8 Nodes x 256, 1 Stats x 256] flat array
    - semantic_context (1024D): Qwen embedding
    """
    if embedding_model is not None:
        unload_embedding()
    if reranker_model is not None:
        unload_reranker()
    
    if not load_vljepa():
        raise HTTPException(503, "VL-JEPA not available")
    
    try:
        hg_state = np.array(json.loads(hypergraph_state), dtype=np.float32)
        sem_ctx = np.array(json.loads(semantic_context), dtype=np.float32)
        
        # Validation for v10.1 (2304D)
        EXPECTED_HG_DIM = 2304 # 9 tokens * 256
        if len(hg_state) != EXPECTED_HG_DIM:
             # Fallback logic validation or padding if legacy client calls
             if len(hg_state) < EXPECTED_HG_DIM:
                 hg_state = np.pad(hg_state, (0, EXPECTED_HG_DIM - len(hg_state)))
             else:
                 hg_state = hg_state[:EXPECTED_HG_DIM]
        
        if len(sem_ctx) != 1024:
            sem_ctx = np.pad(sem_ctx, (0, max(0, 1024 - len(sem_ctx))))[:1024]
        
        # Convert to tensors
        hg_tensor = torch.tensor(hg_state, dtype=torch.float32).unsqueeze(0).to(device)
        sem_tensor = torch.tensor(sem_ctx, dtype=torch.float32).unsqueeze(0).to(device)
        
        # Predict
        with torch.no_grad():
            predicted, action_logits = vljepa_model(hg_tensor, sem_tensor)
            action_probs = F.softmax(action_logits, dim=-1)
        
        return {
            "predicted_embedding": predicted.squeeze().cpu().numpy().tolist(),
            "embedding_dimension": predicted.shape[-1],
            "action_probabilities": action_probs.squeeze().cpu().numpy().tolist(),
            "top_action": int(torch.argmax(action_probs).item()),
            "layer": "L18.5",
            "model": "VL-JEPA v10.1 (Sequence-Based)",
            "paper": "arXiv:2512.10942"
        }
        
    except json.JSONDecodeError:
        raise HTTPException(400, "Invalid JSON")
    except Exception as e:
        raise HTTPException(500, str(e))

@app.post("/analyze_hypergraph")
async def analyze_hypergraph(
    hypergraph_state: str = Form(...)
):
    """Analyze hypergraph state for VL-JEPA compatibility."""
    try:
        hg_state = np.array(json.loads(hypergraph_state), dtype=np.float32)
        analysis = analyze_hypergraph_for_vljepa(hg_state)
        return analysis
    except Exception as e:
        raise HTTPException(500, str(e))

@app.post("/train_step")
async def train_step(
    hypergraph_state: str = Form(..., description="JSON array 2304D sequence"),
    semantic_context: str = Form(..., description="JSON array 1024D semantic"),
    target_text: str = Form(...),
    learning_rate: float = Form(1e-4) # Adaptive LR
):
    """
    L18.5: Single JEPA training step v10.1 (Active Learning).
    
    1. Predicts embedding from Inputs.
    2. Encodes Target Text using Y-Encoder (Ground Truth).
    3. Calculates JEPA Loss.
    4. Backpropagates and updates weights (AdamW).
    """
    if not load_vljepa():
        raise HTTPException(503, "VL-JEPA not available")
    
    try:
        # Load Optimizer (Lazy Init)
        global optimizer
        if optimizer is None:
            optimizer = torch.optim.AdamW(vljepa_model.parameters(), lr=learning_rate)
        
        # Update LR if changed
        for param_group in optimizer.param_groups:
            param_group['lr'] = learning_rate

        hg_state = np.array(json.loads(hypergraph_state), dtype=np.float32)
        sem_ctx = np.array(json.loads(semantic_context), dtype=np.float32)
        
        # Validation for v10.1 (2304D)
        EXPECTED_HG_DIM = 2304 
        if len(hg_state) != EXPECTED_HG_DIM:
             if len(hg_state) < EXPECTED_HG_DIM:
                 hg_state = np.pad(hg_state, (0, EXPECTED_HG_DIM - len(hg_state)))
             else:
                 hg_state = hg_state[:EXPECTED_HG_DIM]
                 
        if len(sem_ctx) != 1024:
            sem_ctx = np.pad(sem_ctx, (0, max(0, 1024 - len(sem_ctx))))[:1024]
        
        # Get target embedding from Y-Encoder (frozen)
        with torch.no_grad():
            target_embedding = y_encoder.encode(target_text, normalize_embeddings=True)
            target_tensor = torch.tensor(target_embedding, dtype=torch.float32).unsqueeze(0).to(device)
        
        # Forward pass
        hg_tensor = torch.tensor(hg_state, dtype=torch.float32).unsqueeze(0).to(device)
        sem_tensor = torch.tensor(sem_ctx, dtype=torch.float32).unsqueeze(0).to(device)
        
        # Optimization Step
        vljepa_model.train()
        optimizer.zero_grad()
        
        predicted, _ = vljepa_model(hg_tensor, sem_tensor)
        
        jepa_loss = JEPALoss()
        loss_dict = jepa_loss(predicted, target_tensor)
        
        # BACKPROPAGATION
        loss_dict["total"].backward()
        torch.nn.utils.clip_grad_norm_(vljepa_model.parameters(), 1.0) # Stability
        optimizer.step()
        
        vljepa_model.eval()
        
        return {
            "status": "weights_updated",
            "loss": {k: float(v) for k, v in loss_dict.items() if k != "total"}, # Serialize
            "total_loss": float(loss_dict["total"]),
            "cosine_similarity": loss_dict["cosine"]
        }
        
    except Exception as e:
        raise HTTPException(500, str(e))

# Global error tracking
last_error = None

@app.get("/health")
def health():
    return {
        "status": "healthy",
        "version": "10.0.0",
        "vljepa_loaded": vljepa_model is not None,
        "y_encoder_loaded": y_encoder is not None,
        "device": str(device) if device else None,
        "last_error": last_error
    }

@app.get("/layers")
def layers():
    vljepa_params = sum(p.numel() for p in vljepa_model.parameters()) if vljepa_model else 0
    
    return {
        "L14": {
            "name": "Qwen3-VL-Embedding",
            "loaded": embedding_model is not None
        },
        "L15": {
            "name": "Matrioshka MRL",
            "dimensions": [64, 128, 256, 512, 1024, 2048]
        },
        "L16": {
            "name": "Qwen3-VL-Reranker",
            "loaded": reranker_model is not None
        },
        "L18_5": {
            "name": "VL-JEPA Paper-Faithful",
            "paper": "arXiv:2512.10942",
            "architecture": {
                "predictor": f"TransformerEncoder ({VLJEPA_CONFIG['num_layers']} layers)",
                "y_encoder": Y_ENCODER_MODEL,
                "hidden_dim": VLJEPA_CONFIG['hidden_dim'],
                "num_heads": VLJEPA_CONFIG['num_heads']
            },
            "params": vljepa_params,
            "loaded": vljepa_model is not None,
            "trainable": True
        }
    }

@app.post("/save_thought")
async def save_thought(
    topic: str = Form(...),
    thought_json: str = Form(...)
):
    """
    Explicit Long-Term Memory Storage.
    Persists a thought to the Azure Bunker (64GB SSD).
    """
    try:
        data = json.loads(thought_json)
        # Use BunkerClient to save (Async/Fail-Safe)
        success = bunker.save_thought(data, topic=topic)
        
        if success:
            return {"status": "queued_for_bunker", "location": f"thoughts/{topic}"}
        else:
            raise HTTPException(500, "Failed to queue thought for bunker")
            
    except json.JSONDecodeError:
        raise HTTPException(400, "Invalid JSON")

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