"""セグメンテーションの評価（`tias-seg eval`）を手元で通すための合成データを生成する.

実データ（TongueSegmentation の `input/`）と同じ構成を作る:

    <out>/input/train/{sample,label}/*.jpg
    <out>/input/test/cropped/{sample,label}/*.jpg
    <out>/input/test_extra/cropped/{sample,label}/*.jpg
    <out>/input/test/fullsize/{sample,label}/*.jpg

撮影画像（fullsize）は青みがかった灰色の背景に舌の代わりの赤い楕円を描いたもの，正解マスクは
楕円の内側を 255 にしたグレースケール JPEG（実データと同じく JPEG のため境界に中間値を含む）．
切り出し済み（cropped）は楕円の外接矩形に余白を付けて切り出したもの．一部の画像には楕円の外に
小さな赤い点を置き，小領域除去が働く場面を作る．ダミーの ONNX モデルは
`scripts/make_dummy_seg_models.py` で作る．

使い方:
    uv run python scripts/make_synthetic_segmentation.py --out data/synthetic_seg --per-set 8
"""

from __future__ import annotations

import argparse
import logging
import shutil
import sys
from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw

logger = logging.getLogger(__name__)

# 出力先がこのスクリプトの生成物であることを示す目印（--overwrite で消すのは目印のある所だけ）．
MARKER_FILE = ".synthetic_segmentation"

# (評価セットのディレクトリ, 撮影画像か)．
SETS: tuple[tuple[str, bool], ...] = (
    ("train", False),
    ("test/cropped", False),
    ("test_extra/cropped", False),
    ("test/fullsize", True),
)

BACKGROUND_RGB = (110, 120, 150)  # 青みがかった灰色（ダミーの ResUNet が背景と判定する色）
TONGUE_RGB = (200, 95, 105)  # 赤みのある舌の色
COLOR_JITTER = 10
PIXEL_NOISE_STD = 5.0
MIN_SIZE = 200
MAX_SIZE = 360
SPOT_PROBABILITY = 0.5  # 楕円の外に小さな赤い点を置く確率
JPEG_QUALITY = 95


