"""
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      = 80      # control-space resolution — increase to 100–120 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_MATCH_DZ = 0.8         # max z-gap for inter-row branch matching

# ── 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. Branch tracking ───────────────────────────────────────────────────────

def track_branches(roots_along_axis):
    """
    Track branches along one axis (fixed d, varying c) using greedy
    nearest-neighbour matching.  Returns a list of tracks; each track is a
    list of length GRID where entry i is the root value at column i, or None
    when the branch does not exist there.

    Sorting-index matching (branch 0 always = branch 0) breaks at folds
    because two adjacent branches coalesce, shifting every higher index by
    one.  Nearest-neighbour tracking follows each physical sheet through the
    fold correctly: the two merging branches each get None past the fold, and
    the surviving sheet keeps its track unbroken.
    """
    n = len(roots_along_axis)
    if n == 0:
        return []

    tracks = [[r] for r in roots_along_axis[0]]

    for i in range(1, n):
        curr = roots_along_axis[i]

        prev_live = [(ti, t[-1]) for ti, t in enumerate(tracks)
                     if t[-1] is not None]
        prev_matched, curr_matched = set(), set()
        matches = {}  # track_idx → curr_root_idx

        cands = sorted(
            [(abs(pv - curr[ci]), ti, ci)
             for ti, pv in prev_live
             for ci in range(len(curr))],
            key=lambda x: x[0]
        )
        for _, ti, ci in cands:
            if ti not in prev_matched and ci not in curr_matched:
                matches[ti] = ci
                prev_matched.add(ti)
                curr_matched.add(ci)

        for ti, t in enumerate(tracks):
            t.append(curr[matches[ti]] if ti in matches else None)

        # Branches that appear for the first time at this column
        for ci in range(len(curr)):
            if ci not in curr_matched:
                tracks.append([None] * i + [curr[ci]])

    return tracks

# ── 3. 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[i][j] = sorted roots at (c_vals[i], d_vals[j])
    roots_grid = [[find_roots(c, d) for d in d_vals] for c in c_vals]

    # Track branches along each row (fixed d, varying c)
    print('  Tracking branches…')
    row_tracks = []
    for j in range(GRID):
        row = [roots_grid[i][j] for i in range(GRID)]
        row_tracks.append(track_branches(row))
    # row_tracks[j] = list of tracks; track[i] = root value at column i

    triangles = []

    for j in range(GRID - 1):
        d0, d1 = d_vals[j], d_vals[j+1]
        tracks_j  = row_tracks[j]
        tracks_j1 = row_tracks[j + 1]

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

            # Collect valid quad-edge segments for both rows at this column
            # A segment is valid when both endpoints of the edge exist
            segs_j  = [(t[i], t[i+1], k)
                       for k, t in enumerate(tracks_j)
                       if t[i] is not None and t[i+1] is not None]
            segs_j1 = [(t[i], t[i+1], k)
                       for k, t in enumerate(tracks_j1)
                       if t[i] is not None and t[i+1] is not None]

            if not segs_j or not segs_j1:
                continue

            # Greedy nearest-neighbour matching between the two row-bands.
            # Each j1-segment is claimed by at most one j-segment so we never
            # emit overlapping quads.
            j1_used = set()
            # Sort j-segments by midpoint so iteration order is deterministic
            for x00, x10, _ in sorted(segs_j, key=lambda s: (s[0] + s[1]) / 2):
                best_k1 = best_x01 = best_x11 = None
                best_dist = MAX_MATCH_DZ
                for x01, x11, k1 in segs_j1:
                    if k1 in j1_used:
                        continue
                    dist = max(abs(x00 - x01), abs(x10 - x11))
                    if dist < best_dist:
                        best_dist, best_k1 = dist, k1
                        best_x01, best_x11 = x01, x11

                if best_k1 is None:
                    continue
                j1_used.add(best_k1)

                p00 = (c0, d0, x00)
                p10 = (c1, d0, x10)
                p11 = (c1, d1, best_x11)
                p01 = (c0, d1, best_x01)
                triangles.append((p00, p10, p11))
                triangles.append((p00, p11, p01))

    return triangles

# ── 4. 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

# ── 5. 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}')

# ── 6. 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.')