import os
import sys
import json
import re
import torch
from transformers import AutoTokenizer

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from model import RecursiveCausalLM, ModelConfig, KVCache

def extract_last_number(text):
    nums = re.findall(r'-?\d*\.?\d+', text.replace(',', ''))
    return float(nums[-1]) if nums else None

def generate_completion(model, tokenizer, prompt, device, max_tokens=100):
    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
    kv_cache = KVCache(model.config, max_batch_size=1, device=device, dtype=torch.float16)
    
    generated_tokens = []
    use_amp = (device.type == "cuda")
    
    with torch.no_grad():
        with torch.amp.autocast(device_type="cuda", enabled=use_amp, dtype=torch.float16):
            logits, _ = model(input_ids, kv_cache=kv_cache)
            
    next_token_logits = logits[0, -1, :]
    
    for _ in range(max_tokens):
        # Temperature 0.1 for maximum deterministic accuracy during evaluation
        next_token_logits = next_token_logits / 0.1
        v, _ = torch.topk(next_token_logits, min(50, next_token_logits.size(-1)))
        next_token_logits[next_token_logits < v[-1]] = -float('Inf')
        probs = torch.softmax(next_token_logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1).item()
        generated_tokens.append(next_token)
        
        if next_token == tokenizer.eos_token_id:
            break
            
        curr_input = torch.tensor([[next_token]], dtype=torch.long, device=device)
        with torch.no_grad():
            with torch.amp.autocast(device_type="cuda", enabled=use_amp, dtype=torch.float16):
                logits, _ = model(curr_input, kv_cache=kv_cache)
                next_token_logits = logits[0, -1, :]
                
    return tokenizer.decode(generated_tokens)

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("=======================================================================")
    print("🎯 THE OFFICIAL PRM MATHEMATICAL EVALUATOR")
    print("=======================================================================\n")
    
    config = ModelConfig(
        vocab_size=50272,
        d_model=768,
        n_iterations=16,
        n_heads=12,
        n_kv_heads=4,
        d_ff=2048,
        max_seq_len=512
    )
    
    model_path = "/home/zeus/micro_llm_200m/uct_target_rl_math.pt"
    dataset_path = "/home/zeus/micro_llm_200m/three_domain_sft_dataset_math_heavy.json"
    tokenizer_path = "/home/zeus/micro_llm_200m/tokenizer"
    
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
    model = RecursiveCausalLM(config).to(device)
    model.load_state_dict(torch.load(model_path, map_location=device, weights_only=False)["model_state_dict"], strict=False)
    model.eval()
    
    # Load dataset and extract math prompts
    with open(dataset_path, "r", encoding="utf-8") as f:
        data = json.load(f)
        
    math_items = [item for item in data if item.get("domain") == "math"]
    print(f"Loaded {len(math_items)} total arithmetic prompts.")
    
    # Test on the first 50 math prompts
    eval_count = min(50, len(math_items))
    print(f"Running rigorous evaluation on {eval_count} test cases...\n")
    
    correct_count = 0
    for idx in range(eval_count):
        item = math_items[idx]
        prompt = item["prompt"]
        gold_response = item["response"]
        
        gold_num = extract_last_number(gold_response)
        
        # Generate completion
        completion = generate_completion(model, tokenizer, prompt, device)
        gen_num = extract_last_number(completion)
        
        is_correct = False
        if gold_num is not None and gen_num is not None and abs(gold_num - gen_num) < 1e-4:
            is_correct = True
            correct_count += 1
            
        status_str = "✅ PASS" if is_correct else "❌ FAIL"
        print(f"Sample {idx+1:2d}/50 | Prompt: '{prompt.strip()}'")
        print(f"        | Generated: '{completion.strip()}'")
        print(f"        | Expected:  {gold_num} | Predicted: {gen_num} | {status_str}")
        print("-" * 75)
        
    accuracy = (correct_count / eval_count) * 100
    print("\n=======================================================================")
    print("📊 FINAL MATHEMATICAL EVALUATION REPORT")
    print("=======================================================================")
    print(f"-> Total Evaluated Prompts:  {eval_count}")
    print(f"-> Total Correct Answers:     {correct_count}")
    print(f"-> Exact Match (EM) Accuracy: {accuracy:.2f}%")
    print("=======================================================================")

if __name__ == '__main__':
    main()