module/text_matcher.py

"""文本场景比对管线：OCR提取 + 语义嵌入 + SSIM结构比对 + 印章/手写检测"""

import time
import cv2
import numpy as np
import torch
from skimage.metrics import structural_similarity as ssim

from module.config import (
    TEXT_SIM_WEIGHT, SSIM_WEIGHT, SEAL_WEIGHT, TEXT_SCENE_THRESHOLD,
    TEXT_SCENE_BGE_THRESHOLD,
)


def _cv_imread(path, flags=cv2.IMREAD_COLOR):
    """支持中文路径的cv2.imread"""
    return cv2.imdecode(np.fromfile(path, dtype=np.uint8), flags)


# ==================== OCR文字提取 ====================

def extract_text(image_path, ocr_engine):
    """使用PaddleOCR提取图片中的全部文字

    Args:
        image_path: 图片路径
        ocr_engine: PaddleOCR引擎实例

    Returns:
        full_text: 拼接后的完整文本
        text_lines: 每行文字的列表（含置信度）
    """
    result = ocr_engine.ocr(image_path, cls=True)

    text_lines = []
    if result and result[0]:
        for line in result[0]:
            text = line[1][0]
            conf = line[1][1]
            text_lines.append({'text': text, 'confidence': float(conf)})

    full_text = ' '.join([t['text'] for t in text_lines])
    return full_text, text_lines


# ==================== 语义嵌入比对 ====================

def compute_text_similarity(text1, text2, bge_tokenizer, bge_model):
    """使用BGE-small-zh计算两段文本的语义相似度

    Args:
        text1, text2: 待比较的两段文本
        bge_tokenizer: BGE分词器实例
        bge_model: BGE模型实例

    Returns:
        similarity: 余弦相似度，范围[0,1]
    """
    if not text1.strip() or not text2.strip():
        return 0.0

    def _encode(text):
        """编码单段文本为归一化向量"""
        encoded = bge_tokenizer(text, padding=True, truncation=True, return_tensors='pt', max_length=512)
        if torch.cuda.is_available():
            encoded = {k: v.cuda() for k, v in encoded.items()}
        with torch.no_grad():
            outputs = bge_model(**encoded)
        cls_embedding = outputs.last_hidden_state[:, 0]
        cls_embedding = torch.nn.functional.normalize(cls_embedding, p=2, dim=1)
        return cls_embedding.cpu().numpy()

    emb1 = _encode(text1)
    emb2 = _encode(text2)

    sim = float(np.dot(emb1[0], emb2[0]))
    return sim


# ==================== SSIM结构比对 ====================

def compute_ssim(image_path1, image_path2):
    """计算两张图片的SSIM结构相似度

    将图片缩放到相同尺寸后转为灰度计算SSIM

    Returns:
        ssim_score: SSIM值，范围[0,1]
    """
    img1 = _cv_imread(image_path1)
    img2 = _cv_imread(image_path2)

    if img1 is None or img2 is None:
        return 0.0

    h = min(img1.shape[0], img2.shape[0])
    w = min(img1.shape[1], img2.shape[1])
    img1 = cv2.resize(img1, (w, h))
    img2 = cv2.resize(img2, (w, h))

    gray1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    gray2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

    score, _ = ssim(gray1, gray2, full=True)
    return float(score)


# ==================== 印章/手写检测 ====================

