import numpy as np
import json, random

try:
    from .mlp.mlp import Mlp
    from .bal.mybal import MyBal
    from .som.process_hand import loadSom, loadMrfSom
    from .proces.processdata import process
except Exception: #ImportError
    from proces.processdata import process
    from mlp.mlp import Mlp
    from bal.mybal import MyBal
    from som.som.som import SOM
    from som.mrf.mrfsom import MRFSOM
    from som.process_hand import loadSom, loadMrfSom


def filt(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}
            # res.append((pos, touch))
        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
            # res.append((pos, touch))
        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"]))
        for i in range(obj["s"]):
            #res.append((obj["p"], obj["t"]))
            pass

    return res



class ComplexModel:

    def __init__(self, filterModel, proprioLeftArmModel, proprioRighttArmModel, touchHandModel, touchForearmModel,
                 leftHandAsoc, rightHandAsoc, leftForearmAsoc, rightForeArmAsoc):
        # filter
        self.filt = filterModel
        # arms SOM models
        self.lSom = proprioLeftArmModel
        self.rSom = proprioRighttArmModel
        # touch MRF models
        self.hand = touchHandModel
        self.forearm = touchForearmModel
        # associators
        self.leftHandAsoc = leftHandAsoc
        self.rightHandAsoc = rightHandAsoc
        self.leftForearmAsoc = leftForearmAsoc
        self.rightForeArmAsoc = rightForeArmAsoc


    def test(self, iCubData):
        res = {}
        processedData = process(iCubData)[0]

        isTouch = sum(processedData["touch"]) != 0
        prediction = self.filt.predict(np.array([processedData["proprio"]]))
        predictTouch = sum(np.array([0 if i < 0.5 else 1 for i in prediction])) != 0
        if isTouch != predictTouch:
            res["filt"] = False
            res["ok"] = False
            return res
        else:
            res["filt"] = True
            res["ok"] = True
            if not isTouch:
                return res

        proprioLeft  = self.lSom.winnerVector(processedData["leftHandPro"])
        proprioRight = self.rSom.winnerVector(processedData["rightHandPro"])

        ok = True

        # hands
        leftHandT  = processedData["leftTouch"][0:9]
        if sum(leftHandT) != 0:
            touchVal = self.hand.winnerVector(leftHandT)
            assocPrediction = self.leftHandAsoc.forward(proprioLeft + proprioRight)
            res["leftHandAssoc"] = np.all(touchVal == assocPrediction)
            ok = ok and res["leftHandAssoc"]

        rightHandT = processedData["rightTouch"][0:9]
        if sum(rightHandT) != 0:
            touchVal = self.hand.winnerVector(rightHandT)
            assocPrediction = self.rightHandAsoc.forward(proprioLeft + proprioRight)
            res["rightHandAssoc"] = np.all(touchVal == assocPrediction)
            ok = ok and res["rightHandAssoc"]

        # forearm
        leftForeT = processedData["leftTouch"][9:32]
        if sum(leftForeT) != 0:
            touchVal = self.forearm.winnerVector(leftForeT)
            assocPrediction = self.leftForearmAsoc.forward(proprioLeft + proprioRight)
            res["leftForeAssoc"] = np.all(touchVal == assocPrediction)
            ok = ok and res["leftForeAssoc"]

        rightForeT = processedData["rightTouch"][9:32]
        if sum(rightForeT) != 0:
            touchVal = self.forearm.winnerVector(rightForeT)
            assocPrediction = self.rightForeArmAsoc.forward(proprioLeft + proprioRight)
            res["rightForeAssoc"] = np.all(touchVal == assocPrediction)
            ok = ok and res["rightForeAssoc"]


        res["ok"] = ok
        return res


if __name__ == "__main__":
    """
    Final test of 50 configurations
    """
    mlp = Mlp([32, 50, 10, 1])
    mlp.load("./mlp/trained/final")

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

    lh = MyBal(216, 70, 49)
    lh.load("./bal/trained/lh1")
    rh = MyBal(216, 70, 49)
    rh.load("./bal/trained/rh1")
    lf = MyBal(216, 70, 108)
    lf.load("./bal/trained/lf1")
    rf = MyBal(216, 70, 108)
    rf.load("./bal/trained/rf1")


    finalModel = ComplexModel(
        filterModel=mlp,
        proprioLeftArmModel=leftSom,
        proprioRighttArmModel=rightSom,
        touchHandModel=handMrf,
        touchForearmModel=forearmMrf,
        leftHandAsoc=lh,
        rightHandAsoc=rh,
        leftForearmAsoc=lf,
        rightForeArmAsoc=rf
    )

    with open('./test_notouch.rawdata') as f:
        testData = json.load(f)
    with open('./test_touch.rawdata') as f:
        testData += json.load(f)
    errNT = []
    for iCubData, i in zip(testData, range(len(testData))):
        res = finalModel.test(iCubData)
        if not res["ok"]:
            errNT.append((i, res))

    for i, bad in errNT:
        print("num.", i, "error", bad)
    print("ok", len(testData)-len(errNT), "bad", len(errNT), "succ", 100 - (len(errNT) / len(testData)) * 100, "%")
