s-canvas/dxf_geometry.py

"""DXF 지오메트리 추출 공통 유틸리티.

모든 구조물 템플릿이 공유하는 DXF 처리 로직:
1. 단위 자동 감지 (mm vs m)
2. 주석/치수/해치 레이어 자동 필터링
3. LINE/LWPOLYLINE/POLYLINE/ARC/SPLINE/CIRCLE 통합 추출
4. 뷰 영역 자동 분할 (평면/정면/측면)

사용법:
    from dxf_geometry import extract_structural_geometry

    result = extract_structural_geometry(dxf_path)
    for layer_name, shapes in result.by_layer.items():
        for shape in shapes:
            # shape.points: [(x, y), ...] (단위 정규화 완료)
            # shape.closed: bool
            # shape.kind: "polyline" | "line" | "arc" | ...
            ...
"""

from __future__ import annotations

import math
import re
from dataclasses import dataclass, field
from pathlib import Path

import ezdxf
import numpy as np


# ---------------------------------------------------------------------------
# 데이터 구조
# ---------------------------------------------------------------------------

@dataclass
class Shape:
    """추출된 단일 지오메트리 요소 (단위: m)."""
    kind: str                # "polyline" | "line" | "arc" | "circle"
    layer: str
    points: list             # [(x, y), ...] — m 단위
    closed: bool = False
    extra: dict = field(default_factory=dict)  # 부가정보 (center, radius, angles 등)

    @property
    def bbox(self) -> tuple[float, float, float, float]:
        """(xmin, ymin, xmax, ymax)"""
        if not self.points:
            return (0, 0, 0, 0)
        arr = np.array(self.points)
        return (float(arr[:, 0].min()), float(arr[:, 1].min()),
                float(arr[:, 0].max()), float(arr[:, 1].max()))

    @property
    def centroid(self) -> tuple[float, float]:
        if not self.points:
            return (0, 0)
        arr = np.array(self.points)
        return (float(arr[:, 0].mean()), float(arr[:, 1].mean()))

    @property
    def area(self) -> float:
        """closed polygon 면적 (shoelace). open이면 0."""
        if not self.closed or len(self.points) < 3:
            return 0.0
        n = len(self.points)
        s = 0.0
        for i in range(n):
            x1, y1 = self.points[i]
            x2, y2 = self.points[(i + 1) % n]
            s += x1 * y2 - x2 * y1
        return abs(s) / 2

    @property
    def length(self) -> float:
        """폴리라인 누적 길이."""
        if len(self.points) < 2:
            return 0.0
        arr = np.array(self.points)
        diffs = np.diff(arr, axis=0)
        return float(np.sum(np.linalg.norm(diffs, axis=1)))


@dataclass
class GeometryResult:
    """DXF에서 추출된 전체 지오메트리."""
    dxf_path: str
    unit_scale: float = 1.0          # 원본 → m 변환 계수 (mm이면 0.001)
    detected_unit: str = "m"         # "mm" | "m"
    shapes: list[Shape] = field(default_factory=list)

    # 분류별 접근 편의 속성
    by_layer: dict[str, list[Shape]] = field(default_factory=dict)
    closed_shapes: list[Shape] = field(default_factory=list)
    open_shapes: list[Shape] = field(default_factory=list)

    # 메타
    total_bounds: tuple[float, float, float, float] = (0, 0, 0, 0)
    raw_text_count: int = 0
    dimension_count: int = 0
    excluded_layers: list[str] = field(default_factory=list)

    def largest_closed(self) -> Shape | None:
        """면적이 가장 큰 closed shape 반환."""
        if not self.closed_shapes:
            return None
        return max(self.closed_shapes, key=lambda s: s.area)

    def longest_polyline(self) -> Shape | None:
        """가장 긴 폴리라인 반환."""
        polys = [s for s in self.shapes if s.kind == "polyline" and len(s.points) > 2]
        if not polys:
            return None
        return max(polys, key=lambda s: s.length)


# ---------------------------------------------------------------------------
# 레이어 필터링
# ---------------------------------------------------------------------------

