from AttributePool import AttributePool
from Capsule import Capsule
from PrimitivesRenderer import Primitives
from Observation import Observation
from MetaLearner import MetaLearner
from Utility import Utility

import copy
import math
import numpy as np
import matplotlib.patches as patches

import time

class CapsuleNetwork:

    def __init__(self, name : str):
        self._primitiveCapsules  : list          = []               # Capsule
        self._semanticCapsules   : list          = []               # Capsule
        self._pixelCapsules      : dict          = {}               # Shape Tuple - Capsule
        self._capsulePrimitive   : dict          = {}               # Primitive Capsule - Primitive Type
        self._semanticLayers     : dict          = {}               # Layer Index - Capsule List
        self._numSemanticLayers  : int           = 0                
        self._attributePool      : AttributePool = AttributePool()
        self._renderer                           = None             # PrimitivesRenderer Instance
        self._metaLearner        : MetaLearner   = MetaLearner()
        self._capsuleCount       : int           = 0
        self._name               : str           = name


        # Adding all Meta Learner Lambdas:
        # 1. Observed Axioms have same $\Omega$ as parent
        self._metaLearner.addLambda(lambda obs, axioms : len(self.findSameParents(list(axioms.keys()))) > 0)
        # 2. Observed Axioms don't have same $\Omega$ as parent
        self._metaLearner.addLambda(lambda obs, axioms : len(self.findSameParents(list(axioms.keys()))) == 0)
        # 3. Parts are tracked from previous scenes
        self._metaLearner.addLambda(lambda obs, axioms : \
            Utility.andElements([Utility.andElements([obsItem.hasPreviousObservation() for obsItem in obsList]) for obsList in axioms.values()]))
        # 4. Parts are NOT tracked from previous scenes
        self._metaLearner.addLambda(lambda obs, axioms : \
            not Utility.andElements([Utility.andElements([obsItem.hasPreviousObservation() for obsItem in obsList]) for obsList in axioms.values()]))
        
        # 5. $\Omega: Z(\vec{\alpha}, \vec{\tilde{\alpha}})$ indicates one attribute mismatch \\ with no entry in memory $\alpha^i >\epsilon$
        # TODO: self._metaLearner.addLambda(lambda obs, axioms : self.agreementOfMostLikelyParent(axioms))
        # 6. $\Omega: Z(\vec{\alpha}, \vec{\tilde{\alpha}})$ indicates attribute mismatch \\ for (position, rotation, size) only
        # TODO:


    def getJSON(self):
        # Only needs to save semantic capsules and data
        semanticData = []
        for layerID in range(self._numSemanticLayers):
            layerCaps = []
            for caps in self._semanticLayers[layerID]:
                layerCaps.append(caps.getJSON())

            semanticData.append({"semanticCapsules" : layerCaps})

        metaLearnerData = self._metaLearner.getJSON()

        return {"semanticLayers" : semanticData, "metaLearner" : metaLearnerData}


    def putJSON(self, data):
        # Only needs to load semantic capsules and data
        semCaps = {}
        for layerData in data["semanticLayers"]:
            for capsData in layerData["semanticCapsules"]:

                capsName = capsData["name"]
                
                obsList = []
                # First Route saved/loaded independently, as it is required to create the
                # Capsule.
                for obsData in capsData["firstRouteObservations"]:
                    obsCaps = self.getCapsuleByName(obsData["name"])
                    obsRoute = obsCaps.getRouteByName(obsData["route"])
                    obsProb = obsData["probability"]

                    attrDict = {}
                    for attrData in obsData["attributes"]:
                        attrDict[obsCaps.getAttributeByName(attrData["attribute"])] = attrData["value"]

                    obsList.append(Observation(obsCaps, obsRoute, [], attrDict, obsProb))
                
                semCaps[capsName] = self.addSemanticCapsule(capsName, obsList, 0)

                # Adding remaining memory
                if "remainingMemory" in capsData:
                    semCaps[capsName].putJSONMemory(capsData["remainingMemory"], self._attributePool, lambda name : self.getCapsuleByName(name))
        
        if "metaLearner" in data:
            self._metaLearner.putJSON(data["metaLearner"])

        return semCaps  # List of Semantic Capsules



    def getName(self):
        return self._name


    def getCapsuleByName(self, name : str):
        for caps in self._primitiveCapsules:
            if caps.getName() == name:
                return caps
        for caps in self._semanticCapsules:
            if caps.getName() == name:
                return caps
        return None


    def agreementOfMostLikelyParent(self, observedAxioms : dict):
        # observedAxioms    # {Observed Axiom (capsule), List of Observations}
        potentialParents = self.findSameParents(list(observedAxioms.keys()))

        maxProbability = 0.0
        maxAgreement = {}
        maxParent = None
        for parent in potentialParents:
            currentProbability, currentAgreement = parent.getMaxAgreement(observedAxioms)
            if currentProbability > maxProbability:
                maxParent = parent
                maxProbability = currentProbability
                maxAgreement = currentAgreement

        return maxAgreement


    def findSameParents(self, capsules : list):
        # capsules  # List of Capsules
        sameParents = []

        for caps in self._semanticCapsules:
            if caps not in capsules:
                if caps.haveSameParent(capsules) is True:
                    sameParents.append(caps)
        
        return sameParents


    def getShapeByPixelCapsule(self, capsule : Capsule):
        for shape, pixelCaps in self._pixelCapsules.items():
            if capsule == pixelCaps:
                return shape
        return (0, 0)

    
    def setRenderer(self, rendererClass):
        # rendererClass     # Class PrimitivesRenderer
        self._renderer = rendererClass(self._attributePool)

    
    def getRenderer(self):
        return self._renderer


    def addPrimitiveCapsule(self, primitive : Primitives, filterShapes : list, additionalTraining : int = 0):
        # filterShapes      # List of Tuples (width, height)
        
        if self._renderer is None:
            print("No Renderer of Type PrimitivesRenderer defined")
            return
        currentCapsule = Capsule(str(primitive), self._capsuleCount)
        self._capsulePrimitive[currentCapsule] = primitive
        self._capsuleCount = self._capsuleCount + 1
        self._renderer.createAttributesForPrimitive(primitive, currentCapsule, self._attributePool)

        for filterShape in filterShapes:
            pixelCapsule = None

            if filterShape not in self._pixelCapsules:
                pixelCapsule = Capsule("PixelLayer-" + str(filterShape[0]) + "-" + str(filterShape[1]), -1)
                self._renderer.createAttributesForPixelLayer(filterShape[0], filterShape[1], pixelCapsule, self._attributePool)
                self._pixelCapsules[filterShape] = pixelCapsule
            else:
                pixelCapsule = self._pixelCapsules[filterShape]

            currentCapsule.addPrimitiveRoute(pixelCapsule, self._renderer, primitive)
        
        self._primitiveCapsules.append(currentCapsule)

        if additionalTraining > 0:
            for i in range(additionalTraining):
                print("-------------------------------- ADDITIONAL TRAINING ROUND " + str(i + 1) + " OF " + str(additionalTraining) + " --------------------------------")
                currentCapsule.continueTraining(True, [0])

        return currentCapsule


    def addSemanticCapsule(self, name : str, fromObservations : list, additionalTraining : int = 0):
        # fromObservations          # List of Observations for one Occurance
        currentCapsule = Capsule(name, self._capsuleCount)
        self._capsuleCount = self._capsuleCount + 1
        currentCapsule.addSemanticRoute(fromObservations, self._attributePool)
        
        maxLayerID = -1
        for obs in fromObservations:
            currentLayer = self.getLayerIndex(obs.getCapsule())
            if currentLayer > maxLayerID:
                maxLayerID = currentLayer

        if self._numSemanticLayers <= maxLayerID + 1:
            self._semanticLayers[self._numSemanticLayers] = []
            self._numSemanticLayers = self._numSemanticLayers + 1

        self._semanticLayers[maxLayerID + 1].append(currentCapsule)
        self._semanticCapsules.append(currentCapsule)

        if additionalTraining > 0:
            for i in range(additionalTraining):
                print("-------------------------------- ADDITIONAL TRAINING ROUND " + str(i + 1) + " OF " + str(additionalTraining) + " --------------------------------")
                currentCapsule.continueTraining(True, [0])

        return currentCapsule


    def addSemanticTraining(self, name : str, fromObservations : list, additionalTraining : int = 0):
        targetCaps = self.findCapsule(name)
        if targetCaps is None:
            return None

        targetCaps.trainSemanticRoute(fromObservations)
        return targetCaps

    
    def addAttribute(self, nameCaps : str, nameAttr : str, fromObservations : list, additionalTraining : int = 0):
        targetCaps = self.findCapsule(nameCaps)
        if targetCaps is None:
            return None

        targetCaps.addSemanticAttribute(fromObservations, nameAttr, self._attributePool)

        for caps in self._semanticCapsules:
            caps.inheritAttributes(self._attributePool)

        return targetCaps


    def addAttributeTraining(self, nameCaps : str, nameAttr : str, fromObservations : list, additionalTraining : int = 0):
        targetCaps = self.findCapsule(nameCaps)
        if targetCaps is None:
            return None

        targetCaps.trainSemanticAttribute(fromObservations, nameAttr)
        return targetCaps


    def findCapsule(self, name : str):
        for caps in self._semanticCapsules:
            if name.lower() in caps.getName().lower():
                return caps
        
        return None


    def getLayerIndex(self, capsule : Capsule):
        for layerID in range(self._numSemanticLayers):
            if capsule in self._semanticLayers[layerID]:
                return layerID
        
        return -1


    def getAttributePool(self):
        return self._attributePool

    
    def clearAllObservations(self):
        for filterShape, capsule in self._pixelCapsules.items():
            capsule.clearObservations()
        for capsule in self._primitiveCapsules:
            capsule.clearObservations()
        for capsule in self._semanticCapsules:
            capsule.clearObservations()
            

    def findObservedAxioms(self, observations : dict):
        # observations      # {Capsule, List of Observations}

        observedAxioms = {}

        # Not the fastest check, but works...
        for caps, obsList in observations.items():
            for obs in obsList:
                foundParent = False
                for checkObs in [x for checkObsList in observations.values() for x in checkObsList] :
                    if checkObs.isParent(obs):
                        foundParent = True
                        break
                if foundParent is False:
                    if caps in observedAxioms:
                        observedAxioms[caps].append(obs) 
                    else:
                        observedAxioms[caps] = [obs] 

        return observedAxioms # {Observed Axioms (Capsule), List of Observations}

    
    def showInput(self, image : list, width : int, height : int, stepSize : int = 1):
        # image             # Linear List of Pixels

        startTime = time.time()
        self.clearAllObservations()

        print("Capsule Network shown an Image of size (" + str(width) + ", " + str(height) + ")")

        offsetLabelX, offsetLabelY, offsetLabelRatio, targetLabelX, targetLabelY, targetLabelSize = self._renderer.getOffsetLabels()
        
        for filterShape, capsule in self._pixelCapsules.items():
            if filterShape[0] <= width and filterShape[1] <= height:

                fsWidth = float(filterShape[0])
                fsHeight = float(filterShape[1])
                fsMaxW = float(max(width, height))

                rads = np.fromfunction(lambda xx, yy : np.sqrt((xx / fsWidth - 0.5) ** 2 + (yy / fsHeight - 0.5) ** 2), (filterShape[0], filterShape[1]), dtype=float)
                angs = np.fromfunction(lambda xx, yy : np.arctan2((xx / fsWidth - 0.5), (yy / fsHeight - 0.5)) % (math.pi * 2), (filterShape[0], filterShape[1]), dtype=float)

                for offsetX in range(0, width - filterShape[0] + 1, stepSize):
                    for offsetY in range(0, height - filterShape[1] + 1, stepSize):
                        attributes = {}
                        for xx in range(filterShape[0]):
                            for yy in range(filterShape[1]):
                                attributes[capsule.getAttributeByName("PixelC-" + str(xx) + "-" + str(yy))] = image[((yy + offsetY) * width + xx + offsetX) * 4]
                                attributes[capsule.getAttributeByName("PixelR-" + str(xx) + "-" + str(yy))] = rads[xx, yy]
                                attributes[capsule.getAttributeByName("PixelA-" + str(xx) + "-" + str(yy))] = angs[xx, yy]
                                attributes[capsule.getAttributeByName("PixelD-" + str(xx) + "-" + str(yy))] = 1.0

                        attributes[capsule.getAttributeByName(offsetLabelX)] = float(offsetX) / fsMaxW
                        attributes[capsule.getAttributeByName(offsetLabelY)] = float(offsetY) / fsMaxW
                        attributes[capsule.getAttributeByName(offsetLabelRatio)] = fsWidth / fsMaxW

                        pixelObs = Observation(capsule, None, [], attributes, 1.0)
                        capsule.addPixelObservation(pixelObs)

        passedTime = time.time() - startTime
        print("Beginning Forward Pass on all Primitive Capsules (Time Passed: " + str(passedTime) + "s)")

        allObs = {}     # Capsule - List Of Observations

        for capsule in self._primitiveCapsules:
            capsule.forwardPass()
            capsule.cleanupObservations(offsetLabelX, offsetLabelY, offsetLabelRatio, targetLabelX, targetLabelY, targetLabelSize)
            allObs[capsule] = capsule.getObservations()

        for layer in range(self._numSemanticLayers):
            passedTime = time.time() - startTime
            print("Beginning Forward Pass on Layer " + str(layer) + " of Semantic Capsules (Time Passed: " + str(passedTime) + "s)")
            for capsule in self._semanticLayers[layer]:
                capsule.forwardPass()
                capsule.cleanupObservations()
                allObs[capsule] = capsule.getObservations()

        recommendation = None
        observedAxioms = self.findObservedAxioms(allObs)
        if len(observedAxioms) > 1:
            recommendation = self._metaLearner.checkResults(allObs, observedAxioms)

        passedTime = time.time() - startTime
        print("Total Time Passed: " + str(passedTime) + "s")

        return allObs, recommendation   # Capsule - List Of Observations, Recommendation String


    def applyOracle(self, oracleDecision : int):
        # TODO: Apply oracle decisions here?
        self._metaLearner.applyOracle(oracleDecision)

        
    def generateImage(self, width : int, height : int, observations : dict, withBackground : bool = False):
        # observations          # Capsule   -   List of Observations

        semantics = {} #  Observation - List of Patches
        texts = []     #  List of (X, Y, Text)
        offsetLabelX, offsetLabelY, offsetLabelRatio, targetLabelX, targetLabelY, targetLabelSize = self._renderer.getOffsetLabels()
        

        # We make a full copy to work on
        obs = {}    # Capsule - List of Observations
        obsMap = {} # Copied Observation - Actual Observation
        for capsule, obsList in observations.items():
            obs[capsule] = []
            for observation in obsList:
                obs[capsule].append(Observation(capsule, observation.getTakenRoute(), observation.getInputObservations(), 
                                                observation.getOutputs(), observation.getProbability()))
                obsMap[obs[capsule][-1]] = observation

        # Generate Semantic Labels for Semantic Capsules
        for capsule, obsList in obs.items():
            if capsule in self._semanticCapsules:
                for observation in obsList:
                    xOffset1 = observation.getOutput(capsule.getAttributeByName(targetLabelX))
                    yOffset1 = observation.getOutput(capsule.getAttributeByName(targetLabelY))
                    
                    xOffset1 = int(xOffset1 * float(max(width, height)))
                    yOffset1 = int(yOffset1 * float(max(width, height)))

                    semantics[obsMap[observation]] = []
    
                    for inObs in observation.getInputObservations():
                        xOffset2 = inObs.getOutput(inObs.getCapsule().getAttributeByName(targetLabelX))
                        yOffset2 = inObs.getOutput(inObs.getCapsule().getAttributeByName(targetLabelY))
                        
                        xOffset2 = int(xOffset2 * float(max(width, height)))
                        yOffset2 = int(yOffset2 * float(max(width, height)))
                        
                        semantics[obsMap[observation]].append(patches.Arrow(xOffset1, yOffset1, xOffset2 - xOffset1, yOffset2 - yOffset1, linewidth = 1, edgecolor = 'r', facecolor = 'none' ))
                        
                    semantics[obsMap[observation]].append(patches.Circle((xOffset1, yOffset1), radius = 1, color = 'r'))    


        for layerIndex in range(self._numSemanticLayers - 1, -1, -1):
            # Iterate the list backwards
            for capsule in self._semanticLayers[layerIndex]:
                if capsule in obs.keys():
                    for observation in obs[capsule]:
                        # Only top level observed symbols that have not yet been generated
                        if not observation.getInputObservations():
                            newObslist = capsule.backwardPass(observation, False)
                            for obsCaps, newObs in newObslist.items():
                                if obsCaps in obs:
                                    obs[obsCaps].extend(newObs)
                                else:
                                    obs[obsCaps] = newObs

                    # Remove the parsed Observations
                    del obs[capsule]


        # Order all observations for the primitive capsules
        capsObsPairs = []
        for capsule, obsList in obs.items():
            if capsule in self._primitiveCapsules:
                for observation in obsList:
                    capsObsPairs.append((capsule, observation))
        capsObsPairs = sorted(capsObsPairs, key=lambda tup: tup[1].getProbability())


        image = np.zeros(width * height * 4)

        for capsule, observation in capsObsPairs:            
            pixelShape = self.getShapeByPixelCapsule(capsule.getPixelLayerInput())

            xOffset = observation.getOutput(capsule.getAttributeByName(targetLabelX))
            yOffset = observation.getOutput(capsule.getAttributeByName(targetLabelY))
            size = observation.getOutput(capsule.getAttributeByName(targetLabelSize))

            observation.setOutput(capsule.getAttributeByName(targetLabelX), 0.5)
            observation.setOutput(capsule.getAttributeByName(targetLabelY), 0.5)
            observation.setOutput(capsule.getAttributeByName(targetLabelSize), size * float(max(width, height)) / float(pixelShape[0]))

            obsPixelLayer = capsule.backwardPass(observation, False)

            pixelObs = list(obsPixelLayer.values())[0][0]
            pixelLay = list(obsPixelLayer.keys())[0]

            newxOffset = int(xOffset * float(max(width, height))) - int(pixelShape[0] / 2)
            newyOffset = int(yOffset * float(max(width, height))) - int(pixelShape[1] / 2)

            # We actually have far more accurate segmentation (including rotation, etc), but its hard to do nicely in matplotlib,
            # so we decided to just box it roughly.
            minX = newxOffset + pixelShape[0]
            maxX = newxOffset
            minY = newyOffset + pixelShape[1]
            maxY = newyOffset

            for xx in range(pixelShape[0]):
                for yy in range(pixelShape[1]):
                    if xx + newxOffset < width and xx + newxOffset >= 0 and yy + newyOffset < height and yy + newyOffset >= 0:
                        depth = pixelObs.getOutput(pixelLay.getAttributeByName("PixelD-" + str(xx) + "-" + str(yy)))
                        
                        if depth < 1.0:
                            minX = min(xx + newxOffset, minX)
                            maxX = max(xx + newxOffset, maxX)
                            minY = min(yy + newyOffset, minY)
                            maxY = max(yy + newyOffset, maxY)

                        if withBackground is True or depth < 1.0:
                            image[((yy + newyOffset) * width + xx + newxOffset) * 4] = pixelObs.getOutput(pixelLay.getAttributeByName("PixelC-" + str(xx) + "-" + str(yy)))

            if observation in obsMap:
                semantics[obsMap[observation]] = [patches.Rectangle((minX, minY), maxX - minX, maxY - minY, linewidth = 1, edgecolor = 'y', facecolor = 'none')]
            else:
                semantics[observation] = [patches.Rectangle((minX, minY), maxX - minX, maxY - minY, linewidth = 1, edgecolor = 'y', facecolor = 'none')]

            observation.setOutput(capsule.getAttributeByName(targetLabelX), xOffset)
            observation.setOutput(capsule.getAttributeByName(targetLabelY), yOffset)
            observation.setOutput(capsule.getAttributeByName(targetLabelSize), size)

            texts.append((minX, minY, capsule.getName()))

        return image, semantics, texts    # Linear List of Pixels


    def producePrimitiveObservations(self, observation : Observation):
        if observation.getCapsule() not in self._primitiveCapsules:
            outputObsList = []
            if not observation.getInputObservations():
                newObsDict = observation.getCapsule().backwardPass(observation, False)
                observation.clearInputObservations()
                for newObsList in newObsDict.values():
                    for newObs in newObsList:
                        outputObsList = outputObsList + self.producePrimitiveObservations(newObs)
            else:
                for newObs in observation.getInputObservations():
                    outputObsList = outputObsList + self.producePrimitiveObservations(newObs)
            return outputObsList    # List of Observations for Primitive Capsules
        else:
            return [observation]    # List of Observation for this Primitive Capsules


    def distance(self, observationA : Observation, observationB : Observation):
        distance = 100.0
        normal1 = [0.0, 0.0]
        normal2 = [0.0, 0.0]

        compareA = self.producePrimitiveObservations(observationA)
        compareB = self.producePrimitiveObservations(observationB)

        for primA in compareA:
            for primB in compareB:
                    testDistance, testNormal1, testNormal2 = self._renderer.getDistance(
                        self._capsulePrimitive[primA.getCapsule()], self._capsulePrimitive[primB.getCapsule()],
                        primA.getOutputs(), primB.getOutputs())
                    if testDistance < distance:
                        distance = testDistance
                        normal1 = testNormal1
                        normal2 = testNormal2

        return distance, normal1, normal2

    

