"""
Butterfly Catastrophe Surface - ASCII STL Generator
Potential: V(x) = x^6 + a*x^4 + b*x^3
Equilibrium surface: dV/dx = 6x^5 + 4a*x^3 + 3b*x^2 = 0
"""

import numpy as np
import os

# ---------------------------------------------------------------------------
# 1. Find equilibrium roots of dV/dx = 0
# ---------------------------------------------------------------------------

def dV(x, a, b):
    return 6*x**5 + 4*a*x**3 + 3*b*x**2

def d2V(x, a, b):
    return 30*x**4 + 12*a*x**2 + 6*b*x

def find_roots(a, b, n_starts=40, x_range=2.2):
    xs = np.linspace(-x_range, x_range, n_starts)
    roots = []
    for x0 in xs:
        x = float(x0)
        for _ in range(100):
            fx  = dV(x, a, b)
            dfx = d2V(x, a, b)
            if abs(dfx) < 1e-14:
                break
            step = fx / dfx
            x -= step
            if abs(step) < 1e-10:
                break
        if abs(dV(x, a, b)) < 1e-7 and abs(x) <= x_range + 0.05:
            if not any(abs(x - r) < 1e-4 for r in roots):
                roots.append(x)
    return sorted(roots)

# ---------------------------------------------------------------------------
# 2. Build mesh
# ---------------------------------------------------------------------------

def build_mesh(grid=40):
    a_vals = np.linspace(-2.5, 1.2, grid)
    b_vals = np.linspace(-2.5, 2.5, grid)

    # Pre-compute roots at every grid point
    roots_grid = [[find_roots(a, b) for b in b_vals] for a in a_vals]

    triangles = []

    for i in range(grid - 1):
        for j in range(grid - 1):
            corners_roots = [
                roots_grid[i  ][j  ],
                roots_grid[i+1][j  ],
                roots_grid[i+1][j+1],
                roots_grid[i  ][j+1],
            ]
            a_c = [a_vals[i], a_vals[i+1], a_vals[i+1], a_vals[i  ]]
            b_c = [b_vals[j], b_vals[j  ], b_vals[j+1], b_vals[j+1]]

            max_branch = max(len(r) for r in corners_roots)

            for branch in range(max_branch):
                pts = []
                for k in range(4):
                    rs = corners_roots[k]
                    if rs:
                        x = rs[branch] if branch < len(rs) else rs[-1]
                    else:
                        pts = None
                        break
                    pts.append((a_c[k], b_c[k], x))

                if pts and len(pts) == 4:
                    p0, p1, p2, p3 = pts
                    triangles.append((p0, p1, p2))
                    triangles.append((p0, p2, p3))

    return triangles

# ---------------------------------------------------------------------------
# 3. Add flat base
# ---------------------------------------------------------------------------

def add_base(triangles, z_base=-2.4):
    a0, a1 = -2.5, 1.2
    b0, b1 = -2.5, 2.5
    zb, zt = z_base, z_base + 0.15

    triangles += [
        ((a0,b0,zt),(a1,b0,zt),(a1,b1,zt)),
        ((a0,b0,zt),(a1,b1,zt),(a0,b1,zt)),
        ((a0,b0,zb),(a1,b1,zb),(a1,b0,zb)),
        ((a0,b0,zb),(a0,b1,zb),(a1,b1,zb)),
    ]
    walls = [
        ((a0,b0),(a1,b0)), ((a1,b0),(a1,b1)),
        ((a1,b1),(a0,b1)), ((a0,b1),(a0,b0)),
    ]
    for (x0,y0),(x1,y1) in walls:
        triangles += [
            ((x0,y0,zb),(x1,y1,zb),(x1,y1,zt)),
            ((x0,y0,zb),(x1,y1,zt),(x0,y0,zt)),
        ]
    return triangles

# ---------------------------------------------------------------------------
# 4. Write ASCII STL
# ---------------------------------------------------------------------------

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, 1])

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('Building butterfly catastrophe mesh (grid=40)…')
    tris = build_mesh(grid=40)
    tris = add_base(tris)
    print('Writing ASCII STL…')
    write_ascii_stl(tris, 'butterfly_catastrophe.stl')
    print('Done.')