module/evaluator.py

import os
import traceback
import numpy as np
from datetime import datetime

from module.file_manager import get_history_images, load_metadata, save_metadata
from module.text_classifier import classify_scene
from module.text_matcher import compare_text_scene, compare_text_scene_cached, compare_text_scene_pure_bge


DINO_PRESCREEN_THRESHOLD = 0.3


def evaluate_scene_consistency_cached(matcher, query_paths, username,
                                       ocr_engine=None, bge_tokenizer=None,
                                       bge_model=None, scene_type_q='complex',
                                       qd_search_fn=None, dinov2_extractor=None,
                                       precomputed_ocr_result=None,
                                      precomputed_bge_vec=None,
                                      qd_cache=None):
    """Qdrant ANN 搜索管线: DINOv2/BGE 预筛 + MASt3R 精确判定 + 早停

    流程 (与 test_5.ipynb 一致):
    文本场景:
      1. 提取查询图 OCR 文本 -> BGE 编码
      2. Qdrant ANN 搜索 (BGE 向量) -> top-K 候选
      3. 对候选逐一 compare_text_scene_pure_bge (纯BGE语义向量)
      4. 早停: 匹配到同场景即跳过剩余
    复杂场景:
      1. 提取查询图 DINOv2 特征
      2. Qdrant ANN 搜索 (DINOv2 向量) -> top-K 候选
      3. DINOv2 预筛: 相似度 < 0.3 直接跳过 MASt3R
      4. 通过预筛的候选 -> matcher.compare() (MASt3R 联合推理)
      5. 早停: 匹配到同场景即跳过剩余
    """
    evaluation_logs = []
    best_match_result = None
    best_match_score = -1

    import time as _t
    for q_path in query_paths:
        if scene_type_q == 'text' and bge_tokenizer is not None:
            _et0 = _t.perf_counter()
            if precomputed_ocr_result is not None:
                full_text = ' '.join(line[1][0] for line in precomputed_ocr_result[0]) if (precomputed_ocr_result and precomputed_ocr_result[0]) else ''
            else:
                from module.text_matcher import extract_text
                full_text, _ = extract_text(q_path, ocr_engine)
            _et1 = _t.perf_counter()
            if precomputed_bge_vec is not None:
                q_vec = precomputed_bge_vec
            else:
                from module.qdrant_manager import _bge_encode
                q_vec = _bge_encode(full_text, bge_tokenizer, bge_model)
            _et2 = _t.perf_counter()
            _qd_cache = qd_cache.get('text', {}) if (qd_cache and isinstance(qd_cache, dict)) else {}
            if not _qd_cache:
                from module.qdrant_manager import qdrant_query_scene
                _qd_cache = qdrant_query_scene(username, 'text')
            _et3 = _t.perf_counter()
            hits = []
            for _pid, _entry in _qd_cache.items():
                _h_path = _entry.get('path', '')
                if os.path.abspath(q_path) == os.path.abspath(_h_path):
                    continue
                if not os.path.exists(_h_path):
                    continue
                _h_vec = _entry.get('_vector', _entry.get('bge_vector', np.zeros(512, dtype=np.float32)))
                _sim = float(np.dot(q_vec, _h_vec))
                if _sim >= 0.3:
                    hits.append((_sim, _pid, _entry))
            hits.sort(key=lambda x: x[0], reverse=True)
            hits = hits[:3]
            print(f'  [TIME-EVAL] text_extract: {_et1-_et0:.3f}s  bge_encode: {_et2-_et1:.3f}s  qdrant_scroll: {_et3-_et2:.3f}s  python_search: {_t.perf_counter()-_et3:.3f}s  n_hits: {len(hits)}')

            for score, pid, entry in hits:
                h_path = entry.get('path', '')
                if os.path.abspath(q_path) == os.path.abspath(h_path):
                    continue
                try:
                    result = compare_text_scene_pure_bge(
                        q_path, entry, ocr_engine, bge_tokenizer, bge_model,
                        precomputed_text=full_text, precomputed_bge=q_vec
                    )
                    log_entry = {
                        'query_image': os.path.basename(q_path),
                        'history_image': entry.get('filename', ''),
                        'query_path': q_path,
                        'history_path': h_path,
                        'scene_type': 'text',
                        'doc_score1': 0.0,
                        'doc_score2': 0.0,
                        'similarity_score': result['similarity_score'],
                        'match_count': result.get('match_count', 0),
                        'raw_match_count': 0,
                        'inlier_ratio': result.get('inlier_ratio', 0.0),
                        'avg_confidence': result.get('avg_confidence', 0.0),
                        'text_similarity': result.get('text_similarity'),
                        'ssim_score': result.get('ssim_score'),
                        'seal_bonus': result.get('seal_bonus'),
                        'is_same_scene': result['is_same_scene'],
                        'gamma_info': [],
                        'bge_search_score': round(score, 4),
                        'query_text': full_text,
                        'history_text': entry.get('ocr_text', ''),
                    }
                except Exception as e:
                    log_entry = {
                        'query_image': os.path.basename(q_path),
                        'history_image': entry.get('filename', ''),
                        'query_path': q_path,
                        'history_path': h_path,
                        'scene_type': 'text',
                        'doc_score1': 0.0, 'doc_score2': 0.0,
                        'similarity_score': 0.0,
                        'match_count': 0, 'raw_match_count': 0,
                        'inlier_ratio': 0.0, 'avg_confidence': 0.0,
                        'is_same_scene': False, 'gamma_info': [],
                        'error': str(e),
                    }
                evaluation_logs.append(log_entry)

                if log_entry['is_same_scene']:
                    if log_entry['similarity_score'] > best_match_score:
                        best_match_score = log_entry['similarity_score']
                        best_match_result = log_entry.copy()
                    print(f"  [早停] {os.path.basename(q_path)} 与 {log_entry['history_image']} 匹配为同场景, 跳过剩余比对")
                    break
                elif log_entry['similarity_score'] > best_match_score:
                    best_match_score = log_entry['similarity_score']
                    best_match_result = log_entry.copy()

        elif scene_type_q == 'complex' and dinov2_extractor is not None:
            from module.qdrant_manager import _dinov2_encode
            q_vec = _dinov2_encode(q_path, dinov2_extractor)
            if not isinstance(q_vec, np.ndarray):
                q_vec = q_vec.cpu().numpy() if hasattr(q_vec, 'cpu') else np.array(q_vec, dtype=np.float32)
            _qd_cache_c = qd_cache.get('complex', {}) if (qd_cache and isinstance(qd_cache, dict)) else {}
            if not _qd_cache_c:
                from module.qdrant_manager import qdrant_query_scene
                _qd_cache_c = qdrant_query_scene(username, 'complex')
            hits = []
            for _pid, _entry in _qd_cache_c.items():
                _h_path = _entry.get('path', '')
                if os.path.abspath(q_path) == os.path.abspath(_h_path):
                    continue
                if not os.path.exists(_h_path):
                    continue
                _h_vec = _entry.get('_vector', np.zeros(384, dtype=np.float32))
                _sim = float(np.dot(q_vec, _h_vec))
                if _sim >= DINO_PRESCREEN_THRESHOLD:
                    hits.append((_sim, _pid, _entry))
            hits.sort(key=lambda x: x[0], reverse=True)
            hits = hits[:3]

            for score, pid, entry in hits:
                h_path = entry.get('path', '')
                if os.path.abspath(q_path) == os.path.abspath(h_path):
                    continue
                dino_sim = score

                if dino_sim < DINO_PRESCREEN_THRESHOLD:
                    log_entry = {
                        'query_image': os.path.basename(q_path),
                        'history_image': entry.get('filename', ''),
                        'query_path': q_path,
                        'history_path': h_path,
                        'scene_type': 'complex',
                        'doc_score1': 0.0, 'doc_score2': 0.0,
                        'similarity_score': round(dino_sim, 4),
                        'match_count': 0, 'raw_match_count': 0,
                        'inlier_ratio': 0.0, 'avg_confidence': 0.0,
                        'dinov2_similarity': round(dino_sim, 4),
                        'is_same_scene': False, 'gamma_info': [],
                        'mast3r_skipped': True,
                        'dino_search_score': round(dino_sim, 4),
                    }
                else:
                    try:
                        result = matcher.compare(q_path, h_path, dinov2_sim_override=dino_sim)
                        log_entry = {
                            'query_image': os.path.basename(q_path),
                            'history_image': entry.get('filename', ''),
                            'query_path': q_path,
                            'history_path': h_path,
                            'scene_type': 'complex',
                            'doc_score1': 0.0, 'doc_score2': 0.0,
                            'similarity_score': result['similarity_score'],
                            'match_count': result['match_count'],
                            'raw_match_count': result['raw_match_count'],
                            'inlier_ratio': result['inlier_ratio'],
                            'avg_confidence': result['avg_confidence'],
                            'mast3r_is_same': result.get('mast3r_is_same'),
                            'dinov2_is_same': result.get('dinov2_is_same'),
                            'dinov2_similarity': result.get('dinov2_similarity'),
                            'is_same_scene': result['is_same_scene'],
                            'gamma_info': result.get('gamma_info', []),
                            'patch_match_info': result.get('patch_match_info'),
                            'mast3r_skipped': False,
                            'dino_search_score': round(dino_sim, 4),
                        }
                    except Exception as e:
                        log_entry = {
                            'query_image': os.path.basename(q_path),
                            'history_image': entry.get('filename', ''),
                            'query_path': q_path,
                            'history_path': h_path,
                            'scene_type': 'complex',
                            'doc_score1': 0.0, 'doc_score2': 0.0,
                            'similarity_score': 0.0,
                            'match_count': 0, 'raw_match_count': 0,
                            'inlier_ratio': 0.0, 'avg_confidence': 0.0,
                            'is_same_scene': False, 'gamma_info': [],
                            'error': str(e),
                        }
                evaluation_logs.append(log_entry)

                if log_entry['is_same_scene']:
                    if log_entry['similarity_score'] > best_match_score:
                        best_match_score = log_entry['similarity_score']
                        best_match_result = log_entry.copy()
                    print(f"  [早停] {os.path.basename(q_path)} 与 {log_entry['history_image']} 匹配为同场景, 跳过剩余比对")
                    break
                elif log_entry['similarity_score'] > best_match_score:
                    best_match_score = log_entry['similarity_score']
                    best_match_result = log_entry.copy()

    any_same = any(log['is_same_scene'] for log in evaluation_logs)
    final_result = not any_same

    return final_result, evaluation_logs, best_match_result


