"""
HOLD Integration Examples

cascade-lattice HOLD is MODEL AGNOSTIC.
These examples show how to plug in ANY model.

The interface is simple:
    hold_data = {
        "action_probs": [0.1, 0.2, 0.7],  # From YOUR model
        "action_names": ["a", "b", "c"],   # Optional
        "expected_values": [1.0, 0.5, 2.0], # Optional
    }
"""

# ═══════════════════════════════════════════════════════════════
# EXAMPLE 1: PyTorch Model
# ═══════════════════════════════════════════════════════════════

def pytorch_example():
    """Using HOLD with a PyTorch model."""
    import torch
    import torch.nn as nn
    from cascade.hold import Hold, HoldMeshServer
    
    # Your PyTorch model (any architecture)
    class YourPolicy(nn.Module):
        def __init__(self, obs_dim, num_actions):
            super().__init__()
            self.net = nn.Sequential(
                nn.Linear(obs_dim, 64),
                nn.ReLU(),
                nn.Linear(64, num_actions),
                nn.Softmax(dim=-1)
            )
        
        def forward(self, x):
            return self.net(x)
    
    model = YourPolicy(obs_dim=10, num_actions=4)
    hold = Hold.get()
    
    # Your inference loop
    obs = torch.randn(10)
    
    with torch.no_grad():
        action_probs = model(obs).numpy()
    
    # HOLD - model agnostic
    resolution = hold.yield_point(
        action_probs=action_probs,
        action_names=["up", "down", "left", "right"],
        brain_id="my_pytorch_model"
    )
    
    # Use resolved action
    action = resolution.action


# ═══════════════════════════════════════════════════════════════
# EXAMPLE 2: JAX/Flax Model  
# ═══════════════════════════════════════════════════════════════

def jax_example():
    """Using HOLD with a JAX model."""
    import jax
    import jax.numpy as jnp
    from cascade.hold import Hold
    
    # Your JAX model
    def policy(params, obs):
        logits = jnp.dot(obs, params['w']) + params['b']
        return jax.nn.softmax(logits)
    
    params = {
        'w': jax.random.normal(jax.random.PRNGKey(0), (10, 4)),
        'b': jnp.zeros(4)
    }
    
    hold = Hold.get()
    obs = jnp.ones(10)
    
    action_probs = policy(params, obs)
    
    # HOLD - same interface regardless of framework
    resolution = hold.yield_point(
        action_probs=np.array(action_probs),  # Convert to numpy
        action_names=["action_0", "action_1", "action_2", "action_3"],
        brain_id="my_jax_model"
    )
    
    action = resolution.action


# ═══════════════════════════════════════════════════════════════
# EXAMPLE 3: Hugging Face Transformers
# ═══════════════════════════════════════════════════════════════

def hf_transformers_example():
    """Using HOLD with Hugging Face models."""
    from transformers import AutoModelForSequenceClassification, AutoTokenizer
    import torch
    from cascade.hold import Hold
    
    # Load any HF model
    model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    
    hold = Hold.get()
    
    text = "The movie was really great!"
    inputs = tokenizer(text, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.softmax(outputs.logits, dim=-1).numpy()[0]
    
    # HOLD at classification decision
    resolution = hold.yield_point(
        action_probs=probs,
        action_names=["negative", "positive"],
        observation={"text": text},
        brain_id="bert-sentiment"
    )
    
    label = resolution.action


# ═══════════════════════════════════════════════════════════════
# EXAMPLE 4: Stable Baselines3 (RL)
# ═══════════════════════════════════════════════════════════════

def sb3_example():
    """Using HOLD with Stable Baselines3 RL agents."""
    from stable_baselines3 import PPO
    from cascade.hold import Hold
    import numpy as np
    
    # Load or create SB3 model
    model = PPO.load("ppo_cartpole")
    
    hold = Hold.get()
    
    obs = np.zeros(4)  # CartPole observation
    
    # Get action distribution
    action, _states = model.predict(obs, deterministic=False)
    
    # Get the actual probabilities from the policy
    obs_tensor = model.policy.obs_to_tensor(obs)[0]
    distribution = model.policy.get_distribution(obs_tensor)
    probs = distribution.distribution.probs.detach().numpy()[0]
    
    # HOLD
    resolution = hold.yield_point(
        action_probs=probs,
        action_names=["left", "right"],
        brain_id="ppo_cartpole"
    )
    
    action = resolution.action


# ═══════════════════════════════════════════════════════════════
# EXAMPLE 5: OpenAI API (LLM)
# ═══════════════════════════════════════════════════════════════

def openai_example():
    """Using HOLD with OpenAI API for tool/function selection."""
    import openai
    from cascade.hold import Hold
    import numpy as np
    
    hold = Hold.get()
    
    # Define available tools
    tools = [
        {"name": "search", "description": "Search the web"},
        {"name": "calculate", "description": "Do math"},
        {"name": "write", "description": "Write text"},
        {"name": "none", "description": "No tool needed"},
    ]
    
    # Get LLM to score tools
    response = openai.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": "What's 2+2?"}],
        # In practice, use function calling or logprobs
    )
    
    # Simulated tool probabilities from LLM
    tool_probs = np.array([0.1, 0.7, 0.1, 0.1])  # Would come from logprobs
    
    # HOLD at tool selection
    resolution = hold.yield_point(
        action_probs=tool_probs,
        action_names=[t["name"] for t in tools],
        observation={"query": "What's 2+2?"},
        brain_id="gpt-4-tool-selector"
    )
    
    selected_tool = tools[resolution.action]


