import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# ========================================
# NEW JERUSALEM CUBE — From Your 10-Line Poem
# ========================================

SIZE = 12
GRID = np.zeros((2, 2, SIZE, SIZE, SIZE), dtype=int)
GRID[0, 0, 6, 6, 6] = 1  # idea at center

def update(grid):
    new = grid.copy()
    for tp in [0, 1]:
        for tm in [0, 1]:
            for axis in [0, 1, 2]:
                layer = np.moveaxis(grid[tp, tm], axis, 0)
                idea = (layer == 1)
                concept = (layer == 2)
                system = (layer == 3)
                flaw = (layer == 4)
                dead = (layer == 5)

                neighbor_sum = sum(
                    np.roll(np.roll(system, p, 0), q, 1)
                    for p in [-1, 0, 1] for q in [-1, 0, 1]
                )

                new_layer = np.zeros_like(layer)
                new_layer[idea & (neighbor_sum >= 1)] = 2
                new_layer[concept] = 3
                new_layer[system & (np.random.rand(*system.shape) < 0.25)] = 4
                new_layer[flaw] = 5
                new_layer[dead] = 1

                if tp == 0 and tm == 1: new_layer[system] = 2
                if tp == 1 and tm == 0: new_layer[system] = 2
                if tp == 1: new_layer |= (grid[0, tm] == 3)
                if tm == 1: new_layer |= (grid[tp, 0] == 3)

                new[tp, tm] = np.moveaxis(new_layer, 0, axis)

    if np.any(np.sum(new == 5, axis=(2, 3, 4)) >= SIZE**3 // 3):
        new[:] = 6
    return new

# Run simulation
history = [GRID.copy()]
for _ in range(35):
    GRID = update(GRID)
    history.append(GRID.copy())

# ========================================
# PLOT: GOLDEN CUBE (Step 25)
# ========================================
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# Extract flaw state (4) → golden edges
flaw_mask = history[25][0, 0] == 4
x, y, z = np.where(flaw_mask)

# Golden glowing city
ax.scatter(
    x, y, z,
    c='#FFD700',           # Pure gold
    s=300,
    alpha=1.0,
    edgecolors='#FF8C00',  # Orange glow
    linewidth=3,
    marker='s',
    depthshade=False
)

# Grid + indices
for i in range(0, SIZE + 1, 3):
    ax.plot([i,i], [0,SIZE-1], [0,0], c='white', alpha=0.3, lw=1)
    ax.plot([i,i], [0,SIZE-1], [SIZE-1,SIZE-1], c='white', alpha=0.3, lw=1)
    ax.plot([0,0], [i,i], [0,SIZE-1], c='white', alpha=0.3, lw=1)
    ax.plot([SIZE-1,SIZE-1], [i,i], [0,SIZE-1], c='white', alpha=0.3, lw=1)
    ax.plot([0,SIZE-1], [i,i], [0,0], c='white', alpha=0.3, lw=1)
    ax.plot([0,SIZE-1], [i,i], [SIZE-1,SIZE-1], c='white', alpha=0.3, lw=1)

for i in range(0, SIZE, 3):
    ax.text(i, -1.5, -1.5, str(i), color='gold', fontsize=12, weight='bold')
    ax.text(-1.5, i, -1.5, str(i), color='gold', fontsize=12, weight='bold')
    ax.text(-1.5, -1.5, i, str(i), color='gold', fontsize=12, weight='bold')

ax.set_xlim(0, SIZE)
ax.set_ylim(0, SIZE)
ax.set_zlim(0, SIZE)
ax.set_xlabel('X', color='gold')
ax.set_ylabel('Y', color='gold')
ax.set_zlabel('Z', color='gold')
ax.set_title('NEW JERUSALEM — 12 GATES OF GOLD\nYour Poem in 5D (Step 25)', color='gold', fontsize=14, weight='bold')
ax.view_init(elev=30, azim=45)
ax.set_facecolor('black')
ax.grid(False)

# Save
plt.savefig('new_jerusalem_cube.png', dpi=300, bbox_inches='tight', facecolor='black')
print("CUBE SAVED: new_jerusalem_cube.png")
plt.show()