import os

import random

import time

from typing import Dict, Optional, Tuple

import numpy as np

# 后视镜折叠场景配置（请按你的4种分辨率填写）

# key: (width, height)  value: (x1, y1, x2, y2) 车辆黑色区域在原图上的像素坐标

MIRROR_FOLD_CAR_BOXES: Dict[Tuple[int, int], Tuple[int, int, int, int]] = {

    # (960, 1088): (x1, y1, x2, y2),

    # (1280, 720): (x1, y1, x2, y2),

    # (1920, 1080): (x1, y1, x2, y2),

    # (1440, 1080): (x1, y1, x2, y2),

}

# 若分辨率不在上表中，可用比例兜底（0-1），None 表示不启用兜底

MIRROR_FOLD_CAR_BOX_RATIOS: Optional[Tuple[float, float, float, float]] = None

# 是否启用后视镜折叠增强

MIRROR_FOLD_ENABLE = True

MIRROR_FOLD_PROB = 1.0  # 1.0=每次都做，0.5=50%概率

MIRROR_FOLD_APPLY_TO_VAL = True  # 验证阶段也应用同样增强

MIRROR_FOLD_AUTO_INFER_CAR_BOX = True

MIRROR_FOLD_DARK_THRESH = 35

# 当既没有分辨率配置，也无法从图像中推断车辆区域时的兜底框（比例）

MIRROR_FOLD_DEFAULT_CAR_BOX_RATIOS: Tuple[float, float, float, float] = (0.42, 0.30, 0.58, 0.72)

# 粉色色块颜色

MIRROR_FOLD_PINK_COLOR_BGR = (255, 0, 255)  # OpenCV/BGR

MIRROR_FOLD_PINK_COLOR_RGB = (255, 0, 255)  # PIL/RGB

# 语义分割标签填充值（按你的数据集语义修改）

FSD_PINK_VALUE = 1  # FSD中粉色区域视为可行驶

RM_PINK_VALUE = 0   # RM中粉色区域视为背景

# Debug保存（预处理后可视化）

MIRROR_FOLD_DEBUG_SAVE = False

MIRROR_FOLD_DEBUG_APPLY_TO_VAL = True

MIRROR_FOLD_DEBUG_DIR = "runs/mirror_fold_debug"

MIRROR_FOLD_DEBUG_MAX = 200

MIRROR_FOLD_DEBUG_EVERY = 1

MIRROR_FOLD_DEBUG_ALPHA = 0.45

_debug_counts = {"det": 0, "fsd": 0, "rm": 0}

 

def _clamp_box(x1: int, y1: int, x2: int, y2: int, w: int, h: int) -> Optional[Tuple[int, int, int, int]]:

    x1 = int(max(0, min(x1, w)))

    x2 = int(max(0, min(x2, w)))

    y1 = int(max(0, min(y1, h)))

    y2 = int(max(0, min(y2, h)))

    if x2 <= x1 or y2 <= y1:

        return None

    return x1, y1, x2, y2

 

def get_car_box_for_shape(width: int, height: int) -> Optional[Tuple[int, int, int, int]]:

    car_box = MIRROR_FOLD_CAR_BOXES.get((width, height))

    if car_box is None and MIRROR_FOLD_CAR_BOX_RATIOS is not None:

        x1r, y1r, x2r, y2r = MIRROR_FOLD_CAR_BOX_RATIOS

        car_box = (int(x1r * width), int(y1r * height), int(x2r * width), int(y2r * height))

    if car_box is None:

        return None

    return _clamp_box(*car_box, w=width, h=height)

 

def build_pink_mask(width: int, height: int, car_box: Tuple[int, int, int, int]) -> Optional[np.ndarray]:

    x1, y1, x2, y2 = _clamp_box(*car_box, w=width, h=height) or (None, None, None, None)

    if x1 is None:

        return None

    mask = np.zeros((height, width), dtype=bool)

    if x1 > 0:

        mask[y1:y2, :x1] = True

    if x2 < width:

        mask[y1:y2, x2:] = True

    return mask

 

def should_apply_mirror_fold() -> bool:

    return MIRROR_FOLD_ENABLE and random.random() < MIRROR_FOLD_PROB

 

def get_debug_save_path(branch: str, img_path: str, suffix: str = "jpg") -> Optional[str]:

    if not MIRROR_FOLD_DEBUG_SAVE:

        return None

    count = _debug_counts.get(branch, 0)

    if count >= MIRROR_FOLD_DEBUG_MAX:

        return None

    if MIRROR_FOLD_DEBUG_EVERY > 1 and (count % MIRROR_FOLD_DEBUG_EVERY) != 0:

        _debug_counts[branch] = count + 1

        return None

    _debug_counts[branch] = count + 1

    base = os.path.splitext(os.path.basename(img_path))[0]

    out_dir = os.path.join(MIRROR_FOLD_DEBUG_DIR, branch)

    os.makedirs(out_dir, exist_ok=True)

    ts = int(time.time() * 1000)

    return os.path.join(out_dir, f"{base}_{count:06d}_{ts}.{suffix}")

 

def infer_car_box_from_image(img: np.ndarray) -> Optional[Tuple[int, int, int, int]]:

    if img is None or img.ndim < 2:

        return None

    if img.ndim == 3:

        gray = img.mean(axis=2)

    else:

        gray = img

    h, w = gray.shape[:2]

    if h < 8 or w < 8:

        return None

    dark = gray < MIRROR_FOLD_DARK_THRESH

    col_count = dark.sum(axis=0)

    center_l = int(w * 0.30)

    center_r = int(w * 0.70)

    center_slice = col_count[center_l:center_r]

    if center_slice.size == 0:

        return None

    peak_rel = int(np.argmax(center_slice))

    peak = center_l + peak_rel

    min_col_dark = max(3, int(h * 0.08))

    if col_count[peak] < min_col_dark:

        return None

    thresh_col = max(2, int(col_count[peak] * 0.35))

    x1 = peak

    x2 = peak

    while x1 > 0 and col_count[x1 - 1] >= thresh_col:

        x1 -= 1

    while x2 < w - 1 and col_count[x2 + 1] >= thresh_col:

        x2 += 1

    if x2 - x1 < 3:

        return None

    row_count = dark[:, x1:x2 + 1].sum(axis=1)

    min_row_dark = max(1, int((x2 - x1 + 1) * 0.2))

    rows = np.where(row_count >= min_row_dark)[0]

    if rows.size == 0:

        return None

    y1 = int(rows.min())

    y2 = int(rows.max()) + 1

    return _clamp_box(x1, y1, x2 + 1, y2, w, h)

 

def get_car_box_with_fallback(width: int, height: int, img: Optional[np.ndarray] = None) -> Optional[Tuple[int, int, int, int]]:

    car_box = get_car_box_for_shape(width, height)

    if car_box is not None:

        return car_box

    if MIRROR_FOLD_AUTO_INFER_CAR_BOX and img is not None:

        car_box = infer_car_box_from_image(img)

        if car_box is not None:

            return car_box

    x1r, y1r, x2r, y2r = MIRROR_FOLD_DEFAULT_CAR_BOX_RATIOS

    return _clamp_box(int(x1r * width), int(y1r * height), int(x2r * width), int(y2r * height), width, height)