def detect_seals_and_handwriting(image_path, min_area_ratio=0.001):
    """检测图片中的红色印章和手写笔迹区域

    印章检测原理：HSV颜色空间中提取红色区域

    Args:
        image_path: 图片路径
        min_area_ratio: 最小区域面积占比，过滤噪声

    Returns:
        seals: 印章区域列表，每项含 {'bbox', 'area', 'area_ratio'}
        seal_total_ratio: 印章总面积占比
    """
    img = _cv_imread(image_path)
    if img is None:
        return [], 0.0

    h, w = img.shape[:2]
    total_area = h * w

    # --- 印章检测：HSV红色区域 ---
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    # 红色在HSV中有两段：0-10 和 160-180
    mask1 = cv2.inRange(hsv, np.array([0, 70, 50]), np.array([10, 255, 255]))
    mask2 = cv2.inRange(hsv, np.array([160, 70, 50]), np.array([180, 255, 255]))
    red_mask = cv2.bitwise_or(mask1, mask2)

    # 形态学操作去噪
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_CLOSE, kernel, iterations=2)
    red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_OPEN, kernel, iterations=1)

    # 查找轮廓
    contours, _ = cv2.findContours(red_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    seals = []
    min_area = total_area * min_area_ratio
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area >= min_area:
            x, y, bw, bh = cv2.boundingRect(cnt)
            seals.append({
                'bbox': (x, y, bw, bh),
                'area': int(area),
                'area_ratio': round(area / total_area, 4),
            })

    seal_total_ratio = sum(s['area_ratio'] for s in seals)
    return seals, round(seal_total_ratio, 4)


def compare_seals(seals1, seals2):
    """比对两张图的印章区域

    比较逻辑：
    - 两张图都有印章 → 比较数量和面积占比差异
    - 一张有一张无 → 印章不匹配
    - 两张都无 → 印章项跳过

    Returns:
        match: bool或None，印章是否匹配（None表示两张都无印章）
        detail: 详细信息字典
    """
    has_seal1 = len(seals1) > 0
    has_seal2 = len(seals2) > 0
    ratio1 = sum(s['area_ratio'] for s in seals1)
    ratio2 = sum(s['area_ratio'] for s in seals2)

    detail = {
        'img1_seal_count': len(seals1),
        'img2_seal_count': len(seals2),
        'img1_seal_ratio': round(ratio1, 4),
        'img2_seal_ratio': round(ratio2, 4),
    }

    if not has_seal1 and not has_seal2:
        return None, detail

    if has_seal1 != has_seal2:
        detail['reason'] = '一张有印章一张无'
        return False, detail

    # 两张都有印章：比较数量和面积
    count_match = len(seals1) == len(seals2)
    ratio_diff = abs(ratio1 - ratio2)
    area_match = ratio_diff < 0.02

    match = count_match and area_match
    detail['count_match'] = count_match
    detail['area_diff'] = round(ratio_diff, 4)
    detail['area_match'] = area_match

    return match, detail


def compare_text_scene_cached(image_path, cached_h, ocr_engine,
                               bge_tokenizer, bge_model,
                               precomputed_text=None, precomputed_bge=None):
    """文本场景缓存比对: 仅实时 OCR/BGE/Seal 查询图，历史图从缓存读取

    Args:
        image_path: 查询图片路径
        cached_h: 预存储的历史图缓存 dict (ocr_text, bge_vector, seal_count, seal_ratio)
        ocr_engine: PaddleOCR引擎
        bge_tokenizer: BGE分词器
        bge_model: BGE模型

    Returns:
        result dict (与 compare_text_scene() 兼容的格式)
    """
    start_time = time.time()

    # 1. OCR 仅查询图 (优先使用预计算结果)
    if precomputed_text is not None:
        text1 = precomputed_text
        lines1 = []
        ocr_time = 0.0
    else:
        text1, lines1 = extract_text(image_path, ocr_engine)
        ocr_time = time.time() - start_time
    text2 = cached_h.get('ocr_text', '')

    # 2. 语义嵌入 (优先使用预计算向量)
    if precomputed_bge is not None:
        emb_q = precomputed_bge
        emb_h = cached_h.get('bge_vector', np.zeros(512, dtype=np.float32))
        text_sim = float(np.dot(emb_q, emb_h)) if text1.strip() else 0.0
        embed_time = 0.0
    elif not text1.strip():
        text_sim = 0.0
        embed_time = 0.0
    else:
        t_embed = time.time()
        def _encode_q(text):
            encoded = bge_tokenizer(text, padding=True, truncation=True,
                                     return_tensors='pt', max_length=512)
            if torch.cuda.is_available():
                encoded = {k: v.cuda() for k, v in encoded.items()}
            with torch.no_grad():
                outputs = bge_model(**encoded)
            cls_embedding = outputs.last_hidden_state[:, 0]
            cls_embedding = torch.nn.functional.normalize(cls_embedding, p=2, dim=1)
            return cls_embedding.cpu().numpy()[0]

        emb_q = _encode_q(text1)
        emb_h = cached_h.get('bge_vector', np.zeros(512, dtype=np.float32))
        text_sim = float(np.dot(emb_q, emb_h))
        embed_time = time.time() - t_embed

    # 3. SSIM 仍需两张图
    t_ssim = time.time()
    ssim_score = compute_ssim(image_path, cached_h['path'])
    ssim_time = time.time() - t_ssim

    # 4. 印章：仅查询图（历史图用缓存）
    t_seal = time.time()
    seals_q, ratio_q = detect_seals_and_handwriting(image_path)
    cnt_h = cached_h.get('seal_count', 0)
    ratio_h = cached_h.get('seal_ratio', 0.0)

    has_q, has_h = len(seals_q) > 0, cnt_h > 0
    if not has_q and not has_h:
        seal_bonus = 0.5
    elif has_q != has_h:
        seal_bonus = 0.0
    else:
        count_match = len(seals_q) == cnt_h
        area_match = abs(ratio_q - ratio_h) < 0.02
        seal_bonus = 1.0 if (count_match and area_match) else 0.0
    seal_time = time.time() - t_seal

    # 5. 综合判定
    final_score = text_sim * TEXT_SIM_WEIGHT + ssim_score * SSIM_WEIGHT + seal_bonus * SEAL_WEIGHT
    is_same = final_score >= TEXT_SCENE_THRESHOLD

    total_time = time.time() - start_time

    result = {
        'is_same_scene': bool(is_same),
        'scene_type': 'text',
        'final_score': round(final_score, 4),
        'text_similarity': round(text_sim, 4),
        'ssim_score': round(ssim_score, 4),
        'seal_bonus': seal_bonus,
        'seal_match': None if (not has_q and not has_h) else (has_q and has_h),
        'seal_detail': {
            'img1_seal_count': len(seals_q), 'img2_seal_count': cnt_h,
            'img1_seal_ratio': ratio_q, 'img2_seal_ratio': ratio_h,
        },
        'img1_text_length': len(text1), 'img2_text_length': len(text2),
        'img1_seal_count': len(seals_q), 'img2_seal_count': cnt_h,
        'img1_seal_ratio': ratio_q, 'img2_seal_ratio': ratio_h,
        'similarity_score': round(final_score, 4),
        'match_count': len(seals_q) + cnt_h,
        'inlier_ratio': round(text_sim, 4),
        'avg_confidence': round(ssim_score, 4),
        'time_breakdown': {
            'ocr': round(ocr_time, 2), 'embedding': round(embed_time, 2),
            'ssim': round(ssim_time, 2), 'seal': round(seal_time, 2),
            'total': round(total_time, 2),
        },
    }

    return result





def compare_text_scene_pure_bge(image_path, cached_h, ocr_engine,
                                 bge_tokenizer, bge_model,
                                 precomputed_text=None, precomputed_bge=None):
    """Pure BGE: only OCR->BGE cosine similarity, no SSIM/seal"""
    start_time = time.time()
    print(f'  [TIME-PURE-BGE] enter, precomputed_text={precomputed_text is not None}, precomputed_bge={precomputed_bge is not None}')

    if precomputed_text is not None:
        text1 = precomputed_text
        ocr_time = 0.0
    else:
        text1, _ = extract_text(image_path, ocr_engine)
        ocr_time = time.time() - start_time

    if precomputed_bge is not None:
        emb_q = precomputed_bge
        emb_h = cached_h.get('_vector', cached_h.get('bge_vector', np.zeros(512, dtype=np.float32)))
        text_sim = float(np.dot(emb_q, emb_h)) if text1.strip() else 0.0
        embed_time = 0.0
    elif not text1.strip():
        text_sim = 0.0
        embed_time = 0.0
    else:
        t_embed = time.time()
        def _encode_q(text):
            encoded = bge_tokenizer(text, padding=True, truncation=True,
                                     return_tensors='pt', max_length=512)
            if torch.cuda.is_available():
                encoded = {k: v.cuda() for k, v in encoded.items()}
            with torch.no_grad():
                outputs = bge_model(**encoded)
            cls_embedding = outputs.last_hidden_state[:, 0]
            cls_embedding = torch.nn.functional.normalize(cls_embedding, p=2, dim=1)
            return cls_embedding.cpu().numpy()[0]
        emb_q = _encode_q(text1)
        emb_h = cached_h.get('_vector', cached_h.get('bge_vector', np.zeros(512, dtype=np.float32)))
        text_sim = float(np.dot(emb_q, emb_h))
        embed_time = time.time() - t_embed

    is_same = text_sim >= TEXT_SCENE_BGE_THRESHOLD
    total_time = time.time() - start_time
    print(f'  [TIME-PURE-BGE] text_sim={text_sim:.4f} is_same={is_same} total={total_time:.3f}s')

    result = {
        'is_same_scene': bool(is_same),
        'scene_type': 'text',
        'final_score': round(text_sim, 4),
        'text_similarity': round(text_sim, 4),
        'ssim_score': 0.0,
        'seal_bonus': 0.0,
        'similarity_score': round(text_sim, 4),
        'match_count': 0,
        'inlier_ratio': round(text_sim, 4),
        'avg_confidence': 0.0,
        'time_breakdown': {
            'ocr': round(ocr_time, 2), 'embedding': round(embed_time, 2),
            'ssim': 0.0, 'seal': 0.0, 'total': round(total_time, 2),
        },
    }
    return result

# ==================== 综合判定 ====================

def compare_text_scene(image_path1, image_path2, ocr_engine, bge_tokenizer, bge_model):
    """文本场景比对主函数

    Args:
        image_path1: 第一张图片路径
        image_path2: 第二张图片路径
        ocr_engine: PaddleOCR引擎实例
        bge_tokenizer: BGE分词器实例
        bge_model: BGE模型实例

    Returns:
        result: dict，含各项分数和最终判定
    """
    start_time = time.time()

    # 1. OCR文字提取
    text1, lines1 = extract_text(image_path1, ocr_engine)
    text2, lines2 = extract_text(image_path2, ocr_engine)
    ocr_time = time.time() - start_time

    # 2. 语义嵌入比对
    text_sim = compute_text_similarity(text1, text2, bge_tokenizer, bge_model)
    embed_time = time.time() - start_time - ocr_time

    # 3. SSIM结构比对
    ssim_score = compute_ssim(image_path1, image_path2)
    ssim_time = time.time() - start_time - ocr_time - embed_time

    # 4. 印章检测与比对
    seals1, seal_ratio1 = detect_seals_and_handwriting(image_path1)
    seals2, seal_ratio2 = detect_seals_and_handwriting(image_path2)
    seal_match, seal_detail = compare_seals(seals1, seals2)
    seal_time = time.time() - start_time - ocr_time - embed_time - ssim_time

    # 5. 综合判定
    if seal_match is True:
        seal_bonus = 1.0
    elif seal_match is False:
        seal_bonus = 0.0
    else:
        seal_bonus = 0.5  # 两张图都没有印章，中性

    final_score = text_sim * TEXT_SIM_WEIGHT + ssim_score * SSIM_WEIGHT + seal_bonus * SEAL_WEIGHT
    is_same = final_score >= TEXT_SCENE_THRESHOLD

    total_time = time.time() - start_time

    result = {
        'is_same_scene': bool(is_same),
        'scene_type': 'text',
        'final_score': round(final_score, 4),
        'text_similarity': round(text_sim, 4),
        'ssim_score': round(ssim_score, 4),
        'seal_bonus': seal_bonus,
        'seal_match': seal_match,
        'seal_detail': seal_detail,
        'img1_text_length': len(text1),
        'img2_text_length': len(text2),
        'img1_seal_count': len(seals1),
        'img2_seal_count': len(seals2),
        'img1_seal_ratio': seal_ratio1,
        'img2_seal_ratio': seal_ratio2,
        'similarity_score': round(final_score, 4),
        'match_count': len(seals1) + len(seals2),  # 文本场景用印章数代替
        'inlier_ratio': round(text_sim, 4),
        'avg_confidence': round(ssim_score, 4),
        'time_breakdown': {
            'ocr': round(ocr_time, 2),
            'embedding': round(embed_time, 2),
            'ssim': round(ssim_time, 2),
            'seal': round(seal_time, 2),
            'total': round(total_time, 2),
        },
    }

    return result