# Dual-Size Base-12 Vajra Generator (144 mm + 120 mm flask-friendly)
# Rodney's dual-size script (2026-10-20) — geometry code verbatim as supplied.
# Completed by Squish: OBJ / PLY / DXF exporters (adapted from the original
# generate_vajra_cad.py) so both sizes export in all four formats as intended.
import struct
import numpy as np

# =====================================================================
# DUAL-SIZE BASE-12 GEOMETRIC PARAMETERS (All dimensions in millimeters)
# =====================================================================
# Size Configs:
# 144 mm version: 1U = 12 mm (Total Length = 144 mm)
# 120 mm version: 1U = 10 mm (Total Length = 120 mm - fits standard foundry flasks)
CONFIGS = {
    "144mm": {"total_length": 144.0, "half_length": 72.0, "hub_r": 12.0, "scale": 1.0},
    "120mm": {"total_length": 120.0, "half_length": 60.0, "hub_r": 10.0, "scale": 10.0 / 12.0}
}

PRONG_COUNT = 6
LOTUS_PETALS = 12
TAPER_ANGLE_DEG = 12.0


def central_body_radius(z, theta, half_len, hub_r):
    az = np.abs(z)
    r = np.zeros_like(az)
    scale = hub_r / 12.0  # Proportional scaling factor

    # 1. Central Hub (Bindu)
    m1 = az < (10.0 * scale)
    if np.any(m1):
        base_sphere = hub_r * np.sqrt(np.clip(1.0 - (az[m1] / (11.2 * scale)) ** 2, 0.2, 1.0))
        r[m1] = base_sphere + (0.25 * scale) * np.cos(12.0 * theta[m1])

    # 2. Triple Beaded Collar Rings
    m2 = (az >= (10.0 * scale)) & (az < (16.0 * scale))
    if np.any(m2):
        zc = az[m2] - (10.0 * scale)
        r[m2] = (9.5 * scale) + (0.35 * zc) + (0.9 * scale) * np.abs(np.sin(np.pi * (zc / (2.0 * scale))))

    # 3. 12-Petal Lotus Base
    m3 = (az >= (16.0 * scale)) & (az < (28.0 * scale))
    if np.any(m3):
        zl = (az[m3] - (16.0 * scale)) / (12.0 * scale)
        envelope = (10.5 * scale) + (3.85 * scale) * np.sin(np.pi * zl)
        petals = (1.15 * scale) * np.sin(np.pi * zl) * (0.5 + 0.5 * np.cos(LOTUS_PETALS * theta[m3]))
        r[m3] = envelope + petals

    # 4. Collar Transition Disc
    m4 = (az >= (28.0 * scale)) & (az < (30.0 * scale))
    if np.any(m4):
        zd = (az[m4] - (28.0 * scale)) / (2.0 * scale)
        r[m4] = (11.5 * scale) * (1.0 - 0.55 * zd)

    # 5. Central Axial Shaft
    m5 = (az >= (30.0 * scale)) & (az < (64.0 * scale))
    if np.any(m5):
        zs = (az[m5] - (30.0 * scale)) / (34.0 * scale)
        twist = (2.0 * np.pi / 3.0) * zs * np.sign(z[m5])
        hex_mod = 1.0 + 0.08 * np.cos(6.0 * (theta[m5] - twist))
        r[m5] = ((4.5 * scale) - (0.8 * scale) * np.sin(np.pi * zs)) * hex_mod

    # 6. Diamond Tip Apex
    m6 = az >= (64.0 * scale)
    if np.any(m6):
        zt = (az[m6] - (64.0 * scale)) / (8.0 * scale)
        tip = np.where(zt < 0.25, (4.5 * scale) + (2.2 * scale) * (zt / 0.25), (6.7 * scale) * (1.0 - (zt - 0.25) / 0.75))
        r[m6] = np.clip(tip * (1.0 + 0.06 * np.cos(6.0 * theta[m6])), 0.0, None)

    return r


