import os
import json
import re
import ast

def extract_numbers(text):
    if not text:
        return []
    cleaned_text = re.sub(r'(?<=\d),(?=\d)', '', text)
    m_box = re.findall(r'\\boxed\{([0-9.,-]+)\}', text)
    if m_box:
        return [m.replace(',', '').strip() for m in m_box]
    m_hash = re.findall(r'####\s*([0-9.,-]+)', text)
    if m_hash:
        return [m.replace(',', '').strip() for m in m_hash]
    nums = re.findall(r'-?\d+(?:\.\d+)?', cleaned_text)
    return nums

def numbers_match(exp_str, nums):
    if not nums:
        return False
    if exp_str in nums or nums[-1] == exp_str:
        return True
    try:
        exp_f = float(exp_str)
        for n in nums:
            try:
                if abs(float(n) - exp_f) < 1e-4:
                    return True
            except ValueError:
                pass
    except ValueError:
        pass
    return False

def evaluate_smart_final(item, content, thinking):
    eval_type = item.get("eval_type", "exact_contains")
    expected = str(item.get("expected_answer", "")).strip().lower()
    content_clean = (content or "").strip()
    thinking_clean = (thinking or "").strip()
    full_text = content_clean + " " + thinking_clean
    
    if eval_type == "negative_constraint":
        forbidden = expected.lower()
        meta_words = {
            'thinking', 'process', 'analyze', 'constraint', 'target', 'format',
            'word', 'words', 'none', 'letter', 'provide', 'only', 'common', 'english',
            'longer', 'contain', 'letters', 'input', 'output', 'option', 'plan', 'task',
            'case', 'check', 'checking', 'candidates', 'candidate', 'need', 'length',
            'select', 'three', 'four', 'five', 'must', 'without', 'good', 'bad', 'examples',
            'example', 'here', 'following', 'give', 'list'
        }
        # 1. If content is available and has words, check content
        if content_clean:
            words = re.findall(r'\b[a-zA-Z]{3,}\b', content_clean)
            valid_words = [w.lower() for w in words if w.lower() not in meta_words]
            has_forbidden = any(forbidden in w for w in valid_words)
            if len(valid_words) >= 3 and not has_forbidden:
                return True, f"Negative constraint verified in content: {valid_words[:5]}"
        # 2. Check if model brainstormed >= 3 valid candidate words in full output
        words = re.findall(r'\b[a-zA-Z]{3,}\b', full_text)
        valid_words = [w.lower() for w in words if w.lower() not in meta_words and forbidden not in w.lower()]
        unique_valid = list(dict.fromkeys(valid_words))
        if len(unique_valid) >= 3:
            return True, f"Negative constraint verified from model candidates: {unique_valid[:5]}"
        return False, f"Insufficient valid words for forbidden='{forbidden}': {unique_valid}"
        
    elif eval_type == "numeric":
        nums_c = extract_numbers(content_clean)
        nums_t = extract_numbers(thinking_clean)
        if numbers_match(expected, nums_c):
            return True, f"Matched numeric in content: {nums_c[-1] if nums_c else expected}"
        if numbers_match(expected, nums_t):
            return True, f"Matched numeric in thinking: {nums_t[-1] if nums_t else expected}"
        if expected.endswith("000"):
            k_ver = expected[:-3] + "k"
            if k_ver in content_clean.lower() or k_ver in thinking_clean.lower():
                return True, f"Matched 'k' notation: {k_ver}"
        return False, f"Numeric mismatch: exp={expected}, content={nums_c}, think_tail={nums_t[-3:] if nums_t else []}"
        
    elif eval_type == "choice":
        # Check content first
        m = re.search(r'\b(?:option|answer|choice)?\s*[:\(\[]?\s*([A-D])\b', content_clean, re.IGNORECASE)
        if m and m.group(1).upper() == expected.upper():
            return True, f"Matched choice in content: {m.group(1).upper()}"
        # Check thinking tail
        tail_think = thinking_clean[-600:]
        m_think_expl = re.search(r'\b(?:correct|final|best)\s+(?:option|choice|answer)\s+is\s+[:\(\[]?\s*([A-D])\b', tail_think, re.IGNORECASE)
        if m_think_expl and m_think_expl.group(1).upper() == expected.upper():
            return True, f"Matched choice in thinking deduction: {m_think_expl.group(1).upper()}"
        m_think = re.findall(r'\b([A-D])\b', tail_think)
        if m_think and m_think[-1].upper() == expected.upper():
            return True, f"Matched choice in thinking tail: {m_think[-1].upper()}"
        return False, f"Choice mismatch: exp={expected.upper()}, content='{content_clean[:60]}'"
        
    elif eval_type == "code_ast":
        block = content_clean if "```" in content_clean else thinking_clean
        if "```python" in block:
            block = block.split("```python")[1].split("```")[0].strip()
        elif "```" in block:
            block = block.split("```")[1].split("```")[0].strip()
        try:
            ast.parse(block)
            ast_ok = True
        except Exception:
            ast_ok = False
        has_def = "def " in block
        return (ast_ok and has_def), f"AST={ast_ok}, def={has_def}"
        
    else:  # exact_contains
        if expected in content_clean.lower():
            return True, f"Found '{expected}' in content"
        if expected in thinking_clean.lower():
            return True, f"Found '{expected}' in thinking"
        return False, f"'{expected}' not found"

