"""
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      = 200     # 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 _emit_quad(triangles, p00, p10, p11, p01):
    triangles.append((p00, p10, p11))
    triangles.append((p00, p11, p01))


def _fold_terminations(tracks, axis_vals):
    """
    Scan tracks along one axis and return a list of fold-termination events.
    Each event is (axis_val, xa, xb) where axis_val is the last valid position,
    and xa < xb are the two branch values that die together at a fold.

    Branches come in pairs at fold lines (two coalesce), so we pair adjacent
    sorted dying values.  Events are indexed by the axis position of the last
    valid step.
    """
    events = []
    n = len(axis_vals) - 1  # number of steps
    for step in range(n):
        dying = sorted(t[step] for t in tracks
                       if t[step] is not None and t[step + 1] is None)
        for k in range(0, len(dying) - 1, 2):
            events.append((axis_vals[step], dying[k], dying[k + 1]))
    return events


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 rows (fixed d, varying c) and columns (fixed c, varying d).
    print('  Tracking branches…')
    row_tracks = [track_branches([roots_grid[i][j] for i in range(GRID)])
                  for j in range(GRID)]
    col_tracks = [track_branches([roots_grid[i][j] for j in range(GRID)])
                  for i in range(GRID)]

    triangles = []

    # ── Main surface quads ────────────────────────────────────────────────────
    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]

            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

            j1_used = set()
            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)
                _emit_quad(triangles,
                           (c0, d0, x00), (c1, d0, x10),
                           (c1, d1, best_x11), (c0, d1, best_x01))

    # ── C-direction fold caps ─────────────────────────────────────────────────
    # The fold line runs at an angle through the (c, d) grid, so the column
    # where two branches die can differ by ±1 between adjacent rows.
    # We collect all fold terminations per row, then match them across the
    # row pair regardless of exact column, connecting the dying edges with
    # (possibly trapezoidal) quads.
    c_terms = [_fold_terminations(row_tracks[j], c_vals) for j in range(GRID)]

    for j in range(GRID - 1):
        d0, d1 = d_vals[j], d_vals[j + 1]
        terms_j  = c_terms[j]
        terms_j1 = c_terms[j + 1]

        j1_used = set()
        for c_j, xa, xb in terms_j:
            mid = (xa + xb) / 2
            best_k = None
            best_dist = 1.0   # max allowed x-midpoint distance between matched pairs
            for k1, (c_j1, xa1, xb1) in enumerate(terms_j1):
                if k1 in j1_used:
                    continue
                dist = abs(mid - (xa1 + xb1) / 2)
                if dist < best_dist:
                    best_dist, best_k = dist, k1
            if best_k is None:
                continue
            j1_used.add(best_k)
            c_j1, xa1, xb1 = terms_j1[best_k]
            # Cap quad: lies at the fold edge, spanning d0→d1 between the two
            # dying branches.  c may differ slightly between the two rows if
            # the fold line is diagonal.
            _emit_quad(triangles,
                       (c_j,  d0, xa),  (c_j,  d0, xb),
                       (c_j1, d1, xb1), (c_j1, d1, xa1))

    # ── D-direction fold caps ─────────────────────────────────────────────────
    d_terms = [_fold_terminations(col_tracks[i], d_vals) for i in range(GRID)]

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

        i1_used = set()
        for d_i, xa, xb in terms_i:
            mid = (xa + xb) / 2
            best_k = None
            best_dist = 1.0
            for k1, (d_i1, xa1, xb1) in enumerate(terms_i1):
                if k1 in i1_used:
                    continue
                dist = abs(mid - (xa1 + xb1) / 2)
                if dist < best_dist:
                    best_dist, best_k = dist, k1
            if best_k is None:
                continue
            i1_used.add(best_k)
            d_i1, xa1, xb1 = terms_i1[best_k]
            _emit_quad(triangles,
                       (c0, d_i,  xa),  (c1, d_i1, xa1),
                       (c1, d_i1, xb1), (c0, d_i,  xb))

    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.')