import itertools
 
import matplotlib.pyplot as plt
import numpy as np
from IPython import embed
 
if __name__ == '__main__':
 
    V = [
        (0, 0.1),  #  dim 1
        (0, 0.05, 0.1),  #  dim 2
        (-0.05, 0, 0.05, 0.1),  #  dim 3
        (-0.1, -0.05, 0, 0.05, 0.1),  #  dim 4
        (-0.01, 0.0, 0.1),  #  dim 5
        (-0.01, 0.0, 0.1),  #  dim 6
    ]
    shape = [len(v) for v in V]
 
    X = np.array(list(itertools.product(*(range(a) for a in shape))))
 
    hparams = [[v[x[i]] for i, v in enumerate(V)] for x in X]
 
    width, height = shape[0::2], shape[1::2]
    buffer = np.zeros((np.product(height), np.product(width)))
 
    def func(xs):
        return np.sum(xs)
 
    def get_idx(idx):
        assert len(idx) % 2 == 0
 
        y, x = 0, 0
        wi, hi = idx[0::2], idx[1::2]
        # x = idx[0] * width[-1] + idx[2] * 1
        # y = idx[1] * height[-1] + idx[3] * 1
 
        # x = wi[-2] * width[-1] + wi[-1] * 1
        # y = hi[-2] * height[-1] + idx[-1] * 1
 
        # x = wi[-3] * width[-2] * width[-1] + wi[-2] * width[-1] + wi[-1] * 1
        # y = hi[-3] * height[-2] * height[-1] + hi[-2] * height[-1] + idx[-1] * 1
        for level in range(len(idx) // 2):
            dt = np.product(width[-level:]) if level > 0 else 1
            x += wi[-(level+1)] * dt
            dt = np.product(height[-level:]) if level > 0 else 1
            y += hi[-(level+1)] * dt
        return y, x
 
    for x, hparam in zip(X, hparams):
        buffer[get_idx(x)] = func(hparam)
 
    plt.imshow(buffer)
 
    for level in range(len(height)-1):
        dt = np.product(height[-(level+1):])
        for j in range(buffer.shape[0]):
            if j % dt == 0:
                plt.axhline(y=j-0.5, color='k', linewidth=level+1)
 
    for level in range(len(width)-1):
        d = np.product(width[-(level+1):])
        for j in range(buffer.shape[1]):
            if j % d == 0:
                plt.axvline(x=j-0.5, color='k', linewidth=level+1)
 
    plt.show()