--- a/scripts/learn-motion-patterns.py	Wed Jun 20 12:04:22 2018 -0400
+++ b/scripts/learn-motion-patterns.py	Wed Jun 20 16:48:20 2018 -0400
@@ -3,15 +3,16 @@
 import sys, argparse
 
 import numpy as np
+import matplotlib.pyplot as plt
 
 from trafficintelligence import ml, utils, storage, moving
 
-parser = argparse.ArgumentParser(description='The program learns prototypes for the motion patterns') #, 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 (the minimum cluster size argument is not used for now as it may change prototypes when assigning other trajectories)''') #, 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 = ['objectfeatures', 'feature', 'object'], default = 'objectfeatures')
+parser.add_argument('-t', dest = 'trajectoryType', help = 'type of trajectories to learn from', choices = ['objectfeature', 'feature', 'object'], default = 'objectfeatures')
 parser.add_argument('--max-nobjectfeatures', dest = 'maxNObjectFeatures', help = 'maximum number of features per object to load', type = int, default = 1)
 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)
@@ -24,25 +25,22 @@
 parser.add_argument('--subsample', dest = 'positionSubsamplingRate', help = 'rate of position subsampling (1 every n positions)', type = int)
 parser.add_argument('--display', dest = 'display', help = 'display trajectories', action = 'store_true')
 parser.add_argument('--save-similarities', dest = 'saveSimilarities', help = 'save computed similarities (in addition to prototypes)', action = 'store_true')
-parser.add_argument('--save-matches', dest = 'saveMatches', help = 'saves the assignments of the objects (not for features) to the prototypes', action = 'store_true')
+parser.add_argument('--save-assignments', dest = 'saveAssignments', help = 'saves the assignments of the objects to the prototypes', action = 'store_true')
 parser.add_argument('--assign', dest = 'assign', help = 'assigns the objects to the prototypes and saves the assignments', action = 'store_true')
 
 args = parser.parse_args()
 
 # use cases
 # 1. learn proto from one file, save in same or another
-# 2. load proto, load objects, update proto, save proto
-# 3. assign objects from one db to proto
-# 4. load objects from several files, save in another -> see metadata: site with view and times
-# 5. keep prototypes, with positions/velocities, in separate db (keep link to original data through filename, type and index)
+# 2. load proto, load objects (from same or other db), update proto matchings, save proto
+# TODO 3. on same dataset, learn and assign trajectories (could be done with min cluster size)
+# TODO? 4. when assigning, allow min cluster size only to avoid assigning to small clusters (but prototypes are not removed even if in small clusters, can be done after assignment with nmatchings)
 
 # TODO add possibility to cluster with velocities
-# TODO add possibilite to load all trajectories and use minclustersize
-# save the objects that match the prototypes
-# write an assignment function for objects
+# TODO add possibility to load all trajectories and use minclustersize
 
 # load trajectories to cluster or assign
-if args.trajectoryType == 'objectfeatures':
+if args.trajectoryType == 'objectfeature':
     trajectoryType = 'feature'
     objectFeatureNumbers = storage.loadObjectFeatureFrameNumbers(args.databaseFilename, objectNumbers = args.nTrajectories)
     featureNumbers = []
@@ -88,9 +86,7 @@
     #         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)
-    prototypeIndices, labels = ml.assignToPrototypeClusters(trajectories, prototypeIndices, similarities, args.minSimilarity, similarityFunc)
-    clusterSizes = ml.computeClusterSizes(labels, prototypeIndices, -1)
-    print(clusterSizes)
+    assignedPrototypeIndices, labels = ml.assignToPrototypeClusters(trajectories, prototypeIndices, similarities, args.minSimilarity, similarityFunc)
 
 if args.learn and not args.assign:
     prototypes = []
@@ -107,39 +103,52 @@
         if args.inputPrototypeDatabaseFilename == args.outputPrototypeDatabaseFilename:
             storage.deleteFromSqlite(args.outputPrototypeDatabaseFilename, 'prototype')
     storage.savePrototypesToSqlite(outputPrototypeDatabaseFilename, prototypes)
+    if args.display:
+        plt.figure()
+        for p in prototypes:
+            p.getMovingObject().plot()
+        plt.axis('equal')
+        plt.show()
 
-if not args.learn and args.assign: # no new prototypes # not save assignments of past prototypes if removes with minClusterSize
-    prototypes = []
+if not args.learn and args.assign: # no modification to prototypes, can work with initialPrototypes
+    clusterSizes = ml.computeClusterSizes(labels, prototypeIndices, -1)
     for i in prototypeIndices:
         nMatchings = clusterSizes[i]-1
         if initialPrototypes[i].nMatchings is None:
             initialPrototypes[i].nMatchings = nMatchings
         else:
             initialPrototypes[i].nMatchings += nMatchings
-        prototypes.append(initialPrototypes[i])
     if args.outputPrototypeDatabaseFilename is None:
         outputPrototypeDatabaseFilename = args.databaseFilename
     else:
         outputPrototypeDatabaseFilename = args.outputPrototypeDatabaseFilename
-    storage.setPrototypeMatchingsInSqlite(outputPrototypeDatabaseFilename, prototypes)
-
-    labelsToProtoIndices = {protoId: i for i, protoId in enumerate(prototypeIndices)}
-    if args.saveMatches:
-        storage.savePrototypeAssignmentsToSqlite(args.databaseFilename, objects, trajectoryType, [labelsToProtoIndices[l] for l in labels], prototypes)
+    storage.setPrototypeMatchingsInSqlite(outputPrototypeDatabaseFilename, initialPrototypes)
+    if args.saveAssignments:
+        if args.trajectoryType == 'objectfeature': # consider that the object is assigned through its longest features
+            objectNumbers = []
+            objectLabels = []
+            for objNum, objFeatureNumbers in objectFeatureNumbers.items():
+                objLabels = []
+                for i, o in enumerate(objects):
+                    if o.getNum() in objFeatureNumbers:
+                        objLabels.append(labels[i+len(initialPrototypes)])
+                objectLabels.append(utils.mostCommon(objLabels))
+                objectNumbers.append(objNum)
+            storage.savePrototypeAssignmentsToSqlite(args.databaseFilename, objectNumbers, 'object', objectLabels, initialPrototypes)
+        else:
+            storage.savePrototypeAssignmentsToSqlite(args.databaseFilename, [obj.getNum() for obj in objects], trajectoryType, labels[len(initialPrototypes):], initialPrototypes)
+    if args.display:
+        plt.figure()
+        for i,o in enumerate(objects):
+            if labels[i+len(initialPrototypes)] < 0:
+                o.plot('kx-')
+            else:
+                o.plot(utils.colors[labels[i+len(initialPrototypes)]])
+        for i,p in enumerate(initialPrototypes):
+            p.getMovingObject().plot(utils.colors[i]+'o')
+        plt.axis('equal')
+        plt.show()
 
 if (args.learn or args.assign) and args.saveSimilarities:
     np.savetxt(utils.removeExtension(args.databaseFilename)+'-prototype-similarities.txt.gz', similarities, '%.4f')
 
-if args.display and args.assign:
-    from matplotlib.pyplot import figure, show, axis
-    figure()
-    for i,o in enumerate(objects):
-        if i not in prototypeIndices:
-            if labels[i] < 0:
-                o.plot('kx')
-            else:
-                o.plot(utils.colors[labels[i]])
-    for i in prototypeIndices:
-            objects[i].plot(utils.colors[i]+'o')
-    axis('equal')
-    show()