garment_3d.py

"""
3D Garment Visualizer — Built FROM 2D Pattern Pieces

Takes the actual 2D sewing pattern pieces from pattern_generator.py,
triangulates them, and wraps each piece onto the correct body region
using cylindrical/cone projections.

Each 3D surface directly corresponds to a 2D pattern piece.
"""
import numpy as np
import plotly.graph_objects as go
from typing import Dict, List, Tuple, Optional

# ── Body constants (cm, Z=0 at floor, Z=170 at head top) ───────────────────
BODY_HEIGHT = 170.0
Z_ANKLE = 7
Z_KNEE = 45
Z_CROTCH = 82
Z_HIP = 88
Z_WAIST = 104
Z_BUST = 120
Z_SHOULDER = 145
Z_NECK = 150
Z_HEAD = 170

_BODY_Z  = [0, 7, 30, 45, 65, 82, 88, 100, 104, 112, 120, 132, 140, 145, 150, 158, 162, 170]
_BODY_RX = [4, 3, 4.5, 4, 6.5, 8, 9.5, 8, 7, 8.5, 10, 9, 8.5, 11.5, 3.5, 3, 5.5, 4.5]
_BODY_RY = [3, 2.5, 4, 4, 6, 7.5, 8.5, 7.5, 6.5, 8, 9, 8, 7.5, 8, 3.5, 3, 5, 4.5]

R_NECK = 5.5
R_WRIST = 3.5

PIECE_COLORS = {
    'Front Bodice': '#5B9BD5', 'Back Bodice': '#4472C4',
    'Sleeve': '#6BB3E0', 'Collar': '#FFFFFF', 'Cuff': '#E8E8E8',
    'Front Skirt': '#C48BB8', 'Back Skirt': '#A86CA8',
    'Front Pant': '#2C5F8A', 'Back Pant': '#1A4570',
    'Waistband': '#333333', 'Hood': '#3A7BD5', 'Pocket': '#7A7A7A',
}


def _point_in_polygon(x, y, polygon):
    n = len(polygon)
    inside = False
    j = n - 1
    for i in range(n):
        xi, yi = polygon[i]
        xj, yj = polygon[j]
        if ((yi > y) != (yj > y)) and (x < (xj - xi) * (y - yi) / (yj - yi + 1e-12) + xi):
            inside = not inside
        j = i
    return inside


def _triangulate_piece(pts_2d, nu=16, nv=16):
    pts = np.array(pts_2d, dtype=float)
    if len(pts) < 3:
        return pts, np.zeros((0, 3), dtype=int)
    if np.allclose(pts[0], pts[-1]):
        pts = pts[:-1]

    x_min, x_max = pts[:, 0].min(), pts[:, 0].max()
    y_min, y_max = pts[:, 1].min(), pts[:, 1].max()
    dx = max(x_max - x_min, 0.1)
    dy = max(y_max - y_min, 0.1)

    us = np.linspace(x_min - dx * 0.01, x_max + dx * 0.01, nu)
    vs = np.linspace(y_min - dy * 0.01, y_max + dy * 0.01, nv)

    grid_idx = -np.ones((nv, nu), dtype=int)
    valid_pts = []
    for i in range(nv):
        for j in range(nu):
            if _point_in_polygon(us[j], vs[i], pts):
                grid_idx[i, j] = len(valid_pts)
                valid_pts.append([us[j], vs[i]])

    if len(valid_pts) < 3:
        return _triangulate_fan(pts)

    tris = []
    for i in range(nv - 1):
        for j in range(nu - 1):
            a, b = grid_idx[i, j], grid_idx[i, j + 1]
            c, d = grid_idx[i + 1, j], grid_idx[i + 1, j + 1]
            if a >= 0 and b >= 0 and c >= 0:
                tris.append([a, b, c])
            if b >= 0 and c >= 0 and d >= 0:
                tris.append([b, d, c])

    if len(tris) == 0:
        return _triangulate_fan(pts)
    return np.array(valid_pts), np.array(tris)


