"""
古文書・くずし字の文字自動切り出しデモ (解像度に応じて baseline / SAHI を自動選択)

- 画像をアップロード → 長辺/imgsz の縮小率で baseline か SAHI を自動判定
- yolov11x-komonjo-char (古文書, デフォルト) / yolov11x-codh-char (典籍) で文字 bbox を検出
- bbox 描画画像 + 座標 JSON を返す

縮小率 > 2.0 (高解像度) → SAHI (タイル分割推論)
それ以下 (低解像度)      → baseline (通常推論)
"""
from __future__ import annotations
import json
import sys
import tempfile
from pathlib import Path

import gradio as gr
import numpy as np
from PIL import Image, ImageDraw
from huggingface_hub import hf_hub_download
from ultralytics import YOLO

# ZeroGPU: 推論関数を @spaces.GPU でラップすると、呼び出し時だけ GPU が割り当てられる。
# spaces 未インストールのローカル環境では no-op デコレータにフォールバックして CPU で動く。
try:
    import spaces
    GPU = spaces.GPU
except Exception:
    def GPU(*args, **kwargs):
        # @GPU でも @GPU(duration=...) でも両対応する no-op
        if len(args) == 1 and callable(args[0]) and not kwargs:
            return args[0]
        def _deco(fn):
            return fn
        return _deco

IMGSZ = 1280
DOWNSCALE_THRESHOLD = 2.0
DEFAULT_MODEL = "komonjo (古文書・書状)"
MODEL_REPOS = {
    "komonjo (古文書・書状)": "nakamura196/yolov11x-komonjo-char",
    "codh (典籍・版本)": "nakamura196/yolov11x-codh-char",
}
HERE = Path(__file__).resolve().parent

# 文字検出モデル (デフォルトは起動時にロード、もう一方は初回選択時に遅延ロード)
_models: dict = {}

def get_model(name: str):
    """(weights_path, YOLO) を返す。初回のみ HF からダウンロードしてロード。"""
    if name not in _models:
        repo = MODEL_REPOS.get(name, MODEL_REPOS[DEFAULT_MODEL])
        print(f"loading weights: {repo} ...")
        weights = hf_hub_download(repo_id=repo, filename="best.pt")
        _models[name] = (weights, YOLO(weights))
        print(f"model ready: {name}")
    return _models[name]

get_model(DEFAULT_MODEL)

# NDL古典籍OCR-Lite の行検出 (RTMDet, optional)。読み込み失敗しても本体は動く
LINE_DETECTOR = None
try:
    sys.path.insert(0, str(HERE / "ndl_line"))
    from rtmdet import RTMDet
    LINE_DETECTOR = RTMDet(
        model_path=str(HERE / "ndl_line" / "rtmdet-s-1280x1280.onnx"),
        class_mapping_path=str(HERE / "ndl_line" / "ndl.yaml"),
        conf_thresold=0.3, iou_threshold=0.4, device="CPU")
    print("NDL line detector ready")
except Exception as e:
    print(f"NDL line detector unavailable: {e}")


def decide_method(w: int, h: int) -> tuple[str, float]:
    downscale = max(w, h) / IMGSZ
    method = "SAHI" if downscale > DOWNSCALE_THRESHOLD else "baseline"
    return method, downscale


def _infer_device():
    """推論デバイスを決定。ZeroGPU では GPU は @spaces.GPU 関数の内側でのみ可視なので、
    モジュール冒頭ではなく推論呼び出し時に判定する。"""
    try:
        import torch
        if torch.cuda.is_available():
            return 0  # ultralytics の device 表記: 0 = 1 枚目の CUDA デバイス
    except Exception:
        pass
    return "cpu"


@GPU(duration=120)
def run_baseline(img_path: str, model_name: str = DEFAULT_MODEL):
    _, model = get_model(model_name)
    r = model.predict(source=img_path, conf=0.25, iou=0.45, imgsz=IMGSZ,
                      max_det=2000, device=_infer_device(), verbose=False, save=False)[0]
    return [{"xyxy": [round(v, 2) for v in b.xyxy[0].tolist()],
             "conf": round(float(b.conf[0]), 4)} for b in r.boxes]


