"""
build_dataset_report.py
对任意图片目录（可含多层子目录）批量预测脸型并生成 HTML 报告。
保留图片原始所属的子目录名作为「分组」，在报告中按分组展示与统计。

用法:
    ./venv/bin/python face/build_dataset_report.py --src <图片目录> [--sample 50] [--seed 42]

示例:
    ./venv/bin/python face/build_dataset_report.py \
        --src face/test_img/脸型测试集合 --sample 50 --name 脸型测试集合

输出:
    static/<slug>_report.html
    static/<slug>_report/images/*.jpg
"""

from __future__ import annotations

import argparse
import html
import random
import re
import shutil
import sys
import unicodedata
from collections import Counter, defaultdict
from datetime import datetime
from pathlib import Path
from typing import Dict, List

import cv2

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from face.face_shape_classifier import classify_from_image  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
MAX_IMAGE_SIDE = 900
JPEG_QUALITY = 88

SHAPE_ORDER = ["圆形脸", "心形脸", "菱形脸", "鹅蛋脸", "方形脸", "长形脸", "瓜子脸"]
SHAPE_COLORS = {
    "圆形脸": "#e67e22",
    "心形脸": "#e74c3c",
    "菱形脸": "#9b59b6",
    "鹅蛋脸": "#27ae60",
    "方形脸": "#2980b9",
    "长形脸": "#16a085",
    "瓜子脸": "#c0392b",
    "检测失败": "#7f8c8d",
}
# 数据集分组名与分类器脸型口径的近似对应（仅用于交叉表高亮参考，非严格标签）
TAXONOMY_EQUIV = {
    "方形脸": "方形脸",
    "长形脸": "长形脸",
    "瓜子脸": "瓜子脸",
    "标准脸": "鹅蛋脸",
    "娃娃脸": "圆形脸",
}

FEATURE_KEYS = [
    "face_width",
    "face_height",
    "jaw_angle",
    "taper_ratio",
    "forehead_ratio",
    "cheekbone_ratio",
    "jaw_ratio",
    "chin_ratio",
    "chin_sharpness",
    "width_uniformity",
    "face_curve_score",
]


def natural_key(text: str):
    parts = re.split(r"(\d+)", text)
    return [int(p) if p.isdigit() else p for p in parts]


def collect_images(src: Path) -> List[Path]:
    return sorted(
        (p for p in src.rglob("*") if p.suffix.lower() in IMAGE_SUFFIXES),
        key=lambda p: natural_key(str(p.relative_to(src))),
    )


def group_of(path: Path, src: Path) -> str:
    """图片相对根目录的父目录名；直接位于根目录则记为「根目录」。"""
    rel = path.relative_to(src).parent
    return str(rel) if str(rel) != "." else "（根目录）"


def ascii_slug(text: str, fallback: str) -> str:
    """生成安全的 ASCII 文件名片段（中文目录名转拼音不可靠，直接编号兜底）。"""
    norm = unicodedata.normalize("NFKD", text).encode("ascii", "ignore").decode()
    norm = re.sub(r"[^A-Za-z0-9_-]+", "_", norm).strip("_")
    return norm or fallback


def stratified_sample(
    images: List[Path], src: Path, total: int, seed: int, min_per_group: int
) -> List[Path]:
    """
    按分组分层抽样：先保证每组至少 min_per_group 张，剩余名额按组大小比例分配。
    小分组（如只有 3 张的梨形脸）在纯随机抽样下几乎必然缺席，分层可保证覆盖。
    """
    rng = random.Random(seed)
    buckets: Dict[str, List[Path]] = defaultdict(list)
    for p in images:
        buckets[group_of(p, src)].append(p)

    groups = sorted(buckets, key=natural_key)
    quota = {g: min(min_per_group, len(buckets[g])) for g in groups}

    remaining = total - sum(quota.values())
    if remaining > 0:
        spare = {g: len(buckets[g]) - quota[g] for g in groups}
        pool = sum(spare.values())
        if pool > 0:
            # 按剩余可选量比例分配，再把取整误差补给最大的分组
            extra = {g: int(remaining * spare[g] / pool) for g in groups}
            for g in sorted(groups, key=lambda g: -spare[g]):
                if sum(extra.values()) >= remaining:
                    break
                if extra[g] < spare[g]:
                    extra[g] += 1
            for g in groups:
                quota[g] += min(extra[g], spare[g])

    chosen: List[Path] = []
    for g in groups:
        chosen.extend(rng.sample(buckets[g], min(quota[g], len(buckets[g]))))
    return chosen