def _triangulate_fan(pts):
    n = len(pts)
    if n < 3:
        return pts, np.zeros((0, 3), dtype=int)
    centroid = pts.mean(axis=0, keepdims=True)
    all_pts = np.vstack([pts, centroid])
    tris = np.array([[i, (i + 1) % n, n] for i in range(n)])
    return all_pts, tris


def _wrap_cylinder(pts_2d, radius, z_bottom, z_top, angle_start, angle_range,
                   x_offset=0.0, y_offset=0.0):
    pts = np.array(pts_2d, dtype=float)
    x_min, x_max = pts[:, 0].min(), pts[:, 0].max()
    y_min, y_max = pts[:, 1].min(), pts[:, 1].max()
    x_norm = (pts[:, 0] - x_min) / max(x_max - x_min, 0.01)
    y_norm = (pts[:, 1] - y_min) / max(y_max - y_min, 0.01)
    theta = angle_start + x_norm * angle_range
    z = z_bottom + y_norm * (z_top - z_bottom)
    x_3d = radius * np.cos(theta) + x_offset
    y_3d = radius * np.sin(theta) + y_offset
    return np.stack([x_3d, y_3d, z], axis=1)


def _wrap_cone(pts_2d, r_top, r_bottom, z_top, z_bottom, angle_start, angle_range):
    pts = np.array(pts_2d, dtype=float)
    x_min, x_max = pts[:, 0].min(), pts[:, 0].max()
    y_min, y_max = pts[:, 1].min(), pts[:, 1].max()
    x_norm = (pts[:, 0] - x_min) / max(x_max - x_min, 0.01)
    y_norm = (pts[:, 1] - y_min) / max(y_max - y_min, 0.01)
    theta = angle_start + x_norm * angle_range
    z = z_top - y_norm * (z_top - z_bottom)
    radius = r_top + y_norm * (r_bottom - r_top)
    x_3d = radius * np.cos(theta)
    y_3d = radius * np.sin(theta)
    return np.stack([x_3d, y_3d, z], axis=1)


def _wrap_sleeve(pts_2d, side='right', shoulder_x=14.0, shoulder_z=142.0,
                 arm_length=35.0, arm_angle_deg=25.0, r_top=5.5, r_bottom=3.5):
    pts = np.array(pts_2d, dtype=float)
    x_min, x_max = pts[:, 0].min(), pts[:, 0].max()
    y_min, y_max = pts[:, 1].min(), pts[:, 1].max()
    x_norm = (pts[:, 0] - x_min) / max(x_max - x_min, 0.01)
    y_norm = (pts[:, 1] - y_min) / max(y_max - y_min, 0.01)
    theta = x_norm * 2 * np.pi
    t_arm = 1.0 - y_norm
    radius = r_top + t_arm * (r_bottom - r_top)
    lx = radius * np.cos(theta)
    ly = -t_arm * arm_length
    lz = radius * np.sin(theta)
    ang = np.radians(arm_angle_deg)
    rx = lx * np.cos(ang) - ly * np.sin(ang)
    rz_local = lx * np.sin(ang) + ly * np.cos(ang)
    sign = 1.0 if side == 'right' else -1.0
    return np.stack([sign * (rx + shoulder_x), lz, rz_local + shoulder_z], axis=1)


def _wrap_pant_leg(pts_2d, side='right', hip_x=5.0, z_top=88.0, z_bottom=7.0,
                   r_top=8.5, r_bottom=4.5, is_front=True):
    pts = np.array(pts_2d, dtype=float)
    x_min, x_max = pts[:, 0].min(), pts[:, 0].max()
    y_min, y_max = pts[:, 1].min(), pts[:, 1].max()
    x_norm = (pts[:, 0] - x_min) / max(x_max - x_min, 0.01)
    y_norm = (pts[:, 1] - y_min) / max(y_max - y_min, 0.01)
    if is_front:
        theta = -np.pi / 2 + x_norm * np.pi
    else:
        theta = np.pi / 2 + x_norm * np.pi
    z = z_bottom + y_norm * (z_top - z_bottom)
    radius = r_bottom + y_norm * (r_top - r_bottom)
    sign = 1.0 if side == 'right' else -1.0
    x_3d = sign * hip_x + radius * np.cos(theta)
    y_3d = radius * np.sin(theta)
    return np.stack([x_3d, y_3d, z], axis=1)


