text: """
  sample_usage.py
  Created on Oct 18 2020 15:05
  @author: Moayed Haji Ali mali18@ku.edu.tr
  
  """
  import numpy as np
  from basic_usage.sketchformer import continuous_embeddings
  import time
  import warnings
  import random
  
  import matplotlib
  matplotlib.use('TkAgg')
  import matplotlib.pyplot as plt
  
  # warnings.filterwarnings("ignore")
  
  
  class Basic_Test:
  
      def __init__(self):
          # prepare sample data from the quickdraw
          filename = "basic_usage/quickdraw_samples/sketchrnn_apple.npz"
          self.apples = np.load(filename, encoding='latin1', allow_pickle=True)
          filename = "basic_usage/quickdraw_samples/sketchrnn_baseball.npz"
          self.baseball = np.load(filename, encoding='latin1', allow_pickle=True)
  
      def performe_test(self, model):
          print("Performing tests:")
          # extract sample embedding of N apples, and M baseballs and observe the distances
          N_apple = N_baseball = 2
  
          apple_sketch, baseball_sketch = [], []
          apple_embedding, baseball_embedding = [], []
  
          # get random N_apple sketches 
          for _ in range(N_apple):
              ind = random.randint(0, len(self.apples['test']) - 1)
              apple_sketch.append(self.apples['test'][ind])
  
          # get random N_baseball sketches 
          for _ in range(N_baseball):
              ind = random.randint(0, len(self.baseball['test']) - 1)
              baseball_sketch.append(self.baseball['test'][ind])
          
          re_con = model.get_re_construction(np.concatenate((apple_sketch, baseball_sketch)))
          
  
          # visulaizing the reconstruction of the sketches 
          for sketch in re_con:
              self.visualize(sketch)
  
  
          embeddings = model.get_embeddings(np.concatenate((apple_sketch, baseball_sketch)))
  
          apple_embedding = embeddings[:N_apple]
          baseball_embedding = embeddings[N_apple:]
          for i, apple_emb1 in enumerate(apple_embedding):
              for j, apple_emb2 in enumerate(apple_embedding):
                  if i > j:
                      print("[Apple {} - Apple {}] embedding vectors norm: ".format(i, j), np.linalg.norm(apple_emb1 - apple_emb2))
  
          for i, base_emb1 in enumerate(baseball_embedding):
              for j, base_emb2 in enumerate(baseball_embedding):
                  if i > j:
                      print("[Baseball {} - Baseball {}] embedding vectors norm: ".format(i, j), np.linalg.norm(base_emb1 - base_emb2))
  
          for i, apple_emb in enumerate(apple_embedding):
              for j, base_emb in enumerate(baseball_embedding):
                      print("[Apple {} - Baseball {}] embedding vectors norm: ".format(i, j), np.linalg.norm(apple_emb - base_emb))
  
          # classify sample apple sketch
          pred_class = model.classify(np.concatenate((apple_sketch, baseball_sketch)))
  
          for i in range(N_apple):
              print("Predicted class for a Apple sketch: ", pred_class[i])
          
          for i in range(N_baseball):
              print("Predicted class for a Baseball sketch: ", pred_class[N_apple + i])
          # classify sample baseball sketch
  
      def visualize(self, sketch):
          X = []
          Y = []
  
          tmp_x, tmp_y = [], []
          sx = sy = 0
          for p in sketch:
              sx += p[0]
              sy += p[1]
              if p[2] == 1:
                  X.append(tmp_x)
                  Y.append(tmp_y)
                  tmp_x, tmp_y = [], []
              else:
                  tmp_x.append(sx)
                  tmp_y.append(-sy)
  
          X.append(tmp_x)
          Y.append(tmp_y)
  
          for x, y in zip(X, Y):
              plt.plot(x, y)
  
          # save the image.
          # plt.savefig("sample.png")
  
          # show the plot
          plt.show()
  
      def pre_trained_model_test(self):
          """peforme tests on the pretrained model
          """
          # obtain the pre-trained model
          sketchformer = continuous_embeddings.get_pretrained_model()
          self.performe_test(sketchformer)
      
      def new_model_test(self):
          """peforme tests on a new model trained on two classes (Apple, Baseball)
          """
          # train a new model
          print("Training a new model")
          MODEL_ID = "my_new_model"
          OUT_DIR = "basic_usage/pre_trained_model"
          sketches_x = np.concatenate((self.apples['train'][:300], self.baseball['train'][:300]))
          sketches_y = np.concatenate((np.zeros(len(self.apples['train'])), np.ones(len(self.baseball['train']))))
          new_model = continuous_embeddings(sketches_x, sketches_y, ['apple', 'baseball'], MODEL_ID, OUT_DIR, resume=False)
          self.performe_test(new_model)
  
          # using the embedding from the checkpoint
          print("Obtain the embeddings from the stored checkpoint of the new model")
  
          resume_model = continuous_embeddings([], [], ['apple', 'baseball'], MODEL_ID, OUT_DIR, resume=True)
          self.performe_test(resume_model)
  
  
  Basic_Test().pre_trained_model_test()
  
  Basic_Test().new_model_test()