@GPU(duration=120)
def run_sahi(img_path: str, model_name: str = DEFAULT_MODEL):
    from sahi import AutoDetectionModel
    from sahi.predict import get_sliced_prediction
    weights, _ = get_model(model_name)
    dev = "cuda:0" if _infer_device() != "cpu" else "cpu"
    det = AutoDetectionModel.from_pretrained(
        model_type="ultralytics", model_path=weights,
        confidence_threshold=0.3, device=dev)
    res = get_sliced_prediction(
        image=img_path, detection_model=det,
        slice_height=1024, slice_width=1024,
        overlap_height_ratio=0.2, overlap_width_ratio=0.2,
        postprocess_match_threshold=0.2, postprocess_class_agnostic=True, verbose=0)
    return [{"xyxy": [round(p.bbox.minx, 2), round(p.bbox.miny, 2),
                      round(p.bbox.maxx, 2), round(p.bbox.maxy, 2)],
             "conf": round(float(p.score.value), 4)}
            for p in res.object_prediction_list]


def detect_lines(image: Image.Image):
    """NDL RTMDet で行 bbox を返す。失敗時は空リスト。"""
    if LINE_DETECTOR is None:
        return []
    try:
        dets = LINE_DETECTOR.detect(np.array(image))
        return [[int(d["box"][0]), int(d["box"][1]), int(d["box"][2]), int(d["box"][3])] for d in dets]
    except Exception as e:
        print(f"line detect error: {e}")
        return []


def _center_in_any_line(box, line_boxes):
    """文字 bbox の中心が、いずれかの行 bbox (文字サイズ分だけ余裕) に入るか。

    tolerance は行の「短辺 (≒文字サイズ)」基準にする。縦書きの行高で 15% などにすると
    縦に巨大な許容になり、行末より下の定規・色見本ノイズを拾ってしまうため。
    """
    cx = (box[0] + box[2]) / 2
    cy = (box[1] + box[3]) / 2
    for lx1, ly1, lx2, ly2 in line_boxes:
        tol = min(lx2 - lx1, ly2 - ly1) * 0.5  # 半文字分
        if lx1 - tol <= cx <= lx2 + tol and ly1 - tol <= cy <= ly2 + tol:
            return True
    return False


