#!/usr/bin/env python3

# Plot a Menger sponge.
#
# The construction of a Menger sponge can be described as follows:
#
# 1. Begin with a cube.
# 2. Divide every face of the cube into 9 squares, like a Rubik's Cube. This
#    will sub-divide the cube into 27 smaller cubes.
# 3. Remove the smaller cube in the middle of each face, and remove the smaller
#    cube in the very center of the larger cube, leaving 20 smaller cubes. This
#    is a level-1 Menger sponge (resembling a Void Cube).
# 4. Repeat steps 2 and 3 for each of the remaining smaller cubes, and continue
#    to iterate.

import argparse

import matplotlib.pyplot as plt

from math import ceil

from matplotlib import pylab
from mpl_toolkits.mplot3d import Axes3D

from linpy import Le, Polyhedron, symbols


x, y, z = symbols('x y z')

_x, _y, _z = x.asdummy(), y.asdummy(), z.asdummy()


def translate(domain, *, dx=0, dy=0, dz=0):
    domain &= Polyhedron([x - _x + dx, y - _y + dy, z - _z + dz])
    domain = domain.project([x, y, z])
    domain = domain.subs({_x: x, _y: y, _z: z})
    return domain


def _menger(domain, size):
    result = domain
    result |= translate(domain, dx=0, dy=size, dz=0)
    result |= translate(domain, dx=0, dy=2*size, dz=0)
    result |= translate(domain, dx=size, dy=0, dz=0)
    result |= translate(domain, dx=size, dy=2*size, dz=0)
    result |= translate(domain, dx=2*size, dy=0, dz=0)
    result |= translate(domain, dx=2*size, dy=size, dz=0)
    result |= translate(domain, dx=2*size, dy=2*size, dz=0)
    result |= translate(domain, dx=0, dy=0, dz=size)
    result |= translate(domain, dx=0, dy=2*size, dz=size)
    result |= translate(domain, dx=2*size, dy=0, dz=size)
    result |= translate(domain, dx=2*size, dy=2*size, dz=size)
    result |= translate(domain, dx=0, dy=0, dz=2*size)
    result |= translate(domain, dx=0, dy=size, dz=2*size)
    result |= translate(domain, dx=0, dy=2*size, dz=2*size)
    result |= translate(domain, dx=size, dy=0, dz=2*size)
    result |= translate(domain, dx=size, dy=2*size, dz=2*size)
    result |= translate(domain, dx=2*size, dy=0, dz=2*size)
    result |= translate(domain, dx=2*size, dy=size, dz=2*size)
    result |= translate(domain, dx=2*size, dy=2*size, dz=2*size)
    return result


def menger(domain, count=1, cut=False):
    size = 1
    for i in range(count):
        domain = _menger(domain, size)
        size *= 3
    if cut:
        domain &= Le(x + y + z, ceil(3 * size / 2))
    return domain

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Compute a Menger sponge.')
    parser.add_argument(
        '-n', '--iterations', type=int, default=2,
        help='number of iterations (default: 2)')
    parser.add_argument(
        '-c', '--cut', action='store_true', default=False,
        help='cut the sponge')
    args = parser.parse_args()
    cube = Le(0, x) & Le(x, 1) & Le(0, y) & Le(y, 1) & Le(0, z) & Le(z, 1)
    fractal = menger(cube, args.iterations, args.cut)
    fig = plt.figure(facecolor='white')
    plot = fig.add_subplot(1, 1, 1, projection='3d', aspect='equal')
    fractal.plot(plot)
    pylab.show()