import os
import sys
import time
import torch
from sentence_transformers import SentenceTransformer
import pandas as pd

# ── 1. Path setup (same pattern as run_ppo-2.py) ────────────────────────────
ROOT = os.path.dirname(__file__)
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)

from model import RouterNetwork
from grpo_trainer import simulate_rag_pipeline   # reuse your existing tool caller

# ── 2. Config ────────────────────────────────────────────────────────────────
CHECKPOINT  = os.path.join(ROOT, "adapters/router_grpo_epoch_16.pt")   # ← change epoch here
EMBED_MODEL = "all-MiniLM-L6-v2"
DEVICE      = "cuda" if torch.cuda.is_available() else "cpu"

ACTION_NAMES = {
    0: "Guidelines  (academic rules / visa / scholarships)",
    1: "Web-Search   (news / announcements)",
    2: "CBNU Database (courses / departments / contacts)",
}

# ── 3. Load embedder + model ─────────────────────────────────────────────────
embedder  = SentenceTransformer(EMBED_MODEL)
embed_dim = embedder.get_sentence_embedding_dimension()
model     = RouterNetwork(embedding_dim=embed_dim, num_classes=3).to(DEVICE)

# Load checkpoint (supports both raw state_dict and wrapped dict)
sd = torch.load(CHECKPOINT, map_location="cpu")
if isinstance(sd, dict) and "model_state_dict" in sd:
    model.load_state_dict(sd["model_state_dict"], strict=False)
else:
    model.load_state_dict(sd, strict=False)

model.eval()
print(f"✅  Loaded checkpoint: {CHECKPOINT}")
print(f"    Device: {DEVICE}\n")

# ── 4. Single-query inference function ───────────────────────────────────────
def route_query(query: str, call_tool: bool = True):
    """
    Route one query through the loaded GRPO checkpoint.

    Args:
        query     : The user question string.
        call_tool : If True, also calls simulate_rag_pipeline and prints the answer.
                    Set False for routing-only / fast testing.

    Returns:
        action_idx (int), action_name (str), context (str), answer (str)
    """
    # Embed
    with torch.no_grad():
        emb = embedder.encode(query, convert_to_tensor=True).to(DEVICE)
        action, log_prob, _ = model.get_action(emb)
        action_idx = int(action.item())

    action_name = ACTION_NAMES[action_idx]
    print(f"Query   : {query}")
    print(f"→ Routed to Action {action_idx}: {action_name}")
    print(f"  log_prob = {log_prob.item():.4f}")

    context, answer = "", ""
    if call_tool:
        print("  Calling tool …")
        context, answer = simulate_rag_pipeline(query, action_idx)
        print(f"  Answer  : {answer[:300]}")

    return action_idx, action_name, context, answer


# ── 5. Main: test with one or more queries ───────────────────────────────────
if __name__ == "__main__":
    test_queries = [
        "What are the graduation requirements for international students?",
        "How do I register for courses next semester?",
        "What is the latest news from CBNU?",
    ]
    test_queries = [
        "2학기 개강일은?",
        "알바하고 싶은데 어떤 조건이 필요한가?",
        "아르바이트 하려는데 조건은?",
        "시간제 취업 방법은?",
        "건강보험 어떻게 가입해?",
        "성적확인 기간은?"
    ]
    
    eval_df = pd.read_csv('ko_eval_queries.csv')
    test_queries = eval_df['question'].tolist()
    results = []
    for idx, row in eval_df.iterrows():
        q = row.get('question')
        ground = row.get('action_label')
        start = time.perf_counter()
        pred_idx, pred_name, context, answer = route_query(q, call_tool=False)
        end = time.perf_counter()
        elapsed = end - start
        results.append({
            'question': q,
            'ground_truth': ground,
            'predicted_label': pred_idx,
            'time_took': elapsed,
        })
        
        
    out_df = pd.DataFrame(results)
    out_df.to_csv('ko_eval_with_timing_results_epoch_16.csv', index=False)
    print(f"Wrote results to eval_with_timing_results.csv ({len(out_df)} rows)")