#!/usr/bin/env python3
"""
Test unified search vs BM25 baseline on 20 questions.
Runs both engines, compares results, reports quality metrics.
"""
import json
import sys
import time
import os

sys.path.insert(0, "/data/knowledge/maritime")
os.environ['OLLAMA_EMBED_URL'] = 'http://127.0.0.1:11435/api/embed'

# Load both engines
from search_engine import search as bm25_search
from search_unified import search as unified_search, _do_classify as classify_query

# 20 test questions covering all 5 types
QUESTIONS = [
    # fact (5)
    {"q": "2019年全国注册船员总数是多少？", "type": "fact", "expected_keywords": ["1659188", "船员"]},
    {"q": "船员适任证书的有效期是多少年？", "type": "fact", "expected_keywords": ["证书", "有效期"]},
    {"q": "琼州海峡船舶定线制是什么时候施行的？", "type": "fact", "expected_keywords": ["琼州海峡", "施行"]},
    {"q": "船舶进出港报告制度的法律依据是什么？", "type": "fact", "expected_keywords": ["进出港", "报告"]},
    {"q": "内河船舶船员培训机构有多少家？", "type": "fact", "expected_keywords": ["培训机构", "内河"]},
    # compare (4)
    {"q": "海上交通安全法与水污染防治法在船舶管理方面的规定有何异同？", "type": "compare", "expected_keywords": ["海上交通安全法", "水污染防治法"]},
    {"q": "对比2019年和2020年中国船员发展报告的主要变化", "type": "compare", "expected_keywords": ["2019", "2020", "船员"]},
    {"q": "国内航行船舶和国际航行船舶的检验要求有什么区别？", "type": "compare", "expected_keywords": ["国内", "国际", "检验"]},
    {"q": "海事局规范性文件与交通运输部规章有何区别？", "type": "compare", "expected_keywords": ["海事局", "交通运输部"]},
    # relation (4)
    {"q": "海事局的上级主管部门是哪个？", "type": "relation", "expected_keywords": ["海事局", "交通运输部"]},
    {"q": "船员适任证书制度依据哪些国际公约？", "type": "relation", "expected_keywords": ["适任证书", "STCW", "公约"]},
    {"q": "防治船舶污染海洋环境管理条例由谁制定？", "type": "relation", "expected_keywords": ["防治", "污染", "条例"]},
    {"q": "船舶检验机构的资质由哪个部门认可？", "type": "relation", "expected_keywords": ["检验", "资质", "认可"]},
    # synthesis (4)
    {"q": "请梳理中国海事法律法规的层级体系", "type": "synthesis", "expected_keywords": ["全国人大", "国务院", "交通运输部", "海事局"]},
    {"q": "总结船员管理方面的主要法规和政策", "type": "synthesis", "expected_keywords": ["船员", "管理", "法规"]},
    {"q": "概述海洋环境保护相关的法律框架", "type": "synthesis", "expected_keywords": ["海洋", "环境", "保护"]},
    {"q": "分析港口管理的法规体系", "type": "synthesis", "expected_keywords": ["港口", "管理"]},
    # temporal (3)
    {"q": "海上交通安全法最新修订是什么时候？", "type": "temporal", "expected_keywords": ["海上交通安全法", "修订"]},
    {"q": "船员培训管理规则的最新版本何时施行？", "type": "temporal", "expected_keywords": ["培训", "施行"]},
    {"q": "水污染防治法经历了哪些修订？", "type": "temporal", "expected_keywords": ["水污染防治法", "修订"]},
]

print("=" * 70)
print("  LambdaRAG Unified Search vs BM25 Baseline Test")
print("  20 questions × 2 engines")
print("=" * 70)