cf51c211525d1752c4ecb72616bec27f7bb15495
<mask token>


class CapsuleNetwork:
    <mask token>
    <mask token>

    def putJSON(self, data):
        semCaps = {}
        for layerData in data['semanticLayers']:
            for capsData in layerData['semanticCapsules']:
                capsName = capsData['name']
                obsList = []
                for obsData in capsData['firstRouteObservations']:
                    obsCaps = self.getCapsuleByName(obsData['name'])
                    obsRoute = obsCaps.getRouteByName(obsData['route'])
                    obsProb = obsData['probability']
                    attrDict = {}
                    for attrData in obsData['attributes']:
                        attrDict[obsCaps.getAttributeByName(attrData[
                            'attribute'])] = attrData['value']
                    obsList.append(Observation(obsCaps, obsRoute, [],
                        attrDict, obsProb))
                semCaps[capsName] = self.addSemanticCapsule(capsName,
                    obsList, 0)
                if 'remainingMemory' in capsData:
                    semCaps[capsName].putJSONMemory(capsData[
                        'remainingMemory'], self._attributePool, lambda
                        name: self.getCapsuleByName(name))
        if 'metaLearner' in data:
            self._metaLearner.putJSON(data['metaLearner'])
        return semCaps
    <mask token>
    <mask token>

    def agreementOfMostLikelyParent(self, observedAxioms: dict):
        potentialParents = self.findSameParents(list(observedAxioms.keys()))
        maxProbability = 0.0
        maxAgreement = {}
        maxParent = None
        for parent in potentialParents:
            currentProbability, currentAgreement = parent.getMaxAgreement(
                observedAxioms)
            if currentProbability > maxProbability:
                maxParent = parent
                maxProbability = currentProbability
                maxAgreement = currentAgreement
        return maxAgreement
    <mask token>
    <mask token>

    def setRenderer(self, rendererClass):
        self._renderer = rendererClass(self._attributePool)
    <mask token>

    def addPrimitiveCapsule(self, primitive: Primitives, filterShapes: list,
        additionalTraining: int=0):
        if self._renderer is None:
            print('No Renderer of Type PrimitivesRenderer defined')
            return
        currentCapsule = Capsule(str(primitive), self._capsuleCount)
        self._capsulePrimitive[currentCapsule] = primitive
        self._capsuleCount = self._capsuleCount + 1
        self._renderer.createAttributesForPrimitive(primitive,
            currentCapsule, self._attributePool)
        for filterShape in filterShapes:
            pixelCapsule = None
            if filterShape not in self._pixelCapsules:
                pixelCapsule = Capsule('PixelLayer-' + str(filterShape[0]) +
                    '-' + str(filterShape[1]), -1)
                self._renderer.createAttributesForPixelLayer(filterShape[0],
                    filterShape[1], pixelCapsule, self._attributePool)
                self._pixelCapsules[filterShape] = pixelCapsule
            else:
                pixelCapsule = self._pixelCapsules[filterShape]
            currentCapsule.addPrimitiveRoute(pixelCapsule, self._renderer,
                primitive)
        self._primitiveCapsules.append(currentCapsule)
        if additionalTraining > 0:
            for i in range(additionalTraining):
                print(
                    '-------------------------------- ADDITIONAL TRAINING ROUND '
                     + str(i + 1) + ' OF ' + str(additionalTraining) +
                    ' --------------------------------')
                currentCapsule.continueTraining(True, [0])
        return currentCapsule

    def addSemanticCapsule(self, name: str, fromObservations: list,
        additionalTraining: int=0):
        currentCapsule = Capsule(name, self._capsuleCount)
        self._capsuleCount = self._capsuleCount + 1
        currentCapsule.addSemanticRoute(fromObservations, self._attributePool)
        maxLayerID = -1
        for obs in fromObservations:
            currentLayer = self.getLayerIndex(obs.getCapsule())
            if currentLayer > maxLayerID:
                maxLayerID = currentLayer
        if self._numSemanticLayers <= maxLayerID + 1:
            self._semanticLayers[self._numSemanticLayers] = []
            self._numSemanticLayers = self._numSemanticLayers + 1
        self._semanticLayers[maxLayerID + 1].append(currentCapsule)
        self._semanticCapsules.append(currentCapsule)
        if additionalTraining > 0:
            for i in range(additionalTraining):
                print(
                    '-------------------------------- ADDITIONAL TRAINING ROUND '
                     + str(i + 1) + ' OF ' + str(additionalTraining) +
                    ' --------------------------------')
                currentCapsule.continueTraining(True, [0])
        return currentCapsule

    def addSemanticTraining(self, name: str, fromObservations: list,
        additionalTraining: int=0):
        targetCaps = self.findCapsule(name)
        if targetCaps is None:
            return None
        targetCaps.trainSemanticRoute(fromObservations)
        return targetCaps

    def addAttribute(self, nameCaps: str, nameAttr: str, fromObservations:
        list, additionalTraining: int=0):
        targetCaps = self.findCapsule(nameCaps)
        if targetCaps is None:
            return None
        targetCaps.addSemanticAttribute(fromObservations, nameAttr, self.
            _attributePool)
        for caps in self._semanticCapsules:
            caps.inheritAttributes(self._attributePool)
        return targetCaps

    def addAttributeTraining(self, nameCaps: str, nameAttr: str,
        fromObservations: list, additionalTraining: int=0):
        targetCaps = self.findCapsule(nameCaps)
        if targetCaps is None:
            return None
        targetCaps.trainSemanticAttribute(fromObservations, nameAttr)
        return targetCaps
    <mask token>
    <mask token>
    <mask token>

    def clearAllObservations(self):
        for filterShape, capsule in self._pixelCapsules.items():
            capsule.clearObservations()
        for capsule in self._primitiveCapsules:
            capsule.clearObservations()
        for capsule in self._semanticCapsules:
            capsule.clearObservations()

    def findObservedAxioms(self, observations: dict):
        observedAxioms = {}
        for caps, obsList in observations.items():
            for obs in obsList:
                foundParent = False
                for checkObs in [x for checkObsList in observations.values(
                    ) for x in checkObsList]:
                    if checkObs.isParent(obs):
                        foundParent = True
                        break
                if foundParent is False:
                    if caps in observedAxioms:
                        observedAxioms[caps].append(obs)
                    else:
                        observedAxioms[caps] = [obs]
        return observedAxioms

    def showInput(self, image: list, width: int, height: int, stepSize: int=1):
        startTime = time.time()
        self.clearAllObservations()
        print('Capsule Network shown an Image of size (' + str(width) +
            ', ' + str(height) + ')')
        (offsetLabelX, offsetLabelY, offsetLabelRatio, targetLabelX,
            targetLabelY, targetLabelSize) = self._renderer.getOffsetLabels()
        for filterShape, capsule in self._pixelCapsules.items():
            if filterShape[0] <= width and filterShape[1] <= height:
                fsWidth = float(filterShape[0])
                fsHeight = float(filterShape[1])
                fsMaxW = float(max(width, height))
                rads = np.fromfunction(lambda xx, yy: np.sqrt((xx / fsWidth -
                    0.5) ** 2 + (yy / fsHeight - 0.5) ** 2), (filterShape[0
                    ], filterShape[1]), dtype=float)
                angs = np.fromfunction(lambda xx, yy: np.arctan2(xx /
                    fsWidth - 0.5, yy / fsHeight - 0.5) % (math.pi * 2), (
                    filterShape[0], filterShape[1]), dtype=float)
                for offsetX in range(0, width - filterShape[0] + 1, stepSize):
                    for offsetY in range(0, height - filterShape[1] + 1,
                        stepSize):
                        attributes = {}
                        for xx in range(filterShape[0]):
                            for yy in range(filterShape[1]):
                                attributes[capsule.getAttributeByName(
                                    'PixelC-' + str(xx) + '-' + str(yy))
                                    ] = image[((yy + offsetY) * width + xx +
                                    offsetX) * 4]
                                attributes[capsule.getAttributeByName(
                                    'PixelR-' + str(xx) + '-' + str(yy))
                                    ] = rads[xx, yy]
                                attributes[capsule.getAttributeByName(
                                    'PixelA-' + str(xx) + '-' + str(yy))
                                    ] = angs[xx, yy]
                                attributes[capsule.getAttributeByName(
                                    'PixelD-' + str(xx) + '-' + str(yy))] = 1.0
                        attributes[capsule.getAttributeByName(offsetLabelX)
                            ] = float(offsetX) / fsMaxW
                        attributes[capsule.getAttributeByName(offsetLabelY)
                            ] = float(offsetY) / fsMaxW
                        attributes[capsule.getAttributeByName(offsetLabelRatio)
                            ] = fsWidth / fsMaxW
                        pixelObs = Observation(capsule, None, [],
                            attributes, 1.0)
                        capsule.addPixelObservation(pixelObs)
        passedTime = time.time() - startTime
        print(
            'Beginning Forward Pass on all Primitive Capsules (Time Passed: ' +
            str(passedTime) + 's)')
        allObs = {}
        for capsule in self._primitiveCapsules:
            capsule.forwardPass()
            capsule.cleanupObservations(offsetLabelX, offsetLabelY,
                offsetLabelRatio, targetLabelX, targetLabelY, targetLabelSize)
            allObs[capsule] = capsule.getObservations()
        for layer in range(self._numSemanticLayers):
            passedTime = time.time() - startTime
            print('Beginning Forward Pass on Layer ' + str(layer) +
                ' of Semantic Capsules (Time Passed: ' + str(passedTime) + 's)'
                )
            for capsule in self._semanticLayers[layer]:
                capsule.forwardPass()
                capsule.cleanupObservations()
                allObs[capsule] = capsule.getObservations()
        recommendation = None
        observedAxioms = self.findObservedAxioms(allObs)
        if len(observedAxioms) > 1:
            recommendation = self._metaLearner.checkResults(allObs,
                observedAxioms)
        passedTime = time.time() - startTime
        print('Total Time Passed: ' + str(passedTime) + 's')
        return allObs, recommendation

    def applyOracle(self, oracleDecision: int):
        self._metaLearner.applyOracle(oracleDecision)

    def generateImage(self, width: int, height: int, observations: dict,
        withBackground: bool=False):
        semantics = {}
        texts = []
        (offsetLabelX, offsetLabelY, offsetLabelRatio, targetLabelX,
            targetLabelY, targetLabelSize) = self._renderer.getOffsetLabels()
        obs = {}
        obsMap = {}
        for capsule, obsList in observations.items():
            obs[capsule] = []
            for observation in obsList:
                obs[capsule].append(Observation(capsule, observation.
                    getTakenRoute(), observation.getInputObservations(),
                    observation.getOutputs(), observation.getProbability()))
                obsMap[obs[capsule][-1]] = observation
        for capsule, obsList in obs.items():
            if capsule in self._semanticCapsules:
                for observation in obsList:
                    xOffset1 = observation.getOutput(capsule.
                        getAttributeByName(targetLabelX))
                    yOffset1 = observation.getOutput(capsule.
                        getAttributeByName(targetLabelY))
                    xOffset1 = int(xOffset1 * float(max(width, height)))
                    yOffset1 = int(yOffset1 * float(max(width, height)))
                    semantics[obsMap[observation]] = []
                    for inObs in observation.getInputObservations():
                        xOffset2 = inObs.getOutput(inObs.getCapsule().
                            getAttributeByName(targetLabelX))
                        yOffset2 = inObs.getOutput(inObs.getCapsule().
                            getAttributeByName(targetLabelY))
                        xOffset2 = int(xOffset2 * float(max(width, height)))
                        yOffset2 = int(yOffset2 * float(max(width, height)))
                        semantics[obsMap[observation]].append(patches.Arrow
                            (xOffset1, yOffset1, xOffset2 - xOffset1, 
                            yOffset2 - yOffset1, linewidth=1, edgecolor='r',
                            facecolor='none'))
                    semantics[obsMap[observation]].append(patches.Circle((
                        xOffset1, yOffset1), radius=1, color='r'))
        for layerIndex in range(self._numSemanticLayers - 1, -1, -1):
            for capsule in self._semanticLayers[layerIndex]:
                if capsule in obs.keys():
                    for observation in obs[capsule]:
                        if not observation.getInputObservations():
                            newObslist = capsule.backwardPass(observation, 
                                False)
                            for obsCaps, newObs in newObslist.items():
                                if obsCaps in obs:
                                    obs[obsCaps].extend(newObs)
                                else:
                                    obs[obsCaps] = newObs
                    del obs[capsule]
        capsObsPairs = []
        for capsule, obsList in obs.items():
            if capsule in self._primitiveCapsules:
                for observation in obsList:
                    capsObsPairs.append((capsule, observation))
        capsObsPairs = sorted(capsObsPairs, key=lambda tup: tup[1].
            getProbability())
        image = np.zeros(width * height * 4)
        for capsule, observation in capsObsPairs:
            pixelShape = self.getShapeByPixelCapsule(capsule.
                getPixelLayerInput())
            xOffset = observation.getOutput(capsule.getAttributeByName(
                targetLabelX))
            yOffset = observation.getOutput(capsule.getAttributeByName(
                targetLabelY))
            size = observation.getOutput(capsule.getAttributeByName(
                targetLabelSize))
            observation.setOutput(capsule.getAttributeByName(targetLabelX), 0.5
                )
            observation.setOutput(capsule.getAttributeByName(targetLabelY), 0.5
                )
            observation.setOutput(capsule.getAttributeByName(
                targetLabelSize), size * float(max(width, height)) / float(
                pixelShape[0]))
            obsPixelLayer = capsule.backwardPass(observation, False)
            pixelObs = list(obsPixelLayer.values())[0][0]
            pixelLay = list(obsPixelLayer.keys())[0]
            newxOffset = int(xOffset * float(max(width, height))) - int(
                pixelShape[0] / 2)
            newyOffset = int(yOffset * float(max(width, height))) - int(
                pixelShape[1] / 2)
            minX = newxOffset + pixelShape[0]
            maxX = newxOffset
            minY = newyOffset + pixelShape[1]
            maxY = newyOffset
            for xx in range(pixelShape[0]):
                for yy in range(pixelShape[1]):
                    if (xx + newxOffset < width and xx + newxOffset >= 0 and
                        yy + newyOffset < height and yy + newyOffset >= 0):
                        depth = pixelObs.getOutput(pixelLay.
                            getAttributeByName('PixelD-' + str(xx) + '-' +
                            str(yy)))
                        if depth < 1.0:
                            minX = min(xx + newxOffset, minX)
                            maxX = max(xx + newxOffset, maxX)
                            minY = min(yy + newyOffset, minY)
                            maxY = max(yy + newyOffset, maxY)
                        if withBackground is True or depth < 1.0:
                            image[((yy + newyOffset) * width + xx +
                                newxOffset) * 4] = pixelObs.getOutput(pixelLay
                                .getAttributeByName('PixelC-' + str(xx) +
                                '-' + str(yy)))
            if observation in obsMap:
                semantics[obsMap[observation]] = [patches.Rectangle((minX,
                    minY), maxX - minX, maxY - minY, linewidth=1, edgecolor
                    ='y', facecolor='none')]
            else:
                semantics[observation] = [patches.Rectangle((minX, minY), 
                    maxX - minX, maxY - minY, linewidth=1, edgecolor='y',
                    facecolor='none')]
            observation.setOutput(capsule.getAttributeByName(targetLabelX),
                xOffset)
            observation.setOutput(capsule.getAttributeByName(targetLabelY),
                yOffset)
            observation.setOutput(capsule.getAttributeByName(
                targetLabelSize), size)
            texts.append((minX, minY, capsule.getName()))
        return image, semantics, texts

    def producePrimitiveObservations(self, observation: Observation):
        if observation.getCapsule() not in self._primitiveCapsules:
            outputObsList = []
            if not observation.getInputObservations():
                newObsDict = observation.getCapsule().backwardPass(observation,
                    False)
                observation.clearInputObservations()
                for newObsList in newObsDict.values():
                    for newObs in newObsList:
                        outputObsList = (outputObsList + self.
                            producePrimitiveObservations(newObs))
            else:
                for newObs in observation.getInputObservations():
                    outputObsList = (outputObsList + self.
                        producePrimitiveObservations(newObs))
            return outputObsList
        else:
            return [observation]

    def distance(self, observationA: Observation, observationB: Observation):
        distance = 100.0
        normal1 = [0.0, 0.0]
        normal2 = [0.0, 0.0]
        compareA = self.producePrimitiveObservations(observationA)
        compareB = self.producePrimitiveObservations(observationB)
        for primA in compareA:
            for primB in compareB:
                testDistance, testNormal1, testNormal2 = (self._renderer.
                    getDistance(self._capsulePrimitive[primA.getCapsule()],
                    self._capsulePrimitive[primB.getCapsule()], primA.
                    getOutputs(), primB.getOutputs()))
                if testDistance < distance:
                    distance = testDistance
                    normal1 = testNormal1
                    normal2 = testNormal2
        return distance, normal1, normal2

