"""
密集堆叠箱子 YOLO 实例分割数据集生成器 v3
==============================================
整体流程：
  1. 在临时画布上【正交拼接】所有箱子 → 无间隙、无黑色底色
  2. 对整张拼图做【一次全局透视变换】→ 统一的俯视畸变
  3. 贴到背景图【中间 60% 区域】
  4. labels: 整个堆垛的 外轮廓多边形（所有箱子角点的凸包）

修复：
  - 拼图画布用第一个箱子纹理铺底，零间隙 → 无黑色
  - label 为整体堆垛多边形（如 3 箱合并为一个 6+ 点多边形）
"""

import cv2
import numpy as np
from pathlib import Path
import random
import os
from concurrent.futures import ThreadPoolExecutor, as_completed


# ============================================================
# 配置
# ============================================================

class Config:
    # ---- 必填 ----
    BOX_DIR = r""            # 箱子正面图文件夹
    BG_DIR = r""             # 背景环境图文件夹

    # ---- 输出 ----
    OUTPUT_DIR = r"./box_dataset"
    NUM_IMAGES = 500
    TRAIN_RATIO = 0.8
    IMG_WIDTH = 1280
    IMG_HEIGHT = 960

    # ---- 画面占比 ----
    STACK_CENTER_RATIO = 0.60    # 堆垛占画面中间的比例(宽高各60%)
    STACK_Y_OFFSET = 0.08        # 堆垛整体下移

    # ---- 堆垛结构 ----
    MIN_COLS = 2
    MAX_COLS = 5
    MIN_BASE_ROWS = 1
    MAX_BASE_ROWS = 4
    MAX_STACK_HEIGHT = 4

    # ---- 箱子尺寸(正交画布上) ----
    BOX_W_MIN = 80
    BOX_W_MAX = 180
    BOX_H_MIN = 100
    BOX_H_MAX = 240

    # ---- 箱子间距(设为0=零间隙，无黑色缝隙) ----
    GAP_X = 0
    GAP_Y = 0

    # ---- 全局透视 (上边/下边) ----
    PERSPECTIVE_TOP_RATIO = 1.05
    PERSPECTIVE_VAR = 0.02

    # ---- 箱子外观 ----
    BRIGHTNESS_RANGE = (0.80, 1.25)
    SATURATION_RANGE = (0.75, 1.20)
    HUE_SHIFT_RANGE = (-12, 12)

    # ---- 背景 ----
    BG_BLUR_RANGE = (0, 3)

    # ---- 多线程 ----
    WORKERS = 4


# ============================================================
# 工具
# ============================================================

def load_images(directory: str) -> list:
    exts = ('.jpg', '.jpeg', '.png', '.bmp', '.webp')
    imgs = []
    p = Path(directory)
    if not p.exists():
        return imgs
    for f in sorted(p.iterdir()):
        if f.suffix.lower() in exts:
            img = cv2.imread(str(f))
            if img is not None:
                imgs.append(img)
    return imgs


def adjust_color(img, brightness=1.0, saturation=1.0, hue_shift=0):
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float32)
    hsv[:, :, 0] = (hsv[:, :, 0] + hue_shift) % 180
    hsv[:, :, 1] = np.clip(hsv[:, :, 1] * saturation, 0, 255)
    hsv[:, :, 2] = np.clip(hsv[:, :, 2] * brightness, 0, 255)
    return cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)


# ============================================================
# 堆垛布局
# ============================================================