def make_fullsize(rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray]:
    """撮影画像を模した画像と正解マスクを 1 組作る.

    Args:
        rng: 乱数生成器．

    Returns:
        （RGB 画像 `(H, W, 3)` uint8, 正解マスク `(H, W)` bool）．
    """
    width = int(rng.integers(MIN_SIZE, MAX_SIZE + 1))
    height = int(rng.integers(MIN_SIZE, MAX_SIZE + 1))
    ew = int(rng.integers(width // 4, width // 2))
    eh = int(rng.integers(height // 4, height // 2))
    cx = int(rng.integers(ew // 2 + 5, width - ew // 2 - 5))
    cy = int(rng.integers(eh // 2 + 5, height - eh // 2 - 5))
    mask_img = Image.new("L", (width, height), 0)
    ImageDraw.Draw(mask_img).ellipse(
        (cx - ew // 2, cy - eh // 2, cx + ew // 2, cy + eh // 2), fill=255
    )
    mask = np.asarray(mask_img) > 0  # (H, W)

    # 背景は横方向に少し明るさが変わる（一様な色だけにしない）．
    gradient = np.linspace(-15.0, 15.0, width)[None, :, None]
    background = np.array(BACKGROUND_RGB, dtype=np.float64) + gradient
    tongue = np.array(TONGUE_RGB, dtype=np.float64) + rng.integers(
        -COLOR_JITTER, COLOR_JITTER + 1, size=3
    )
    pixels = np.where(mask[..., None], tongue, background)
    pixels = pixels + rng.normal(0.0, PIXEL_NOISE_STD, size=(height, width, 3))

    if rng.random() < SPOT_PROBABILITY:
        # 舌と同じ色の小さな点（小領域除去で消えるべき誤検出の種）．正解マスクには含めない．
        spot = Image.new("L", (width, height), 0)
        sx = int(rng.integers(3, width - 8))
        sy = int(rng.integers(3, height - 8))
        ImageDraw.Draw(spot).ellipse((sx, sy, sx + 4, sy + 4), fill=255)
        spot_mask = (np.asarray(spot) > 0) & ~mask
        pixels[spot_mask] = tongue
    return np.clip(pixels, 0, 255).astype(np.uint8), mask


def crop_around(
    rgb: np.ndarray, mask: np.ndarray, rng: np.random.Generator
) -> tuple[np.ndarray, np.ndarray]:
    """正解マスクの外接矩形に余白を付けて切り出す（切り出し済みの画像を模す）.

    Args:
        rgb: RGB 画像（`(H, W, 3)`, uint8）．
        mask: 正解マスク（`(H, W)`, bool）．
        rng: 乱数生成器．

    Returns:
        （切り出した RGB 画像, 切り出した正解マスク）．
    """
    ys, xs = np.nonzero(mask)
    height, width = mask.shape
    mx = int(rng.integers(2, max(3, (xs.max() - xs.min()) // 8)))
    my = int(rng.integers(2, max(3, (ys.max() - ys.min()) // 8)))
    x0, x1 = max(0, xs.min() - mx), min(width, xs.max() + 1 + mx)
    y0, y1 = max(0, ys.min() - my), min(height, ys.max() + 1 + my)
    return rgb[y0:y1, x0:x1], mask[y0:y1, x0:x1]


def generate(out_dir: Path, per_set: int, seed: int) -> int:
    """合成データを生成する.

    Args:
        out_dir: 出力先（空または存在しないディレクトリ）．
        per_set: 評価セットあたりの枚数．
        seed: 乱数 seed．

    Returns:
        生成した組の数．
    """
    rng = np.random.default_rng(seed)
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / MARKER_FILE).write_text(f"seed={seed}\nper_set={per_set}\n", encoding="utf-8")
    serial = 0
    for set_dir, fullsize in SETS:
        sample_dir = out_dir / "input" / set_dir / "sample"
        label_dir = out_dir / "input" / set_dir / "label"
        sample_dir.mkdir(parents=True, exist_ok=True)
        label_dir.mkdir(parents=True, exist_ok=True)
        for _ in range(per_set):
            rgb, mask = make_fullsize(rng)
            if not fullsize:
                rgb, mask = crop_around(rgb, mask, rng)
            name = f"seg_{serial:06d}.jpg"
            Image.fromarray(rgb).save(sample_dir / name, quality=JPEG_QUALITY)
            Image.fromarray((mask * 255).astype(np.uint8)).save(
                label_dir / name, quality=JPEG_QUALITY
            )
            serial += 1
        logger.info("%s: %d 組", set_dir, per_set)
    return serial


def _prepare_out_dir(out_dir: Path, overwrite: bool) -> None:
    if not out_dir.exists():
        return
    if not out_dir.is_dir():
        raise ValueError(f"出力先がディレクトリではない: {out_dir}")
    if not any(out_dir.iterdir()):
        return
    if not overwrite:
        raise ValueError(f"出力先が空でない: {out_dir}（作り直すには --overwrite を付ける）")
    if not (out_dir / MARKER_FILE).is_file():
        raise ValueError(
            f"出力先に合成データの目印 {MARKER_FILE} が無いため削除しない: {out_dir}"
            "（合成データ以外を誤って消さないため）"
        )
    logger.info("既存の合成データを削除する: %s", out_dir)
    shutil.rmtree(out_dir)


def _parse_args(argv: list[str] | None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="セグメンテーションの評価用の合成データを生成する")
    parser.add_argument(
        "--out",
        type=Path,
        default=Path("data/synthetic_seg"),
        help="出力先（既定: data/synthetic_seg）",
    )
    parser.add_argument("--per-set", type=int, default=8, help="評価セットあたりの枚数（既定: 8）")
    parser.add_argument("--seed", type=int, default=0, help="乱数 seed（既定: 0）")
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="既存の合成データを削除して作り直す（目印ファイルのある出力先のみ．models/ も消える）",
    )
    args = parser.parse_args(argv)
    if args.per_set < 1:
        parser.error("--per-set は 1 以上を指定する")
    return args


def main(argv: list[str] | None = None) -> int:
    """エントリポイント.

    Args:
        argv: コマンドライン引数（None なら `sys.argv[1:]`）．

    Returns:
        終了コード（成功 0・入力の問題 1）．
    """
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
    )
    args = _parse_args(argv)
    try:
        _prepare_out_dir(args.out, args.overwrite)
    except ValueError as e:
        logger.error("%s", e)
        return 1
    count = generate(args.out, args.per_set, args.seed)
    logger.info("合成データを生成した: %d 組 -> %s", count, args.out / "input")
    return 0


if __name__ == "__main__":
    sys.exit(main())