def _group_into_columns(boxes):
    """文字 bbox を縦の列にグループ化し、読み順 (右列→左列、列内は上→下) で返す。"""
    if not boxes:
        return []
    def cx(b): return (b["xyxy"][0] + b["xyxy"][2]) / 2
    def cy(b): return (b["xyxy"][1] + b["xyxy"][3]) / 2
    widths = sorted(b["xyxy"][2] - b["xyxy"][0] for b in boxes)
    med_w = widths[len(widths) // 2] or 1
    items = sorted(boxes, key=lambda b: -cx(b))  # 右から
    columns, cur = [], [items[0]]
    for b in items[1:]:
        col_mean = sum(cx(x) for x in cur) / len(cur)
        if abs(cx(b) - col_mean) > med_w * 1.3:
            columns.append(cur); cur = [b]
        else:
            cur.append(b)
    columns.append(cur)
    for c in columns:
        c.sort(key=cy)  # 列内 上→下
    return columns


def _rounded_card(crop: Image.Image, cell: int):
    """文字 crop を白い角丸カード (cell×cell) に中央配置して返す (RGBA)。"""
    card = Image.new("RGBA", (cell, cell), (0, 0, 0, 0))
    mask = Image.new("L", (cell, cell), 0)
    ImageDraw.Draw(mask).rounded_rectangle([0, 0, cell - 1, cell - 1], radius=cell // 8, fill=255)
    white = Image.new("RGBA", (cell, cell), (255, 255, 255, 255))
    card.paste(white, (0, 0), mask)
    inner = int(cell * 0.84)
    c = crop.copy()
    c.thumbnail((inner, inner), Image.LANCZOS)
    ox = (cell - c.size[0]) // 2
    oy = (cell - c.size[1]) // 2
    card.paste(c.convert("RGB"), (ox, oy))
    # 角丸の枠線
    ImageDraw.Draw(card).rounded_rectangle([0, 0, cell - 1, cell - 1], radius=cell // 8,
                                           outline=(180, 165, 140, 255), width=1)
    return card


def _group_by_lines(boxes, line_boxes):
    """NDL 行 bbox を使って文字を列に割当てる。原文の列構造に忠実。読み順は行を右→左。"""
    def cx(b): return (b["xyxy"][0] + b["xyxy"][2]) / 2
    def cy(b): return (b["xyxy"][1] + b["xyxy"][3]) / 2
    # 行を右→左に並べる
    lines_sorted = sorted(line_boxes, key=lambda l: -((l[0] + l[2]) / 2))
    groups = [[] for _ in lines_sorted]
    leftover = []
    for b in boxes:
        bx, by = cx(b), cy(b)
        best, best_d = -1, None
        for i, (lx1, ly1, lx2, ly2) in enumerate(lines_sorted):
            if lx1 <= bx <= lx2 and ly1 <= by <= ly2:
                best = i; break
            # 含まれない場合は中心距離で最寄り行を候補に
            ccx, ccy = (lx1 + lx2) / 2, (ly1 + ly2) / 2
            d = (bx - ccx) ** 2 + (by - ccy) ** 2
            if best_d is None or d < best_d:
                best_d, best_near = d, i
        if best >= 0:
            groups[best].append(b)
        else:
            groups[best_near].append(b)
    cols = [sorted(g, key=cy) for g in groups if g]
    return cols


def _reading_order(boxes, line_boxes=None):
    """読み順 (右列→左列・列内上→下) に並べ替えて 1 次元化。"""
    cols = _group_by_lines(boxes, line_boxes) if line_boxes else _group_into_columns(boxes)
    return [b for col in cols for b in col]


def make_char_montage(image: Image.Image, boxes, target_w: int = 1500,
                      cell: int = 96, gap: int = 8, pad_ratio: float = 0.14, line_boxes=None):
    """検出文字を「正方形セル + padding」の字典風グリッドで並べる。

    各文字は正方セルの中に、アスペクト維持 + 余白付きで中央配置。読み順に左→右・上→下。
    列構造は無視して均等なグリッドにするので、一字ずつ均等で読みやすい。
    """
    if not boxes:
        return None
    W, H = image.size
    ordered = _reading_order(boxes, line_boxes)

    crops = []
    for b in ordered:
        x1, y1, x2, y2 = [int(v) for v in b["xyxy"]]
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(W, x2), min(H, y2)
        if x2 <= x1 or y2 <= y1:
            continue
        crops.append(image.crop((x1, y1, x2, y2)))
    if not crops:
        return None
    n_chars = len(crops)

    margin = 28
    header_h = 64
    pitch = cell + gap
    cols = max(1, (target_w - margin * 2 + gap) // pitch)
    rows = (n_chars + cols - 1) // cols
    grid_w = margin * 2 + cols * pitch - gap
    total_h = margin + header_h + rows * pitch - gap + margin

    # 和紙風グラデーション背景
    top = np.array([248, 243, 230], dtype=np.float32)
    bot = np.array([235, 226, 205], dtype=np.float32)
    grad = (top[None, :] * (1 - np.linspace(0, 1, total_h))[:, None]
            + bot[None, :] * np.linspace(0, 1, total_h)[:, None]).astype(np.uint8)
    bg = Image.fromarray(np.repeat(grad[:, None, :], grid_w, axis=1))
    draw = ImageDraw.Draw(bg, "RGBA")

    # ヘッダー (右寄せ)
    from PIL import ImageFont
    # 日本語グリフを持つフォントを優先的に探す。HF Space(Linux) では packages.txt の
    # fonts-noto-cjk が入る。DejaVu は日本語を持たない(豆腐になる)ので最後の保険のみ。
    fpath = None
    for p in [str(HERE / "fonts" / "NotoSansJP-Regular.ttf"),          # 同梱した場合に最優先
              "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc",  # packages.txt: fonts-noto-cjk
              "/System/Library/Fonts/ヒラギノ角ゴシック W6.ttc",       # macOS ローカル開発
              "/System/Library/Fonts/Hiragino Sans GB.ttc",
              "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"]:  # 最後の保険(日本語なし)
        if Path(p).exists():
            fpath = p; break
    f_title = ImageFont.truetype(fpath, 30) if fpath else ImageFont.load_default()
    f_sub = ImageFont.truetype(fpath, 16) if fpath else ImageFont.load_default()
    title_t = "切り出し文字"
    sub_t = f"{n_chars} 字"
    tw = draw.textlength(title_t, font=f_title)
    sw = draw.textlength(sub_t, font=f_sub)
    draw.text((grid_w - margin - tw, 16), title_t, fill=(70, 55, 35), font=f_title)
    draw.text((grid_w - margin - sw, 50), sub_t, fill=(120, 105, 80), font=f_sub)
    draw.line([(margin, header_h + 6), (grid_w - margin, header_h + 6)], fill=(200, 185, 160), width=2)

    pad_in = max(4, int(cell * pad_ratio))
    radius = cell // 9
    y0 = margin + header_h
    for i, crop in enumerate(crops):
        r, c_idx = divmod(i, cols)
        x = margin + c_idx * pitch
        y = y0 + r * pitch
        # 影
        draw.rounded_rectangle([x + 2, y + 3, x + cell + 1, y + cell + 2], radius=radius, fill=(0, 0, 0, 38))
        # 白い正方カード
        card = Image.new("RGBA", (cell, cell), (255, 255, 255, 255))
        mask = Image.new("L", (cell, cell), 0)
        ImageDraw.Draw(mask).rounded_rectangle([0, 0, cell - 1, cell - 1], radius=radius, fill=255)
        # 文字を padding 内に「長辺基準でリサイズ (拡大も許可)」して均一サイズに、中央配置
        fit = cell - 2 * pad_in
        cw, ch = crop.size
        s = fit / max(cw, ch)  # 小さい文字は拡大、大きい文字は縮小して揃える
        nw, nh = max(1, round(cw * s)), max(1, round(ch * s))
        cc = crop.convert("RGB").resize((nw, nh), Image.LANCZOS)
        card.paste(cc, ((cell - nw) // 2, (cell - nh) // 2))
        bg.paste(card, (x, y), mask)
        ImageDraw.Draw(bg, "RGBA").rounded_rectangle(
            [x, y, x + cell - 1, y + cell - 1], radius=radius, outline=(185, 170, 145, 255), width=1)
    return bg


def segment(image: Image.Image, mode: str = "自動 (解像度で判定)",
            show_lines: bool = False, filter_outside: bool = False,
            model_name: str = DEFAULT_MODEL):
    if image is None:
        return None, None, None, "画像をアップロードしてください。"
    image = image.convert("RGB")
    W, H = image.size
    auto_method, downscale = decide_method(W, H)

    # mode に応じて手法を決定
    if mode.startswith("baseline"):
        method = "baseline"
        chosen_note = "手動選択"
    elif mode.startswith("SAHI"):
        method = "SAHI"
        chosen_note = "手動選択"
    else:
        method = auto_method
        chosen_note = "自動判定"

    # 一時ファイルに保存して推論
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tf:
        image.save(tf.name)
        tmp_path = tf.name

    boxes = run_sahi(tmp_path, model_name) if method == "SAHI" else run_baseline(tmp_path, model_name)

    # 行検出 (表示 or フィルタのどちらかが有効なら実行)
    n_lines = 0
    line_boxes = []
    if show_lines or filter_outside:
        line_boxes = detect_lines(image)
        n_lines = len(line_boxes)

    # 行領域外の文字検出を除外
    n_before = len(boxes)
    n_excluded = 0
    if filter_outside and line_boxes:
        kept = [b for b in boxes if _center_in_any_line(b["xyxy"], line_boxes)]
        n_excluded = n_before - len(kept)
        boxes = kept

    # 描画
    vis = image.copy()
    d = ImageDraw.Draw(vis, "RGBA")
    lw = max(2, min(W, H) // 600)
    # 行 bbox を先に青で (背面)
    for lb in line_boxes:
        x1, y1, x2, y2 = lb
        d.rectangle([x1, y1, x2, y2], outline=(40, 120, 255, 230), width=lw + 1)
    # 文字 bbox を緑で (前面)
    for b in boxes:
        x1, y1, x2, y2 = b["xyxy"]
        d.rectangle([x1, y1, x2, y2], outline=(0, 200, 0, 230), width=lw)

    # JSON ファイル
    result = {
        "image_size": {"width": W, "height": H},
        "model": MODEL_REPOS.get(model_name, MODEL_REPOS[DEFAULT_MODEL]),
        "downscale_ratio": round(downscale, 3),
        "method": method,
        "method_selection": chosen_note,
        "auto_recommended": auto_method,
        "n_chars": len(boxes),
        "boxes": boxes,
    }
    out_json = Path(tempfile.gettempdir()) / "segmentation_result.json"
    out_json.write_text(json.dumps(result, ensure_ascii=False, indent=2))

    # 切り出し文字を並べた montage (行検出済みなら列構造を忠実に再現)
    montage = make_char_montage(image, boxes, line_boxes=line_boxes if line_boxes else None)

    line_info = f"\n- NDL 行検出 (青枠): **{n_lines} 行**" if (show_lines or filter_outside) else ""
    filter_info = f"\n- 行領域外を除外: **{n_excluded} 個** を除去 ({n_before}→{len(boxes)})" if filter_outside else ""
    info = (f"### 結果\n"
            f"- 画像サイズ: **{W} × {H}** / 縮小率 (長辺/{IMGSZ}): **{downscale:.2f}x**\n"
            f"- 使用モデル: **{model_name}**\n"
            f"- 使用手法: **{method}** ({chosen_note})\n"
            f"- 自動判定の推奨: **{auto_method}** (※判定は画像サイズのみで決定、行検出は不使用)\n"
            f"- 検出文字数 (緑枠): **{len(boxes)}**{line_info}{filter_info}\n\n"
            f"※ 同じ画像で手法を切り替えると、SAHI の効果 (特に高解像度) を比較できます。\n")
    return vis, montage, str(out_json), info


CUSTOM_CSS = """
.hero {text-align:center; padding: 18px 12px 6px;}
.hero h1 {font-size: 2.0rem; margin: 0 0 6px; font-weight: 800; letter-spacing: .02em;}
.hero p {font-size: 1.02rem; color: var(--body-text-color-subdued); margin: 0;}
.badges {text-align:center; margin: 4px 0 10px;}
.badges img {display:inline; margin: 0 3px; vertical-align: middle;}
#run-btn {font-size: 1.05rem; padding: 12px;}
footer {visibility: hidden;}
"""

THEME = gr.themes.Soft(
    primary_hue="emerald",
    secondary_hue="blue",
    font=[gr.themes.GoogleFont("Noto Sans JP"), "sans-serif"],
)

with gr.Blocks(title="くずし字 文字切り出し") as demo:
    gr.HTML(
        '<div class="hero">'
        '<h1>📜 くずし字 文字切り出し</h1>'
        '<p>古文書・くずし字・典籍の画像から、文字を 1 字ずつ自動で切り出します</p>'
        '</div>'
        '<div class="badges">'
        '<img src="https://img.shields.io/badge/model-komonjo(古文書)%20%2F%20codh(典籍)-success">'
        '<img src="https://img.shields.io/badge/解像度で-baseline%2FSAHI%20自動選択-blue">'
        '<img src="https://img.shields.io/badge/license-app%20MIT-lightgrey">'
        '</div>'
    )
    with gr.Row():
        with gr.Column(scale=1):
            inp = gr.Image(type="pil", label="① 史料画像をアップロード", height=340)
            model_sel = gr.Radio(
                choices=list(MODEL_REPOS),
                value=DEFAULT_MODEL,
                label="② モデル",
                info="自筆書状・文書など連綿の強い古文書は komonjo、版本・古典籍は codh が向きます",
            )
            mode = gr.Radio(
                choices=["baseline (通常推論)", "SAHI (タイル分割)", "自動 (解像度で判定)"],
                value="baseline (通常推論)",
                label="③ 手法",
                info="まず baseline を試し、取りこぼしが多ければ SAHI に切替を。SAHI は高解像度の連綿で有効な場合があります",
            )
            with gr.Row():
                show_lines = gr.Checkbox(value=False, label="行 bbox も表示 (青)")
                filter_outside = gr.Checkbox(value=True, label="行外ノイズ除去")
            btn = gr.Button("④ 文字を切り出す", variant="primary", elem_id="run-btn")
        with gr.Column(scale=2):
            out_img = gr.Image(type="pil", label="検出結果 (緑枠 = 検出文字 / 青枠 = NDL 行)", height=380)
            out_info = gr.Markdown()
            out_file = gr.File(label="座標データ (JSON)")
    # 切り出し文字一覧 (最もインパクトのある出力なので全幅で大きく表示)
    out_montage = gr.Image(type="pil", label="✂️ 切り出した文字一覧 (列構造を保った読み順)")

    gr.Examples(
        examples=[
            ["examples/toji_hyakugo_komonjo_CC-BY.jpg", "komonjo (古文書・書状)", "baseline (通常推論)", False, True],
            ["examples/ito_hirobumi_letter_NDL_PDM.jpg", "komonjo (古文書・書状)", "baseline (通常推論)", False, True],
            ["examples/taketori_kuzushiji_NDL_PDM.jpg", "codh (典籍・版本)", "baseline (通常推論)", False, True],
        ],
        inputs=[inp, model_sel, mode, show_lines, filter_outside],
        label="サンプル (中世古文書・近代書状・くずし字版本)",
    )
    gr.Markdown(
        "**サンプル出典**: "
        "「紀伊郡司解」京都府立京都学・歴彩館 東寺百合文書WEB から (CC BY 2.1 JP) ／ "
        "「伊藤博文書簡 岩倉具視宛」「竹取物語」国立国会図書館デジタルコレクション (Public Domain Mark)"
    )

    with gr.Accordion("このデモについて / 注意点", open=False):
        gr.Markdown(
            "**仕組み**: 画像の縮小率 (長辺 ÷ 1280) が 2.0 を超える高解像度では、画像をタイル分割して "
            "原寸に近い形で推論する **SAHI** を自動選択します。縮小で文字が潰れるのを防げます。"
            "低解像度では通常推論 (baseline)。\n\n"
            "**オプション**: 「行 bbox 表示」は NDL古典籍OCR-Lite の行検出 (RTMDet) を青で重ね、"
            "「行→文字」の構造を見せます。「行外ノイズ除去」は行領域の外に出た検出 (定規・色見本への誤検出) を捨てます。\n\n"
            "**モデル**: 「komonjo」は東京大学史料編纂所 電子くずし字字典データベース由来の古文書データで学習し、"
            "御内書のような連綿草書で codh 版を大きく上回ります。「codh」は CODH 古典籍くずし字データセットで学習し、"
            "版本・古典籍向きです。\n\n"
            "**注意**: いずれも日本のくずし字で学習しているため、漢籍の整版・楷書や近代活字など "
            "学習分布から外れた資料では検出が大きく落ちます。検出は完璧ではなく、最終的な文字同定には人の眼が必要です。\n\n"
            "**サンプル出典**: 「紀伊郡司解」京都府立京都学・歴彩館 東寺百合文書WEB から (CC BY 2.1 JP) / "
            "「伊藤博文書簡 岩倉具視宛」「竹取物語」国立国会図書館デジタルコレクション (Public Domain Mark)。\n\n"
            "**使用モデル**: [`nakamura196/yolov11x-komonjo-char`](https://huggingface.co/nakamura196/yolov11x-komonjo-char) "
            "(kuzushiji-db-derived-rights-reserved: 非営利の研究・教育利用可) / "
            "[`nakamura196/yolov11x-codh-char`](https://huggingface.co/nakamura196/yolov11x-codh-char) (MIT) / "
            "[NDL古典籍OCR-Lite](https://github.com/ndl-lab/ndlkotenocr-lite) (CC BY 4.0)。アプリ本体のコードは MIT。"
        )

    btn.click(segment, inputs=[inp, mode, show_lines, filter_outside, model_sel],
              outputs=[out_img, out_montage, out_file, out_info])

if __name__ == "__main__":
    demo.launch(theme=THEME, css=CUSTOM_CSS)