"""
Ejemplos de uso del modelo Metamatemático NLE v7.0
====================================================

Requisitos:
    pip install torch>=2.5.1 torch-geometric>=2.7.0 huggingface-hub>=0.31

Ejecutar:
    python example.py
"""

from model import MetamaticoPredictor, ActorCriticNetwork, encode_query, encode_goal, empty_graph_data, ACTIONS
import torch


# ─── Ejemplo 1: Uso rápido con from_pretrained ────────────────────────────────

print("=" * 60)
print("Ejemplo 1 — Uso rápido")
print("=" * 60)

model = MetamaticoPredictor.from_pretrained()

result = model.predict("Prove that the square root of 2 is irrational")
print(f"Query    : Prove that sqrt(2) is irrational")
print(f"Action   : {result['action']}")
print(f"Confidence: {result['confidence']:.1%}")
print(f"Probs    : {result['probabilities']}")
print()


# ─── Ejemplo 2: Predicciones en lote ─────────────────────────────────────────

print("=" * 60)
print("Ejemplo 2 — Predicciones en lote")
print("=" * 60)

queries = [
    "Prove that there are infinitely many prime numbers",
    "What is the definition of a category?",
    "Compute the integral of e^x from 0 to 1",
    "Is the group Z/6Z isomorphic to Z/2Z × Z/3Z?",
    "What is 2 + 2?",
    "Hello, who are you?",
    "Formalize the first isomorphism theorem in Lean 4",
    "Prove by induction that sum of first n naturals is n(n+1)/2",
]

results = model.predict_batch(queries)
print(f"{'Action':<12} {'Conf':>6}  Query")
print("-" * 60)
for q, r in zip(queries, results):
    print(f"{r['action']:<12} {r['confidence']:>5.0%}  {q[:55]}")
print()


# ─── Ejemplo 3: Con goal Lean 4 ───────────────────────────────────────────────

print("=" * 60)
print("Ejemplo 3 — Con goal Lean 4")
print("=" * 60)

lean_queries = [
    ("Is this valid?",      "⊢ ∀ n : ℕ, n + 0 = n"),
    ("How do I prove this?","⊢ ∀ a b : ℤ, a * b = b * a"),
    ("Lean goal check",     "⊢ ∃ x : ℝ, x ^ 2 = 2"),
]

for query, goal in lean_queries:
    r = model.predict(query, lean_goal=goal)
    print(f"Query : {query}")
    print(f"Goal  : {goal}")
    print(f"Action: {r['action']}  (conf={r['confidence']:.1%})")
    print()


# ─── Ejemplo 4: Uso directo con PyTorch ───────────────────────────────────────

print("=" * 60)
print("Ejemplo 4 — Forward pass directo con PyTorch")
print("=" * 60)

from huggingface_hub import hf_hub_download

weights_path = hf_hub_download(
    repo_id="metamatematico/Metamatematico",
    filename="neural_agent.pt"
)

net = ActorCriticNetwork()
net.load_state_dict(torch.load(weights_path, map_location="cpu", weights_only=True))
net.eval()

graph   = empty_graph_data()
q_emb   = encode_query("prove the fundamental theorem of algebra").unsqueeze(0)
goal_emb = encode_goal("⊢ ∀ p : Polynomial ℂ, p.degree ≥ 1 → ∃ z : ℂ, p.eval z = 0").unsqueeze(0)

with torch.no_grad():
    logits, value = net(graph, q_emb, goal_emb=goal_emb)
    probs = torch.softmax(logits, dim=-1)

print(f"Logits : {logits.squeeze().tolist()}")
print(f"Probs  : {probs.squeeze().tolist()}")
print(f"Value  : {value.item():.4f}")
print(f"Action : {ACTIONS[int(probs.argmax())]}")
print()


# ─── Ejemplo 5: Cargar best checkpoint ───────────────────────────────────────

print("=" * 60)
print("Ejemplo 5 — Cargar best.pt (mejor checkpoint supervisado)")
print("=" * 60)

best_path = hf_hub_download(
    repo_id="metamatematico/Metamatematico",
    filename="best.pt"
)
model_best = MetamaticoPredictor(best_path)
r = model_best.predict("Show that every finite group of prime order is cyclic")
print(f"Action: {r['action']}  (conf={r['confidence']:.1%})")