--- a/scripts/learn-motion-patterns.py	Wed Jul 04 17:39:39 2018 -0400
+++ b/scripts/learn-motion-patterns.py	Thu Jul 05 17:06:40 2018 -0400
@@ -5,20 +5,20 @@
 import numpy as np
 import matplotlib.pyplot as plt
 
-from trafficintelligence import ml, utils, storage, moving
+from trafficintelligence import ml, utils, storage, moving, processing
 
-parser = argparse.ArgumentParser(description='''The program clusters trajectories, each cluster being represented by a trajectory. It can either work on the same dataset (database) or different ones, but only does learning or assignment at a time to avoid issues (the minimum cluster size argument is not used for now as it may change prototypes when assigning other trajectories)''') #, epilog = ''
+parser = argparse.ArgumentParser(description='''The program clusters trajectories, each cluster being represented by a trajectory. It can either work on the same dataset (database) or different ones, but only does learning or assignment at a time to avoid issues''') #, epilog = ''
 #parser.add_argument('--cfg', dest = 'configFilename', help = 'name of the configuration file')
 parser.add_argument('-d', dest = 'databaseFilename', help = 'name of the Sqlite database file', required = True)
 parser.add_argument('-o', dest = 'outputPrototypeDatabaseFilename', help = 'name of the Sqlite database file to save prototypes')
 parser.add_argument('-i', dest = 'inputPrototypeDatabaseFilename', help = 'name of the Sqlite database file for prototypes to start the algorithm with')
-parser.add_argument('-t', dest = 'trajectoryType', help = 'type of trajectories to learn from', choices = ['objectfeature', 'feature', 'object'], default = 'objectfeature')
-parser.add_argument('--max-nobjectfeatures', dest = 'maxNObjectFeatures', help = 'maximum number of features per object to load', type = int, default = 1)
+parser.add_argument('-t', dest = 'trajectoryType', help = 'type of trajectories to process', choices = ['feature', 'object'], default = 'feature')
+parser.add_argument('--nfeatures-per-object', dest = 'nLongestFeaturesPerObject', help = 'maximum number of features per object to load', type = int)
 parser.add_argument('-n', dest = 'nTrajectories', help = 'number of the object or feature trajectories to load', type = int, default = None)
 parser.add_argument('-e', dest = 'epsilon', help = 'distance for the similarity of trajectory points', type = float, required = True)
 parser.add_argument('--metric', dest = 'metric', help = 'metric for the similarity of trajectory points', default = 'cityblock') # default is manhattan distance
 parser.add_argument('-s', dest = 'minSimilarity', help = 'minimum similarity to put a trajectory in a cluster', type = float, required = True)
-parser.add_argument('-c', dest = 'minClusterSize', help = 'minimum cluster size', type = int, default = 0)
+#parser.add_argument('-c', dest = 'minClusterSize', help = 'minimum cluster size', type = int, default = 0)
 parser.add_argument('--learn', dest = 'learn', help = 'learn', action = 'store_true')
 parser.add_argument('--optimize', dest = 'optimizeCentroid', help = 'recompute centroid at each assignment', action = 'store_true')
 parser.add_argument('--random', dest = 'randomInitialization', help = 'random initialization of clustering algorithm', action = 'store_true')
@@ -40,57 +40,41 @@
 # TODO add possibility to cluster with velocities
 # TODO add possibility to load all trajectories and use minclustersize
 
-# load trajectories to cluster or assign
-objects = storage.loadTrajectoriesFromSqlite(args.databaseFilename, args.trajectoryType, args.nTrajectories, timeStep = args.positionSubsamplingRate)
-trajectories = [o.getPositions().asArray().T for o in objects]
+if args.learn and args.assign:
+    print('Cannot learn and assign simultaneously')
+    sys.exit(0)
+
+objects = storage.loadTrajectoriesFromSqlite(args.databaseFilename, args.trajectoryType, args.nTrajectories, timeStep = args.positionSubsamplingRate, nLongestFeaturesPerObject = args.nLongestFeaturesPerObject)
+if args.trajectoryType == 'object' and args.nLongestFeaturesPerObject is not None:
+    objectsWithFeatures = objects
+    objects = [f for o in objectsWithFeatures for f in o.getFeatures()]
+    prototypeType = 'feature'
+else:
+    prototypeType = args.trajectoryType
 
 # load initial prototypes, if any    
 if args.inputPrototypeDatabaseFilename is not None:
     initialPrototypes = storage.loadPrototypesFromSqlite(args.inputPrototypeDatabaseFilename, True)
-    trajectories = [p.getMovingObject().getPositions().asArray().T for p in initialPrototypes]+trajectories
-    if len(initialPrototypes) > 0:
-        initialPrototypeIndices = list(range(len(initialPrototypes)))
-    else:
-        initialPrototypeIndices = None
 else:
     initialPrototypes = []
-    initialPrototypeIndices = None
 
 lcss = utils.LCSS(metric = args.metric, epsilon = args.epsilon)
+similarityFunc = lambda x,y : lcss.computeNormalized(x, y)
+nTrajectories = len(initialPrototypes)+len(objects)
 if args.similaritiesFilename is not None:
     similarities = np.loadtxt(args.similaritiesFilename)
