"""
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

Control parameters (print base): c (horizontal), d (depth)
State variable (print height):    x

With a = -3, the bifurcation set in the (c, d) plane forms the characteristic
butterfly shape: a self-intersecting loop passing through (c=3, d=0), enclosing
a region with 5 equilibria ("butterfly pocket"), surrounded by a 3-root region
with outer fold wings, and a single-root region outside.  This is structurally
different from the cusp catastrophe and cannot be seen with c = d = 0 fixed.
"""

import numpy as np
import os

A_FIXED = -3.0      # butterfly unfolding parameter (must be negative)

# ── Tuning ──────────────────────────────────────────────────────────────────
GRID      = 50      # control-space resolution — increase to 80–100 for print
C_RANGE   = (-2.0, 7.0)   # range of control parameter c
D_RANGE   = (-6.0, 6.0)   # range of control parameter d
X_RANGE   = 2.5            # half-width of root search window
MAX_EDGE_DZ = 0.6          # max z-jump per quad edge (filters branch mismatches)

# ── 1. Root finding ──────────────────────────────────────────────────────────

def dV(x, c, d):
    return 6*x**5 + 4*A_FIXED*x**3 + 2*c*x + d

def d2V(x, c, d):
    return 30*x**4 + 12*A_FIXED*x**2 + 2*c

def find_roots(c, d, n_starts=80):
    """Return sorted real roots of dV/dx = 0 for the given (c, d)."""
    xs = np.linspace(-X_RANGE, X_RANGE, n_starts)
    roots = []
    for x0 in xs:
        x = float(x0)
        for _ in range(200):
            fx  = dV(x, c, d)
            if abs(fx) < 1e-12:
                break
            dfx = d2V(x, c, d)
            if abs(dfx) < 1e-14:
                break
            step = fx / dfx
            x -= step
            if abs(x) > 2.0 * X_RANGE:   # diverged — abandon
                break
            if abs(step) < 1e-10:
                break
        if abs(dV(x, c, d)) < 1e-7 and abs(x) <= X_RANGE + 0.15:
            if not any(abs(x - r) < 1e-4 for r in roots):
                roots.append(x)
    return sorted(roots)

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

def build_mesh():
    c_vals = np.linspace(*C_RANGE, GRID)
    d_vals = np.linspace(*D_RANGE, GRID)

    print(f'  Computing roots on {GRID}×{GRID} grid…')
    roots_grid = [[find_roots(c, d) for d in d_vals] for c in c_vals]

    triangles = []

    for i in range(GRID - 1):
        for j in range(GRID - 1):
            c0, c1 = c_vals[i], c_vals[i+1]
            d0, d1 = d_vals[j], d_vals[j+1]

            r00 = roots_grid[i  ][j  ]
            r10 = roots_grid[i+1][j  ]
            r11 = roots_grid[i+1][j+1]
            r01 = roots_grid[i  ][j+1]

            # Only draw branch k when ALL four corners have it.
            # The rs[-1] fallback used previously created false triangles at
            # fold edges by connecting unrelated branches — this is the main
            # cause of the messy, self-intersecting geometry.
            n = min(len(r00), len(r10), len(r11), len(r01))

            for k in range(n):
                p0 = (c0, d0, r00[k])
                p1 = (c1, d0, r10[k])
                p2 = (c1, d1, r11[k])
                p3 = (c0, d1, r01[k])

                # Reject quads where any edge has a large z-jump.  A large
                # jump indicates a branch-index mismatch near a fold line
                # (sorted order is preserved within a branch but can
                # "swap" across folds where adjacent branches coalesce).
                if any(abs(a[2] - b[2]) > MAX_EDGE_DZ
                       for a, b in [(p0,p1),(p1,p2),(p2,p3),(p3,p0)]):
                    continue

                triangles.append((p0, p1, p2))
                triangles.append((p0, p2, p3))

    return triangles

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

def add_base(triangles, z_base=-2.8):
    c0, c1 = C_RANGE
    d0, d1 = D_RANGE
    zb, zt = z_base, z_base + 0.15

    triangles += [
        ((c0,d0,zt),(c1,d0,zt),(c1,d1,zt)),
        ((c0,d0,zt),(c1,d1,zt),(c0,d1,zt)),
        ((c0,d0,zb),(c1,d1,zb),(c1,d0,zb)),
        ((c0,d0,zb),(c0,d1,zb),(c1,d1,zb)),
    ]
    for (x0,y0),(x1,y1) in [
        ((c0,d0),(c1,d0)), ((c1,d0),(c1,d1)),
        ((c1,d1),(c0,d1)), ((c0,d1),(c0,d0)),
    ]:
        triangles += [
            ((x0,y0,zb),(x1,y1,zb),(x1,y1,zt)),
            ((x0,y0,zb),(x1,y1,zt),(x0,y0,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)
    length = np.linalg.norm(n)
    return n / length if length > 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 mesh (grid={GRID})…')
    tris = build_mesh()
    tris = add_base(tris)
    print('Writing ASCII STL…')
    write_ascii_stl(tris, 'butterfly_catastrophe.stl')
    print('Done.')