def analyze(
    src: Path, sample: int, seed: int, img_dir: Path, min_per_group: int
) -> List[Dict]:
    all_images = collect_images(src)
    if not all_images:
        raise SystemExit(f"目录中没有图片: {src}")

    if sample and sample < len(all_images):
        if min_per_group > 0:
            chosen = stratified_sample(all_images, src, sample, seed, min_per_group)
        else:
            chosen = random.Random(seed).sample(all_images, sample)
        chosen.sort(key=lambda p: natural_key(str(p.relative_to(src))))
    else:
        chosen = all_images

    mode = f"分层抽样，每组至少 {min_per_group} 张" if min_per_group > 0 else "纯随机抽样"
    print(f"共发现 {len(all_images)} 张图片，本次测试 {len(chosen)} 张（{mode}，seed={seed}）\n")

    if img_dir.exists():
        shutil.rmtree(img_dir)
    img_dir.mkdir(parents=True)

    group_slugs: Dict[str, str] = {}
    rows: List[Dict] = []

    for idx, path in enumerate(chosen, 1):
        group = group_of(path, src)
        if group not in group_slugs:
            group_slugs[group] = ascii_slug(group, f"g{len(group_slugs) + 1}")
        out_name = f"{group_slugs[group]}_{idx:03d}.jpg"

        item = {
            "index": idx,
            "group": group,
            "file": path.name,
            "rel_path": str(path.relative_to(src)),
            # 相对 static/ 的路径（报告 HTML 也放在 static/ 根下）
            "img_src": f"{img_dir.relative_to(ROOT / 'static').as_posix()}/{out_name}",
            "ok": False,
            "predicted": None,
            "display": None,
            "confidence": None,
            "score": None,
            "top3": [],
            "features": {},
            "error": None,
        }

        try:
            result = classify_from_image(path, return_details=True, return_annotated=True)
            annotated = result["annotated"]
            h, w = annotated.shape[:2]
            if max(h, w) > MAX_IMAGE_SIDE:
                scale = MAX_IMAGE_SIDE / max(h, w)
                annotated = cv2.resize(
                    annotated, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA
                )
            cv2.imwrite(str(img_dir / out_name), annotated, [int(cv2.IMWRITE_JPEG_QUALITY), JPEG_QUALITY])

            item.update(
                {
                    "ok": True,
                    "predicted": result["face_shape"],
                    "display": result["display"],
                    "confidence": result["confidence"],
                    "score": result["details"]["ranked"][0][1],
                    "top3": result["details"]["ranked"][:3],
                    "features": {k: result["features"][k] for k in FEATURE_KEYS},
                }
            )
        except Exception as exc:  # noqa: BLE001 - 报告需要汇总所有失败样本
            img = cv2.imread(str(path))
            if img is not None:
                h, w = img.shape[:2]
                if max(h, w) > MAX_IMAGE_SIDE:
                    scale = MAX_IMAGE_SIDE / max(h, w)
                    img = cv2.resize(img, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)
                cv2.imwrite(str(img_dir / out_name), img, [int(cv2.IMWRITE_JPEG_QUALITY), JPEG_QUALITY])
            item["error"] = str(exc)

        rows.append(item)
        print(f"[{idx:3d}/{len(chosen)}] [{group}] {path.name} -> {item['display'] or 'ERR: ' + str(item['error'])}")

    return rows


def bar_chart(counter: Counter) -> str:
    if not counter:
        return "<p class='muted'>无数据</p>"
    total = sum(counter.values())
    parts = []
    order = [s for s in SHAPE_ORDER if counter.get(s)] + [
        s for s in counter if s not in SHAPE_ORDER
    ]
    for shape in order:
        n = counter[shape]
        color = SHAPE_COLORS.get(shape, "#7f8c8d")
        pct = n / total * 100
        parts.append(
            f"<div class='bar-row'><span class='bar-label'>{html.escape(shape)}</span>"
            f"<div class='bar-track'><div class='bar-fill' style='width:{pct:.1f}%;background:{color}'></div></div>"
            f"<span class='bar-num'>{n}（{pct:.0f}%）</span></div>"
        )
    return "".join(parts)


def fmt_feat(key: str, value: float) -> str:
    if key in {"face_width", "face_height"}:
        return f"{value:.1f}px"
    if key == "jaw_angle":
        return f"{value:.1f}°"
    return f"{value:.3f}"


