import numpy as np
import matplotlib.pyplot as plt
import itertools

handPos = [
    (1,1),
    (1,3),
    (1,5),
    (2,6),
    (3,0),
    (3,3),
    (3,4),
    (4,3),
    (4,4),
]



def loadSom(rows, cols, folderName):
    model = [0] * rows
    for i in range(rows):
        model[i] = np.loadtxt("./data/"+folderName+"/"+str(i)+".txt")
    return np.array(model)


def showSom(som):
    somR, somC, _ = som.shape
    f, ax = plt.subplots(somC, somR)
    for r, c in itertools.product(range(somR), range(somC)):
        heatMap = np.ones((6,7)) * -0.5
        weight = som[r][c]
        # hand
        for i in range(9):
            x, y = handPos[i][0], handPos[i][1]
            heatMap[x][y] = np.floor(weight[i] * 10) / 10

        ax[r][c].set_yticklabels([])
        ax[r][c].set_xticklabels([])
        ax[r][c].imshow(heatMap)
    plt.show()

if __name__ == "__main__":
    """
    Visualization of trained MRF-SOM for hand
    """
    som = loadSom(7, 7, "h_touch")
    showSom(som)
