import json, random
import numpy as np
import matplotlib.pyplot as plt

try:
    from .mybal import MyBal
except Exception:
    from mybal import MyBal


def filterAmbiguousTouch(labeld):
    res = []
    seen = {}
    for pos, touch in labeld:
        tmp = np.array(pos)
        idx1 = np.argmax(tmp)
        tmp[idx1] = 0
        idx2 = np.argmax(tmp)
        key = idx1 * 1000 + idx2
        val = np.argmax(touch)

        if key not in seen:
            seen[key] = {val: {"p": pos, "t": touch, "s": 1}, "max": 1, "maxP": val}
        elif val in seen[key]:
            seen[key][val]["s"] += 1
            numb = seen[key][val]["s"]
            if numb > seen[key]["max"]:
                seen[key]["max"] = numb
                seen[key]["maxP"] = val
        else:
            seen[key][val] = {"p": pos, "t": touch, "s": 1}

    for key in seen.keys():
        obj = seen[key]
        obj = obj[obj["maxP"]]
        res.append((obj["p"], obj["t"]))
        #res.append((obj["p"], obj["t"]))
        for i in range(obj["s"]):
            # res.append((obj["p"], obj["t"]))
            pass

    return res


if __name__ == "__main__":
    """
    Training associators on transformed data in ./trainingdata/[trainingName].somdata
    """
    trainingName = "lh"
    trainingName = "rh"
    # trainingName = "lf"
    # trainingName = "rf"
    finalTraining = False
    print("final" if finalTraining else "test", "part", trainingName)

    with open("./trainingdata/" + trainingName + ".somdata") as f:
        data = json.load(f)


    labeled = [(np.array(i['pos']), np.array(i["touch"])) for i in data]
    labeled = filterAmbiguousTouch(labeled)
    random.shuffle(labeled)

    if finalTraining:
        line = len(labeled)
    else:
        line = int(len(labeled) * 0.75)
    trainig = labeled[0:line]
    test = labeled[line:]

    inp, hid, out = len(labeled[0][0]), 110, len(labeled[0][1])
    model = MyBal(inp, hid, out, alpha=0.25)
    res = model.train(trainig, epoch=1000)

    plt.show()
    replot = plt.subplot(311)
    replot.set_title("mean squared error")
    replot.plot([i for i, _, _ in res])

    errplot = plt.subplot(313)
    errplot.set_title("classification error [%]")
    errplot.set_ylim([-5, 105])
    errplot.plot([i for _, (i, _), _ in res])

    plt.show()

    err = 0
    for X, y in test:
        prediction = model.forward(X)
        if not np.all(prediction == y):
            err += 1

    if len(test) != 0:
        print("test: succ", 100 - ((err / len(test))*100),"%")

    if finalTraining:
        model.save("./trained/" + trainingName)
        model = MyBal(inp, hid, out)
        model.load("./trained/" + trainingName)
        err = 0
        for X, y in trainig:
            prediction = model.forward(X)
            if not np.all(prediction == y):
                err += 1

        print("test on all data: succ", 100 - ((err / max(1, len(trainig))) * 100), "%")

