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


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(somR, somC)
    for r, c in itertools.product(range(somR), range(somC)):
        tmp = np.append(som[r][c], [0])
        heatMap = np.reshape(tmp, (6, 4))

        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 forearm
    """
    som = loadSom(12, 9, "lavepred4")
    showSom(som)