<mask token>


class CapsuleNetwork:

    def __init__(self, name: str):
        self._primitiveCapsules: list = []
        self._semanticCapsules: list = []
        self._pixelCapsules: dict = {}
        self._capsulePrimitive: dict = {}
        self._semanticLayers: dict = {}
        self._numSemanticLayers: int = 0
        self._attributePool: AttributePool = AttributePool()
        self._renderer = None
        self._metaLearner: MetaLearner = MetaLearner()
        self._capsuleCount: int = 0
        self._name: str = name
        self._metaLearner.addLambda(lambda obs, axioms: len(self.
            findSameParents(list(axioms.keys()))) > 0)
        self._metaLearner.addLambda(lambda obs, axioms: len(self.
            findSameParents(list(axioms.keys()))) == 0)
        self._metaLearner.addLambda(lambda obs, axioms: Utility.andElements
            ([Utility.andElements([obsItem.hasPreviousObservation() for
            obsItem in obsList]) for obsList in axioms.values()]))
        self._metaLearner.addLambda(lambda obs, axioms: not Utility.
            andElements([Utility.andElements([obsItem.
            hasPreviousObservation() for obsItem in obsList]) for obsList in
            axioms.values()]))

    def getJSON(self):
        semanticData = []
        for layerID in range(self._numSemanticLayers):
            layerCaps = []
            for caps in self._semanticLayers[layerID]:
                layerCaps.append(caps.getJSON())
            semanticData.append({'semanticCapsules': layerCaps})
        metaLearnerData = self._metaLearner.getJSON()
        return {'semanticLayers': semanticData, 'metaLearner': metaLearnerData}

    def putJSON(self, data):
        semCaps = {}
        for layerData in data['semanticLayers']:
            for capsData in layerData['semanticCapsules']:
                capsName = capsData['name']
                obsList = []
                for obsData in capsData['firstRouteObservations']:
                    obsCaps = self.getCapsuleByName(obsData['name'])
                    obsRoute = obsCaps.getRouteByName(obsData['route'])
                    obsProb = obsData['probability']
                    attrDict = {}
                    for attrData in obsData['attributes']:
                        attrDict[obsCaps.getAttributeByName(attrData[
                            'attribute'])] = attrData['value']
                    obsList.append(Observation(obsCaps, obsRoute, [],
                        attrDict, obsProb))
                semCaps[capsName] = self.addSemanticCapsule(capsName,
                    obsList, 0)
                if 'remainingMemory' in capsData:
                    semCaps[capsName].putJSONMemory(capsData[
                        'remainingMemory'], self._attributePool, lambda
                        name: self.getCapsuleByName(name))
        if 'metaLearner' in data:
            self._metaLearner.putJSON(data['metaLearner'])
        return semCaps

    def getName(self):
        return self._name
    <mask token>

    def agreementOfMostLikelyParent(self, observedAxioms: dict):
        potentialParents = self.findSameParents(list(observedAxioms.keys()))
        maxProbability = 0.0
        maxAgreement = {}
        maxParent = None
        for parent in potentialParents:
            currentProbability, currentAgreement = parent.getMaxAgreement(
                observedAxioms)
            if currentProbability > maxProbability:
                maxParent = parent
                maxProbability = currentProbability
                maxAgreement = currentAgreement
        return maxAgreement
    <mask token>
    <mask token>

    def setRenderer(self, rendererClass):
        self._renderer = rendererClass(self._attributePool)
    <mask token>

    def addPrimitiveCapsule(self, primitive: Primitives, filterShapes: list,
        additionalTraining: int=0):
        if self._renderer is None:
            print('No Renderer of Type PrimitivesRenderer defined')
            return
        currentCapsule = Capsule(str(primitive), self._capsuleCount)
        self._capsulePrimitive[currentCapsule] = primitive
        self._capsuleCount = self._capsuleCount + 1
        self._renderer.createAttributesForPrimitive(primitive,
            currentCapsule, self._attributePool)
        for filterShape in filterShapes:
            pixelCapsule = None
            if filterShape not in self._pixelCapsules:
                pixelCapsule = Capsule('PixelLayer-' + str(filterShape[0]) +
                    '-' + str(filterShape[1]), -1)
                self._renderer.createAttributesForPixelLayer(filterShape[0],
                    filterShape[1], pixelCapsule, self._attributePool)
                self._pixelCapsules[filterShape] = pixelCapsule
            else:
                pixelCapsule = self._pixelCapsules[filterShape]
            currentCapsule.addPrimitiveRoute(pixelCapsule, self._renderer,
                primitive)
        self._primitiveCapsules.append(currentCapsule)
        if additionalTraining > 0:
            for i in range(additionalTraining):
                print(
                    '-------------------------------- ADDITIONAL TRAINING ROUND '
                     + str(i + 1) + ' OF ' + str(additionalTraining) +
                    ' --------------------------------')
                currentCapsule.continueTraining(True, [0])
        return currentCapsule

    def addSemanticCapsule(self, name: str, fromObservations: list,
        additionalTraining: int=0):
        currentCapsule = Capsule(name, self._capsuleCount)
        self._capsuleCount = self._capsuleCount + 1
        currentCapsule.addSemanticRoute(fromObservations, self._attributePool)
        maxLayerID = -1
        for obs in fromObservations:
            currentLayer = self.getLayerIndex(obs.getCapsule())
            if currentLayer > maxLayerID:
                maxLayerID = currentLayer
        if self._numSemanticLayers <= maxLayerID + 1:
            self._semanticLayers[self._numSemanticLayers] = []
            self._numSemanticLayers = self._numSemanticLayers + 1
        self._semanticLayers[maxLayerID + 1].append(currentCapsule)
        self._semanticCapsules.append(currentCapsule)
        if additionalTraining > 0:
            for i in range(additionalTraining):
                print(
                    '-------------------------------- ADDITIONAL TRAINING ROUND '
                     + str(i + 1) + ' OF ' + str(additionalTraining) +
                    ' --------------------------------')
                currentCapsule.continueTraining(True, [0])
        return currentCapsule

    def addSemanticTraining(self, name: str, fromObservations: list,
        additionalTraining: int=0):
        targetCaps = self.findCapsule(name)
        if targetCaps is None:
            return None
        targetCaps.trainSemanticRoute(fromObservations)
        return targetCaps

    def addAttribute(self, nameCaps: str, nameAttr: str, fromObservations:
        list, additionalTraining: int=0):
        targetCaps = self.findCapsule(nameCaps)
        if targetCaps is None:
            return None
        targetCaps.addSemanticAttribute(fromObservations, nameAttr, self.
            _attributePool)
        for caps in self._semanticCapsules:
            caps.inheritAttributes(self._attributePool)
        return targetCaps

    def addAttributeTraining(self, nameCaps: str, nameAttr: str,
        fromObservations: list, additionalTraining: int=0):
        targetCaps = self.findCapsule(nameCaps)
        if targetCaps is None:
            return None
        targetCaps.trainSemanticAttribute(fromObservations, nameAttr)
        return targetCaps
    <mask token>

    def getLayerIndex(self, capsule: Capsule):
        for layerID in range(self._numSemanticLayers):
            if capsule in self._semanticLayers[layerID]:
                return layerID
        return -1

    def getAttributePool(self):
        return self._attributePool

    def clearAllObservations(self):
        for filterShape, capsule in self._pixelCapsules.items():
            capsule.clearObservations()
        for capsule in self._primitiveCapsules:
            capsule.clearObservations()
        for capsule in self._semanticCapsules:
            capsule.clearObservations()

    def findObservedAxioms(self, observations: dict):
        observedAxioms = {}
        for caps, obsList in observations.items():
            for obs in obsList:
                foundParent = False
                for checkObs in [x for checkObsList in observations.values(
                    ) for x in checkObsList]:
                    if checkObs.isParent(obs):
                        foundParent = True
                        break
                if foundParent is False:
                    if caps in observedAxioms:
                        observedAxioms[caps].append(obs)
                    else:
                        observedAxioms[caps] = [obs]
        return observedAxioms

    def showInput(self, image: list, width: int, height: int, stepSize: int=1):
        startTime = time.time()
        self.clearAllObservations()
        print('Capsule Network shown an Image of size (' + str(width) +
            ', ' + str(height) + ')')
        (offsetLabelX, offsetLabelY, offsetLabelRatio, targetLabelX,
            targetLabelY, targetLabelSize) = self._renderer.getOffsetLabels()
        for filterShape, capsule in self._pixelCapsules.items():
            if filterShape[0] <= width and filterShape[1] <= height:
                fsWidth = float(filterShape[0])
                fsHeight = float(filterShape[1])
                fsMaxW = float(max(width, height))
                rads = np.fromfunction(lambda xx, yy: np.sqrt((xx / fsWidth -
                    0.5) ** 2 + (yy / fsHeight - 0.5) ** 2), (filterShape[0
                    ], filterShape[1]), dtype=float)
                angs = np.fromfunction(lambda xx, yy: np.arctan2(xx /
                    fsWidth - 0.5, yy / fsHeight - 0.5) % (math.pi * 2), (
                    filterShape[0], filterShape[1]), dtype=float)
                for offsetX in range(0, width - filterShape[0] + 1, stepSize):
                    for offsetY in range(0, height - filterShape[1] + 1,
                        stepSize):
                        attributes = {}
                        for xx in range(filterShape[0]):
                            for yy in range(filterShape[1]):
                                attributes[capsule.getAttributeByName(
                                    'PixelC-' + str(xx) + '-' + str(yy))
                                    ] = image[((yy + offsetY) * width + xx +
                                    offsetX) * 4]
                                attributes[capsule.getAttributeByName(
                                    'PixelR-' + str(xx) + '-' + str(yy))
                                    ] = rads[xx, yy]
                                attributes[capsule.getAttributeByName(
                                    'PixelA-' + str(xx) + '-' + str(yy))
                                    ] = angs[xx, yy]
                                attributes[capsule.getAttributeByName(
                                    'PixelD-' + str(xx) + '-' + str(yy))] = 1.0
                        attributes[capsule.getAttributeByName(offsetLabelX)
                            ] = float(offsetX) / fsMaxW
                        attributes[capsule.getAttributeByName(offsetLabelY)
                            ] = float(offsetY) / fsMaxW
                        attributes[capsule.getAttributeByName(offsetLabelRatio)
                            ] = fsWidth / fsMaxW
                        pixelObs = Observation(capsule, None, [],
                            attributes, 1.0)
                        capsule.addPixelObservation(pixelObs)
        passedTime = time.time() - startTime
        print(
            'Beginning Forward Pass on all Primitive Capsules (Time Passed: ' +
            str(passedTime) + 's)')
        allObs = {}
        for capsule in self._primitiveCapsules:
            capsule.forwardPass()
            capsule.cleanupObservations(offsetLabelX, offsetLabelY,
                offsetLabelRatio, targetLabelX, targetLabelY, targetLabelSize)
            allObs[capsule] = capsule.getObservations()
        for layer in range(self._numSemanticLayers):
            passedTime = time.time() - startTime
            print('Beginning Forward Pass on Layer ' + str(layer) +
                ' of Semantic Capsules (Time Passed: ' + str(passedTime) + 's)'
                )
            for capsule in self._semanticLayers[layer]:
                capsule.forwardPass()
                capsule.cleanupObservations()
                allObs[capsule] = capsule.getObservations()
        recommendation = None
        observedAxioms = self.findObservedAxioms(allObs)
        if len(observedAxioms) > 1:
            recommendation = self._metaLearner.checkResults(allObs,
                observedAxioms)
        passedTime = time.time() - startTime
        print('Total Time Passed: ' + str(passedTime) + 's')
        return allObs, recommendation

    def applyOracle(self, oracleDecision: int):
        self._metaLearner.applyOracle(oracleDecision)

    def generateImage(self, width: int, height: int, observations: dict,
        withBackground: bool=False):
        semantics = {}
        texts = []
        (offsetLabelX, offsetLabelY, offsetLabelRatio, targetLabelX,
            targetLabelY, targetLabelSize) = self._renderer.getOffsetLabels()
        obs = {}
        obsMap = {}
        for capsule, obsList in observations.items():
            obs[capsule] = []
            for observation in obsList:
                obs[capsule].append(Observation(capsule, observation.
                    getTakenRoute(), observation.getInputObservations(),
                    observation.getOutputs(), observation.getProbability()))
                obsMap[obs[capsule][-1]] = observation
        for capsule, obsList in obs.items():
            if capsule in self._semanticCapsules:
                for observation in obsList:
                    xOffset1 = observation.getOutput(capsule.
                        getAttributeByName(targetLabelX))
                    yOffset1 = observation.getOutput(capsule.
                        getAttributeByName(targetLabelY))
                    xOffset1 = int(xOffset1 * float(max(width, height)))
                    yOffset1 = int(yOffset1 * float(max(width, height)))
                    semantics[obsMap[observation]] = []
                    for inObs in observation.getInputObservations():
                        xOffset2 = inObs.getOutput(inObs.getCapsule().
                            getAttributeByName(targetLabelX))
                        yOffset2 = inObs.getOutput(inObs.getCapsule().
                            getAttributeByName(targetLabelY))
                        xOffset2 = int(xOffset2 * float(max(width, height)))
                        yOffset2 = int(yOffset2 * float(max(width, height)))
                        semantics[obsMap[observation]].append(patches.Arrow
                            (xOffset1, yOffset1, xOffset2 - xOffset1, 
                            yOffset2 - yOffset1, linewidth=1, edgecolor='r',
                            facecolor='none'))
                    semantics[obsMap[observation]].append(patches.Circle((
                        xOffset1, yOffset1), radius=1, color='r'))
        for layerIndex in range(self._numSemanticLayers - 1, -1, -1):
            for capsule in self._semanticLayers[layerIndex]:
                if capsule in obs.keys():
                    for observation in obs[capsule]:
                        if not observation.getInputObservations():
                            newObslist = capsule.backwardPass(observation, 
                                False)
                            for obsCaps, newObs in newObslist.items():
                                if obsCaps in obs:
                                    obs[obsCaps].extend(newObs)
                                else:
                                    obs[obsCaps] = newObs
                    del obs[capsule]
        capsObsPairs = []
        for capsule, obsList in obs.items():
            if capsule in self._primitiveCapsules:
                for observation in obsList:
                    capsObsPairs.append((capsule, observation))
        capsObsPairs = sorted(capsObsPairs, key=lambda tup: tup[1].
            getProbability())
        image = np.zeros(width * height * 4)
        for capsule, observation in capsObsPairs:
            pixelShape = self.getShapeByPixelCapsule(capsule.
                getPixelLayerInput())
            xOffset = observation.getOutput(capsule.getAttributeByName(
                targetLabelX))
            yOffset = observation.getOutput(capsule.getAttributeByName(
                targetLabelY))
            size = observation.getOutput(capsule.getAttributeByName(
                targetLabelSize))
            observation.setOutput(capsule.getAttributeByName(targetLabelX), 0.5
                )
            observation.setOutput(capsule.getAttributeByName(targetLabelY), 0.5
                )
            observation.setOutput(capsule.getAttributeByName(
                targetLabelSize), size * float(max(width, height)) / float(
                pixelShape[0]))
            obsPixelLayer = capsule.backwardPass(observation, False)
            pixelObs = list(obsPixelLayer.values())[0][0]
            pixelLay = list(obsPixelLayer.keys())[0]
            newxOffset = int(xOffset * float(max(width, height))) - int(
                pixelShape[0] / 2)
            newyOffset = int(yOffset * float(max(width, height))) - int(
                pixelShape[1] / 2)
            minX = newxOffset + pixelShape[0]
            maxX = newxOffset
            minY = newyOffset + pixelShape[1]
            maxY = newyOffset
            for xx in range(pixelShape[0]):
                for yy in range(pixelShape[1]):
                    if (xx + newxOffset < width and xx + newxOffset >= 0 and
                        yy + newyOffset < height and yy + newyOffset >= 0):
                        depth = pixelObs.getOutput(pixelLay.
                            getAttributeByName('PixelD-' + str(xx) + '-' +
                            str(yy)))
                        if depth < 1.0:
                            minX = min(xx + newxOffset, minX)
                            maxX = max(xx + newxOffset, maxX)
                            minY = min(yy + newyOffset, minY)
                            maxY = max(yy + newyOffset, maxY)
                        if withBackground is True or depth < 1.0:
                            image[((yy + newyOffset) * width + xx +
                                newxOffset) * 4] = pixelObs.getOutput(pixelLay
                                .getAttributeByName('PixelC-' + str(xx) +
                                '-' + str(yy)))
            if observation in obsMap:
                semantics[obsMap[observation]] = [patches.Rectangle((minX,
                    minY), maxX - minX, maxY - minY, linewidth=1, edgecolor
                    ='y', facecolor='none')]
            else:
                semantics[observation] = [patches.Rectangle((minX, minY), 
                    maxX - minX, maxY - minY, linewidth=1, edgecolor='y',
                    facecolor='none')]
            observation.setOutput(capsule.getAttributeByName(targetLabelX),
                xOffset)
            observation.setOutput(capsule.getAttributeByName(targetLabelY),
                yOffset)
            observation.setOutput(capsule.getAttributeByName(
                targetLabelSize), size)
            texts.append((minX, minY, capsule.getName()))
        return image, semantics, texts

    def producePrimitiveObservations(self, observation: Observation):
        if observation.getCapsule() not in self._primitiveCapsules:
            outputObsList = []
            if not observation.getInputObservations():
                newObsDict = observation.getCapsule().backwardPass(observation,
                    False)
                observation.clearInputObservations()
                for newObsList in newObsDict.values():
                    for newObs in newObsList:
                        outputObsList = (outputObsList + self.
                            producePrimitiveObservations(newObs))
            else:
                for newObs in observation.getInputObservations():
                    outputObsList = (outputObsList + self.
                        producePrimitiveObservations(newObs))
            return outputObsList
        else:
            return [observation]

    def distance(self, observationA: Observation, observationB: Observation):
        distance = 100.0
        normal1 = [0.0, 0.0]
        normal2 = [0.0, 0.0]
        compareA = self.producePrimitiveObservations(observationA)
        compareB = self.producePrimitiveObservations(observationB)
        for primA in compareA:
            for primB in compareB:
                testDistance, testNormal1, testNormal2 = (self._renderer.
                    getDistance(self._capsulePrimitive[primA.getCapsule()],
                    self._capsulePrimitive[primB.getCapsule()], primA.
                    getOutputs(), primB.getOutputs()))
                if testDistance < distance:
                    distance = testDistance
                    normal1 = testNormal1
                    normal2 = testNormal2
        return distance, normal1, normal2