def build_vajra_mesh(size_key):
    cfg = CONFIGS[size_key]
    half_l = cfg["half_length"]
    hub_r = cfg["hub_r"]
    scale = cfg["scale"]

    # Central Body Mesh
    n_z, n_theta = 360, 144
    z_vals = np.linspace(-half_l, half_l, n_z)
    theta_vals = np.linspace(0.0, 2.0 * np.pi, n_theta, endpoint=False)
    Z, THETA = np.meshgrid(z_vals, theta_vals, indexing="ij")
    R = central_body_radius(Z, THETA, half_l, hub_r)

    X = R * np.cos(THETA)
    Y = R * np.sin(THETA)
    X[0, :], Y[0, :], Z[0, :] = 0.0, 0.0, -half_l
    X[-1, :], Y[-1, :], Z[-1, :] = 0.0, 0.0, half_l

    all_verts = [np.column_stack((X.ravel(), Y.ravel(), Z.ravel()))]
    v_offset = len(all_verts[0])
    faces = []

    for i in range(n_z - 1):
        for j in range(n_theta):
            jn = (j + 1) % n_theta
            v00, v01 = i * n_theta + j, i * n_theta + jn
            v10, v11 = (i + 1) * n_theta + j, (i + 1) * n_theta + jn
            if i == 0:
                faces.append([v00, v10, v11])
            elif i == n_z - 2:
                faces.append([v00, v10, v01])
            else:
                faces.extend([[v00, v10, v11], [v00, v11, v01]])

    all_faces = [np.array(faces, dtype=np.int32)]

    # Outer Prongs
    z_start = 27.5 * scale
    z_end = 66.0 * scale
    r_max = 17.15 * scale
    n_t, n_phi = 100, 32
    t = np.linspace(0.0, 1.0, n_t)
    phi = np.linspace(0.0, 2.0 * np.pi, n_phi, endpoint=False)

    r_center = np.where(
        t < 0.32,
        (8.5 * scale) + (r_max - 8.5 * scale) * np.sin((np.pi / 2.0) * (t / 0.32)),
        r_max - np.tan(np.radians(TAPER_ANGLE_DEG)) * ((t - 0.32) * (z_end - z_start)) - (3.8 * scale) * ((t - 0.32) / 0.68) ** 2
    )
    tube_r = (2.4 * scale) * (1.0 - 0.28 * t)
    tube_r[0], tube_r[-1] = 0.0, 0.0

    for z_sign in (-1.0, 1.0):
        z_center = z_sign * (z_start + (z_end - z_start) * t)
        dr, dz = np.gradient(r_center), np.gradient(z_center)
        norm = np.sqrt(dr**2 + dz**2) + 1e-12
        nr, nz = -dz / norm, dr / norm

        T_g, PHI_g = np.meshgrid(np.arange(n_t), phi, indexing="ij")
        tr = tube_r[T_g] * (1.0 + 0.08 * np.cos(4.0 * PHI_g))
        local_r = r_center[T_g] + tr * np.cos(PHI_g) * nr[T_g]
        local_p = tr * np.sin(PHI_g)
        Zp = z_center[T_g] + tr * np.cos(PHI_g) * nz[T_g]

        for k in range(PRONG_COUNT):
            ang = k * (2.0 * np.pi / PRONG_COUNT)
            Xp = local_r * np.cos(ang) - local_p * np.sin(ang)
            Yp = local_r * np.sin(ang) + local_p * np.cos(ang)

            p_verts = np.column_stack((Xp.ravel(), Yp.ravel(), Zp.ravel()))
            p_faces = []
            for i in range(n_t - 1):
                for j in range(n_phi):
                    jn = (j + 1) % n_phi
                    v00, v01 = i * n_phi + j, i * n_phi + jn
                    v10, v11 = (i + 1) * n_phi + j, (i + 1) * n_phi + jn
                    if z_sign > 0:
                        p_faces.extend([[v00, v10, v11], [v00, v11, v01]])
                    else:
                        p_faces.extend([[v00, v11, v10], [v00, v01, v11]])

            all_verts.append(p_verts)
            all_faces.append(np.array(p_faces, dtype=np.int32) + v_offset)
            v_offset += len(p_verts)

    return np.vstack(all_verts), np.vstack(all_faces)