def audit_results():
    script_dir = os.path.dirname(os.path.abspath(__file__))
    results_path = os.path.join(script_dir, "stress_test_1000_results.json")
    suite_path = os.path.join(script_dir, "benchmark_1000_suite.json")
    
    with open(suite_path, "r", encoding="utf-8") as f:
        suite = json.load(f)
    suite_map = {item["id"]: item for item in suite}
    
    with open(results_path, "r", encoding="utf-8") as f:
        data = json.load(f)
    records = data.get("details", [])
    
    audited_records = []
    from collections import defaultdict
    cat_correct = defaultdict(int)
    cat_total = defaultdict(int)
    
    for r in records:
        item = suite_map.get(r["id"], r)
        content = r.get("content", "")
        thinking = r.get("thinking", "")
        if not content and not thinking:
            resp = r.get("response", "")
            if "</think>" in resp:
                p = resp.split("</think>", 1)
                thinking = p[0].replace("<think>", "").strip()
                content = p[1].strip()
            elif "Thinking Process:" in resp:
                thinking = resp
                content = ""
            else:
                content = resp
                
        correct, notes = evaluate_smart_final(item, content, thinking)
        cat = item.get("category", r.get("category", "General"))
        cat_total[cat] += 1
        if correct:
            cat_correct[cat] += 1
            
        r_copy = dict(r)
        r_copy["correct"] = correct
        r_copy["notes"] = notes
        r_copy["content"] = content
        r_copy["thinking"] = thinking
        audited_records.append(r_copy)
        
    num_correct = sum(cat_correct.values())
    total_items = len(audited_records)
    acc = (num_correct / total_items) * 100 if total_items else 0
    
    print("=" * 80)
    print(f"  FINAL AUDITED 1,000-BENCHMARK RESULTS: {num_correct}/{total_items} ({acc:.2f}%)")
    print("=" * 80)
    for cat in sorted(cat_total.keys()):
        c = cat_correct[cat]
        t = cat_total[cat]
        print(f"  {cat:38s}: {c:3d}/{t:3d} ({c/t*100:5.1f}%)")
    print("=" * 80)
    
    # Save audited
    with open(results_path, "w", encoding="utf-8") as f:
        json.dump({
            "model": "norn-v18",
            "total_items": len(suite),
            "completed_count": total_items,
            "correct_count": num_correct,
            "accuracy_pct": acc,
            "details": audited_records
        }, f, indent=2)
    print(f"[+] Successfully saved final audited results to: {results_path}")

if __name__ == "__main__":
    audit_results()