"""合成データセットを生成する（手元で実データと同じコードパスを流すため）.

実データと同じ構成 `<out>/images/<label>_<severity>/*.png` と `<out>/manifest.csv` を作る．
画像は黒背景 (0, 0, 0) の RGB PNG に，クラスごとに色の異なる楕円（舌の代わり）を描いたもの．

使い方:
    uv run python scripts/make_synthetic_dataset.py --out data/synthetic --per-dir 10 --seed 0
"""

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

from tias_judge_model.data.classes import ALL_DIR_NAMES, parse_dir_name
from tias_judge_model.data.manifest import build_manifest_records, write_manifest

logger = logging.getLogger(__name__)

# 出力先がこのスクリプトの生成物であることを示す目印．
# --overwrite で消してよいのは目印のあるディレクトリだけにし，実データを誤って消さないようにする．
MARKER_FILE = ".synthetic_dataset"

# クラスごとの基準色（RGB）．学習 smoke テストで学習可能なよう，互いに十分離しておく．
BASE_COLORS: dict[str, tuple[int, int, int]] = {
    "pale": (235, 205, 200),
    "light_red": (225, 140, 140),
    "red": (215, 40, 45),
    "dark_red": (120, 20, 30),
    "purple": (110, 50, 140),
}

# mild は strong より「正常（淡紅）」寄りの薄い色にする（境界例を模すため）．
MILD_BLEND_TO_NORMAL = 0.25

MIN_SIZE = 96
MAX_SIZE = 192
COLOR_JITTER = 8
PIXEL_NOISE_STD = 4.0


def _class_color(label: str, severity: str) -> np.ndarray:
    base = np.array(BASE_COLORS[label], dtype=np.float64)
    if severity == "mild":
        normal = np.array(BASE_COLORS["light_red"], dtype=np.float64)
        base = (1.0 - MILD_BLEND_TO_NORMAL) * base + MILD_BLEND_TO_NORMAL * normal
    return base


def make_image(color: np.ndarray, rng: np.random.Generator) -> Image.Image:
    """黒背景に楕円を描いた RGB 画像を 1 枚作る.

    Args:
        color: 楕円の基準色（`(3,)`, float64, 値域 0〜255）．
        rng: 乱数生成器．

    Returns:
        RGB 画像（サイズは `MIN_SIZE`〜`MAX_SIZE` の範囲でランダム）．
    """
    width = int(rng.integers(MIN_SIZE, MAX_SIZE + 1))
    height = int(rng.integers(MIN_SIZE, MAX_SIZE + 1))
    margin_x = int(rng.integers(0, width // 10 + 1))
    margin_y = int(rng.integers(0, height // 10 + 1))

    mask_img = Image.new("L", (width, height), 0)
    ImageDraw.Draw(mask_img).ellipse(
        (margin_x, margin_y, width - 1 - margin_x, height - 1 - margin_y), fill=255
    )
    mask = np.asarray(mask_img) > 0  # (H, W), bool

    jittered = color + rng.integers(-COLOR_JITTER, COLOR_JITTER + 1, size=3)
    noise = rng.normal(0.0, PIXEL_NOISE_STD, size=(height, width, 3))
    tongue = np.clip(jittered + noise, 1, 255)  # 舌領域が背景の黒 (0,0,0) と紛れないよう 1 以上
    pixels = np.zeros((height, width, 3), dtype=np.uint8)  # (H, W, 3)
    pixels[mask] = tongue[mask].astype(np.uint8)
    return Image.fromarray(pixels)


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

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

    Returns:
        生成した画像の枚数．
    """
    rng = np.random.default_rng(seed)
    images_dir = out_dir / "images"
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / MARKER_FILE).write_text(f"seed={seed}\nper_dir={per_dir}\n", encoding="utf-8")

    serial = 0
    for dir_name in ALL_DIR_NAMES:
        label, severity = parse_dir_name(dir_name)
        color = _class_color(label, severity)
        class_dir = images_dir / dir_name
        class_dir.mkdir(parents=True, exist_ok=True)
        for _ in range(per_dir):
            make_image(color, rng).save(class_dir / f"syn_{serial:06d}.png")
            serial += 1
        logger.info("%s/: %d 枚", dir_name, per_dir)

    records = build_manifest_records(images_dir)
    write_manifest(records, out_dir / "manifest.csv")
    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"), help="出力先（既定: data/synthetic）"
    )
    parser.add_argument(
        "--per-dir", type=int, default=10, help="1 ディレクトリあたりの枚数（既定: 10）"
    )
    parser.add_argument("--seed", type=int, default=0, help="乱数 seed（既定: 0）")
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="既存の合成データを削除して作り直す（目印ファイルのある出力先のみ）",
    )
    args = parser.parse_args(argv)
    if args.per_dir < 1:
        parser.error("--per-dir は 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_dir, args.seed)
    logger.info("合成データを生成した: %d 枚 -> %s", count, args.out)
    return 0


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