diff --git a/home/alegen/ipv6tests/train.py b/home/alegen/repos/nmap/luis/ipv6tests/train.py
index e4cd9eb..6fdf210 100755
--- a/home/alegen/ipv6tests/train.py
+++ b/home/alegen/repos/nmap/luis/ipv6tests/train.py
@@ -24,7 +24,10 @@ The output is a .model file.
   -g, --group=GROUP_FILENAME  read groups from GROUP_FILENAME (required).
   -h, --help                  show this help.
   -s, --set=SET_FILENAME      use the set of features in SET_FILENAME (required).
-  --scale                     scale feature vectors to the range [0, 1].\
+  --scale                     scale feature vectors to the range [0, 1].
+  --imputations               number of different imputation sets to be generated with mice (default 20).
+  --iterations                number of iterations per imputation set (default 50).
+  --merge                     merge nmap fingerprint databased (groups file) with previous mice output if available (default 1).
 """ % sys.argv[0]
 
 class options (object):
@@ -69,15 +72,18 @@ def scale(features):
                 features[j, i] = (features[j, i] - mn) / denom
     return features, s_min, s_max
 
-def prepare_features(groups, do_scale):
+def prepare_features(groups, imp_methods, do_scale):
     """Impute and scale features, and assign them back to groups."""
     feature_list = []
+    k_group = 0
     for group in groups:
+        k_group += 1
         for features in group.features:
-            feature_list.append(features)
+            feature_list.append([k_group] + features)
 
     feature_matrix = np.vstack(feature_list)
-    feature_matrix = impute.impute(feature_matrix)
+    feature_matrix = impute.impute_mice(feature_matrix, ['skip'] + imp_methods, **imp_args)
+    feature_matrix = np.delete(feature_matrix, 0, 1) # remove first column which is group identifier
     if do_scale:
         feature_matrix, s_min, s_max = scale(feature_matrix)
         scale_params = zip(s_min, s_max)
@@ -106,8 +112,9 @@ def train(groups, cost=1):
     # Re-train to get an accuracy estimate.
     param.cross_validation = True
     param.nr_fold = 5
-    acc = ll.train(prob, param)
-    print >> sys.stderr, "Accuracy", acc
+    print "Accuracy estimates"
+    for i in range(5):
+        print ll.train(prob, param)
     return model
 
 def mean(a):
@@ -174,7 +181,7 @@ def save_model(f, feature_names, groups, model, scale_params = None):
     print >>f, "end liblinear"
 
 def main(set_filename, group_filename, cost, do_scale):
-    feature_names = parse.parse_feature_set_file(set_filename)
+    feature_names, imp_methods = parse.parse_feature_set_file(set_filename)
 
     groups = parse.parse_groups_file(group_filename)
     for group in groups:
@@ -183,15 +190,18 @@ def main(set_filename, group_filename, cost, do_scale):
             features = vectorize.vectorize(feature_names, rs)
             group.features.append(features)
 
-    scale_params = prepare_features(groups, do_scale)
+    scale_params = prepare_features(groups, imp_methods, do_scale)
 
     print >> sys.stderr, "Training."
     model = train(groups, cost)
 
-    save_model(sys.stdout, feature_names, groups, model, scale_params)
+    model_file = open('nmap.model', 'w')
+    save_model(model_file, feature_names, groups, model, scale_params)
+    model_file.close()
 
 if __name__ == "__main__":
-    opts, args = getopt.gnu_getopt(sys.argv[1:], "c:g:hs:", ["group=", "help", "set=", "scale"])
+    opts, args = getopt.gnu_getopt(sys.argv[1:], "c:g:hs:", ["group=", "help", "set=", "scale", "imputations=", "iterations=", "merge="])
+    imp_args = {}
     for o, a in opts:
         if o == "-c":
             options.cost = float(a)
@@ -204,6 +214,12 @@ if __name__ == "__main__":
             options.set_filename = a
         elif o == "--scale":
             options.do_scale = True
+        elif o == "--imputations":
+            imp_args["imputations"] = int(a)
+        elif o == "--iterations":
+            imp_args["iterations"] = int(a)
+        elif o == "--merge":
+            imp_args["merge"] = bool(a)
 
     if options.set_filename is None:
         usage(sys.stderr)