From 2b4edbda8cf29454aec6008088e50684101a37c1 Mon Sep 17 00:00:00 2001
From: g <g@g>
Date: Sat, 20 Apr 2019 17:56:26 -0400
Subject: [PATCH] fixes

---
 utils/similarity/__init__.py | 13 ++-----------
 utils/similarity/plot.py     | 33 ++++++++++++++++++++++-----------
 2 files changed, 24 insertions(+), 22 deletions(-)

diff --git a/utils/similarity/__init__.py b/utils/similarity/__init__.py
index dc48cf579..5a1473b46 100644
--- a/utils/similarity/__init__.py
+++ b/utils/similarity/__init__.py
@@ -7,15 +7,6 @@ except ImportError:
     fastText = None
 
 if fastText is None:
-    print("In order to use slither-simil, you need to install fastText 0.2.0:")
-    print("$ pip3 install https://github.com/facebookresearch/fastText/archive/0.2.0.zip --user")
+    print("ERROR: in order to use slither-simil, you need to install fastText 0.2.0:")
+    print("$ pip3 install https://github.com/facebookresearch/fastText/archive/0.2.0.zip --user\n")
     sys.exit(-1)
-
-try:
-    import sklearn
-except ImportError:
-    sklearn = None
-
-if sklearn is None:
-    print("In order to use plot mode in slither-simil, you need to install sklearn:")
-    print("$ pip3 install sklearn --user")
diff --git a/utils/similarity/plot.py b/utils/similarity/plot.py
index 84fcd69b9..d6bd75e33 100644
--- a/utils/similarity/plot.py
+++ b/utils/similarity/plot.py
@@ -3,19 +3,30 @@ import sys
 import traceback
 import operator
 import numpy as np
+import tqdm
 import random
 
-from sklearn import decomposition
-import matplotlib.pyplot as plt
-
+try:
+    from sklearn import decomposition
+    import matplotlib.pyplot as plt
+except ImportError:
+    decomposition = None
+    plt = None
+    
 from fastText import load_model
 from .cache import load_cache
 
-logger = logging.getLogger("crytic-pred")
+logger = logging.getLogger("Slither-simil")
 
 def plot(args):
 
+    if decomposition is None or plt is None:
+        print("ERROR: In order to use plot mode in slither-simil, you need to install sklearn and matplotlib:")
+        print("$ pip3 install sklearn matplotlib --user")
+        sys.exit(-1)
+
     try:
+
         model = args.model
         model = load_model(model)
         filename = args.filename
@@ -29,32 +40,32 @@ def plot(args):
             logger.error('The plot mode requieres contract, fname and input parameters.')
             sys.exit(-1)
 
+        logger.info('Loading data..')
         cache = load_cache(infile, model, ext=ext, solc=solc)
-        #save_cache("cache.npz", cache)
 
         data = list()
         fs = list()
+
+        logger.info('Procesing data..')
         for (f,c,n),y in cache.items():
             if c == contract and n == fname:
                 fs.append(f)
                 data.append(y)
-            #r[x] = similarity(fvector, y)
-
        
         data = np.array(data)
         pca = decomposition.PCA(n_components=2)
         tdata = pca.fit_transform(data)
+
+        logger.info('Plotting data..')
         plt.figure()
         assert(len(tdata) == len(fs))
         for ([x,y],l) in zip(tdata, fs):
             x = random.gauss(0, 0.01) + x
             y = random.gauss(0, 0.01) + y
             plt.scatter(x, y, c='blue')
-            plt.text(x-0.001,y+0.001, l.split("_")[1].replace(".sol.ast.compact.json",""))
+            #plt.text(x-0.001,y+0.001, l.split("_")[1].replace(".sol.ast.compact.json",""))
 
-        plt.show()
-        #r = sorted(r.items(), key=operator.itemgetter(1), reverse=True)
-        #for x,score in r[:10]:
+        plt.savefig('plot.png', bbox_inches='tight')
  
     except Exception:
         logger.error('Error in %s' % args.filename)