<mask token>


class CapsuleNetwork:

    def __init__(self, name: str):
        self._primitiveCapsules: list = []
        self._semanticCapsules: list = []
        self._pixelCapsules: dict = {}
        self._capsulePrimitive: dict = {}
        self._semanticLayers: dict = {}
        self._numSemanticLayers: int = 0
        self._attributePool: AttributePool = AttributePool()
        self._renderer = None
        self._metaLearner: MetaLearner = MetaLearner()
        self._capsuleCount: int = 0
        self._name: str = name
        self._metaLearner.addLambda(lambda obs, axioms: len(self.
            findSameParents(list(axioms.keys()))) > 0)
        self._metaLearner.addLambda(lambda obs, axioms: len(self.
            findSameParents(list(axioms.keys()))) == 0)
        self._metaLearner.addLambda(lambda obs, axioms: Utility.andElements
            ([Utility.andElements([obsItem.hasPreviousObservation() for
            obsItem in obsList]) for obsList in axioms.values()]))
        self._metaLearner.addLambda(lambda obs, axioms: not Utility.
            andElements([Utility.andElements([obsItem.
            hasPreviousObservation() for obsItem in obsList]) for obsList in
            axioms.values()]))

    def getJSON(self):
        semanticData = []
        for layerID in range(self._numSemanticLayers):
            layerCaps = []
            for caps in self._semanticLayers[layerID]:
                layerCaps.append(caps.getJSON())
            semanticData.append({'semanticCapsules': layerCaps})
        metaLearnerData = self._metaLearner.getJSON()
        return {'semanticLayers': semanticData, 'metaLearner': metaLearnerData}

    def putJSON(self, data):
        semCaps = {}
        for layerData in data['semanticLayers']:
            for capsData in layerData['semanticCapsules']:
                capsName = capsData['name']
                obsList = []
                for obsData in capsData['firstRouteObservations']:
                    obsCaps = self.getCapsuleByName(obsData['name'])
                    obsRoute = obsCaps.getRouteByName(obsData['route'])
                    obsProb = obsData['probability']
                    attrDict = {}
                    for attrData in obsData['attributes']:
                        attrDict[obsCaps.getAttributeByName(attrData[
                            'attribute'])] = attrData['value']
                    obsList.append(Observation(obsCaps, obsRoute, [],
                        attrDict, obsProb))
                semCaps[capsName] = self.addSemanticCapsule(capsName,
                    obsList, 0)
                if 'remainingMemory' in capsData:
                    semCaps[capsName].putJSONMemory(capsData[
                        'remainingMemory'], self._attributePool, lambda
                        name: self.getCapsuleByName(name))
        if 'metaLearner' in data:
            self._metaLearner.putJSON(data['metaLearner'])
        return semCaps

    def getName(self):
        return self._name
    <mask token>

    def agreementOfMostLikelyParent(self, observedAxioms: dict):
        potentialParents = self.findSameParents(list(observedAxioms.keys()))
        maxProbability = 0.0
        maxAgreement = {}
        maxParent = None
        for parent in potentialParents:
            currentProbability, currentAgreement = parent.getMaxAgreement(
                observedAxioms)
            if currentProbability > maxProbability:
                maxParent = parent
                maxProbability = currentProbability
                maxAgreement = currentAgreement
        return maxAgreement

    def findSameParents(self, capsules: list):
        sameParents = []
        for caps in self._semanticCapsules:
            if caps not in capsules:
                if caps.haveSameParent(capsules) is True:
                    sameParents.append(caps)
        return sameParents
    <mask token>

    def setRenderer(self, rendererClass):
        self._renderer = rendererClass(self._attributePool)
    <mask token>

    def addPrimitiveCapsule(self, primitive: Primitives, filterShapes: list,
        additionalTraining: int=0):
        if self._renderer is None:
            print('No Renderer of Type PrimitivesRenderer defined')
            return
        currentCapsule = Capsule(str(primitive), self._capsuleCount)
        self._capsulePrimitive[currentCapsule] = primitive
        self._capsuleCount = self._capsuleCount + 1
        self._renderer.createAttributesForPrimitive(primitive,
            currentCapsule, self._attributePool)
        for filterShape in filterShapes:
            pixelCapsule = None
            if filterShape not in self._pixelCapsules:
                pixelCapsule = Capsule('PixelLayer-' + str(filterShape[0]) +
                    '-' + str(filterShape[1]), -1)
                self._renderer.createAttributesForPixelLayer(filterShape[0],
                    filterShape[1], pixelCapsule, self._attributePool)
                self._pixelCapsules[filterShape] = pixelCapsule
            else:
                pixelCapsule = self._pixelCapsules[filterShape]
            currentCapsule.addPrimitiveRoute(pixelCapsule, self._renderer,
                primitive)
        self._primitiveCapsules.append(currentCapsule)
        if additionalTraining > 0:
            for i in range(additionalTraining):
                print(
                    '-------------------------------- ADDITIONAL TRAINING ROUND '
                     + str(i + 1) + ' OF ' + str(additionalTraining) +
                    ' --------------------------------')
                currentCapsule.continueTraining(True, [0])
        return currentCapsule

    def addSemanticCapsule(self, name: str, fromObservations: list,
        additionalTraining: int=0):
        currentCapsule = Capsule(name, self._capsuleCount)
        self._capsuleCount = self._capsuleCount + 1
        currentCapsule.addSemanticRoute(fromObservations, self._attributePool)
        maxLayerID = -1
        for obs in fromObservations:
            currentLayer = self.getLayerIndex(obs.getCapsule())
            if currentLayer > maxLayerID:
                maxLayerID = currentLayer
        if self._numSemanticLayers <= maxLayerID + 1:
            self._semanticLayers[self._numSemanticLayers] = []
            self._numSemanticLayers = self._numSemanticLayers + 1
        self._semanticLayers[maxLayerID + 1].append(currentCapsule)
        self._semanticCapsules.append(currentCapsule)
        if additionalTraining > 0:
            for i in range(additionalTraining):
                print(
                    '-------------------------------- ADDITIONAL TRAINING ROUND '
                     + str(i + 1) + ' OF ' + str(additionalTraining) +
                    ' --------------------------------')
                currentCapsule.continueTraining(True, [0])
        return currentCapsule

    def addSemanticTraining(self, name: str, fromObservations: list,
        additionalTraining: int=0):
        targetCaps = self.findCapsule(name)
        if targetCaps is None:
            return None
        targetCaps.trainSemanticRoute(fromObservations)
        return targetCaps

    def addAttribute(self, nameCaps: str, nameAttr: str, fromObservations:
        list, additionalTraining: int=0):
        targetCaps = self.findCapsule(nameCaps)
        if targetCaps is None:
            return None
        targetCaps.addSemanticAttribute(fromObservations, nameAttr, self.
            _attributePool)
        for caps in self._semanticCapsules:
            caps.inheritAttributes(self._attributePool)
        return targetCaps

    def addAttributeTraining(self, nameCaps: str, nameAttr: str,
        fromObservations: list, additionalTraining: int=0):
        targetCaps = self.findCapsule(nameCaps)
        if targetCaps is None:
            return None
        targetCaps.trainSemanticAttribute(fromObservations, nameAttr)
        return targetCaps

    def findCapsule(self, name: str):
        for caps in self._semanticCapsules:
            if name.lower() in caps.getName().lower():
                return caps
        return None

    def getLayerIndex(self, capsule: Capsule):
        for layerID in range(self._numSemanticLayers):
            if capsule in self._semanticLayers[layerID]:
                return layerID
        return -1

    def getAttributePool(self):
        return self._attributePool

    def clearAllObservations(self):
        for filterShape, capsule in self._pixelCapsules.items():
            capsule.clearObservations()
        for capsule in self._primitiveCapsules:
            capsule.clearObservations()
        for capsule in self._semanticCapsules:
            capsule.clearObservations()

    def findObservedAxioms(self, observations: dict):
        observedAxioms = {}
        for caps, obsList in observations.items():
            for obs in obsList:
                foundParent = False
                for checkObs in [x for checkObsList in observations.values(
                    ) for x in checkObsList]:
                    if checkObs.isParent(obs):
                        foundParent = True
                        break
                if foundParent is False:
                    if caps in observedAxioms:
                        observedAxioms[caps].append(obs)
                    else:
                        observedAxioms[caps] = [obs]
        return observedAxioms

    def showInput(self, image: list, width: int, height: int, stepSize: int=1):
        startTime = time.time()
        self.clearAllObservations()
        print('Capsule Network shown an Image of size (' + str(width) +
            ', ' + str(height) + ')')
        (offsetLabelX, offsetLabelY, offsetLabelRatio, targetLabelX,
            targetLabelY, targetLabelSize) = self._renderer.getOffsetLabels()
        for filterShape, capsule in self._pixelCapsules.items():
            if filterShape[0] <= width and filterShape[1] <= height:
                fsWidth = float(filterShape[0])
                fsHeight = float(filterShape[1])
                fsMaxW = float(max(width, height))
                rads = np.fromfunction(lambda xx, yy: np.sqrt((xx / fsWidth -
                    0.5) ** 2 + (yy / fsHeight - 0.5) ** 2), (filterShape[0
                    ], filterShape[1]), dtype=float)
                angs = np.fromfunction(lambda xx, yy: np.arctan2(xx /
                    fsWidth - 0.5, yy / fsHeight - 0.5) % (math.pi * 2), (
                    filterShape[0], filterShape[1]), dtype=float)
                for offsetX in range(0, width - filterShape[0] + 1, stepSize):
                    for offsetY in range(0, height - filterShape[1] + 1,
                        stepSize):
                        attributes = {}
                        for xx in range(filterShape[0]):
                            for yy in range(filterShape[1]):
                                attributes[capsule.getAttributeByName(
                                    'PixelC-' + str(xx) + '-' + str(yy))
                                    ] = image[((yy + offsetY) * width + xx +
                                    offsetX) * 4]
                                attributes[capsule.getAttributeByName(
                                    'PixelR-' + str(xx) + '-' + str(yy))
                                    ] = rads[xx, yy]
                                attributes[capsule.getAttributeByName(
                                    'PixelA-' + str(xx) + '-' + str(yy))
                                    ] = angs[xx, yy]
                                attributes[capsule.getAttributeByName(
                                    'PixelD-' + str(xx) + '-' + str(yy))] = 1.0
                        attributes[capsule.getAttributeByName(offsetLabelX)
                            ] = float(offsetX) / fsMaxW
                        attributes[capsule.getAttributeByName(offsetLabelY)
                            ] = float(offsetY) / fsMaxW
                        attributes[capsule.getAttributeByName(offsetLabelRatio)
                            ] = fsWidth / fsMaxW
                        pixelObs = Observation(capsule, None, [],
                            attributes, 1.0)
                        capsule.addPixelObservation(pixelObs)
        passedTime = time.time() - startTime
        print(
            'Beginning Forward Pass on all Primitive Capsules (Time Passed: ' +
            str(passedTime) + 's)')
        allObs = {}
        for capsule in self._primitiveCapsules:
            capsule.forwardPass()
            capsule.cleanupObservations(offsetLabelX, offsetLabelY,
                offsetLabelRatio, targetLabelX, targetLabelY, targetLabelSize)
            allObs[capsule] = capsule.getObservations()
        for layer in range(self._numSemanticLayers):
            passedTime = time.time() - startTime
            print('Beginning Forward Pass on Layer ' + str(layer) +
                ' of Semantic Capsules (Time Passed: ' + str(passedTime) + 's)'
                )
            for capsule in self._semanticLayers[layer]:
                capsule.forwardPass()
                capsule.cleanupObservations()
                allObs[capsule] = capsule.getObservations()
        recommendation = None
        observedAxioms = self.findObservedAxioms(allObs)
        if len(observedAxioms) > 1:
            recommendation = self._metaLearner.checkResults(allObs,
                observedAxioms)
        passedTime = time.time() - startTime
        print('Total Time Passed: ' + str(passedTime) + 's')
        return allObs, recommendation

    def applyOracle(self, oracleDecision: int):
        self._metaLearner.applyOracle(oracleDecision)

    def generateImage(self, width: int, height: int, observations: dict,
        withBackground: bool=False):
        semantics = {}
        texts = []
        (offsetLabelX, offsetLabelY, offsetLabelRatio, targetLabelX,
            targetLabelY, targetLabelSize) = self._renderer.getOffsetLabels()
        obs = {}
        obsMap = {}
        for capsule, obsList in observations.items():
            obs[capsule] = []
            for observation in obsList:
                obs[capsule].append(Observation(capsule, observation.
                    getTakenRoute(), observation.getInputObservations(),
                    observation.getOutputs(), observation.getProbability()))
                obsMap[obs[capsule][-1]] = observation
        for capsule, obsList in obs.items():
            if capsule in self._semanticCapsules:
                for observation in obsList:
                    xOffset1 = observation.getOutput(capsule.
                        getAttributeByName(targetLabelX))
                    yOffset1 = observation.getOutput(capsule.
                        getAttributeByName(targetLabelY))
                    xOffset1 = int(xOffset1 * float(max(width, height)))
                    yOffset1 = int(yOffset1 * float(max(width, height)))
                    semantics[obsMap[observation]] = []
                    for inObs in observation.getInputObservations():
                        xOffset2 = inObs.getOutput(inObs.getCapsule().
                            getAttributeByName(targetLabelX))
                        yOffset2 = inObs.getOutput(inObs.getCapsule().
                            getAttributeByName(targetLabelY))
                        xOffset2 = int(xOffset2 * float(max(width, height)))
                        yOffset2 = int(yOffset2 * float(max(width, height)))
                        semantics[obsMap[observation]].append(patches.Arrow
                            (xOffset1, yOffset1, xOffset2 - xOffset1, 
                            yOffset2 - yOffset1, linewidth=1, edgecolor='r',
                            facecolor='none'))
                    semantics[obsMap[observation]].append(patches.Circle((
                        xOffset1, yOffset1), radius=1, color='r'))
        for layerIndex in range(self._numSemanticLayers - 1, -1, -1):
            for capsule in self._semanticLayers[layerIndex]:
                if capsule in obs.keys():
                    for observation in obs[capsule]:
                        if not observation.getInputObservations():
                            newObslist = capsule.backwardPass(observation, 
                                False)
                            for obsCaps, newObs in newObslist.items():
                                if obsCaps in obs:
                                    obs[obsCaps].extend(newObs)
                                else:
                                    obs[obsCaps] = newObs
                    del obs[capsule]
        capsObsPairs = []
        for capsule, obsList in obs.items():
            if capsule in self._primitiveCapsules:
                for observation in obsList:
                    capsObsPairs.append((capsule, observation))
        capsObsPairs = sorted(capsObsPairs, key=lambda tup: tup[1].
            getProbability())
        image = np.zeros(width * height * 4)
        for capsule, observation in capsObsPairs:
            pixelShape = self.getShapeByPixelCapsule(capsule.
                getPixelLayerInput())
            xOffset = observation.getOutput(capsule.getAttributeByName(
                targetLabelX))
            yOffset = observation.getOutput(capsule.getAttributeByName(
                targetLabelY))
            size = observation.getOutput(capsule.getAttributeByName(
                targetLabelSize))
            observation.setOutput(capsule.getAttributeByName(targetLabelX), 0.5
                )
            observation.setOutput(capsule.getAttributeByName(targetLabelY), 0.5
                )
            observation.setOutput(capsule.getAttributeByName(
                targetLabelSize), size * float(max(width, height)) / float(
                pixelShape[0]))
            obsPixelLayer = capsule.backwardPass(observation, False)
            pixelObs = list(obsPixelLayer.values())[0][0]
            pixelLay = list(obsPixelLayer.keys())[0]
            newxOffset = int(xOffset * float(max(width, height))) - int(
                pixelShape[0] / 2)
            newyOffset = int(yOffset * float(max(width, height))) - int(
                pixelShape[1] / 2)
            minX = newxOffset + pixelShape[0]
            maxX = newxOffset
            minY = newyOffset + pixelShape[1]
            maxY = newyOffset
            for xx in range(pixelShape[0]):
                for yy in range(pixelShape[1]):
                    if (xx + newxOffset < width and xx + newxOffset >= 0 and
                        yy + newyOffset < height and yy + newyOffset >= 0):
                        depth = pixelObs.getOutput(pixelLay.
                            getAttributeByName('PixelD-' + str(xx) + '-' +
                            str(yy)))
                        if depth < 1.0:
                            minX = min(xx + newxOffset, minX)
                            maxX = max(xx + newxOffset, maxX)
                            minY = min(yy + newyOffset, minY)
                            maxY = max(yy + newyOffset, maxY)
                        if withBackground is True or depth < 1.0:
                            image[((yy + newyOffset) * width + xx +
                                newxOffset) * 4] = pixelObs.getOutput(pixelLay
                                .getAttributeByName('PixelC-' + str(xx) +
                                '-' + str(yy)))
            if observation in obsMap:
                semantics[obsMap[observation]] = [patches.Rectangle((minX,
                    minY), maxX - minX, maxY - minY, linewidth=1, edgecolor
                    ='y', facecolor='none')]
            else:
                semantics[observation] = [patches.Rectangle((minX, minY), 
                    maxX - minX, maxY - minY, linewidth=1, edgecolor='y',
                    facecolor='none')]
            observation.setOutput(capsule.getAttributeByName(targetLabelX),
                xOffset)
            observation.setOutput(capsule.getAttributeByName(targetLabelY),
                yOffset)
            observation.setOutput(capsule.getAttributeByName(
                targetLabelSize), size)
            texts.append((minX, minY, capsule.getName()))
        return image, semantics, texts

    def producePrimitiveObservations(self, observation: Observation):
        if observation.getCapsule() not in self._primitiveCapsules:
            outputObsList = []
            if not observation.getInputObservations():
                newObsDict = observation.getCapsule().backwardPass(observation,
                    False)
                observation.clearInputObservations()
                for newObsList in newObsDict.values():
                    for newObs in newObsList:
                        outputObsList = (outputObsList + self.
                            producePrimitiveObservations(newObs))
            else:
                for newObs in observation.getInputObservations():
                    outputObsList = (outputObsList + self.
                        producePrimitiveObservations(newObs))
            return outputObsList
        else:
            return [observation]

    def distance(self, observationA: Observation, observationB: Observation):
        distance = 100.0
        normal1 = [0.0, 0.0]
        normal2 = [0.0, 0.0]
        compareA = self.producePrimitiveObservations(observationA)
        compareB = self.producePrimitiveObservations(observationB)
        for primA in compareA:
            for primB in compareB:
                testDistance, testNormal1, testNormal2 = (self._renderer.
                    getDistance(self._capsulePrimitive[primA.getCapsule()],
                    self._capsulePrimitive[primB.getCapsule()], primA.
                    getOutputs(), primB.getOutputs()))
                if testDistance < distance:
                    distance = testDistance
                    normal1 = testNormal1
                    normal2 = testNormal2
        return distance, normal1, normal2

