#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
シンプルなDeepSeek-OCR-2 PDFアプリケーション
PDFファイルを選択してOCR処理を実行します

事前インストールが必要なパッケージ (Windows):
pip install --no-user transformers==4.46.3 tokenizers==0.20.3 einops addict easydict gitpython PyMuPDF
pip install --no-user https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.15/flash_attn-2.8.3+cu126torch2.9-cp312-cp312-win_amd64.whl

flash-attnはプリビルドwhlをダウンロードしてインストール（環境に合わせて選択）:
https://github.com/mjun0812/flash-attention-prebuild-wheels/releases
"""

# DeepSeek-OCR-2 PDF・画像OCRアプリケーション

import warnings
warnings.filterwarnings('ignore')

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
os.environ['TRANSFORMERS_VERBOSITY'] = 'error'

import sys
import logging

logging.getLogger('transformers').setLevel(logging.CRITICAL)
logging.getLogger('transformers.generation').setLevel(logging.CRITICAL)
logging.getLogger('transformers.modeling_utils').setLevel(logging.CRITICAL)

import tkinter as tk
from tkinter import filedialog
from transformers import AutoModel, AutoTokenizer
import torch
from pathlib import Path
import fitz
import tempfile
import re
from PIL import Image, ImageDraw, ImageFont
import cv2
import numpy as np

SUPPORTED_IMAGE_FORMATS = ['.png', '.jpg', '.jpeg', '.bmp', '.tiff']
RESULT_FILENAME = "result.md"
PDF_ZOOM_FACTOR = 150 / 72
DISPLAY_WIDTH = 1024
COORD_NORMALIZE_MAX = 999

MODE_CONFIGS = {
    "1": {"base_size": 512, "image_size": 512, "crop_mode": False, "name": "Tiny", "tokens": "64"},
    "2": {"base_size": 640, "image_size": 640, "crop_mode": False, "name": "Small", "tokens": "100"},
    "3": {"base_size": 1024, "image_size": 1024, "crop_mode": False, "name": "Base", "tokens": "256"},
    "4": {"base_size": 1280, "image_size": 1280, "crop_mode": False, "name": "Large", "tokens": "400"},
    "5": {"base_size": 1024, "image_size": 640, "crop_mode": True, "name": "Gundam", "tokens": "n×100+256 (動的)"}
}

def print_separator():
    print("=" * 60)

def get_japanese_font(size=48):
    font_paths = [
        "C:/Windows/Fonts/msgothic.ttc",
        "C:/Windows/Fonts/meiryo.ttc",
        "/System/Library/Fonts/ヒラギノ角ゴシック W3.ttc",
        "/usr/share/fonts/truetype/fonts-japanese-gothic.ttf",
    ]
    for font_path in font_paths:
        if os.path.exists(font_path):
            return ImageFont.truetype(font_path, size)
    return ImageFont.load_default()

def parse_ocr_result(ocr_text):
    results = []

    pattern = r'<\|ref\|>(.*?)<\|/ref\|><\|det\|>\[\[(.*?)\]\]<\|/det\|>'
    matches = re.findall(pattern, ocr_text, re.DOTALL)

    for text, coords in matches:
        text = text.strip()
        coords_list = [int(float(x)) for x in coords.split(',')]
        if len(coords_list) == 4:
            results.append({
                'text': text,
                'box': coords_list
            })

    return results

def load_and_resize_image(image_path, target_width=DISPLAY_WIDTH):
    img = Image.open(image_path).convert('RGB')
    orig_width, orig_height = img.size

    scale = target_width / orig_width
    new_height = int(orig_height * scale)
    img_resized = img.resize((target_width, new_height), Image.Resampling.LANCZOS)

    return img_resized, orig_width, orig_height, scale

def visualize_result(image_path, ocr_text, output_path):
    img_resized, orig_width, orig_height, display_scale = load_and_resize_image(image_path)

    draw = ImageDraw.Draw(img_resized)
    font = get_japanese_font(48)

    parsed_results = parse_ocr_result(ocr_text)

    if len(parsed_results) == 0:
        img_resized.save(output_path)
        return output_path

    coord_to_orig_scale_x = orig_width / COORD_NORMALIZE_MAX
    coord_to_orig_scale_y = orig_height / COORD_NORMALIZE_MAX

    for idx, item in enumerate(parsed_results):
        if item['box']:
            x1_norm, y1_norm, x2_norm, y2_norm = item['box']

            x1_orig = x1_norm * coord_to_orig_scale_x
            y1_orig = y1_norm * coord_to_orig_scale_y
            x2_orig = x2_norm * coord_to_orig_scale_x
            y2_orig = y2_norm * coord_to_orig_scale_y

            x1_display = int(x1_orig * display_scale)
            y1_display = int(y1_orig * display_scale)
            x2_display = int(x2_orig * display_scale)
            y2_display = int(y2_orig * display_scale)

            draw.rectangle([x1_display, y1_display, x2_display, y2_display], outline=(255, 0, 0), width=5)

            text_y = y1_display - 60 if y1_display > 60 else y2_display + 10
            draw.text((x1_display, text_y), item['text'], font=font, fill=(255, 0, 0))

    img_resized.save(output_path)
    return output_path

def run_ocr_inference(image_path, config, model, tokenizer, prompt):
    import sys
    from io import StringIO

    class TeeOutput:
        def __init__(self, *streams):
            self.streams = streams
        def write(self, data):
            for stream in self.streams:
                stream.write(data)
        def flush(self):
            for stream in self.streams:
                stream.flush()

    class SuppressOutput:
        def write(self, data):
            pass
        def flush(self):
            pass

    old_stdout = sys.stdout
    old_stderr = sys.stderr
    captured_out = StringIO()
    sys.stdout = TeeOutput(old_stdout, captured_out)
    sys.stderr = SuppressOutput()

    model.infer(
        tokenizer,
        prompt=prompt,
        image_file=image_path,
        output_path=".",
        base_size=config['base_size'],
        image_size=config['image_size'],
        crop_mode=config['crop_mode'],
        save_results=True,
        test_compress=False
    )

    sys.stdout = old_stdout
    sys.stderr = old_stderr
    output = captured_out.getvalue()

    filtered_lines = []
    for line in output.split('\n'):
        if line.startswith('BASE:') or line.startswith('PATCHES:') or line.startswith('====='):
            continue
        if line.strip():
            filtered_lines.append(line)

    return '\n'.join(filtered_lines)

def read_result_file():
    try:
        with open(RESULT_FILENAME, 'r', encoding='utf-8') as f:
            content = f.read()
            return content
    except FileNotFoundError:
        return None

def execute_ocr_and_visualize(image_path, filename_base, config, model, tokenizer, prompt, original_path):
    print(f"  [1/3] 座標情報取得中（Grounding）...")
    prompt_grounding = "<image>\n<|grounding|>OCR this image."
    output_grounding = run_ocr_inference(image_path, config, model, tokenizer, prompt_grounding)
    result_grounding = read_result_file()
    if result_grounding is None:
        result_grounding = output_grounding

    detected_boxes = parse_ocr_result(result_grounding)
    print(f"  検出されたテキスト領域数: {len(detected_boxes)}")

    if len(detected_boxes) == 0:
        print(f"  [2/3] テキスト抽出中（Free OCR - 全体）...")
        prompt_text = "<image>\nFree OCR."
        output_text = run_ocr_inference(image_path, config, model, tokenizer, prompt_text)
        result_text = read_result_file()
        if result_text is None:
            result_text = output_text

        vis_output = f"visualized_{filename_base}"
        visualize_result(image_path, result_grounding, vis_output)

        return {
            'filename': filename_base,
            'text': result_text,
            'visual': vis_output,
            'original': original_path,
            'boxes': []
        }

    original_image = cv2.imread(image_path)
    original_height, original_width = original_image.shape[:2]

    scale_x = original_width / COORD_NORMALIZE_MAX
    scale_y = original_height / COORD_NORMALIZE_MAX

    print(f"  [2/3] 各領域の高精度テキスト認識中（{len(detected_boxes)}個）...")

    refined_results = []
    margin = 5

    for idx, obj in enumerate(detected_boxes, 1):
        x1_norm, y1_norm, x2_norm, y2_norm = obj['box']

        x1_orig = int(x1_norm * scale_x)
        y1_orig = int(y1_norm * scale_y)
        x2_orig = int(x2_norm * scale_x)
        y2_orig = int(y2_norm * scale_y)

        x1_crop = max(0, x1_orig - margin)
        y1_crop = max(0, y1_orig - margin)
        x2_crop = min(original_width, x2_orig + margin)
        y2_crop = min(original_height, y2_orig + margin)

        cropped_image = original_image[y1_crop:y2_crop, x1_crop:x2_crop]

        crop_height, crop_width = cropped_image.shape[:2]

        if crop_width < 10 or crop_height < 10:
            refined_results.append({
                'text': obj['text'],
                'box': obj['box']
            })
            continue

        temp_crop_path = f"temp_crop_{idx}.png"
        cv2.imwrite(temp_crop_path, cropped_image)

        try:
            prompt_free = "<image>\nFree OCR."
            ocr_output = run_ocr_inference(temp_crop_path, config, model, tokenizer, prompt_free)
            refined_text = read_result_file()
            if refined_text is None:
                refined_text = ocr_output
            refined_text = refined_text.strip()

            if not refined_text:
                refined_text = obj['text']

            refined_results.append({
                'text': refined_text,
                'box': obj['box']
            })

        finally:
            if os.path.exists(temp_crop_path):
                os.remove(temp_crop_path)

    print(f"  [3/3] 可視化中...")

    combined_text = '\n'.join([r['text'] for r in refined_results])

    combined_grounding = ''
    for r in refined_results:
        coords_str = ','.join(map(str, r['box']))
        combined_grounding += f"<|ref|>{r['text']}<|/ref|><|det|>[[{coords_str}]]<|/det|>\n"

    vis_output = f"visualized_{filename_base}"
    visualize_result(image_path, combined_grounding, vis_output)

    return {
        'filename': filename_base,
        'text': combined_text,
        'visual': vis_output,
        'original': original_path,
        'boxes': refined_results
    }

def process_file(file_path, config, model, tokenizer, prompt):
    file_ext = Path(file_path).suffix.lower()

    if file_ext in SUPPORTED_IMAGE_FORMATS:
        result_dict = execute_ocr_and_visualize(
            file_path,
            Path(file_path).name,
            config,
            model,
            tokenizer,
            prompt,
            file_path
        )
        return [result_dict]
    elif file_ext == '.pdf':
        doc = fitz.open(file_path)
        total_pages = len(doc)
        print(f"  ページ数: {total_pages}")
        results = []

        for page_num in range(total_pages):
            page = doc[page_num]
            mat = fitz.Matrix(PDF_ZOOM_FACTOR, PDF_ZOOM_FACTOR)
            pix = page.get_pixmap(matrix=mat)
            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
                tmp_path = tmp_file.name
            pix.save(tmp_path)

            filename_base = f"{Path(file_path).stem}_page{page_num + 1}.png"
            result_dict = execute_ocr_and_visualize(
                tmp_path,
                filename_base,
                config,
                model,
                tokenizer,
                prompt,
                file_path
            )
            result_dict['filename'] = f"{Path(file_path).name} Page {page_num + 1}/{total_pages}"

            os.unlink(tmp_path)
            results.append(result_dict)

        doc.close()
        return results
    else:
        print(f"未対応形式: {file_path}")
        return []

def select_files_dialog():
    root = tk.Tk()
    root.withdraw()
    root.update()
    user_files = filedialog.askopenfilenames(
        title="ファイルを選択してください（複数選択可）",
        filetypes=[
            ("対応ファイル", "*.pdf *.png *.jpg *.jpeg *.bmp *.tiff"),
            ("すべてのファイル", "*.*")
        ]
    )
    root.destroy()
    return user_files

def display_visualization(all_results):
    print()
    print_separator()
    print("【ビジュアル表示】")
    print_separator()
    for result in all_results:
        print(f"\n{result['filename']}")
        print(f"  元のファイル: {result['original']}")
        print(f"  保存先: {result['visual']}")

        # バウンディングボックスの座標とテキストを表示
        if 'boxes' in result and len(result['boxes']) > 0:
            print(f"  バウンディングボックス数: {len(result['boxes'])}")
            for idx, box_info in enumerate(result['boxes'], 1):
                print(f"    [{idx}] 座標: {box_info['box']}, テキスト: {box_info['text']}")

        img = cv2.imread(result['visual'])
        if img is None:
            print(f"  画像読み込みエラー")
            continue

        window_name = f"OCR結果: {result['filename']}"
        cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)

        screen_height = 900
        img_height, img_width = img.shape[:2]
        if img_height > screen_height:
            scale = screen_height / img_height
            new_width = int(img_width * scale)
            img = cv2.resize(img, (new_width, screen_height))

        cv2.imshow(window_name, img)
        print(f"  OpenCVウィンドウで表示中（任意のキーで次へ）")
        cv2.waitKey(0)
        cv2.destroyWindow(window_name)

    cv2.destroyAllWindows()

print_separator()
print("DeepSeek-OCR-2: 画像・PDFからMarkdown変換")
print("\n【モデル選択】")
print("1. Tiny (512×512, 64トークン) - 最速")
print("2. Small (640×640, 100トークン) - 高速")
print("3. Base (1024×1024, 256トークン) - 推奨・バランス型")
print("4. Large (1280×1280, 400トークン) - 高精度")
print("5. Gundam (動的解像度) - 最高精度")

choice = input("\nモード選択 (1-5, デフォルト=3): ").strip() or "3"
config = MODE_CONFIGS.get(choice, MODE_CONFIGS["3"])
print(f"\n選択: {config['name']}モード")

print("\n【モデルについて】")
print("""DeepSeek-OCR-2は2つのコンポーネントで構成:
  1. DeepEncoder V2: 画像を人間に近い順序で読み取りビジョントークンに変換
  2. デコーダ: ビジョントークンからテキストを復元（3Bパラメータ）
  特徴: 複雑なレイアウトへの対応強化、OmniDocBenchで前バージョンから精度向上""")

print("\nモデルロード中...")
if not torch.cuda.is_available():
    print_separator()
    print("【エラー】NVIDIA GPUが検出されませんでした")
    print("DeepSeek-OCR-2はNVIDIA GPU専用モデルです")
    print_separator()
    raise RuntimeError("GPU required for DeepSeek-OCR-2")

print(f"デバイス: cuda")

tokenizer = AutoTokenizer.from_pretrained('deepseek-ai/DeepSeek-OCR-2', trust_remote_code=True)
model = AutoModel.from_pretrained(
    'deepseek-ai/DeepSeek-OCR-2',
    _attn_implementation='flash_attention_2',
    trust_remote_code=True,
    use_safetensors=True
).eval().cuda().to(torch.bfloat16)
print("ロード完了\n")

prompt = "<image>\nFree OCR."

print()
print_separator()
print("【ファイル選択】")
print_separator()
print("複数の画像・PDFファイルをアップロード可能です")
print("対応形式: PDF, PNG, JPG, JPEG, BMP, TIFF")
print(f"\nOCR方式: 3段階処理")
print("  1. Grounding（バウンディングボックス検出）")
print("  2. Free OCR（各領域の高精度テキスト抽出）")
print("  3. 可視化\n")

user_files = select_files_dialog()

if user_files:
    all_results = []

    for file_path in user_files:
        print(f"\n処理中: {file_path}")
        file_results = process_file(file_path, config, model, tokenizer, prompt)
        all_results.extend(file_results)

    print()
    print_separator()
    print("【OCR結果一覧】")
    print_separator()
    for result in all_results:
        print(f"\n--- {result['filename']} ---")
        print(result['text'])
        print_separator()

    print("\n【検証データ】")
    total_files = len(user_files)
    total_pages = len(all_results)
    print(f"  処理ファイル数: {total_files}")
    print(f"  総ページ/画像数: {total_pages}")
    print(f"  モード: {config['name']} (base_size={config['base_size']}, image_size={config['image_size']}, crop_mode={config['crop_mode']})")
    print(f"  ビジョントークン数: {config['tokens']}")

    display_visualization(all_results)
else:
    print("ファイル未選択")

print("\n処理完了")