def card(item: Dict) -> str:
    group_tag = f"<span class='group-tag'>{html.escape(item['group'])}</span>"
    if not item["ok"]:
        return f"""
        <article class="card error">
          <a class="img-link" href="{html.escape(item['img_src'])}" target="_blank">
            <img src="{html.escape(item['img_src'])}" alt="{html.escape(item['file'])}" loading="lazy"/>
          </a>
          <div class="body">
            <div class="meta"><h3>{html.escape(item['file'])}</h3>{group_tag}</div>
            <p class="badge bad">检测失败</p>
            <p class="muted">{html.escape(item['error'] or '')}</p>
          </div>
        </article>"""

    color = SHAPE_COLORS.get(item["predicted"], "#34495e")
    top3 = "".join(
        f"<li><span>{html.escape(name)}</span><b>{score:.1f}</b></li>" for name, score in item["top3"]
    )
    feat_html = "".join(
        f"<tr><td>{html.escape(k)}</td><td>{html.escape(fmt_feat(k, v))}</td></tr>"
        for k, v in item["features"].items()
    )
    return f"""
    <article class="card">
      <a class="img-link" href="{html.escape(item['img_src'])}" target="_blank" title="点击查看大图标注">
        <img src="{html.escape(item['img_src'])}" alt="{html.escape(item['file'])}" loading="lazy"/>
      </a>
      <div class="body">
        <div class="meta"><h3>{html.escape(item['file'])}</h3>{group_tag}</div>
        <p class="path muted">{html.escape(item['rel_path'])}</p>
        <p class="badge" style="background:{color}">{html.escape(item['display'])}</p>
        <p class="conf">匹配度 {item['score']:.1f} · 置信度 {item['confidence']:.3f}</p>
        <ul class="scores">{top3}</ul>
        <details>
          <summary>标注特征数值</summary>
          <table>{feat_html}</table>
        </details>
      </div>
    </article>"""


CSS = """
:root { --bg:#f3efe6; --ink:#1c1915; --muted:#6b645a; --card:#fffdf8; --line:#e2d8c8; --accent:#0f6b5c; }
* { box-sizing: border-box; }
body {
  margin:0; font-family:"PingFang SC","Noto Sans SC","Segoe UI",sans-serif; color:var(--ink);
  background: radial-gradient(1200px 600px at 10% -10%, #ffe8c8 0%, transparent 55%),
              radial-gradient(900px 500px at 100% 0%, #d9f2ea 0%, transparent 50%), var(--bg);
}
header { padding:40px 24px 20px; max-width:1320px; margin:0 auto; }
header h1 { margin:0 0 8px; font-size:clamp(1.8rem,3vw,2.4rem); }
header p { margin:4px 0; color:var(--muted); }
.legend-box { max-width:1320px; margin:0 auto 20px; padding:0 24px; }
.legend-box .inner { background:var(--card); border:1px solid var(--line); border-radius:14px;
  padding:14px 16px; font-size:.9rem; line-height:1.55; }
.legend-box code { background:#efe7da; padding:1px 6px; border-radius:4px; font-size:.84rem; }
.swatch { display:inline-block; width:10px; height:10px; border-radius:2px; margin-right:4px; vertical-align:middle; }
.stats { display:grid; grid-template-columns:repeat(auto-fit,minmax(280px,1fr)); gap:16px;
  max-width:1320px; margin:0 auto 28px; padding:0 24px; }
.stat { background:var(--card); border:1px solid var(--line); border-radius:16px; padding:16px 18px; }
.stat h2 { margin:0 0 12px; font-size:1rem; }
.bar-row { display:grid; grid-template-columns:72px 1fr 80px; gap:8px; align-items:center; margin:6px 0; font-size:.86rem; }
.bar-track { height:8px; background:#efe7da; border-radius:999px; overflow:hidden; }
.bar-fill { height:100%; border-radius:999px; }
.bar-num { color:var(--muted); text-align:right; }
section { max-width:1320px; margin:0 auto 36px; padding:0 24px; }
section h2 { margin:0 0 14px; font-size:1.3rem; border-left:4px solid var(--accent); padding-left:10px; }
section h2 small { color:var(--muted); font-weight:400; font-size:.8rem; margin-left:8px; }
.grid { display:grid; grid-template-columns:repeat(auto-fill,minmax(270px,1fr)); gap:16px; }
.card { background:var(--card); border:1px solid var(--line); border-radius:18px; overflow:hidden;
  display:flex; flex-direction:column; box-shadow:0 8px 24px rgba(60,40,10,.05); }
.card.error { opacity:.9; }
.img-link { display:block; }
.card img { width:100%; aspect-ratio:3/4; object-fit:cover; background:#ddd; display:block; }
.card .body { padding:14px; }
.meta { display:flex; justify-content:space-between; align-items:baseline; gap:8px; }
.meta h3 { margin:0; font-size:.95rem; word-break:break-all; }
.group-tag { font-size:.72rem; color:var(--accent); background:#e7f6f2; padding:2px 8px;
  border-radius:999px; white-space:nowrap; }
.path { font-size:.72rem; margin:4px 0 0; word-break:break-all; }
.badge { display:inline-block; margin:10px 0 4px; color:#fff; padding:6px 10px; border-radius:999px;
  font-weight:600; font-size:.92rem; }
.badge.bad { background:#c0392b; }
.conf { margin:0 0 8px; color:var(--muted); font-size:.85rem; }
.scores { list-style:none; padding:0; margin:0 0 8px; }
.scores li { display:flex; justify-content:space-between; padding:4px 0; border-bottom:1px dashed var(--line); font-size:.86rem; }
details { margin-top:8px; }
summary { cursor:pointer; color:var(--accent); font-size:.86rem; }
table { width:100%; border-collapse:collapse; margin-top:8px; font-size:.8rem; }
td { padding:3px 0; border-bottom:1px solid var(--line); }
td:last-child { text-align:right; font-variant-numeric:tabular-nums; }
.muted { color:var(--muted); }
.cross-wrap { overflow-x:auto; background:var(--card); border:1px solid var(--line);
  border-radius:16px; padding:14px 16px; }
table.cross { border-collapse:collapse; width:100%; font-size:.88rem; }
table.cross th, table.cross td { padding:7px 10px; text-align:center; border-bottom:1px solid var(--line);
  white-space:nowrap; }
table.cross thead th { background:#efe7da; font-weight:600; position:sticky; top:0; }
table.cross th.rowh { text-align:left; font-weight:600; }
table.cross th.rowh small { color:var(--muted); font-weight:400; }
table.cross td.num { font-variant-numeric:tabular-nums; }
table.cross td.hit { background:#d8f3e4; color:#0f6b5c; font-weight:700; font-variant-numeric:tabular-nums; }
table.cross td.zero { color:#cfc6b6; }
footer { max-width:1320px; margin:0 auto; padding:8px 24px 40px; color:var(--muted); font-size:.85rem; }
"""