# 구조물 지오메트리로 간주하지 않는 레이어 이름 패턴
# (대소문자 무시, 부분 매칭)
# 주의: MZ-HIDL(숨김선)은 도면에 따라 주 구조물이 있을 수 있어 제외 대상에서 뺌
_EXCLUDE_LAYER_PATTERNS = [
    # 치수선
    r"^DIM$", r"-DIM$", r"^CS-DIM", r"^MZ-DIM", r"치수선",
    # 해치/패턴
    r"^HATCH$", r"-HATCH$", r"^CS-PATT", r"^CZ-PATT", r"^MZ-HATCH",
    r"^CS-HATCH", r"^PATT-", r"해치",
    # 텍스트 (레이어명이 명확히 TEXT 전용인 경우만)
    r"^CS-TEXT$", r"^CZ-TEXT$", r"^CX-BORD", r"^TEXT$", r"^문자$",
    r"^CZ-TEX[0-9]", r"^텍스트$",
    # 지시선/리더
    r"^CS-LEAT$", r"^CS-LEAL$", r"^CS-LEA$", r"^지시선$", r"^LEADER$",
    # 표/프레임
    r"^CS-TABL",
    # 중심선 (완전히 중심선 전용 레이어만)
    r"^중심선\(", r"중심선$",
    # 기타 주석
    r"^Defpoints$", r"^심볼$", r"^SYMB$", r"^AA-",
]


def is_excluded_layer(layer_name: str) -> bool:
    """레이어명이 주석/치수/해치 등 비구조 레이어인지 확인."""
    if not layer_name:
        return False
    return any(re.search(pat, layer_name, re.IGNORECASE) for pat in _EXCLUDE_LAYER_PATTERNS)


# ---------------------------------------------------------------------------
# 단위 자동 감지
# ---------------------------------------------------------------------------

def detect_unit_scale(doc) -> tuple[float, str]:
    """DXF에서 단위를 자동 감지하여 m로 변환할 scale 반환.

    한국 토목도면은 $INSUNITS가 잘못 설정되는 경우가 많으므로
    DIMENSION 값을 최우선 판단 기준으로 사용.

    판단 순서:
    1. DIMENSION 값 분포 (최우선) — 치수 값이 설계 의도를 반영
    2. $INSUNITS header variable (보조 확인)
    3. 전체 bbox 크기 (폴백)

    Returns:
        (scale, unit_name): 원본 × scale = m. unit_name은 "mm"/"m"/"cm"
    """
    msp = doc.modelspace()

    # 1) DIMENSION 값 분포 분석 (가장 신뢰할 만한 신호)
    dim_values = []
    for e in msp:
        if e.dxftype() == "DIMENSION":
            try:
                m = e.dxf.get("actual_measurement", None)
                if m is not None and m > 0:
                    dim_values.append(float(m))
            except Exception:
                pass

    if len(dim_values) >= 3:
        median_val = float(np.median(dim_values))
        # 토목 구조물 치수 관례:
        # - m 단위: 중앙값이 보통 0.5~100 범위
        # - mm 단위: 중앙값이 보통 100~100,000 범위
        # - cm 단위: 중앙값이 50~10,000 범위 (드뭄)
        if median_val >= 100:
            return 0.001, "mm"
        elif median_val >= 1 and median_val < 100:
            return 1.0, "m"

    # 2) $INSUNITS header (DIMENSION이 없거나 모호할 때)
    # 0=unitless, 1=inch, 2=feet, 4=mm, 5=cm, 6=m
    # inch/feet는 한국 토목도면에서 거의 없음 → 무시하고 bbox로 판단
    _INSUNITS_TO_SCALE = {4: (0.001, "mm"), 5: (0.01, "cm"), 6: (1.0, "m")}
    try:
        insunits = int(doc.header.get("$INSUNITS", 0))
        if insunits in _INSUNITS_TO_SCALE:
            return _INSUNITS_TO_SCALE[insunits]
    except Exception:
        pass

    # 3) Bbox 크기로 추정
    try:
        # 엔티티 bbox 직접 계산 (header extents는 부정확할 수 있음)
        all_x = []
        all_y = []
        for e in msp:
            if e.dxftype() in ("LWPOLYLINE", "LINE", "CIRCLE", "ARC"):
                try:
                    if e.dxftype() == "LWPOLYLINE":
                        for p in e.get_points():
                            all_x.append(p[0])
                            all_y.append(p[1])
                    elif e.dxftype() == "LINE":
                        all_x.extend([e.dxf.start.x, e.dxf.end.x])
                        all_y.extend([e.dxf.start.y, e.dxf.end.y])
                    elif e.dxftype() in ("CIRCLE", "ARC"):
                        all_x.append(e.dxf.center.x)
                        all_y.append(e.dxf.center.y)
                except Exception:
                    pass
                # 성능: 처음 1000개만 샘플
                if len(all_x) > 1000:
                    break

        if all_x and all_y:
            diag = math.sqrt((max(all_x) - min(all_x)) ** 2 + (max(all_y) - min(all_y)) ** 2)
            # 토목구조물 전체 크기: 수 m ~ 수백 m
            # mm이면 대각선 수천 ~ 수십만
            if diag > 2000:
                return 0.001, "mm"
            elif diag < 500:
                return 1.0, "m"
    except Exception:
        pass

    # 기본값: 토목도면은 대부분 mm
    return 0.001, "mm"


