import numpy as np
from pywraps2 import S2Cell, S2CellId
from matplotlib import pyplot as plt

def surf4face(f, l):
    m = 2**l; m1 = m + 1; m2 = m - 1
    k = range(0, 2**S2CellId.kMaxLevel, 2**(S2CellId.kMaxLevel - l))
    x = np.empty((m1,m1)); y = np.empty((m1,m1)); z = np.empty((m1,m1));
    for i in range(m):
        for j in range(m):
            c = S2Cell(S2CellId.FromFaceIJ(f, k[i], k[j]).parent(l))
            v = c.GetVertex(0)
            x[i][j] = v.x(); y[i][j] = v.y(); z[i][j] = v.z();
            if i == m2:
                v = c.GetVertex(1)
                x[m][j] = v.x(); y[m][j] = v.y(); z[m][j] = v.z();
            if j == m2:
                v = c.GetVertex(3)
                x[i][m] = v.x(); y[i][m] = v.y(); z[i][m] = v.z();
            if i == m2 and j == m2:
                v = c.GetVertex(2)
                x[m][m] = v.x(); y[m][m] = v.y(); z[m][m] = v.z();
    return (x, y, z)

fig = plt.figure()
ax = fig.add_subplot(projection='3d')

for f in range(6):
    s = surf4face(f, 2)
    ax.plot_surface(*s, color='blue')
    ax.plot_wireframe(*s, color='green')

fig.savefig('s2sphere.svg')