def export_stl(filename, verts, faces):
    v0, v1, v2 = verts[faces[:, 0]], verts[faces[:, 1]], verts[faces[:, 2]]
    normals = np.cross(v1 - v0, v2 - v0)
    normals /= np.linalg.norm(normals, axis=1, keepdims=True) + 1e-12
    dt = np.dtype([("normal", "<f4", (3,)), ("v0", "<f4", (3,)), ("v1", "<f4", (3,)), ("v2", "<f4", (3,)), ("attr", "<u2")])
    data = np.zeros(len(faces), dtype=dt)
    data["normal"], data["v0"], data["v1"], data["v2"] = normals, v0, v1, v2
    with open(filename, "wb") as f:
        f.write(b"Base-12 Vajra Dual Size Export".ljust(80, b" "))
        f.write(struct.pack("<I", len(faces)))
        data.tofile(f)


# =====================================================================
# Additional exporters (from the original generate_vajra_cad.py),
# adapted to the dual-size configs so each size exports in every format.
# =====================================================================
def export_obj(filename, verts, faces):
    with open(filename, "w", encoding="utf-8") as f:
        f.write("# Base-12 Meditation Vajra\n")
        f.write("o Base12_Vajra\n")
        for v in verts:
            f.write(f"v {v[0]:.5f} {v[1]:.5f} {v[2]:.5f}\n")
        for face in faces + 1:
            f.write(f"f {face[0]} {face[1]} {face[2]}\n")


def export_ply(filename, verts, faces):
    with open(filename, "w", encoding="utf-8") as f:
        f.write("ply\nformat ascii 1.0\n")
        f.write(f"element vertex {len(verts)}\n")
        f.write("property float x\nproperty float y\nproperty float z\n")
        f.write(f"element face {len(faces)}\n")
        f.write("property list uchar int vertex_indices\nend_header\n")
        for v in verts:
            f.write(f"{v[0]:.5f} {v[1]:.5f} {v[2]:.5f}\n")
        for face in faces:
            f.write(f"3 {face[0]} {face[1]} {face[2]}\n")


def export_dxf(filename, size_key):
    cfg = CONFIGS[size_key]
    half_l, hub_r = cfg["half_length"], cfg["hub_r"]
    z_vals = np.linspace(-half_l, half_l, 400)
    r_vals = central_body_radius(z_vals, np.zeros_like(z_vals), half_l, hub_r)

    def dxf_line(x1, y1, x2, y2, layer="PROFILE"):
        return (
            f"0\nLINE\n8\n{layer}\n"
            f"10\n{x1:.4f}\n20\n{y1:.4f}\n30\n0.0\n"
            f"11\n{x2:.4f}\n21\n{y2:.4f}\n31\n0.0\n"
        )

    with open(filename, "w", encoding="utf-8") as f:
        f.write("0\nSECTION\n2\nENTITIES\n")
        f.write(dxf_line(-half_l, 0.0, half_l, 0.0, layer="CENTERLINE"))
        for i in range(len(z_vals) - 1):
            f.write(dxf_line(z_vals[i], r_vals[i], z_vals[i + 1], r_vals[i + 1]))
            f.write(dxf_line(z_vals[i], -r_vals[i], z_vals[i + 1], -r_vals[i + 1]))
        f.write("0\nENDSEC\n0\nEOF\n")


if __name__ == "__main__":
    for size in ["144mm", "120mm"]:
        print(f"Generating {size} Vajra (Length: {CONFIGS[size]['total_length']} mm)...")
        v, f = build_vajra_mesh(size)
        print(f" -> Mesh: {len(v):,} vertices | {len(f):,} triangles")
        export_stl(f"base12_vajra_{size}.stl", v, f)
        print(f" -> Saved: base12_vajra_{size}.stl")
        export_obj(f"base12_vajra_{size}.obj", v, f)
        print(f" -> Saved: base12_vajra_{size}.obj")
        export_ply(f"base12_vajra_{size}.ply", v, f)
        print(f" -> Saved: base12_vajra_{size}.ply")
        export_dxf(f"base12_vajra_{size}_profile.dxf", size)
        print(f" -> Saved: base12_vajra_{size}_profile.dxf")
    print("All exports complete!")