# ---------------------------------------------------------------------------
# 메인 추출 함수
# ---------------------------------------------------------------------------

def extract_structural_geometry(
    dxf_path: str | Path,
    exclude_layers: bool = True,
    include_open: bool = True,
    min_points: int = 2,
    unit_override: str | None = None,
    explode_blocks: bool = False,
    max_block_depth: int = 4,
) -> GeometryResult:
    """DXF에서 구조물 지오메트리를 추출.

    Args:
        dxf_path: DXF 파일 경로
        exclude_layers: True면 주석/치수/해치 레이어 제외
        include_open: False면 closed 지오메트리만
        min_points: 최소 점 개수
        unit_override: "mm" | "m" 중 하나면 자동 감지 무시
        explode_blocks: True면 INSERT(블록 참조)를 virtual_entities로 재귀 확장.
                        평면도에 구조물이 블록으로 배치된 경우 필수.
        max_block_depth: 중첩 블록 재귀 최대 깊이

    Returns:
        GeometryResult
    """
    doc = ezdxf.readfile(str(dxf_path))
    msp = doc.modelspace()

    # 단위 감지
    if unit_override == "mm":
        unit_scale, unit_name = 0.001, "mm"
    elif unit_override == "m":
        unit_scale, unit_name = 1.0, "m"
    else:
        unit_scale, unit_name = detect_unit_scale(doc)

    result = GeometryResult(
        dxf_path=str(dxf_path),
        unit_scale=unit_scale,
        detected_unit=unit_name,
    )

    excluded_set = set()

    def _process_entity(entity, inherited_layer: str | None = None, depth: int = 0):
        """단일 엔티티 처리. INSERT면 explode_blocks에 따라 재귀 확장."""
        etype = entity.dxftype()
        # 블록 내부 엔티티의 layer가 "0"이면 INSERT의 레이어를 상속
        raw_layer = getattr(entity.dxf, "layer", "")
        layer = inherited_layer if inherited_layer and raw_layer in ("", "0") else raw_layer

        # 레이어 필터
        if exclude_layers and is_excluded_layer(layer):
            excluded_set.add(layer)
            return

        # INSERT 재귀 확장
        if etype == "INSERT":
            if not explode_blocks or depth >= max_block_depth:
                return
            try:
                for sub in entity.virtual_entities():
                    _process_entity(sub, inherited_layer=layer, depth=depth + 1)
            except Exception:
                pass
            return

        # 엔티티별 추출
        try:
            shape = _extract_entity(entity, etype, unit_scale)
        except Exception:
            return

        if shape is None:
            return
        if len(shape.points) < min_points:
            return
        if not include_open and not shape.closed:
            return

        shape.layer = layer
        result.shapes.append(shape)

        # 레이어별 집계
        result.by_layer.setdefault(layer, []).append(shape)

        if shape.closed:
            result.closed_shapes.append(shape)
        else:
            result.open_shapes.append(shape)

    for entity in msp:
        _process_entity(entity)

    # 메타 카운트
    for e in msp:
        et = e.dxftype()
        if et in ("TEXT", "MTEXT"):
            result.raw_text_count += 1
        elif et == "DIMENSION":
            result.dimension_count += 1

    result.excluded_layers = sorted(excluded_set)

    # 전체 bbox
    if result.shapes:
        all_pts = []
        for s in result.shapes:
            all_pts.extend(s.points)
        arr = np.array(all_pts)
        result.total_bounds = (
            float(arr[:, 0].min()), float(arr[:, 1].min()),
            float(arr[:, 0].max()), float(arr[:, 1].max()),
        )

    return result