def evaluate_scene_consistency(matcher, query_paths, history_paths,
                               ocr_engine=None, bge_tokenizer=None, bge_model=None):
    evaluation_logs = []
    best_match_result = None
    best_match_score = -1

    for q_path in query_paths:
        for h_path in history_paths:
            if os.path.abspath(q_path) == os.path.abspath(h_path):
                continue

            try:
                scene_type1, doc_score1, detail1 = 'complex', 0.0, {}
                scene_type2, doc_score2, detail2 = 'complex', 0.0, {}
                if ocr_engine is not None:
                    scene_type1, doc_score1, detail1, _ = classify_scene(q_path, ocr_engine)
                    scene_type2, doc_score2, detail2, _ = classify_scene(h_path, ocr_engine)

                if scene_type1 == 'text' and scene_type2 == 'text' and bge_tokenizer is not None:
                    result = compare_text_scene(
                        q_path, h_path, ocr_engine, bge_tokenizer, bge_model
                    )
                    log_entry = {
                        'query_image': os.path.basename(q_path),
                        'history_image': os.path.basename(h_path),
                        'query_path': q_path,
                        'history_path': h_path,
                        'scene_type': 'text',
                        'doc_score1': doc_score1,
                        'doc_score2': doc_score2,
                        'similarity_score': result['similarity_score'],
                        'match_count': result.get('match_count', 0),
                        'raw_match_count': 0,
                        'inlier_ratio': result.get('inlier_ratio', 0.0),
                        'avg_confidence': result.get('avg_confidence', 0.0),
                        'text_similarity': result.get('text_similarity'),
                        'ssim_score': result.get('ssim_score'),
                        'seal_bonus': result.get('seal_bonus'),
                        'is_same_scene': result['is_same_scene'],
                        'gamma_info': [],
                    }
                else:
                    result = matcher.compare(q_path, h_path)
                    log_entry = {
                        'query_image': os.path.basename(q_path),
                        'history_image': os.path.basename(h_path),
                        'query_path': q_path,
                        'history_path': h_path,
                        'scene_type': 'complex',
                        'doc_score1': doc_score1,
                        'doc_score2': doc_score2,
                        'similarity_score': result['similarity_score'],
                        'match_count': result['match_count'],
                        'raw_match_count': result['raw_match_count'],
                        'inlier_ratio': result['inlier_ratio'],
                        'avg_confidence': result['avg_confidence'],
                        'mast3r_is_same': result.get('mast3r_is_same'),
                        'dinov2_is_same': result.get('dinov2_is_same'),
                        'dinov2_similarity': result.get('dinov2_similarity'),
                        'is_same_scene': result['is_same_scene'],
                        'gamma_info': result.get('gamma_info', []),
                        'patch_match_info': result.get('patch_match_info'),
                    }
            except Exception as e:
                log_entry = {
                    'query_image': os.path.basename(q_path),
                    'history_image': os.path.basename(h_path),
                    'query_path': q_path,
                    'history_path': h_path,
                    'scene_type': 'unknown',
                    'doc_score1': 0.0, 'doc_score2': 0.0,
                    'similarity_score': 0.0,
                    'match_count': 0, 'raw_match_count': 0,
                    'inlier_ratio': 0.0, 'avg_confidence': 0.0,
                    'is_same_scene': False, 'gamma_info': [],
                    'error': str(e),
                    'traceback': traceback.format_exc(),
                }
            evaluation_logs.append(log_entry)

            if log_entry['is_same_scene']:
                if log_entry['similarity_score'] > best_match_score:
                    best_match_score = log_entry['similarity_score']
                    best_match_result = log_entry.copy()
                print(f"  [早停] {os.path.basename(q_path)} 与 {log_entry['history_image']} 匹配为同场景, 跳过剩余比对")
                break
            elif log_entry['similarity_score'] > best_match_score:
                best_match_score = log_entry['similarity_score']
                best_match_result = log_entry.copy()

    any_same = any(log['is_same_scene'] for log in evaluation_logs)
    final_result = not any_same

    return final_result, evaluation_logs, best_match_result


def update_metadata(user_dir, saved_info, per_image_results):
    metadata = load_metadata(user_dir)
    for original_name, save_path in saved_info:
        img_result = per_image_results.get(save_path, '未知')
        img_meta = {
            'filename': os.path.basename(save_path),
            'path': os.path.abspath(save_path),
            'upload_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
            'original_name': original_name,
            'evaluation_result': img_result,
        }
        metadata['images'].append(img_meta)
    save_metadata(user_dir, metadata)