def cross_table(by_group: Dict[str, List[Dict]]) -> str:
    """原始分组 × 预测脸型 交叉表，对角线（口径对应的格子）高亮。"""
    cols = SHAPE_ORDER + ["检测失败"]
    head = "".join(f"<th>{html.escape(c)}</th>" for c in cols)
    body = []
    for group, items in sorted(by_group.items(), key=lambda kv: natural_key(kv[0])):
        counts = Counter(i["predicted"] if i["ok"] else "检测失败" for i in items)
        equiv = TAXONOMY_EQUIV.get(group)
        cells = []
        for c in cols:
            n = counts.get(c, 0)
            if n == 0:
                cells.append("<td class='zero'>·</td>")
                continue
            cls = "hit" if c == equiv else "num"
            cells.append(f"<td class='{cls}'>{n}</td>")
        label = html.escape(group)
        if equiv:
            label += f" <small>≈{html.escape(equiv)}</small>"
        body.append(f"<tr><th class='rowh'>{label}</th>{''.join(cells)}<th>{len(items)}</th></tr>")
    return (
        "<div class='cross-wrap'><table class='cross'>"
        f"<thead><tr><th>原始分组 \\ 预测</th>{head}<th>合计</th></tr></thead>"
        f"<tbody>{''.join(body)}</tbody></table></div>"
    )