def _wrap_hood(pts_2d, z_base=145.0, z_top=172.0, r_base=12.0, r_top=7.0):
    pts = np.array(pts_2d, dtype=float)
    x_min, x_max = pts[:, 0].min(), pts[:, 0].max()
    y_min, y_max = pts[:, 1].min(), pts[:, 1].max()
    x_norm = (pts[:, 0] - x_min) / max(x_max - x_min, 0.01)
    y_norm = (pts[:, 1] - y_min) / max(y_max - y_min, 0.01)
    theta = x_norm * np.pi
    z = z_base + y_norm * (z_top - z_base)
    radius = r_base + y_norm * (r_top - r_base)
    radius *= (1.0 - 0.3 * y_norm ** 2)
    x_3d = radius * np.cos(theta)
    y_3d = radius * np.sin(theta) - 3
    return np.stack([x_3d, y_3d, z], axis=1)


def _match_piece_type(name):
    name_lower = name.lower()
    if 'front bodice' in name_lower: return 'front_bodice'
    if 'back bodice' in name_lower: return 'back_bodice'
    if 'sleeve' in name_lower: return 'sleeve'
    if 'collar' in name_lower: return 'collar'
    if 'cuff' in name_lower: return 'cuff'
    if 'front skirt' in name_lower: return 'front_skirt'
    if 'back skirt' in name_lower: return 'back_skirt'
    if 'front pant' in name_lower: return 'front_pant'
    if 'back pant' in name_lower: return 'back_pant'
    if 'waistband' in name_lower: return 'waistband'
    if 'hood' in name_lower: return 'hood'
    if 'pocket' in name_lower: return 'pocket'
    return 'unknown'


def _get_piece_color(name):
    for key, color in PIECE_COLORS.items():
        if key.lower() in name.lower():
            return color
    return '#888888'


def _make_body_trace():
    theta = np.linspace(0, 2 * np.pi, 48)
    z_arr = np.array(_BODY_Z, float)
    rx = np.array(_BODY_RX, float)
    ry = np.array(_BODY_RY, float)
    Z = np.outer(z_arr, np.ones(48))
    X = np.outer(rx, np.cos(theta))
    Y = np.outer(ry, np.sin(theta))
    return go.Surface(
        x=X, y=Y, z=Z,
        surfacecolor=np.ones_like(X) * 0.8,
        colorscale=[[0, "#E8D0B0"], [1, "#E8D0B0"]],
        opacity=0.20, showscale=False, name="Body",
        hoverinfo="skip",
        lighting=dict(ambient=0.9, diffuse=0.1, specular=0, roughness=1.0),
    )


def create_3d_figure(analysis: Dict, pattern_pieces: Optional[List[Dict]] = None) -> go.Figure:
    garment_type = analysis.get("garment_type", "shirt").lower()
    measurements = analysis.get("measurements", {})
    features = analysis.get("features", {})

    fig = go.Figure()
    fig.add_trace(_make_body_trace())

    if pattern_pieces is None or len(pattern_pieces) == 0:
        _setup_layout(fig, garment_type)
        return fig

    for piece in pattern_pieces:
        name = piece['name']
        pts_2d = piece['points']
        piece_type = _match_piece_type(name)
        color = _get_piece_color(name)

        mesh_pts, tris = _triangulate_piece(pts_2d)
        if len(tris) == 0:
            continue

        traces = _wrap_piece(piece_type, mesh_pts, tris, name, color,
                             measurements, features, garment_type)
        for trace in traces:
            fig.add_trace(trace)

    _setup_layout(fig, garment_type)
    return fig


