#!/usr/bin/python3
"""
Butterfly Catastrophe Surface — ASCII STL Generator

Potential:   V(x) = x^6 + a*x^4 + c*x^2 + d*x   (a = -3 fixed, b = 0)
Equilibrium: dV/dx = 6x^5 - 12x^3 + 2c*x + d = 0

Rearranged as a single-valued function:
    d(x, c) = -(6x^5 - 12x^3 + 2c*x)
            =  -6x^5 + 12x^3 - 2c*x

The print base is the (x, c) plane; height is d.  Every grid point maps to
exactly one vertex — no root finding, no branch tracking, no holes.

The fold lines of the bifurcation set appear as ridges where the surface has
zero gradient in x: ∂d/∂x = -30x^4 + 36x^2 - 2c = 0, i.e. c = 18x^2 - 15x^4.
"""

import numpy as np
import os

A_FIXED = -3.0

# ── Tuning ──────────────────────────────────────────────────────────────────
GRID    = 400          # grid resolution — higher = smoother ridges
X_RANGE = (-4.5, 4.5)  # state variable x
C_RANGE = (-1.5, 7.5)  # control parameter c
D_SCALE = 0.4          # scale factor applied to the computed d height;
                       # reduce if the model is too tall for your printer

# ── 1. Height function ────────────────────────────────────────────────────────

def d_surface(x, c):
    """d value on the equilibrium manifold: dV/dx = 0 solved for d."""
    return (-6*x**5 + 12*x**3 - 2*c*x) * D_SCALE

# ── 2. Build mesh ─────────────────────────────────────────────────────────────

def build_mesh():
    x_vals = np.linspace(*X_RANGE, GRID)
    c_vals = np.linspace(*C_RANGE, GRID)

    # Pre-compute the full height field in one vectorised call
    X, C = np.meshgrid(x_vals, c_vals, indexing='ij')  # (GRID, GRID)
    D = (-6*X**5 + 12*X**3 - 2*C*X) * D_SCALE

    triangles = []
    for i in range(GRID - 1):
        for j in range(GRID - 1):
            p00 = (x_vals[i],   c_vals[j],   D[i,   j  ])
            p10 = (x_vals[i+1], c_vals[j],   D[i+1, j  ])
            p11 = (x_vals[i+1], c_vals[j+1], D[i+1, j+1])
            p01 = (x_vals[i],   c_vals[j+1], D[i,   j+1])
            triangles.append((p00, p10, p11))
            triangles.append((p00, p11, p01))

    return triangles

# ── 3. Flat base ──────────────────────────────────────────────────────────────

def add_base(triangles):
    x0, x1 = X_RANGE
    c0, c1 = C_RANGE
    X, C = np.meshgrid(np.linspace(x0, x1, GRID),
                       np.linspace(c0, c1, GRID), indexing='ij')
    D = (-6*X**5 + 12*X**3 - 2*C*X) * D_SCALE
    zb = D.min() - 0.15
    zt = zb + 0.15

    # Top and bottom faces of base slab
    triangles += [
        ((x0,c0,zt),(x1,c0,zt),(x1,c1,zt)),
        ((x0,c0,zt),(x1,c1,zt),(x0,c1,zt)),
        ((x0,c0,zb),(x1,c1,zb),(x1,c0,zb)),
        ((x0,c0,zb),(x0,c1,zb),(x1,c1,zb)),
    ]
    for (ax,ay),(bx,by) in [
        ((x0,c0),(x1,c0)), ((x1,c0),(x1,c1)),
        ((x1,c1),(x0,c1)), ((x0,c1),(x0,c0)),
    ]:
        triangles += [
            ((ax,ay,zb),(bx,by,zb),(bx,by,zt)),
            ((ax,ay,zb),(bx,by,zt),(ax,ay,zt)),
        ]
    return triangles

# ── 4. ASCII STL output ───────────────────────────────────────────────────────

def normal(v0, v1, v2):
    a = np.subtract(v1, v0)
    b = np.subtract(v2, v0)
    n = np.cross(a, b)
    L = np.linalg.norm(n)
    return n / L if L > 1e-14 else np.array([0.0, 0.0, 1.0])

def write_ascii_stl(triangles, filename):
    with open(filename, 'w') as f:
        f.write('solid butterfly_catastrophe\n')
        for tri in triangles:
            v0, v1, v2 = [np.array(v, dtype=float) for v in tri]
            nx, ny, nz = normal(v0, v1, v2)
            f.write(f'  facet normal {nx:.6e} {ny:.6e} {nz:.6e}\n')
            f.write('    outer loop\n')
            for v in (v0, v1, v2):
                f.write(f'      vertex {v[0]:.6e} {v[1]:.6e} {v[2]:.6e}\n')
            f.write('    endloop\n')
            f.write('  endfacet\n')
        f.write('endsolid butterfly_catastrophe\n')
    size_kb = os.path.getsize(filename) / 1024
    print(f'  Triangles : {len(triangles):,}')
    print(f'  File size : {size_kb:.0f} KB  →  {filename}')

# ── 5. Main ───────────────────────────────────────────────────────────────────

if __name__ == '__main__':
    print(f'Building butterfly catastrophe surface (grid={GRID})…')
    tris = build_mesh()
    tris = add_base(tris)
    print('Writing ASCII STL…')
    write_ascii_stl(tris, 'butterfly_catastrophe.stl')
    print('Done.')