import numpy as np
import itertools
from scipy.linalg import expm, logm
import json

def find_rigid_permutations(D, z=[]):
    return sum((find_rigid_permutations(D, z+[i]) for i in range(len(D)) if np.allclose(
        D[z+[i]]@D[i],
        D[:len(z)+1]@D[len(z)])),
    []) if len(z)!=len(D)else[np.eye(len(D))[z]]

sticker_coords = np.array([s for s in itertools.product(range(-2,3), repeat=3) 
                           if (np.abs(s)==2).sum()==1])
colors = 120 + 120 * np.round(sticker_coords / 2)
stickers_in_slice = np.sum(sticker_coords[:, 0] < 0)
view = np.linalg.qr(np.random.randn(3, 3))[0]

global_perms = find_rigid_permutations(sticker_coords)

slice_perms = find_rigid_permutations(sticker_coords[:stickers_in_slice])
move = np.eye(len(sticker_coords))


print(slice_perms[3])
move[:stickers_in_slice, :stickers_in_slice] = slice_perms[3]

def np2js(arr):
    real_arr = np.block([[arr.real, arr.imag], [-arr.imag, arr.real]])
    return json.dumps(real_arr.tolist())

def cube(key_actions):
    ret = f"""
    <!DOCTYPE html><html><body><script>

    let mul = (A, B) => A.map((row, i) => B[0].map((_, j) =>
        row.reduce((acc, _, n) => acc + A[i][n] * B[n][j], 0)))

    var state = {np2js(np.eye(len(sticker_coords)))}
    const coords = {np2js(sticker_coords @ view * 14 + 35)}
    var moves = [state]
    document.addEventListener("keypress", (event) => {{
    """
    for i, matrix in enumerate(key_actions):
        ret += f"""
        if (event.key == {i}) {{
            moves = (new Array(10).fill( {np2js(expm(.1 * logm(matrix)))})).concat( moves);
        }}
        """
    ret += """
    });
    setInterval(step, 20);

    function step() {
        if (!moves.length) return;
        state = mul(state, moves.pop());
        const locations = mul(state, coords);
        document.body.innerHTML= `
    """
    for i, color in enumerate(colors):
        ret += f"""
            <div style='
                position: absolute;
                left: ${{locations[{i}][1]}}px;
                top: ${{locations[{i}][2]}}px;
                z-index: ${{Math.round(locations[{i}][0])}};
                color: rgb({color[0]} {color[1]} {color[2]});
            '>
                &#x2B24;
            </div>
        """
    return ret + "`;}</script></body></html>"

with open("output.html", "w") as static_site:
    generators = [move, global_perms[15], global_perms[9]]
    static_site.write(cube(generators))