def _wrap_piece(piece_type, mesh_pts, tris, name, color,
                measurements, features, garment_type):
    traces = []
    bust = measurements.get('bust', 92)
    waist = measurements.get('waist', 74)
    hip = measurements.get('hip', 96)
    sleeve_len = measurements.get('sleeve_length', 60)
    skirt_len = measurements.get('skirt_length', 55)
    pant_len = measurements.get('pant_length', 100)

    r_torso = bust / (2 * np.pi) + 1.5
    r_waist = waist / (2 * np.pi) + 1.5
    r_hip = hip / (2 * np.pi) + 1.5

    if piece_type == 'front_bodice':
        pts_3d = _wrap_cylinder(mesh_pts, r_torso, Z_WAIST, Z_SHOULDER,
                                -np.pi / 2, np.pi)
        traces.append(_make_mesh3d(pts_3d, tris, color, name))

    elif piece_type == 'back_bodice':
        pts_3d = _wrap_cylinder(mesh_pts, r_torso, Z_WAIST, Z_SHOULDER,
                                np.pi / 2, np.pi)
        traces.append(_make_mesh3d(pts_3d, tris, color, name))

    elif piece_type == 'sleeve':
        shoulder_x = r_torso + 1.0
        for side in ['right', 'left']:
            pts_3d = _wrap_sleeve(mesh_pts, side=side, shoulder_x=shoulder_x,
                                  shoulder_z=Z_SHOULDER - 2,
                                  arm_length=sleeve_len * 0.55,
                                  arm_angle_deg=25, r_top=5.5, r_bottom=3.5)
            traces.append(_make_mesh3d(pts_3d, tris, color,
                                       f"{name} ({'R' if side == 'right' else 'L'})"))

    elif piece_type == 'collar':
        pts_3d = _wrap_cylinder(mesh_pts, R_NECK + 1.0, Z_NECK - 2, Z_NECK + 3,
                                -np.pi, 2 * np.pi)
        traces.append(_make_mesh3d(pts_3d, tris, color, name))

    elif piece_type == 'cuff':
        for side in ['right', 'left']:
            sign = 1.0 if side == 'right' else -1.0
            wrist_x = sign * (r_torso + sleeve_len * 0.3 + 2)
            wrist_z = Z_SHOULDER - sleeve_len * 0.5
            pts_3d = _wrap_cylinder(mesh_pts, R_WRIST + 0.5,
                                    wrist_z - 3, wrist_z,
                                    0, 2 * np.pi, x_offset=wrist_x)
            traces.append(_make_mesh3d(pts_3d, tris, color,
                                       f"{name} ({'R' if side == 'right' else 'L'})"))

    elif piece_type == 'front_skirt':
        hem_z = max(Z_WAIST - skirt_len, Z_ANKLE)
        flare = measurements.get('flare', 5)
        r_hem = r_hip + flare * 0.5
        pts_3d = _wrap_cone(mesh_pts, r_waist, r_hem, Z_WAIST, hem_z,
                            -np.pi / 2, np.pi)
        traces.append(_make_mesh3d(pts_3d, tris, color, name))

    elif piece_type == 'back_skirt':
        hem_z = max(Z_WAIST - skirt_len, Z_ANKLE)
        flare = measurements.get('flare', 5)
        r_hem = r_hip + flare * 0.5
        pts_3d = _wrap_cone(mesh_pts, r_waist, r_hem, Z_WAIST, hem_z,
                            np.pi / 2, np.pi)
        traces.append(_make_mesh3d(pts_3d, tris, color, name))

    elif piece_type == 'front_pant':
        hem_z = max(Z_HIP - pant_len * 0.85, Z_ANKLE)
        thigh_circ = measurements.get('thigh', 56)
        ankle_circ = measurements.get('ankle', 24)
        r_thigh = thigh_circ / (2 * np.pi) + 1
        r_ankle_r = ankle_circ / (2 * np.pi) + 1
        for side in ['right', 'left']:
            pts_3d = _wrap_pant_leg(mesh_pts, side=side, hip_x=r_hip * 0.4,
                                    z_top=Z_HIP, z_bottom=hem_z,
                                    r_top=r_thigh, r_bottom=r_ankle_r, is_front=True)
            traces.append(_make_mesh3d(pts_3d, tris, color,
                                       f"{name} ({'R' if side == 'right' else 'L'})"))

    elif piece_type == 'back_pant':
        hem_z = max(Z_HIP - pant_len * 0.85, Z_ANKLE)
        thigh_circ = measurements.get('thigh', 56)
        ankle_circ = measurements.get('ankle', 24)
        r_thigh = thigh_circ / (2 * np.pi) + 1
        r_ankle_r = ankle_circ / (2 * np.pi) + 1
        for side in ['right', 'left']:
            pts_3d = _wrap_pant_leg(mesh_pts, side=side, hip_x=r_hip * 0.4,
                                    z_top=Z_HIP, z_bottom=hem_z,
                                    r_top=r_thigh, r_bottom=r_ankle_r, is_front=False)
            traces.append(_make_mesh3d(pts_3d, tris, color,
                                       f"{name} ({'R' if side == 'right' else 'L'})"))

    elif piece_type == 'waistband':
        wb_h = measurements.get('waistband_height', 4)
        pts_3d = _wrap_cylinder(mesh_pts, r_waist + 0.5, Z_WAIST, Z_WAIST + wb_h,
                                -np.pi, 2 * np.pi)
        traces.append(_make_mesh3d(pts_3d, tris, color, name))

    elif piece_type == 'hood':
        pts_3d_r = _wrap_hood(mesh_pts, z_base=Z_SHOULDER, z_top=Z_HEAD + 2,
                              r_base=r_torso, r_top=7)
        traces.append(_make_mesh3d(pts_3d_r, tris, color, f"{name} (R)"))
        pts_3d_l = pts_3d_r.copy()
        pts_3d_l[:, 0] = -pts_3d_l[:, 0]
        traces.append(_make_mesh3d(pts_3d_l, tris, color, f"{name} (L)"))

    elif piece_type == 'pocket':
        pocket_z = (Z_WAIST + Z_BUST) / 2
        pts_3d = _wrap_cylinder(mesh_pts, r_torso + 0.3,
                                pocket_z - 8, pocket_z, -0.3, 0.6)
        traces.append(_make_mesh3d(pts_3d, tris, color, name))

    return traces