-if args.similaritiesFilename is None or similarities.shape[0] != len(trajectories) or similarities.shape[1] != len(trajectories):
-    similarities = -np.ones((len(trajectories), len(trajectories)))
-similarityFunc = lambda x,y : lcss.computeNormalized(x, y)
-# the next line can be called again without reinitializing similarities
-if args.learn:
-    prototypeIndices = ml.prototypeCluster(trajectories, similarities, args.minSimilarity, similarityFunc, args.optimizeCentroid, args.randomInitialization, initialPrototypeIndices)
-else:
-    prototypeIndices = initialPrototypeIndices
+if args.similaritiesFilename is None or similarities.shape[0] != nTrajectories or similarities.shape[1] != nTrajectories:
+    similarities = -np.ones((nTrajectories, nTrajectories))
 
-if args.assign: # TODO don't touch initial prototypes if not from same db as trajectories
-    #if not args.learn and args.minClusterSize >= 1: # allow only 
-    #   print('Warning: you did not learn the prototypes and you are using minimum cluster size of {}, which may lead to removing prototypes and assigning them to others'.format(args.minClusterSize))
-    # if args.minClusterSize >= 1:
-    #     if initialPrototypeIndices is None:
-    #         prototypeIndices, labels = ml.assignToPrototypeClusters(trajectories, prototypeIndices, similarities, args.minSimilarity, similarityFunc, args.minClusterSize)
-    #     else:
-    #         print('Not assigning with non-zero minimum cluster size and initial prototypes (would remove initial prototypes based on other trajectories')
-    # else:
-    #     prototypeIndices, labels = ml.assignToPrototypeClusters(trajectories, prototypeIndices, similarities, args.minSimilarity, similarityFunc)
-    assignedPrototypeIndices, labels = ml.assignToPrototypeClusters(trajectories, prototypeIndices, similarities, args.minSimilarity, similarityFunc)
+prototypeIndices, labels = processing.learnAssignMotionPatterns(args.learn, args.assign, objects, similarities, args.minSimilarity, similarityFunc, 0, args.optimizeCentroid, args.randomInitialization, False, initialPrototypes)
 
-if args.learn and not args.assign:
+if args.learn:# and not args.assign:
     prototypes = []
-    if args.trajectoryType == 'objectfeature':
-        trajectoryType = 'feature'
-    else:
-        trajectoryType = args.trajectoryType
     for i in prototypeIndices:
         if i<len(initialPrototypes):
             prototypes.append(initialPrototypes[i])
         else:
-            prototypes.append(moving.Prototype(args.databaseFilename, objects[i-len(initialPrototypes)].getNum(), trajectoryType))
+            prototypes.append(moving.Prototype(args.databaseFilename, objects[i-len(initialPrototypes)].getNum(), prototypeType))
 
     if args.outputPrototypeDatabaseFilename is None:
         outputPrototypeDatabaseFilename = args.databaseFilename
@@ -106,10 +90,10 @@
         plt.axis('equal')
         plt.show()
 
-if not args.learn and args.assign: # no modification to prototypes, can work with initialPrototypes
+if args.assign: # not args.learn and  no modification to prototypes, can work with initialPrototypes
     clusterSizes = ml.computeClusterSizes(labels, prototypeIndices, -1)
     for i in prototypeIndices:
-        nMatchings = clusterSizes[i]-1
+        nMatchings = clusterSizes[i]-1 # external prototypes
         if initialPrototypes[i].nMatchings is None:
             initialPrototypes[i].nMatchings = nMatchings
         else:
@@ -120,16 +104,22 @@
         outputPrototypeDatabaseFilename = args.outputPrototypeDatabaseFilename
     storage.setPrototypeMatchingsInSqlite(outputPrototypeDatabaseFilename, initialPrototypes)
     if args.saveAssignments:
-        if args.trajectoryType == 'objectfeature': # consider that the object is assigned through its longest features
+        if args.trajectoryType == 'object' and args.nLongestFeaturesPerObject is not None:
+            # consider that the object is assigned through its longest features
+            # issues are inconsistencies in the number of matchings per prototype and display (will display features, not objects)
             objectNumbers = []
             objectLabels = []
-            for objNum, objFeatureNumbers in objectFeatureNumbers.items():
+            i = 0
+            for obj in objectsWithFeatures:
                 objLabels = []
-                for i, o in enumerate(objects):
-                    if o.getNum() in objFeatureNumbers:
+                for f in obj.getFeatures():
+                    if f == objects[i]:
                         objLabels.append(labels[i+len(initialPrototypes)])
+                        i += 1
+                    else:
+                        print('Issue with obj {} and feature {} (trajectory {})'.format(obj.getNum(), f.getNum(), i))
                 objectLabels.append(utils.mostCommon(objLabels))
-                objectNumbers.append(objNum)
+                objectNumbers.append(obj.getNum())
             storage.savePrototypeAssignmentsToSqlite(args.databaseFilename, objectNumbers, 'object', objectLabels, initialPrototypes)
         else:
             storage.savePrototypeAssignmentsToSqlite(args.databaseFilename, [obj.getNum() for obj in objects], args.trajectoryType, labels[len(initialPrototypes):], initialPrototypes)