import os
import sys
import time
import traceback
import pandas as pd

# Ensure local package imports work
ROOT = os.path.dirname(__file__)
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)

import run_inference as ri


def predict_and_time(input_csv: str, output_csv: str):
    df = pd.read_csv(input_csv)
    results = []

    for idx, row in df.iterrows():
        q = row.get('question')
        ground = row.get('action_label')
        try:
            start = time.perf_counter()
            # fast routing-only prediction using loaded embedder + model
            with ri.torch.no_grad():
                emb = ri.embedder.encode(q, convert_to_tensor=True).to(ri.DEVICE)
                action, log_prob, _ = ri.model.get_action(emb)
                pred_idx = int(action.item())
            end = time.perf_counter()
            elapsed = end - start
        except Exception:
            pred_idx = None
            elapsed = None
            traceback.print_exc()

        results.append({
            'question': q,
            'ground_truth': ground,
            'predicted_label': pred_idx,
            'time_took': elapsed,
        })

    out_df = pd.DataFrame(results)
    out_df.to_csv(output_csv, index=False)
    print(f"Wrote results to {output_csv} ({len(out_df)} rows)")


if __name__ == '__main__':
    inp = os.path.join(ROOT, 'eval_queries.csv')
    out = os.path.join(ROOT, 'eval_with_timing_results.csv')
    print(f"Input: {inp}\nOutput: {out}\nDevice: {ri.DEVICE}")
    predict_and_time(inp, out)