#!/usr/bin/env python3
"""
Deform original Junction Box STL: only stretch the span between corner posts.
Corner posts, wall thickness, screw holes, radii stay ORIGINAL size.
"""
import struct, numpy as np, os

ORIG = "/app/projects/esp32-box/original"
OUT = "/app/projects/esp32-box/exports"
os.makedirs(OUT, exist_ok=True)

def load_stl(path):
    with open(path, "rb") as f:
        header = f.read(80)
        n_tri = struct.unpack("<I", f.read(4))[0]
        verts = np.zeros((n_tri * 3, 3), dtype=np.float32)
        for i in range(n_tri):
            normal = f.read(12)
            for j in range(3):
                data = f.read(12)
                verts[i*3+j] = np.frombuffer(data, dtype=np.float32)
            attr = f.read(2)
    return verts, n_tri

def write_stl(verts, n_tri, path):
    with open(path, "wb") as f:
        f.write(b"\x00" * 80)
        f.write(struct.pack("<I", n_tri))
        for i in range(n_tri):
            f.write(struct.pack("<fff", 0.0, 0.0, 0.0))
            for j in range(3):
                v = verts[i*3+j]
                f.write(struct.pack("<fff", v[0], v[1], v[2]))
            f.write(b"\x00\x00")

def piecewise_scale(vals, orig_min, orig_max, new_min, new_max, 
                    corner_start, corner_end_frac):
    """
    Piecewise linear: corner regions keep original, middle stretches.
    corner_start: distance from each end that stays unchanged
    Returns new values.
    """
    orig_len = orig_max - orig_min
    new_len = new_max - new_min
    
    cs = corner_start
    ce = orig_len - corner_start  # corner end position
    
    # Original middle span
    orig_mid = ce - cs
    new_mid = new_len - 2 * corner_start
    
    scale = new_mid / orig_mid if orig_mid > 0 else 1.0
    
    result = np.copy(vals)
    
    # Left corner (below cs): no change
    # Middle (cs to ce): scale
    # Right corner (above ce): offset
    
    in_left = vals <= orig_min + cs
    in_mid = (vals > orig_min + cs) & (vals < orig_max - cs)
    in_right = vals >= orig_max - cs
    
    result[in_left] = vals[in_left]
    result[in_mid] = new_min + cs + (vals[in_mid] - orig_min - cs) * scale
    result[in_right] = new_min + new_len - (orig_max - vals[in_right])
    
    return result

# ==========================================
# BASE: 60x60x17 -> 80x70x32
# ==========================================
print("=== BASE ===")
verts, n = load_stl(os.path.join(ORIG, "Box.stl"))
v = verts.copy()

# X: corner 12mm each side
v[:, 0] = piecewise_scale(v[:, 0], 0, 60, 0, 80, corner_start=12.0, corner_end_frac=0)
# Y: corner 12mm each side
v[:, 1] = piecewise_scale(v[:, 1], 0, 60, 0, 70, corner_start=12.0, corner_end_frac=0)
# Z: floor 3mm + top 3mm unchanged, middle stretches
v[:, 2] = piecewise_scale(v[:, 2], 0, 17, 0, 32, corner_start=3.0, corner_end_frac=0)

write_stl(v, n, os.path.join(OUT, "base.stl"))
# Verify
print("  X: %.1f-%.1f -> %.1f-%.1f" % (0, 60, v[:,0].min(), v[:,0].max()))
print("  Y: %.1f-%.1f -> %.1f-%.1f" % (0, 60, v[:,1].min(), v[:,1].max()))
print("  Z: %.1f-%.1f -> %.1f-%.1f" % (0, 17, v[:,2].min(), v[:,2].max()))
print("  -> base.stl")

# ==========================================
# LID: 60x60x10 -> 80x70x18
# ==========================================
print("\n=== LID ===")
verts, n = load_stl(os.path.join(ORIG, "Lid.stl"))
v = verts.copy()

v[:, 0] = piecewise_scale(v[:, 0], 0, 60, 0, 80, corner_start=12.0, corner_end_frac=0)
v[:, 1] = piecewise_scale(v[:, 1], 0, 60, 0, 70, corner_start=12.0, corner_end_frac=0)
v[:, 2] = piecewise_scale(v[:, 2], 0, 10, 0, 18, corner_start=3.0, corner_end_frac=0)

write_stl(v, n, os.path.join(OUT, "lid.stl"))
print("  Z: 0-10 -> %.1f-%.1f" % (v[:,2].min(), v[:,2].max()))
print("  -> lid.stl")

# ==========================================
# WINDOWS: only X/Y stretch, Z stays original
# ==========================================
print("\n=== WINDOWS ===")
win_specs = [
    ("Window 20x60mm - liso.stl", "window_smooth.stl"),
    ("Window 20x60mm.stl", "window_keystone_dc.stl"),
    ("Window - Flex Gland.stl", "window_usb.stl"),
]
for src, dst in win_specs:
    verts, n = load_stl(os.path.join(ORIG, src))
    v = verts.copy()
    # Windows fit between corner posts: only stretch X/Y
    v[:, 0] = piecewise_scale(v[:, 0], 0, 60, 0, 80, corner_start=8.0, corner_end_frac=0)
    v[:, 1] = piecewise_scale(v[:, 1], 0, 60, 0, 70, corner_start=8.0, corner_end_frac=0)
    # Z stays same (no change needed)
    write_stl(v, n, os.path.join(OUT, dst))
    print("  X: 0-60 -> %.1f-%.1f" % (v[:,0].min(), v[:,0].max()))
    print("  -> %s" % dst)

# ==========================================
# SCREW: only stretch to fit wider posts
# ==========================================
print("\n=== SCREW ===")
verts, n = load_stl(os.path.join(ORIG, "Screw.stl"))
v = verts.copy()
# Screw scales slightly in X/Y to match wider posts
v[:, 0] = piecewise_scale(v[:, 0], 0, 60, 0, 80, corner_start=6.0, corner_end_frac=0)
v[:, 1] = piecewise_scale(v[:, 1], 0, 60, 0, 70, corner_start=6.0, corner_end_frac=0)
write_stl(v, n, os.path.join(OUT, "screw.stl"))
print("  -> screw.stl")

print("\nDONE: all parts deformed to 80x70mm footprint.")