# ═══════════════════════════════════════════════════════════════
# EXAMPLE 6: Custom Model with Mesh Visualization
# ═══════════════════════════════════════════════════════════════

def custom_with_mesh():
    """Full example with mesh server."""
    import numpy as np
    import time
    from cascade.hold import Hold, HoldMeshServer
    
    # Start mesh server for 3D visualization
    server = HoldMeshServer(port=8766)
    server.start()
    
    hold = Hold.get()
    
    # Your custom decision function (could be anything)
    def my_model_decide(state):
        # This is YOUR model - any logic
        probs = np.random.dirichlet(np.ones(5))  # Random for demo
        values = np.random.randn(5) * 0.5
        return probs, values
    
    action_names = ["attack", "defend", "heal", "flee", "wait"]
    
    try:
        for step in range(100):
            state = {"health": 50, "enemies": 3}
            
            probs, values = my_model_decide(state)
            
            # Build hold data (the universal interface)
            hold_data = {
                "action_probs": probs.tolist(),
                "action_names": action_names,
                "expected_values": values.tolist(),
                "selected_action": int(np.argmax(probs)),
                "value": float(values[np.argmax(probs)]),
                "step": step,
                "model_card": {
                    "name": "My Custom Model",
                    "state": state,
                }
            }
            
            # Update mesh visualization
            server.update(hold_data)
            
            # Optionally yield to HOLD for human intervention
            if step % 10 == 0:  # Every 10 steps, allow intervention
                resolution = hold.yield_point(
                    action_probs=probs,
                    action_names=action_names,
                    value=float(values[np.argmax(probs)]),
                    brain_id="my_custom_model"
                )
                action = resolution.action
            else:
                action = np.argmax(probs)
            
            print(f"Step {step}: {action_names[action]}")
            time.sleep(0.1)
            
    finally:
        server.stop()


# ═══════════════════════════════════════════════════════════════
# THE MINIMAL INTERFACE
# ═══════════════════════════════════════════════════════════════

def minimal_example():
    """
    The absolute minimum to use HOLD.
    
    HOLD doesn't care what model you use.
    It just needs action_probs (a probability distribution).
    """
    import numpy as np
    from cascade.hold import Hold
    
    # 1. Get the HOLD singleton
    hold = Hold.get()
    
    # 2. Your model produces probabilities (any shape, any source)
    action_probs = np.array([0.1, 0.3, 0.4, 0.2])
    
    # 3. HOLD yields at decision point
    resolution = hold.yield_point(action_probs=action_probs)
    
    # 4. Use the resolved action
    action = resolution.action  # int
    was_override = resolution.was_override  # bool
    
    return action


if __name__ == "__main__":
    print("HOLD is model-agnostic.")
    print()
    print("Required: action_probs (numpy array or list)")
    print("Optional: action_names, value, observation, brain_id")
    print()
    print("Run custom_with_mesh() to see it in action with 3D visualization.")
