import csv
import json
import time
from pathlib import Path
import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer

# Resolve file paths relative to script location
TESTS_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = TESTS_DIR.parent

MODEL_PATH = PROJECT_ROOT / "model.onnx"
TOK_PATH = PROJECT_ROOT / "tokenizer.json"
DATASET_PATH = TESTS_DIR / "test_dataset.csv"

LABEL_MAP = {"0": "low", "1": "medium", "2": "hard"}

def main():
    if not MODEL_PATH.exists() or not TOK_PATH.exists():
        raise FileNotFoundError(f"Model artifacts not found in {PROJECT_ROOT}")
    if not DATASET_PATH.exists():
        raise FileNotFoundError(f"Test dataset not found at {DATASET_PATH}")
        
    print(f"Loading tokenizer from {TOK_PATH}...")
    tokenizer = Tokenizer.from_file(str(TOK_PATH))
    
    print(f"Loading ONNX session from {MODEL_PATH}...")
    session = ort.InferenceSession(str(MODEL_PATH), providers=["CPUExecutionProvider"])
    
    # Read test dataset
    queries = []
    targets = []
    with open(DATASET_PATH, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            queries.append(row["query"])
            targets.append(row["label"])
            
    print(f"Loaded {len(queries)} test queries from {DATASET_PATH.name}.")
    
    predictions = []
    latencies = []
    
    # Warmup session
    encoded_warmup = tokenizer.encode("Warmup query")
    w_ids = np.array([[encoded_warmup.ids[0]]], dtype=np.int64)
    w_mask = np.array([[1]], dtype=np.int64)
    _ = session.run(None, {"input_ids": w_ids, "attention_mask": w_mask})
    
    for q in queries:
        t0 = time.perf_counter()
        encoded = tokenizer.encode(q)
        input_ids = np.array([encoded.ids], dtype=np.int64)
        attention_mask = np.array([encoded.attention_mask], dtype=np.int64)
        
        outputs = session.run(None, {"input_ids": input_ids, "attention_mask": attention_mask})
        logits = outputs[0][0]
        
        exp_l = np.exp(logits - np.max(logits))
        probs = exp_l / np.sum(exp_l)
        pred_idx = int(np.argmax(probs))
        t1 = time.perf_counter()
        
        latencies.append((t1 - t0) * 1000.0)
        predictions.append(LABEL_MAP[str(pred_idx)])

    # Compute Metrics
    labels = ["low", "medium", "hard"]
    conf_matrix = {t: {p: 0 for p in labels} for t in labels}
    
    correct = 0
    for target, pred in zip(targets, predictions):
        conf_matrix[target][pred] += 1
        if target == pred:
            correct += 1
            
    accuracy = correct / len(targets)
    
    metrics = {}
    f1_scores = []
    for label in labels:
        tp = conf_matrix[label][label]
        fp = sum(conf_matrix[other][label] for other in labels if other != label)
        fn = sum(conf_matrix[label][other] for other in labels if other != label)
        support = sum(conf_matrix[label].values())
        
        precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
        
        f1_scores.append(f1)
        metrics[label] = {
            "precision": precision,
            "recall": recall,
            "f1_score": f1,
            "support": support
        }
        
    macro_f1 = float(np.mean(f1_scores))
    
    p50_lat = float(np.percentile(latencies, 50))
    mean_lat = float(np.mean(latencies))
    p90_lat = float(np.percentile(latencies, 90))
    p99_lat = float(np.percentile(latencies, 99))
    qps = 1000.0 / mean_lat if mean_lat > 0 else 0.0
    
    print("\n" + "="*55)
    print("HELD-OUT EVALUATION RESULTS (301 SAMPLES)")
    print("="*55)
    print(f"Overall Accuracy: {accuracy*100:.2f}%")
    print(f"Macro F1-Score:   {macro_f1:.4f}")
    print("\nPer-Class Breakdown:")
    for label in labels:
        m = metrics[label]
        print(f"  [{label.upper():<6}] Precision: {m['precision']*100:6.2f}% | Recall: {m['recall']*100:6.2f}% | F1: {m['f1_score']:.4f} | Support: {m['support']}")
        
    print("\nConfusion Matrix (Rows=Actual, Cols=Predicted):")
    print(f"{'Actual \\ Pred':<15} {'low':<8} {'medium':<8} {'hard':<8}")
    for t in labels:
        row_str = f"{t:<15} " + " ".join(f"{conf_matrix[t][p]:<8}" for p in labels)
        print(row_str)
        
    print("\nLatency Profile (CPU):")
    print(f"  p50: {p50_lat:.2f} ms")
    print(f"  Mean: {mean_lat:.2f} ms")
    print(f"  p90: {p90_lat:.2f} ms")
    print(f"  p99: {p99_lat:.2f} ms")
    print(f"  Throughput: {qps:.1f} queries/sec")

if __name__ == "__main__":
    main()