def _extract_entity(entity, etype: str, scale: float) -> Shape | None:
    """개별 DXF 엔티티 → Shape 변환."""
    if etype == "LWPOLYLINE":
        pts = [(p[0] * scale, p[1] * scale) for p in entity.get_points()]
        if len(pts) < 2:
            return None
        return Shape(kind="polyline", layer="", points=pts, closed=entity.closed)

    elif etype == "POLYLINE":
        pts = [(v.dxf.location.x * scale, v.dxf.location.y * scale) for v in entity.vertices]
        if len(pts) < 2:
            return None
        return Shape(kind="polyline", layer="", points=pts, closed=entity.is_closed)

    elif etype == "LINE":
        s = entity.dxf.start
        e = entity.dxf.end
        return Shape(
            kind="line", layer="",
            points=[(s.x * scale, s.y * scale), (e.x * scale, e.y * scale)],
            closed=False,
        )

    elif etype == "ARC":
        c = entity.dxf.center
        r = entity.dxf.radius * scale
        sa = math.radians(entity.dxf.start_angle)
        ea = math.radians(entity.dxf.end_angle)
        if ea < sa:
            ea += 2 * math.pi
        n = max(8, int((ea - sa) / math.radians(5)))
        pts = []
        for i in range(n + 1):
            t = sa + (ea - sa) * i / n
            pts.append((c.x * scale + r * math.cos(t), c.y * scale + r * math.sin(t)))
        return Shape(
            kind="arc", layer="", points=pts, closed=False,
            extra={"center": (c.x * scale, c.y * scale), "radius": r,
                   "start_angle": sa, "end_angle": ea},
        )

    elif etype == "CIRCLE":
        c = entity.dxf.center
        r = entity.dxf.radius * scale
        n = 32
        pts = []
        for i in range(n + 1):
            t = 2 * math.pi * i / n
            pts.append((c.x * scale + r * math.cos(t), c.y * scale + r * math.sin(t)))
        return Shape(
            kind="circle", layer="", points=pts, closed=True,
            extra={"center": (c.x * scale, c.y * scale), "radius": r},
        )

    elif etype == "SPLINE":
        try:
            # 제어점으로 근사
            pts = [(pt[0] * scale, pt[1] * scale) for pt in entity.control_points]
            if len(pts) < 2:
                return None
            return Shape(
                kind="polyline", layer="", points=pts, closed=entity.closed,
            )
        except Exception:
            return None

    elif etype == "ELLIPSE":
        # 타원 → 폴리라인 샘플링
        try:
            c = entity.dxf.center
            # flattening으로 점 목록 얻기
            pts = []
            for pt in entity.flattening(distance=0.1):
                pts.append((pt.x * scale, pt.y * scale))
            if len(pts) < 2:
                return None
            return Shape(kind="polyline", layer="", points=pts, closed=True)
        except Exception:
            return None

    return None


# ---------------------------------------------------------------------------
# 뷰 영역 분할 (평면/정면/측면)
# ---------------------------------------------------------------------------