results = []
for i, q in enumerate(QUESTIONS):
    print(f"\n{'─' * 70}")
    print(f"Q{i+1} [{q['type']}] {q['q']}")

    # Classify
    classified = classify_query(q['q'])
    classify_correct = classified == q['type']

    # BM25 baseline
    t0 = time.time()
    try:
        bm25_results = bm25_search(q['q'], top_k=5)
    except Exception as e:
        bm25_results = []
        print(f"  BM25 ERROR: {e}")
    bm25_time = time.time() - t0

    # Unified
    t0 = time.time()
    try:
        unified_results = unified_search(q['q'], top_k=5)
    except Exception as e:
        unified_results = []
        print(f"  UNIFIED ERROR: {e}")
    unified_time = time.time() - t0

    # Quality check: do results contain expected keywords?
    def check_relevance(search_results, keywords):
        if not search_results:
            return 0, []
        all_text = ' '.join([r.get('text', '') + ' ' + r.get('source', '') for r in search_results])
        found = [kw for kw in keywords if kw in all_text]
        return len(found) / max(len(keywords), 1), found

    bm25_score, bm25_found = check_relevance(bm25_results, q['expected_keywords'])
    unified_score, unified_found = check_relevance(unified_results, q['expected_keywords'])

    # Get unified metadata
    u_type = unified_results[0].get('_query_type', '?') if unified_results else '?'
    u_strategy = unified_results[0].get('_strategy', '?') if unified_results else '?'

    print(f"  Classify: predicted={classified} expected={q['type']} {'✓' if classify_correct else '✗'}")
    print(f"  BM25:    {bm25_time:.2f}s | relevance={bm25_score:.0%} | found={bm25_found}")
    print(f"  Unified: {unified_time:.2f}s | relevance={unified_score:.0%} | found={unified_found}")
    print(f"  Route:   {u_type} → {u_strategy}")

    # Check: unified sources different from BM25?
    bm25_sources = set(r.get('source', '')[:30] for r in bm25_results[:3])
    unified_sources = set(r.get('source', '')[:30] for r in unified_results[:3])
    sources_differ = bm25_sources != unified_sources

    record = {
        'id': i + 1,
        'question': q['q'],
        'expected_type': q['type'],
        'classified_type': classified,
        'classify_correct': classify_correct,
        'bm25_time': round(bm25_time, 2),
        'unified_time': round(unified_time, 2),
        'bm25_relevance': round(bm25_score, 2),
        'unified_relevance': round(unified_score, 2),
        'bm25_results_count': len(bm25_results),
        'unified_results_count': len(unified_results),
        'sources_differ': sources_differ,
        'unified_route': f"{u_type} → {u_strategy}",
        'winner': 'unified' if unified_score > bm25_score else ('tie' if unified_score == bm25_score else 'bm25'),
    }
    results.append(record)

# Summary
print(f"\n{'=' * 70}")
print("  SUMMARY")
print(f"{'=' * 70}")

# Classification accuracy
correct = sum(1 for r in results if r['classify_correct'])
print(f"\nClassification accuracy: {correct}/{len(results)} ({correct/len(results)*100:.0f}%)")

# Relevance comparison
bm25_avg = sum(r['bm25_relevance'] for r in results) / len(results)
unified_avg = sum(r['unified_relevance'] for r in results) / len(results)
print(f"\nRelevance (keyword match):")
print(f"  BM25 avg:    {bm25_avg:.0%}")
print(f"  Unified avg: {unified_avg:.0%}")

# Winner counts
wins = {'unified': 0, 'bm25': 0, 'tie': 0}
for r in results:
    wins[r['winner']] += 1
print(f"\nWins: Unified={wins['unified']} BM25={wins['bm25']} Tie={wins['tie']}")

# By type
print(f"\nBy question type:")
for qtype in ['fact', 'compare', 'relation', 'synthesis', 'temporal']:
    type_results = [r for r in results if r['expected_type'] == qtype]
    if not type_results:
        continue
    b_avg = sum(r['bm25_relevance'] for r in type_results) / len(type_results)
    u_avg = sum(r['unified_relevance'] for r in type_results) / len(type_results)
    wins_u = sum(1 for r in type_results if r['winner'] == 'unified')
    wins_b = sum(1 for r in type_results if r['winner'] == 'bm25')
    print(f"  {qtype:10s}: BM25={b_avg:.0%} Unified={u_avg:.0%} | U wins={wins_u} B wins={wins_b}")

# Timing
bm25_avg_time = sum(r['bm25_time'] for r in results) / len(results)
unified_avg_time = sum(r['unified_time'] for r in results) / len(results)
print(f"\nAvg search time:")
print(f"  BM25:    {bm25_avg_time:.2f}s")
print(f"  Unified: {unified_avg_time:.2f}s")

# Sources differ
differ_count = sum(1 for r in results if r['sources_differ'])
print(f"\nSources differ in {differ_count}/{len(results)} questions")

# Issues
print(f"\nPotential issues:")
for r in results:
    if r['unified_results_count'] == 0:
        print(f"  Q{r['id']}: Unified returned 0 results!")
    if not r['classify_correct']:
        print(f"  Q{r['id']}: Misclassified as {r['classified_type']} (expected {r['expected_type']})")
    if r['unified_relevance'] < r['bm25_relevance']:
        print(f"  Q{r['id']}: Unified worse than BM25 ({r['unified_relevance']:.0%} vs {r['bm25_relevance']:.0%})")

# Save results
with open('/data/knowledge/maritime/test_unified_results.json', 'w') as f:
    json.dump(results, f, ensure_ascii=False, indent=2)
print(f"\nResults saved to test_unified_results.json")