def build_html(rows: List[Dict], name: str, src: Path, seed: int, img_dir_name: str) -> str:
    overall = Counter(r["predicted"] if r["ok"] else "检测失败" for r in rows)
    by_group: Dict[str, List[Dict]] = defaultdict(list)
    for r in rows:
        by_group[r["group"]].append(r)

    group_stats = "".join(
        f"<div class='stat'><h2>{html.escape(g)} <span class='muted'>（{len(items)} 张）</span></h2>"
        f"{bar_chart(Counter(i['predicted'] if i['ok'] else '检测失败' for i in items))}</div>"
        for g, items in sorted(by_group.items(), key=lambda kv: natural_key(kv[0]))
    )

    sections = "".join(
        f"<section><h2>{html.escape(g)}<small>{len(items)} 张</small></h2>"
        f"<div class='grid'>{''.join(card(i) for i in items)}</div></section>"
        for g, items in sorted(by_group.items(), key=lambda kv: natural_key(kv[0]))
    )

    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    n_ok = sum(1 for r in rows if r["ok"])
    n_kind = len([s for s in SHAPE_ORDER if overall.get(s)])

    return f"""<!DOCTYPE html>
<html lang="zh-CN">
<head>
<meta charset="UTF-8"/>
<meta name="viewport" content="width=device-width, initial-scale=1"/>
<title>{html.escape(name)} — 脸型分类报告</title>
<style>{CSS}</style>
</head>
<body>
<header>
  <h1>{html.escape(name)} — 脸型分类报告</h1>
  <p>z 分数原型匹配分类 · 照片上标注 face_width / face_height 及各比例特征</p>
  <p>生成时间：{html.escape(now)} · 抽样 {len(rows)} 张（随机种子 {seed}）· 成功 {n_ok} 张 ·
     覆盖 {n_kind} 种脸型 · 共 {len(by_group)} 个原始分组 · 点击图片看大图</p>
  <p class="muted">来源目录：{html.escape(str(src))}</p>
</header>

<div class="legend-box">
  <div class="inner">
    <b>图上标注说明</b><br/>
    <span class="swatch" style="background:#00dc78"></span><code>face_width</code> 颧骨宽度&nbsp;
    <span class="swatch" style="background:#28b4ff"></span><code>face_height</code> 额头顶→下巴&nbsp;
    <span class="swatch" style="background:#ff5a00"></span><code>jaw_angle</code> 下巴到左右下颌角夹角&nbsp;
    <span class="swatch" style="background:#ffc828"></span><code>taper_ratio</code> 额头→下巴收窄&nbsp;
    <span class="swatch" style="background:#28a0ff"></span><code>forehead / jaw / chin ratio</code> 各级宽度比&nbsp;
    <span class="swatch" style="background:#b4ff50"></span><code>face_curve_score</code> 下颌中点→下巴&nbsp;
    右侧柱状条示意 <code>width_uniformity</code>；左上角是完整数值图例。
  </div>
</div>

<section>
  <h2>原始分组 × 预测脸型 对照<small>数据集分组本身是脸型标签，但命名口径与分类器不同</small></h2>
  {cross_table(by_group)}
  <p class="muted" style="margin-top:10px;font-size:.85rem">
    绿色格子表示预测结果与该分组的对应口径一致（标准脸≈鹅蛋脸、娃娃脸≈圆形脸，方形/长形/瓜子同名直接对应）。
    「梨形脸」「混合脸」在分类器的 7 分类里没有对应项，不作一致性判断。
  </p>
</section>

<div class="stats">
  <div class="stat"><h2>总体脸型分布（{len(rows)} 张）</h2>{bar_chart(overall)}</div>
  {group_stats}
</div>

{sections}

<footer>
  分类实现：face/face_shape_classifier.py · 报告生成：face/build_dataset_report.py ·
  图片目录：static/{html.escape(img_dir_name)}/
</footer>
</body>
</html>
"""


def main() -> None:
    ap = argparse.ArgumentParser(description="批量脸型预测并生成 HTML 报告")
    ap.add_argument("--src", required=True, help="图片根目录（可含子目录）")
    ap.add_argument("--sample", type=int, default=50, help="随机抽样张数，0 表示全部")
    ap.add_argument("--seed", type=int, default=42, help="随机种子")
    ap.add_argument("--name", default=None, help="报告标题，默认取目录名")
    ap.add_argument("--slug", default="dataset", help="输出文件名前缀（ASCII）")
    ap.add_argument(
        "--min-per-group",
        type=int,
        default=2,
        help="分层抽样时每个分组至少抽几张，0 表示纯随机抽样",
    )
    args = ap.parse_args()

    src = Path(args.src).expanduser().resolve()
    if not src.is_dir():
        raise SystemExit(f"目录不存在: {src}")

    name = args.name or src.name
    img_dir_name = f"{args.slug}_report"
    img_dir = ROOT / "static" / img_dir_name / "images"
    out_html = ROOT / "static" / f"{args.slug}_report.html"

    rows = analyze(src, args.sample, args.seed, img_dir, args.min_per_group)
    out_html.write_text(
        build_html(rows, name, src, args.seed, img_dir_name), encoding="utf-8"
    )

    overall = Counter(r["predicted"] if r["ok"] else "检测失败" for r in rows)
    total = sum(overall.values())
    print(f"\n写入 {out_html}")
    print("=== 总体脸型分布 ===")
    for shape in SHAPE_ORDER + ["检测失败"]:
        n = overall.get(shape, 0)
        if n:
            print(f"  {shape}: {n:3d}  ({n / total * 100:4.1f}%)  {'#' * n}")


if __name__ == "__main__":
    main()