<mask token>


class CapsuleNetwork:

    def __init__(self, name: str):
        self._primitiveCapsules: list = []
        self._semanticCapsules: list = []
        self._pixelCapsules: dict = {}
        self._capsulePrimitive: dict = {}
        self._semanticLayers: dict = {}
        self._numSemanticLayers: int = 0
        self._attributePool: AttributePool = AttributePool()
        self._renderer = None
        self._metaLearner: MetaLearner = MetaLearner()
        self._capsuleCount: int = 0
        self._name: str = name
        self._metaLearner.addLambda(lambda obs, axioms: len(self.
            findSameParents(list(axioms.keys()))) > 0)
        self._metaLearner.addLambda(lambda obs, axioms: len(self.
            findSameParents(list(axioms.keys()))) == 0)
        self._metaLearner.addLambda(lambda obs, axioms: Utility.andElements
            ([Utility.andElements([obsItem.hasPreviousObservation() for
            obsItem in obsList]) for obsList in axioms.values()]))
        self._metaLearner.addLambda(lambda obs, axioms: not Utility.
            andElements([Utility.andElements([obsItem.
            hasPreviousObservation() for obsItem in obsList]) for obsList in
            axioms.values()]))

    def getJSON(self):
        semanticData = []
        for layerID in range(self._numSemanticLayers):
            layerCaps = []
            for caps in self._semanticLayers[layerID]:
                layerCaps.append(caps.getJSON())
            semanticData.append({'semanticCapsules': layerCaps})
        metaLearnerData = self._metaLearner.getJSON()
        return {'semanticLayers': semanticData, 'metaLearner': metaLearnerData}

    def putJSON(self, data):
        semCaps = {}
        for layerData in data['semanticLayers']:
            for capsData in layerData['semanticCapsules']:
                capsName = capsData['name']
                obsList = []
                for obsData in capsData['firstRouteObservations']:
                    obsCaps = self.getCapsuleByName(obsData['name'])
                    obsRoute = obsCaps.getRouteByName(obsData['route'])
                    obsProb = obsData['probability']
                    attrDict = {}
                    for attrData in obsData['attributes']:
                        attrDict[obsCaps.getAttributeByName(attrData[
                            'attribute'])] = attrData['value']
                    obsList.append(Observation(obsCaps, obsRoute, [],
                        attrDict, obsProb))
                semCaps[capsName] = self.addSemanticCapsule(capsName,
                    obsList, 0)
                if 'remainingMemory' in capsData:
                    semCaps[capsName].putJSONMemory(capsData[
                        'remainingMemory'], self._attributePool, lambda
                        name: self.getCapsuleByName(name))
        if 'metaLearner' in data:
            self._metaLearner.putJSON(data['metaLearner'])
        return semCaps

    def getName(self):
        return self._name

    def getCapsuleByName(self, name: str):
        for caps in self._primitiveCapsules:
            if caps.getName() == name:
                return caps
        for caps in self._semanticCapsules:
            if caps.getName() == name:
                return caps
        return None

    def agreementOfMostLikelyParent(self, observedAxioms: dict):
        potentialParents = self.findSameParents(list(observedAxioms.keys()))
        maxProbability = 0.0
        maxAgreement = {}
        maxParent = None
        for parent in potentialParents:
            currentProbability, currentAgreement = parent.getMaxAgreement(
                observedAxioms)
            if currentProbability > maxProbability:
                maxParent = parent
                maxProbability = currentProbability
                maxAgreement = currentAgreement
        return maxAgreement

    def findSameParents(self, capsules: list):
        sameParents = []
        for caps in self._semanticCapsules:
            if caps not in capsules:
                if caps.haveSameParent(capsules) is True:
                    sameParents.append(caps)
        return sameParents
    <mask token>

    def setRenderer(self, rendererClass):
        self._renderer = rendererClass(self._attributePool)
    <mask token>

    def addPrimitiveCapsule(self, primitive: Primitives, filterShapes: list,
        additionalTraining: int=0):
        if self._renderer is None:
            print('No Renderer of Type PrimitivesRenderer defined')
            return
        currentCapsule = Capsule(str(primitive), self._capsuleCount)
        self._capsulePrimitive[currentCapsule] = primitive
        self._capsuleCount = self._capsuleCount + 1
        self._renderer.createAttributesForPrimitive(primitive,
            currentCapsule, self._attributePool)
        for filterShape in filterShapes:
            pixelCapsule = None
            if filterShape not in self._pixelCapsules:
                pixelCapsule = Capsule('PixelLayer-' + str(filterShape[0]) +
                    '-' + str(filterShape[1]), -1)
                self._renderer.createAttributesForPixelLayer(filterShape[0],
                    filterShape[1], pixelCapsule, self._attributePool)
                self._pixelCapsules[filterShape] = pixelCapsule
            else:
                pixelCapsule = self._pixelCapsules[filterShape]
            currentCapsule.addPrimitiveRoute(pixelCapsule, self._renderer,
                primitive)
        self._primitiveCapsules.append(currentCapsule)
        if additionalTraining > 0:
            for i in range(additionalTraining):
                print(
                    '-------------------------------- ADDITIONAL TRAINING ROUND '
                     + str(i + 1) + ' OF ' + str(additionalTraining) +
                    ' --------------------------------')
                currentCapsule.continueTraining(True, [0])
        return currentCapsule

    def addSemanticCapsule(self, name: str, fromObservations: list,
        additionalTraining: int=0):
        currentCapsule = Capsule(name, self._capsuleCount)
        self._capsuleCount = self._capsuleCount + 1
        currentCapsule.addSemanticRoute(fromObservations, self._attributePool)
        maxLayerID = -1
        for obs in fromObservations:
            currentLayer = self.getLayerIndex(obs.getCapsule())
            if currentLayer > maxLayerID:
                maxLayerID = currentLayer
        if self._numSemanticLayers <= maxLayerID + 1:
            self._semanticLayers[self._numSemanticLayers] = []
            self._numSemanticLayers = self._numSemanticLayers + 1
        self._semanticLayers[maxLayerID + 1].append(currentCapsule)
        self._semanticCapsules.append(currentCapsule)
        if additionalTraining > 0:
            for i in range(additionalTraining):
                print(
                    '-------------------------------- ADDITIONAL TRAINING ROUND '
                     + str(i + 1) + ' OF ' + str(additionalTraining) +
                    ' --------------------------------')
                currentCapsule.continueTraining(True, [0])
        return currentCapsule

    def addSemanticTraining(self, name: str, fromObservations: list,
        additionalTraining: int=0):
        targetCaps = self.findCapsule(name)
        if targetCaps is None:
            return None
        targetCaps.trainSemanticRoute(fromObservations)
        return targetCaps

    def addAttribute(self, nameCaps: str, nameAttr: str, fromObservations:
        list, additionalTraining: int=0):
        targetCaps = self.findCapsule(nameCaps)
        if targetCaps is None:
            return None
        targetCaps.addSemanticAttribute(fromObservations, nameAttr, self.
            _attributePool)
        for caps in self._semanticCapsules:
            caps.inheritAttributes(self._attributePool)
        return targetCaps

    def addAttributeTraining(self, nameCaps: str, nameAttr: str,
        fromObservations: list, additionalTraining: int=0):
        targetCaps = self.findCapsule(nameCaps)
        if targetCaps is None:
            return None
        targetCaps.trainSemanticAttribute(fromObservations, nameAttr)
        return targetCaps

    def findCapsule(self, name: str):
        for caps in self._semanticCapsules:
            if name.lower() in caps.getName().lower():
                return caps
        return None

    def getLayerIndex(self, capsule: Capsule):
        for layerID in range(self._numSemanticLayers):
            if capsule in self._semanticLayers[layerID]:
                return layerID
        return -1

    def getAttributePool(self):
        return self._attributePool

    def clearAllObservations(self):
        for filterShape, capsule in self._pixelCapsules.items():
            capsule.clearObservations()
        for capsule in self._primitiveCapsules:
            capsule.clearObservations()
        for capsule in self._semanticCapsules:
            capsule.clearObservations()

    def findObservedAxioms(self, observations: dict):
        observedAxioms = {}
        for caps, obsList in observations.items():
            for obs in obsList:
                foundParent = False
                for checkObs in [x for checkObsList in observations.values(
                    ) for x in checkObsList]:
                    if checkObs.isParent(obs):
                        foundParent = True
                        break
                if foundParent is False:
                    if caps in observedAxioms:
                        observedAxioms[caps].append(obs)
                    else:
                        observedAxioms[caps] = [obs]
        return observedAxioms

    def showInput(self, image: list, width: int, height: int, stepSize: int=1):
        startTime = time.time()
        self.clearAllObservations()
        print('Capsule Network shown an Image of size (' + str(width) +
            ', ' + str(height) + ')')
        (offsetLabelX, offsetLabelY, offsetLabelRatio, targetLabelX,
            targetLabelY, targetLabelSize) = self._renderer.getOffsetLabels()
        for filterShape, capsule in self._pixelCapsules.items():
            if filterShape[0] <= width and filterShape[1] <= height:
                fsWidth = float(filterShape[0])
                fsHeight = float(filterShape[1])
                fsMaxW = float(max(width, height))
                rads = np.fromfunction(lambda xx, yy: np.sqrt((xx / fsWidth -
                    0.5) ** 2 + (yy / fsHeight - 0.5) ** 2), (filterShape[0
                    ], filterShape[1]), dtype=float)
                angs = np.fromfunction(lambda xx, yy: np.arctan2(xx /
                    fsWidth - 0.5, yy / fsHeight - 0.5) % (math.pi * 2), (
                    filterShape[0], filterShape[1]), dtype=float)
                for offsetX in range(0, width - filterShape[0] + 1, stepSize):
                    for offsetY in range(0, height - filterShape[1] + 1,
                        stepSize):
                        attributes = {}
                        for xx in range(filterShape[0]):
                            for yy in range(filterShape[1]):
                                attributes[capsule.getAttributeByName(
                                    'PixelC-' + str(xx) + '-' + str(yy))
                                    ] = image[((yy + offsetY) * width + xx +
                                    offsetX) * 4]
                                attributes[capsule.getAttributeByName(
                                    'PixelR-' + str(xx) + '-' + str(yy))
                                    ] = rads[xx, yy]
                                attributes[capsule.getAttributeByName(
                                    'PixelA-' + str(xx) + '-' + str(yy))
                                    ] = angs[xx, yy]
                                attributes[capsule.getAttributeByName(
                                    'PixelD-' + str(xx) + '-' + str(yy))] = 1.0
                        attributes[capsule.getAttributeByName(offsetLabelX)
                            ] = float(offsetX) / fsMaxW
                        attributes[capsule.getAttributeByName(offsetLabelY)
                            ] = float(offsetY) / fsMaxW
                        attributes[capsule.getAttributeByName(offsetLabelRatio)
                            ] = fsWidth / fsMaxW
                        pixelObs = Observation(capsule, None, [],
                            attributes, 1.0)
                        capsule.addPixelObservation(pixelObs)
        passedTime = time.time() - startTime
        print(
            'Beginning Forward Pass on all Primitive Capsules (Time Passed: ' +
            str(passedTime) + 's)')
        allObs = {}
        for capsule in self._primitiveCapsules:
            capsule.forwardPass()
            capsule.cleanupObservations(offsetLabelX, offsetLabelY,
                offsetLabelRatio, targetLabelX, targetLabelY, targetLabelSize)
            allObs[capsule] = capsule.getObservations()
        for layer in range(self._numSemanticLayers):
            passedTime = time.time() - startTime
            print('Beginning Forward Pass on Layer ' + str(layer) +
                ' of Semantic Capsules (Time Passed: ' + str(passedTime) + 's)'
                )
            for capsule in self._semanticLayers[layer]:
                capsule.forwardPass()
                capsule.cleanupObservations()
                allObs[capsule] = capsule.getObservations()
        recommendation = None
        observedAxioms = self.findObservedAxioms(allObs)
        if len(observedAxioms) > 1:
            recommendation = self._metaLearner.checkResults(allObs,
                observedAxioms)
        passedTime = time.time() - startTime
        print('Total Time Passed: ' + str(passedTime) + 's')
        return allObs, recommendation

    def applyOracle(self, oracleDecision: int):
        self._metaLearner.applyOracle(oracleDecision)

    def generateImage(self, width: int, height: int, observations: dict,
        withBackground: bool=False):
        semantics = {}
        texts = []
        (offsetLabelX, offsetLabelY, offsetLabelRatio, targetLabelX,
            targetLabelY, targetLabelSize) = self._renderer.getOffsetLabels()
        obs = {}
        obsMap = {}
        for capsule, obsList in observations.items():
            obs[capsule] = []
            for observation in obsList:
                obs[capsule].append(Observation(capsule, observation.
                    getTakenRoute(), observation.getInputObservations(),
                    observation.getOutputs(), observation.getProbability()))
                obsMap[obs[capsule][-1]] = observation
        for capsule, obsList in obs.items():
            if capsule in self._semanticCapsules:
                for observation in obsList:
                    xOffset1 = observation.getOutput(capsule.
                        getAttributeByName(targetLabelX))
                    yOffset1 = observation.getOutput(capsule.
                        getAttributeByName(targetLabelY))
                    xOffset1 = int(xOffset1 * float(max(width, height)))
                    yOffset1 = int(yOffset1 * float(max(width, height)))
                    semantics[obsMap[observation]] = []
                    for inObs in observation.getInputObservations():
                        xOffset2 = inObs.getOutput(inObs.getCapsule().
                            getAttributeByName(targetLabelX))
                        yOffset2 = inObs.getOutput(inObs.getCapsule().
                            getAttributeByName(targetLabelY))
                        xOffset2 = int(xOffset2 * float(max(width, height)))
                        yOffset2 = int(yOffset2 * float(max(width, height)))
                        semantics[obsMap[observation]].append(patches.Arrow
                            (xOffset1, yOffset1, xOffset2 - xOffset1, 
                            yOffset2 - yOffset1, linewidth=1, edgecolor='r',
                            facecolor='none'))
                    semantics[obsMap[observation]].append(patches.Circle((
                        xOffset1, yOffset1), radius=1, color='r'))
        for layerIndex in range(self._numSemanticLayers - 1, -1, -1):
            for capsule in self._semanticLayers[layerIndex]:
                if capsule in obs.keys():
                    for observation in obs[capsule]:
                        if not observation.getInputObservations():
                            newObslist = capsule.backwardPass(observation, 
                                False)
                            for obsCaps, newObs in newObslist.items():
                                if obsCaps in obs:
                                    obs[obsCaps].extend(newObs)
                                else:
                                    obs[obsCaps] = newObs
                    del obs[capsule]
        capsObsPairs = []
        for capsule, obsList in obs.items():
            if capsule in self._primitiveCapsules:
                for observation in obsList:
                    capsObsPairs.append((capsule, observation))
        capsObsPairs = sorted(capsObsPairs, key=lambda tup: tup[1].
            getProbability())
        image = np.zeros(width * height * 4)
        for capsule, observation in capsObsPairs:
            pixelShape = self.getShapeByPixelCapsule(capsule.
                getPixelLayerInput())
            xOffset = observation.getOutput(capsule.getAttributeByName(
                targetLabelX))
            yOffset = observation.getOutput(capsule.getAttributeByName(
                targetLabelY))
            size = observation.getOutput(capsule.getAttributeByName(
                targetLabelSize))
            observation.setOutput(capsule.getAttributeByName(targetLabelX), 0.5
                )
            observation.setOutput(capsule.getAttributeByName(targetLabelY), 0.5
                )
            observation.setOutput(capsule.getAttributeByName(
                targetLabelSize), size * float(max(width, height)) / float(
                pixelShape[0]))
            obsPixelLayer = capsule.backwardPass(observation, False)
            pixelObs = list(obsPixelLayer.values())[0][0]
            pixelLay = list(obsPixelLayer.keys())[0]
            newxOffset = int(xOffset * float(max(width, height))) - int(
                pixelShape[0] / 2)
            newyOffset = int(yOffset * float(max(width, height))) - int(
                pixelShape[1] / 2)
            minX = newxOffset + pixelShape[0]
            maxX = newxOffset
            minY = newyOffset + pixelShape[1]
            maxY = newyOffset
            for xx in range(pixelShape[0]):
                for yy in range(pixelShape[1]):
                    if (xx + newxOffset < width and xx + newxOffset >= 0 and
                        yy + newyOffset < height and yy + newyOffset >= 0):
                        depth = pixelObs.getOutput(pixelLay.
                            getAttributeByName('PixelD-' + str(xx) + '-' +
                            str(yy)))
                        if depth < 1.0:
                            minX = min(xx + newxOffset, minX)
                            maxX = max(xx + newxOffset, maxX)
                            minY = min(yy + newyOffset, minY)
                            maxY = max(yy + newyOffset, maxY)
                        if withBackground is True or depth < 1.0:
                            image[((yy + newyOffset) * width + xx +
                                newxOffset) * 4] = pixelObs.getOutput(pixelLay
                                .getAttributeByName('PixelC-' + str(xx) +
                                '-' + str(yy)))
            if observation in obsMap:
                semantics[obsMap[observation]] = [patches.Rectangle((minX,
                    minY), maxX - minX, maxY - minY, linewidth=1, edgecolor
                    ='y', facecolor='none')]
            else:
                semantics[observation] = [patches.Rectangle((minX, minY), 
                    maxX - minX, maxY - minY, linewidth=1, edgecolor='y',
                    facecolor='none')]
            observation.setOutput(capsule.getAttributeByName(targetLabelX),
                xOffset)
            observation.setOutput(capsule.getAttributeByName(targetLabelY),
                yOffset)
            observation.setOutput(capsule.getAttributeByName(
                targetLabelSize), size)
            texts.append((minX, minY, capsule.getName()))
        return image, semantics, texts

    def producePrimitiveObservations(self, observation: Observation):
        if observation.getCapsule() not in self._primitiveCapsules:
            outputObsList = []
            if not observation.getInputObservations():
                newObsDict = observation.getCapsule().backwardPass(observation,
                    False)
                observation.clearInputObservations()
                for newObsList in newObsDict.values():
                    for newObs in newObsList:
                        outputObsList = (outputObsList + self.
                            producePrimitiveObservations(newObs))
            else:
                for newObs in observation.getInputObservations():
                    outputObsList = (outputObsList + self.
                        producePrimitiveObservations(newObs))
            return outputObsList
        else:
            return [observation]

    def distance(self, observationA: Observation, observationB: Observation):
        distance = 100.0
        normal1 = [0.0, 0.0]
        normal2 = [0.0, 0.0]
        compareA = self.producePrimitiveObservations(observationA)
        compareB = self.producePrimitiveObservations(observationB)
        for primA in compareA:
            for primB in compareB:
                testDistance, testNormal1, testNormal2 = (self._renderer.
                    getDistance(self._capsulePrimitive[primA.getCapsule()],
                    self._capsulePrimitive[primB.getCapsule()], primA.
                    getOutputs(), primB.getOutputs()))
                if testDistance < distance:
                    distance = testDistance
                    normal1 = testNormal1
                    normal2 = testNormal2
        return distance, normal1, normal2