def build_stack_layout(cfg):
    """
    正交坐标系中构建多层堆垛。
    返回: boxes, comp_w, comp_h
    """
    n_cols = random.randint(cfg.MIN_COLS, cfg.MAX_COLS)
    n_base_rows = random.randint(cfg.MIN_BASE_ROWS, cfg.MAX_BASE_ROWS)

    col_widths = [random.randint(cfg.BOX_W_MIN, cfg.BOX_W_MAX)
                  for _ in range(n_cols)]
    row_heights = [random.randint(cfg.BOX_H_MIN, cfg.BOX_H_MAX)
                   for _ in range(n_base_rows)]

    # 零间距计算坐标
    col_x = [0]
    for w in col_widths:
        col_x.append(col_x[-1] + w)
    row_y = [0]
    for h in row_heights:
        row_y.append(row_y[-1] + h)

    boxes = []
    for r in range(n_base_rows):
        for c in range(n_cols):
            boxes.append({
                'row': r, 'col': c, 'layer': 0,
                'x1': col_x[c],
                'y1': row_y[r],
                'x2': col_x[c + 1],
                'y2': row_y[r + 1],
            })

    # 上层随机叠高
    for layer in range(1, cfg.MAX_STACK_HEIGHT):
        for r in range(n_base_rows):
            for c in range(n_cols):
                if random.random() < 0.5:
                    continue
                below = [b for b in boxes
                         if b['row'] == r and b['col'] == c
                         and b['layer'] == layer - 1]
                if not below:
                    continue
                base = below[0]
                shrink = random.uniform(0.82, 0.96)
                bw, bh = base['x2'] - base['x1'], base['y2'] - base['y1']
                nw = max(20, int(bw * shrink))
                nh = max(20, int(bh * shrink))
                nx1 = base['x1'] + (bw - nw) // 2
                ny1 = base['y1'] - nh
                boxes.append({
                    'row': r, 'col': c, 'layer': layer,
                    'x1': nx1, 'y1': ny1,
                    'x2': nx1 + nw, 'y2': ny1 + nh,
                })

    # 获取画布范围
    xs = [b['x1'] for b in boxes] + [b['x2'] for b in boxes]
    ys = [b['y1'] for b in boxes] + [b['y2'] for b in boxes]
    comp_w = int(max(xs) - min(xs) + 2)
    comp_h = int(max(ys) - min(ys) + 2)
    dx = min(xs) - 1
    dy = min(ys) - 1

    for b in boxes:
        b['x1'] -= dx
        b['x2'] -= dx
        b['y1'] -= dy
        b['y2'] -= dy

    return boxes, comp_w, comp_h


# ============================================================
# 拼图画布（无黑色底色）
# ============================================================

