import numpy as np
import json

try:
    from som.som.som import SOM
    from som.mrf.mrfsom import MRFSOM
except Exception: #ImportError
    from som.som import SOM
    from mrf.mrfsom import MRFSOM


def loadSom(dim, rows, cols, folderName):
    model = SOM(dim, rows, cols)
    for i in range(rows):
        model.weights[i] = np.loadtxt(folderName+"/"+str(i)+".txt")
    return model


def loadMrfSom(dim, rows, cols, folderName):
    model = MRFSOM(dim, rows, cols)
    for i in range(rows):
        model.weights[i] = np.loadtxt(folderName+"/"+str(i)+".txt")
    return model


def toOneHot(x, y, maxI, maxJ):
    res =[]
    for i in range(maxI):
        for j in range(maxJ):
            if x == i and y == j:
                res.append(1)
            else:
                res.append(0)
    return res

if __name__ == "__main__":
    """
    Transform ../data.data to 4 somdata files containg transformed configurations and touches
    new files are in som_data, for association training we need to copy them to bal/trainingdata
    """

    leftSom  = loadSom(16, 9, 12, "./trained/left_hand")
    rightSom = loadSom(16, 9, 12, "./trained/right_hand")
    handMrf = loadMrfSom(9, 7, 7, "./trained/h_touch")
    forearmMrf = loadMrfSom(23, 12, 9, "./trained/fore_touch")

    # load data
    f = open("../data.data")
    rawData = json.load(f)#[0:13]
    touch = np.array([i["touch"] for i in rawData])
    left = np.array([i["leftHandPro"] for i in rawData])
    right = np.array([i["rightHandPro"] for i in rawData])

    # process data
    processedLh = []
    processedRh = []
    processedLf = []
    processedRf = []
    for l, r, t in zip(left, right, touch):
        if sum(t) == 0:
            continue
        # left prop
        a = leftSom.winnerVector(l)

        # right prop
        b = rightSom.winnerVector(r)

        # left hand
        cropT = t[0:9]
        if sum(cropT) != 0:
            lh = handMrf.winnerVector(cropT)
            processedLh.append({"pos":a+b, "touch": lh})

        # right hand
        cropT = t[32:41]
        if sum(cropT) != 0:
            rh = handMrf.winnerVector(cropT)
            processedRh.append({"pos":a+b, "touch": rh})

        # left fore
        cropT = t[9:32]
        if sum(cropT) != 0:
            lf = forearmMrf.winnerVector(cropT)
            processedLf.append({"pos":a+b, "touch": lf})

        # right fore
        cropT = t[41:64]
        if sum(cropT) != 0:
            rf = forearmMrf.winnerVector(cropT)
            processedRf.append({"pos":a+b, "touch": rf})

    # save data
    with open('./som_data/lh.somdata', "w") as out:
        out.write(str(json.dumps(processedLh, indent=4)))

    with open('./som_data/rh.somdata', "w") as out:
        out.write(str(json.dumps(processedRh, indent=4)))

    with open('./som_data/lf.somdata', "w") as out:
        out.write(str(json.dumps(processedLf, indent=4)))

    with open('./som_data/rf.somdata', "w") as out:
        out.write(str(json.dumps(processedRf, indent=4)))
    print("done")