def _make_mesh3d(pts_3d, tris, color, name):
    return go.Mesh3d(
        x=pts_3d[:, 0], y=pts_3d[:, 1], z=pts_3d[:, 2],
        i=tris[:, 0], j=tris[:, 1], k=tris[:, 2],
        color=color, opacity=0.85, name=name,
        flatshading=False,
        lighting=dict(ambient=0.55, diffuse=0.8, specular=0.2, roughness=0.6),
        lightposition=dict(x=100, y=200, z=300),
        hoverinfo="name",
    )


def _setup_layout(fig, garment_type):
    fig.update_layout(
        scene=dict(
            xaxis=dict(showgrid=False, showticklabels=False, title="",
                       zeroline=False, range=[-40, 40]),
            yaxis=dict(showgrid=False, showticklabels=False, title="",
                       zeroline=False, range=[-40, 40]),
            zaxis=dict(showgrid=False, showticklabels=False, title="",
                       zeroline=False, range=[0, 180]),
            aspectmode="data",
            camera=dict(eye=dict(x=1.8, y=1.2, z=0.5),
                        center=dict(x=0, y=0, z=-0.1)),
            bgcolor="rgba(245,245,248,1)",
        ),
        margin=dict(l=0, r=0, t=35, b=0),
        height=550,
        title=dict(text=f"3D Preview: {garment_type.title()} (from pattern pieces)",
                   font=dict(size=14)),
        paper_bgcolor="#fafafa",
        showlegend=True,
        legend=dict(font=dict(size=10), bgcolor="rgba(255,255,255,0.8)"),
    )