from AttributePool import AttributePool
from Capsule import Capsule
from PrimitivesRenderer import Primitives
from Observation import Observation
from MetaLearner import MetaLearner
from Utility import Utility

import copy
import math
import numpy as np
import matplotlib.patches as patches

import time

class CapsuleNetwork:

    def __init__(self, name : str):
        self._primitiveCapsules  : list          = []               # Capsule
        self._semanticCapsules   : list          = []               # Capsule
        self._pixelCapsules      : dict          = {}               # Shape Tuple - Capsule
        self._capsulePrimitive   : dict          = {}               # Primitive Capsule - Primitive Type
        self._semanticLayers     : dict          = {}               # Layer Index - Capsule List
        self._numSemanticLayers  : int           = 0                
        self._attributePool      : AttributePool = AttributePool()
        self._renderer                           = None             # PrimitivesRenderer Instance
        self._metaLearner        : MetaLearner   = MetaLearner()
        self._capsuleCount       : int           = 0
        self._name               : str           = name


        # Adding all Meta Learner Lambdas:
        # 1. Observed Axioms have same $\Omega$ as parent
        self._metaLearner.addLambda(lambda obs, axioms : len(self.findSameParents(list(axioms.keys()))) > 0)
        # 2. Observed Axioms don't have same $\Omega$ as parent
        self._metaLearner.addLambda(lambda obs, axioms : len(self.findSameParents(list(axioms.keys()))) == 0)
        # 3. Parts are tracked from previous scenes
        self._metaLearner.addLambda(lambda obs, axioms : \
            Utility.andElements([Utility.andElements([obsItem.hasPreviousObservation() for obsItem in obsList]) for obsList in axioms.values()]))
        # 4. Parts are NOT tracked from previous scenes
        self._metaLearner.addLambda(lambda obs, axioms : \
            not Utility.andElements([Utility.andElements([obsItem.hasPreviousObservation() for obsItem in obsList]) for obsList in axioms.values()]))
        
        # 5. $\Omega: Z(\vec{\alpha}, \vec{\tilde{\alpha}})$ indicates one attribute mismatch \\ with no entry in memory $\alpha^i >\epsilon$
        # TODO: self._metaLearner.addLambda(lambda obs, axioms : self.agreementOfMostLikelyParent(axioms))
        # 6. $\Omega: Z(\vec{\alpha}, \vec{\tilde{\alpha}})$ indicates attribute mismatch \\ for (position, rotation, size) only
        # TODO:


    def getJSON(self):
        # Only needs to save semantic capsules and data
        semanticData = []
        for layerID in range(self._numSemanticLayers):
            layerCaps = []
            for caps in self._semanticLayers[layerID]:
                layerCaps.append(caps.getJSON())

            semanticData.append({"semanticCapsules" : layerCaps})

        metaLearnerData = self._metaLearner.getJSON()

        return {"semanticLayers" : semanticData, "metaLearner" : metaLearnerData}


    def putJSON(self, data):
        # Only needs to load semantic capsules and data
        semCaps = {}
        for layerData in data["semanticLayers"]:
            for capsData in layerData["semanticCapsules"]:

                capsName = capsData["name"]
                
                obsList = []
                # First Route saved/loaded independently, as it is required to create the
                # Capsule.
                for obsData in capsData["firstRouteObservations"]:
                    obsCaps = self.getCapsuleByName(obsData["name"])
                    obsRoute = obsCaps.getRouteByName(obsData["route"])
                    obsProb = obsData["probability"]

                    attrDict = {}
                    for attrData in obsData["attributes"]:
                        attrDict[obsCaps.getAttributeByName(attrData["attribute"])] = attrData["value"]

                    obsList.append(Observation(obsCaps, obsRoute, [], attrDict, obsProb))
                
                semCaps[capsName] = self.addSemanticCapsule(capsName, obsList, 0)

                # Adding remaining memory
                if "remainingMemory" in capsData:
                    semCaps[capsName].putJSONMemory(capsData["remainingMemory"], self._attributePool, lambda name : self.getCapsuleByName(name))
        
        if "metaLearner" in data:
            self._metaLearner.putJSON(data["metaLearner"])

        return semCaps  # List of Semantic Capsules



    def getName(self):
        return self._name


    def getCapsuleByName(self, name : str):
        for caps in self._primitiveCapsules:
            if caps.getName() == name:
                return caps
        for caps in self._semanticCapsules:
            if caps.getName() == name:
                return caps
        return None


    def agreementOfMostLikelyParent(self, observedAxioms : dict):
        # observedAxioms    # {Observed Axiom (capsule), List of Observations}
        potentialParents = self.findSameParents(list(observedAxioms.keys()))

        maxProbability = 0.0
        maxAgreement = {}
        maxParent = None
        for parent in potentialParents:
            currentProbability, currentAgreement = parent.getMaxAgreement(observedAxioms)
            if currentProbability > maxProbability:
                maxParent = parent
                maxProbability = currentProbability
                maxAgreement = currentAgreement

        return maxAgreement


    def findSameParents(self, capsules : list):
        # capsules  # List of Capsules
        sameParents = []

        for caps in self._semanticCapsules:
            if caps not in capsules:
                if caps.haveSameParent(capsules) is True:
                    sameParents.append(caps)
        
        return sameParents


    def getShapeByPixelCapsule(self, capsule : Capsule):
        for shape, pixelCaps in self._pixelCapsules.items():
            if capsule == pixelCaps:
                return shape
        return (0, 0)

    
    def setRenderer(self, rendererClass):
        # rendererClass     # Class PrimitivesRenderer
        self._renderer = rendererClass(self._attributePool)

    
    def getRenderer(self):
        return self._renderer


    def addPrimitiveCapsule(self, primitive : Primitives, filterShapes : list, additionalTraining : int = 0):
        # filterShapes      # List of Tuples (width, height)
        
        if self._renderer is None:
            print("No Renderer of Type PrimitivesRenderer defined")
            return
        currentCapsule = Capsule(str(primitive), self._capsuleCount)
        self._capsulePrimitive[currentCapsule] = primitive
        self._capsuleCount = self._capsuleCount + 1
        self._renderer.createAttributesForPrimitive(primitive, currentCapsule, self._attributePool)

        for filterShape in filterShapes:
            pixelCapsule = None

            if filterShape not in self._pixelCapsules:
                pixelCapsule = Capsule("PixelLayer-" + str(filterShape[0]) + "-" + str(filterShape[1]), -1)
                self._renderer.createAttributesForPixelLayer(filterShape[0], filterShape[1], pixelCapsule, self._attributePool)
                self._pixelCapsules[filterShape] = pixelCapsule
            else:
                pixelCapsule = self._pixelCapsules[filterShape]

            currentCapsule.addPrimitiveRoute(pixelCapsule, self._renderer, primitive)
        
        self._primitiveCapsules.append(currentCapsule)

        if additionalTraining > 0:
            for i in range(additionalTraining):
                print("-------------------------------- ADDITIONAL TRAINING ROUND " + str(i + 1) + " OF " + str(additionalTraining) + " --------------------------------")
                currentCapsule.continueTraining(True, [0])

        return currentCapsule


    def addSemanticCapsule(self, name : str, fromObservations : list, additionalTraining : int = 0):
        # fromObservations          # List of Observations for one Occurance
        currentCapsule = Capsule(name, self._capsuleCount)
        self._capsuleCount = self._capsuleCount + 1
        currentCapsule.addSemanticRoute(fromObservations, self._attributePool)
        
        maxLayerID = -1
        for obs in fromObservations:
            currentLayer = self.getLayerIndex(obs.getCapsule())
            if currentLayer > maxLayerID:
                maxLayerID = currentLayer

        if self._numSemanticLayers <= maxLayerID + 1:
            self._semanticLayers[self._numSemanticLayers] = []
            self._numSemanticLayers = self._numSemanticLayers + 1

        self._semanticLayers[maxLayerID + 1].append(currentCapsule)
        self._semanticCapsules.append(currentCapsule)

        if additionalTraining > 0:
            for i in range(additionalTraining):
                print("-------------------------------- ADDITIONAL TRAINING ROUND " + str(i + 1) + " OF " + str(additionalTraining) + " --------------------------------")
                currentCapsule.continueTraining(True, [0])

        return currentCapsule


    def addSemanticTraining(self, name : str, fromObservations : list, additionalTraining : int = 0):
        targetCaps = self.findCapsule(name)
        if targetCaps is None:
            return None

        targetCaps.trainSemanticRoute(fromObservations)
        return targetCaps

    
    def addAttribute(self, nameCaps : str, nameAttr : str, fromObservations : list, additionalTraining : int = 0):
        targetCaps = self.findCapsule(nameCaps)
        if targetCaps is None:
            return None

        targetCaps.addSemanticAttribute(fromObservations, nameAttr, self._attributePool)

        for caps in self._semanticCapsules:
            caps.inheritAttributes(self._attributePool)

        return targetCaps


    def addAttributeTraining(self, nameCaps : str, nameAttr : str, fromObservations : list, additionalTraining : int = 0):
        targetCaps = self.findCapsule(nameCaps)
        if targetCaps is None:
            return None

        targetCaps.trainSemanticAttribute(fromObservations, nameAttr)
        return targetCaps


    def findCapsule(self, name : str):
        for caps in self._semanticCapsules:
            if name.lower() in caps.getName().lower():
                return caps
        
        return None


    def getLayerIndex(self, capsule : Capsule):
        for layerID in range(self._numSemanticLayers):
            if capsule in self._semanticLayers[layerID]:
                return layerID
        
        return -1


    def getAttributePool(self):
        return self._attributePool

    
    def clearAllObservations(self):
        for filterShape, capsule in self._pixelCapsules.items():
            capsule.clearObservations()
        for capsule in self._primitiveCapsules:
            capsule.clearObservations()
        for capsule in self._semanticCapsules:
            capsule.clearObservations()
            

    def findObservedAxioms(self, observations : dict):
        # observations      # {Capsule, List of Observations}

        observedAxioms = {}

        # Not the fastest check, but works...
        for caps, obsList in observations.items():
            for obs in obsList:
                foundParent = False
                for checkObs in [x for checkObsList in observations.values() for x in checkObsList] :
                    if checkObs.isParent(obs):
                        foundParent = True
                        break
                if foundParent is False:
                    if caps in observedAxioms:
                        observedAxioms[caps].append(obs) 
                    else:
                        observedAxioms[caps] = [obs] 

        return observedAxioms # {Observed Axioms (Capsule), List of Observations}

    
    def showInput(self, image : list, width : int, height : int, stepSize : int = 1):
        # image             # Linear List of Pixels

        startTime = time.time()
        self.clearAllObservations()

        print("Capsule Network shown an Image of size (" + str(width) + ", " + str(height) + ")")

        offsetLabelX, offsetLabelY, offsetLabelRatio, targetLabelX, targetLabelY, targetLabelSize = self._renderer.getOffsetLabels()
        
        for filterShape, capsule in self._pixelCapsules.items():
            if filterShape[0] <= width and filterShape[1] <= height:

                fsWidth = float(filterShape[0])
                fsHeight = float(filterShape[1])
                fsMaxW = float(max(width, height))

                rads = np.fromfunction(lambda xx, yy : np.sqrt((xx / fsWidth - 0.5) ** 2 + (yy / fsHeight - 0.5) ** 2), (filterShape[0], filterShape[1]), dtype=float)
                angs = np.fromfunction(lambda xx, yy : np.arctan2((xx / fsWidth - 0.5), (yy / fsHeight - 0.5)) % (math.pi * 2), (filterShape[0], filterShape[1]), dtype=float)

                for offsetX in range(0, width - filterShape[0] + 1, stepSize):
                    for offsetY in range(0, height - filterShape[1] + 1, stepSize):
                        attributes = {}
                        for xx in range(filterShape[0]):
                            for yy in range(filterShape[1]):
                                attributes[capsule.getAttributeByName("PixelC-" + str(xx) + "-" + str(yy))] = image[((yy + offsetY) * width + xx + offsetX) * 4]
                                attributes[capsule.getAttributeByName("PixelR-" + str(xx) + "-" + str(yy))] = rads[xx, yy]
                                attributes[capsule.getAttributeByName("PixelA-" + str(xx) + "-" + str(yy))] = angs[xx, yy]
                                attributes[capsule.getAttributeByName("PixelD-" + str(xx) + "-" + str(yy))] = 1.0

                        attributes[capsule.getAttributeByName(offsetLabelX)] = float(offsetX) / fsMaxW
                        attributes[capsule.getAttributeByName(offsetLabelY)] = float(offsetY) / fsMaxW
                        attributes[capsule.getAttributeByName(offsetLabelRatio)] = fsWidth / fsMaxW

                        pixelObs = Observation(capsule, None, [], attributes, 1.0)
                        capsule.addPixelObservation(pixelObs)

        passedTime = time.time() - startTime
        print("Beginning Forward Pass on all Primitive Capsules (Time Passed: " + str(passedTime) + "s)")

        allObs = {}     # Capsule - List Of Observations

        for capsule in self._primitiveCapsules:
            capsule.forwardPass()
            capsule.cleanupObservations(offsetLabelX, offsetLabelY, offsetLabelRatio, targetLabelX, targetLabelY, targetLabelSize)
            allObs[capsule] = capsule.getObservations()

        for layer in range(self._numSemanticLayers):
            passedTime = time.time() - startTime
            print("Beginning Forward Pass on Layer " + str(layer) + " of Semantic Capsules (Time Passed: " + str(passedTime) + "s)")
            for capsule in self._semanticLayers[layer]:
                capsule.forwardPass()
                capsule.cleanupObservations()
                allObs[capsule] = capsule.getObservations()

        recommendation = None
        observedAxioms = self.findObservedAxioms(allObs)
        if len(observedAxioms) > 1:
            recommendation = self._metaLearner.checkResults(allObs, observedAxioms)

        passedTime = time.time() - startTime
        print("Total Time Passed: " + str(passedTime) + "s")

        return allObs, recommendation   # Capsule - List Of Observations, Recommendation String


    def applyOracle(self, oracleDecision : int):
        # TODO: Apply oracle decisions here?
        self._metaLearner.applyOracle(oracleDecision)

        
    def generateImage(self, width : int, height : int, observations : dict, withBackground : bool = False):
        # observations          # Capsule   -   List of Observations

        semantics = {} #  Observation - List of Patches
        texts = []     #  List of (X, Y, Text)
        offsetLabelX, offsetLabelY, offsetLabelRatio, targetLabelX, targetLabelY, targetLabelSize = self._renderer.getOffsetLabels()
        

        # We make a full copy to work on
        obs = {}    # Capsule - List of Observations
        obsMap = {} # Copied Observation - Actual Observation
        for capsule, obsList in observations.items():
            obs[capsule] = []
            for observation in obsList:
                obs[capsule].append(Observation(capsule, observation.getTakenRoute(), observation.getInputObservations(), 
                                                observation.getOutputs(), observation.getProbability()))
                obsMap[obs[capsule][-1]] = observation

        # Generate Semantic Labels for Semantic Capsules
        for capsule, obsList in obs.items():
            if capsule in self._semanticCapsules:
                for observation in obsList:
                    xOffset1 = observation.getOutput(capsule.getAttributeByName(targetLabelX))
                    yOffset1 = observation.getOutput(capsule.getAttributeByName(targetLabelY))
                    
                    xOffset1 = int(xOffset1 * float(max(width, height)))
                    yOffset1 = int(yOffset1 * float(max(width, height)))

                    semantics[obsMap[observation]] = []
    
                    for inObs in observation.getInputObservations():
                        xOffset2 = inObs.getOutput(inObs.getCapsule().getAttributeByName(targetLabelX))
                        yOffset2 = inObs.getOutput(inObs.getCapsule().getAttributeByName(targetLabelY))
                        
                        xOffset2 = int(xOffset2 * float(max(width, height)))
                        yOffset2 = int(yOffset2 * float(max(width, height)))
                        
                        semantics[obsMap[observation]].append(patches.Arrow(xOffset1, yOffset1, xOffset2 - xOffset1, yOffset2 - yOffset1, linewidth = 1, edgecolor = 'r', facecolor = 'none' ))
                        
                    semantics[obsMap[observation]].append(patches.Circle((xOffset1, yOffset1), radius = 1, color = 'r'))    


        for layerIndex in range(self._numSemanticLayers - 1, -1, -1):
            # Iterate the list backwards
            for capsule in self._semanticLayers[layerIndex]:
                if capsule in obs.keys():
                    for observation in obs[capsule]:
                        # Only top level observed symbols that have not yet been generated
                        if not observation.getInputObservations():
                            newObslist = capsule.backwardPass(observation, False)
                            for obsCaps, newObs in newObslist.items():
                                if obsCaps in obs:
                                    obs[obsCaps].extend(newObs)
                                else:
                                    obs[obsCaps] = newObs

                    # Remove the parsed Observations
                    del obs[capsule]


        # Order all observations for the primitive capsules
        capsObsPairs = []
        for capsule, obsList in obs.items():
            if capsule in self._primitiveCapsules:
                for observation in obsList:
                    capsObsPairs.append((capsule, observation))
        capsObsPairs = sorted(capsObsPairs, key=lambda tup: tup[1].getProbability())


        image = np.zeros(width * height * 4)

        for capsule, observation in capsObsPairs:            
            pixelShape = self.getShapeByPixelCapsule(capsule.getPixelLayerInput())

            xOffset = observation.getOutput(capsule.getAttributeByName(targetLabelX))
            yOffset = observation.getOutput(capsule.getAttributeByName(targetLabelY))
            size = observation.getOutput(capsule.getAttributeByName(targetLabelSize))

            observation.setOutput(capsule.getAttributeByName(targetLabelX), 0.5)
            observation.setOutput(capsule.getAttributeByName(targetLabelY), 0.5)
            observation.setOutput(capsule.getAttributeByName(targetLabelSize), size * float(max(width, height)) / float(pixelShape[0]))

            obsPixelLayer = capsule.backwardPass(observation, False)

            pixelObs = list(obsPixelLayer.values())[0][0]
            pixelLay = list(obsPixelLayer.keys())[0]

            newxOffset = int(xOffset * float(max(width, height))) - int(pixelShape[0] / 2)
            newyOffset = int(yOffset * float(max(width, height))) - int(pixelShape[1] / 2)

            # We actually have far more accurate segmentation (including rotation, etc), but its hard to do nicely in matplotlib,
            # so we decided to just box it roughly.
            minX = newxOffset + pixelShape[0]
            maxX = newxOffset
            minY = newyOffset + pixelShape[1]
            maxY = newyOffset

            for xx in range(pixelShape[0]):
                for yy in range(pixelShape[1]):
                    if xx + newxOffset < width and xx + newxOffset >= 0 and yy + newyOffset < height and yy + newyOffset >= 0:
                        depth = pixelObs.getOutput(pixelLay.getAttributeByName("PixelD-" + str(xx) + "-" + str(yy)))
                        
                        if depth < 1.0:
                            minX = min(xx + newxOffset, minX)
                            maxX = max(xx + newxOffset, maxX)
                            minY = min(yy + newyOffset, minY)
                            maxY = max(yy + newyOffset, maxY)

                        if withBackground is True or depth < 1.0:
                            image[((yy + newyOffset) * width + xx + newxOffset) * 4] = pixelObs.getOutput(pixelLay.getAttributeByName("PixelC-" + str(xx) + "-" + str(yy)))

            if observation in obsMap:
                semantics[obsMap[observation]] = [patches.Rectangle((minX, minY), maxX - minX, maxY - minY, linewidth = 1, edgecolor = 'y', facecolor = 'none')]
            else:
                semantics[observation] = [patches.Rectangle((minX, minY), maxX - minX, maxY - minY, linewidth = 1, edgecolor = 'y', facecolor = 'none')]

            observation.setOutput(capsule.getAttributeByName(targetLabelX), xOffset)
            observation.setOutput(capsule.getAttributeByName(targetLabelY), yOffset)
            observation.setOutput(capsule.getAttributeByName(targetLabelSize), size)

            texts.append((minX, minY, capsule.getName()))

        return image, semantics, texts    # Linear List of Pixels


    def producePrimitiveObservations(self, observation : Observation):
        if observation.getCapsule() not in self._primitiveCapsules:
            outputObsList = []
            if not observation.getInputObservations():
                newObsDict = observation.getCapsule().backwardPass(observation, False)
                observation.clearInputObservations()
                for newObsList in newObsDict.values():
                    for newObs in newObsList:
                        outputObsList = outputObsList + self.producePrimitiveObservations(newObs)
            else:
                for newObs in observation.getInputObservations():
                    outputObsList = outputObsList + self.producePrimitiveObservations(newObs)
            return outputObsList    # List of Observations for Primitive Capsules
        else:
            return [observation]    # List of Observation for this Primitive Capsules


    def distance(self, observationA : Observation, observationB : Observation):
        distance = 100.0
        normal1 = [0.0, 0.0]
        normal2 = [0.0, 0.0]

        compareA = self.producePrimitiveObservations(observationA)
        compareB = self.producePrimitiveObservations(observationB)

        for primA in compareA:
            for primB in compareB:
                    testDistance, testNormal1, testNormal2 = self._renderer.getDistance(
                        self._capsulePrimitive[primA.getCapsule()], self._capsulePrimitive[primB.getCapsule()],
                        primA.getOutputs(), primB.getOutputs())
                    if testDistance < distance:
                        distance = testDistance
                        normal1 = testNormal1
                        normal2 = testNormal2

        return distance, normal1, normal2

    