def split_views_by_y(result: GeometryResult, n_views: int = 2) -> list[GeometryResult]:
    """Y 좌표 분포로 도면을 n개 뷰로 분할.

    토목도면은 보통 평면(위)/정면(아래) 또는 측면/단면을 수직 배치.
    """
    if not result.shapes or n_views < 2:
        return [result]

    # Y 좌표 집계
    y_vals = []
    for s in result.shapes:
        arr = np.array(s.points)
        y_vals.append(arr[:, 1].mean())
    y_vals = np.array(y_vals)

    # 분위수로 분할
    thresholds = np.quantile(y_vals, [i / n_views for i in range(1, n_views)])

    views = [[] for _ in range(n_views)]
    for i, s in enumerate(result.shapes):
        bucket = n_views - 1
        for ti, t in enumerate(thresholds):
            if y_vals[i] <= t:
                bucket = ti
                break
        views[bucket].append(s)

    # 각 뷰를 GeometryResult로 래핑
    out = []
    for shapes in views:
        v = GeometryResult(
            dxf_path=result.dxf_path,
            unit_scale=result.unit_scale,
            detected_unit=result.detected_unit,
        )
        v.shapes = shapes
        for s in shapes:
            v.by_layer.setdefault(s.layer, []).append(s)
            if s.closed:
                v.closed_shapes.append(s)
            else:
                v.open_shapes.append(s)
        if shapes:
            all_pts = []
            for s in shapes:
                all_pts.extend(s.points)
            arr = np.array(all_pts)
            v.total_bounds = (
                float(arr[:, 0].min()), float(arr[:, 1].min()),
                float(arr[:, 0].max()), float(arr[:, 1].max()),
            )
        out.append(v)

    return out


# ---------------------------------------------------------------------------
# 편의 함수
# ---------------------------------------------------------------------------

def extract_all(dxf_paths: list[str], **kwargs) -> GeometryResult:
    """여러 DXF 파일을 모두 파싱하여 단일 GeometryResult로 반환."""
    combined = GeometryResult(dxf_path=";".join(dxf_paths))

    first_scale = None
    for p in dxf_paths:
        try:
            r = extract_structural_geometry(p, **kwargs)
            if first_scale is None:
                first_scale = r.unit_scale
                combined.unit_scale = r.unit_scale
                combined.detected_unit = r.detected_unit
            combined.shapes.extend(r.shapes)
            for layer, shapes in r.by_layer.items():
                combined.by_layer.setdefault(layer, []).extend(shapes)
            combined.closed_shapes.extend(r.closed_shapes)
            combined.open_shapes.extend(r.open_shapes)
            combined.raw_text_count += r.raw_text_count
            combined.dimension_count += r.dimension_count
            combined.excluded_layers.extend(r.excluded_layers)
        except Exception as e:
            print(f"  추출 실패 ({p}): {e}")

    if combined.shapes:
        all_pts = []
        for s in combined.shapes:
            all_pts.extend(s.points)
        arr = np.array(all_pts)
        combined.total_bounds = (
            float(arr[:, 0].min()), float(arr[:, 1].min()),
            float(arr[:, 0].max()), float(arr[:, 1].max()),
        )

    return combined


if __name__ == "__main__":
    # 샘플 테스트
    import sys
    from pathlib import Path

    paths = sys.argv[1:]
    if not paths:
        base = Path("Gate_Sample")
        paths = [
            str(base / "12995740-M40-001 여수로 수문 설치도(1／2).dxf"),
            str(base / "12995740-M40-002 여수로 수문 설치도(2／2).dxf"),
        ]

    for p in paths:
        print(f"\n=== {Path(p).name} ===")
        r = extract_structural_geometry(p)
        print(f"  단위: {r.detected_unit} (scale={r.unit_scale})")
        print(f"  총 지오메트리: {len(r.shapes)}개")
        print(f"    closed: {len(r.closed_shapes)}, open: {len(r.open_shapes)}")
        print("  레이어별:")
        for layer, shapes in sorted(r.by_layer.items(), key=lambda x: -len(x[1]))[:10]:
            print(f"    {layer}: {len(shapes)}개")
        b = r.total_bounds
        print(f"  bbox: ({b[0]:.2f}, {b[1]:.2f}) ~ ({b[2]:.2f}, {b[3]:.2f}) m")
        print(f"  제외된 레이어: {r.excluded_layers[:5]}")
        if r.largest_closed():
            lc = r.largest_closed()
            print(f"  최대 closed: {lc.layer} ({lc.area:.2f} m², {len(lc.points)}pts)")