def build_composition(boxes_ortho, comp_w, comp_h, box_imgs, cfg):
    """
    用第一个箱子纹理铺满整张画布，再在上面画各个箱子。
    全程零间隙，无黑色出现。
    """
    # 拿一个箱子纹理铺满画布（消除黑色底色）
    base_tex = random.choice(box_imgs).copy()
    base_tex = adjust_color(
        base_tex,
        brightness=random.uniform(*cfg.BRIGHTNESS_RANGE),
        saturation=random.uniform(*cfg.SATURATION_RANGE),
        hue_shift=int(random.uniform(*cfg.HUE_SHIFT_RANGE)),
    )
    comp = cv2.resize(base_tex, (comp_w, comp_h), interpolation=cv2.INTER_AREA)

    # 按层排序（底层先画）
    boxes_sorted = sorted(boxes_ortho, key=lambda b: (b['layer'], b['y1']))

    for bi in boxes_sorted:
        x1 = int(max(0, round(bi['x1'])))
        y1 = int(max(0, round(bi['y1'])))
        x2 = int(min(comp_w, round(bi['x2'])))
        y2 = int(min(comp_h, round(bi['y2'])))
        bw, bh = x2 - x1, y2 - y1
        if bw < 3 or bh < 3:
            continue

        tex = random.choice(box_imgs).copy()
        tex = adjust_color(
            tex,
            brightness=random.uniform(*cfg.BRIGHTNESS_RANGE),
            saturation=random.uniform(*cfg.SATURATION_RANGE),
            hue_shift=int(random.uniform(*cfg.HUE_SHIFT_RANGE)),
        )
        tex = cv2.resize(tex, (bw, bh), interpolation=cv2.INTER_AREA)
        comp[y1:y2, x1:x2] = tex

    # 箱间细微阴影（只画底层底部）
    for bi in boxes_ortho:
        if bi['layer'] != 0:
            continue
        x1 = int(max(0, round(bi['x1'])))
        x2 = int(min(comp_w, round(bi['x2'])))
        y2 = int(min(comp_h, round(bi['y2'])))
        y1 = int(max(0, round(bi['y1'])))
        if y2 - y1 < 20:
            continue
        shadow_h = max(1, (y2 - y1) // 30)
        for si in range(shadow_h):
            sy = y2 - shadow_h + si
            a = 0.20 * (1 - si / shadow_h)
            row = comp[sy, x1:x2].astype(np.float32) * (1 - a)
            comp[sy, x1:x2] = np.clip(row, 0, 255).astype(np.uint8)

    return comp


# ============================================================
# 单张生成
# ============================================================

def generate_one(idx, box_imgs, bg_imgs, cfg, out_img_dir, out_lbl_dir):
    img_w, img_h = cfg.IMG_WIDTH, cfg.IMG_HEIGHT

    # ---- 1. 正交布局 ----
    boxes_ortho, comp_w, comp_h = build_stack_layout(cfg)
    if not boxes_ortho:
        return f"box_{idx:05d}.jpg", 0

    # ---- 2. 拼图（零黑色） ----
    comp = build_composition(boxes_ortho, comp_w, comp_h, box_imgs, cfg)

    # ---- 3. 全局透视 warp ----
    top_ratio = cfg.PERSPECTIVE_TOP_RATIO + random.uniform(
        -cfg.PERSPECTIVE_VAR, cfg.PERSPECTIVE_VAR)

    stack_w = int(img_w * cfg.STACK_CENTER_RATIO)
    stack_h = int(img_h * cfg.STACK_CENTER_RATIO)
    mid_x = img_w / 2.0
    offset_y = int(img_h * cfg.STACK_Y_OFFSET)
    top_y = int(img_h / 2 - stack_h / 2 + offset_y)
    bot_y = top_y + stack_h

    src_corners = np.float32([[0, 0], [comp_w, 0],
                              [comp_w, comp_h], [0, comp_h]])
    dst_corners = np.float32([
        [mid_x - (stack_w / 2) * top_ratio, top_y],
        [mid_x + (stack_w / 2) * top_ratio, top_y],
        [mid_x + stack_w / 2, bot_y],
        [mid_x - stack_w / 2, bot_y],
    ])

    M = cv2.getPerspectiveTransform(src_corners, dst_corners)

    warped = cv2.warpPerspective(
        comp, M, (img_w, img_h),
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=(0, 0, 0),
    )
    mask = cv2.warpPerspective(
        np.ones((comp_h, comp_w), dtype=np.uint8) * 255,
        M, (img_w, img_h),
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=0,
    )

    # ---- 4. 贴到背景 ----
    bg = random.choice(bg_imgs).copy()
    bg = cv2.resize(bg, (img_w, img_h))
    bk = random.randint(*cfg.BG_BLUR_RANGE)
    if bk >= 3:
        if bk % 2 == 0:
            bk += 1
        bg = cv2.GaussianBlur(bg, (bk, bk), 0)

    canvas = bg.copy()
    valid = mask > 0
    for ch in range(3):
        canvas[:, :, ch] = np.where(valid, warped[:, :, ch],
                                     canvas[:, :, ch])

    # ---- 5. 标签：整体轮廓多边形 ----
    # 收集所有箱子角点在最终图像中的坐标
    all_corners_final = []
    for bi in boxes_ortho:
        corners_comp = np.float32([
            [bi['x1'], bi['y1']],
            [bi['x2'], bi['y1']],
            [bi['x2'], bi['y2']],
            [bi['x1'], bi['y2']],
        ])
        ones = np.ones((4, 1), dtype=np.float32)
        pts_h = np.hstack([corners_comp, ones])
        pts_f = (M @ pts_h.T).T
        pts_f[:, 0] /= pts_f[:, 2]
        pts_f[:, 1] /= pts_f[:, 2]
        all_corners_final.append(pts_f[:, :2])

    all_pts = np.vstack(all_corners_final)  # (N, 2)

    # 凸包 → 整个堆垛的外轮廓
    hull = cv2.convexHull(all_pts.astype(np.float32))
    hull = hull.squeeze(1)  # (N, 2)

    # 简化（每隔1个点取一个，避免过多点）
    if len(hull) > 12:
        keep = np.arange(0, len(hull), 2)
        hull = hull[keep]

    norm_parts = []
    for pt in hull:
        norm_parts.append(f"{pt[0] / img_w:.6f}")
        norm_parts.append(f"{pt[1] / img_h:.6f}")
    label_line = "0 " + " ".join(norm_parts)

    # ---- 保存 ----
    fname = f"box_{idx:05d}.jpg"
    cv2.imwrite(os.path.join(out_img_dir, fname), canvas,
                [cv2.IMWRITE_JPEG_QUALITY, 95])

    lbl_path = os.path.join(out_lbl_dir, f"box_{idx:05d}.txt")
    with open(lbl_path, 'w') as f:
        f.write(label_line + "\n")

    return fname, 1  # 每张图1个标签(整体)


# ============================================================
# 主函数
# ============================================================

def main():
    c = Config

    if not c.BOX_DIR or not Path(c.BOX_DIR).exists():
        raise FileNotFoundError("请设置 BOX_DIR")
    if not c.BG_DIR or not Path(c.BG_DIR).exists():
        raise FileNotFoundError("请设置 BG_DIR")

    box_imgs = load_images(c.BOX_DIR)
    bg_imgs = load_images(c.BG_DIR)

    if len(box_imgs) < 2:
        raise ValueError(f"箱子纹理太少({len(box_imgs)})，至少2张")
    if len(bg_imgs) < 1:
        raise ValueError("背景图为空")

    print("=" * 55)
    print(f"[纹理] 箱子: {len(box_imgs)}张  |  背景: {len(bg_imgs)}张")
    print(f"[输出] {c.IMG_WIDTH}x{c.IMG_HEIGHT}  |  {c.NUM_IMAGES}张")
    print(f"[布局] {c.MIN_COLS}-{c.MAX_COLS}列 x "
          f"{c.MIN_BASE_ROWS}-{c.MAX_BASE_ROWS}行 | "
          f"最高{c.MAX_STACK_HEIGHT}层")
    print(f"[标签] 整体堆垛轮廓多边形（凸包）")
    print("=" * 55)

    out = Path(c.OUTPUT_DIR)
    for split in ['train', 'val']:
        (out / 'images' / split).mkdir(parents=True, exist_ok=True)
        (out / 'labels' / split).mkdir(parents=True, exist_ok=True)

    n_train = int(c.NUM_IMAGES * c.TRAIN_RATIO)
    n_val = c.NUM_IMAGES - n_train

    print(f"\n[生成] train={n_train}  val={n_val}")

    jobs = []
    for i in range(c.NUM_IMAGES):
        split = 'train' if i < n_train else 'val'
        jobs.append((i, str(out / 'images' / split),
                     str(out / 'labels' / split)))

    done = 0
    with ThreadPoolExecutor(max_workers=c.WORKERS) as pool:
        futures = {
            pool.submit(generate_one, idx, box_imgs, bg_imgs, c, idir, ldir): idx
            for idx, idir, ldir in jobs
        }
        for fut in as_completed(futures):
            fname, _ = fut.result()
            done += 1
            if done % 50 == 0 or done == c.NUM_IMAGES:
                print(f"  [{done}/{c.NUM_IMAGES}] {fname}")

    yaml_text = f"""# YOLO Instance Segmentation - Box Stack
path: {Path(c.OUTPUT_DIR).absolute()}
train: images/train
val: images/val
nc: 1
names:
  0: box_stack
"""
    with open(out / 'data.yaml', 'w') as f:
        f.write(yaml_text)

    print(f"\n{'=' * 55}")
    print(f"完成! train={n_train}  val={n_val}")
    print(f"每图1个整体堆垛标签")
    print(f"输出: {out.absolute()}")
    print("=" * 55)


if __name__ == '__main__':
    main()