'''
  This script aims to extract k-means clustering from a 
  a priori trained k-means.
  '''
  
  import argparse
  import numpy as np
  import pickle
  from data import read_file, index_by_id, write_line
  import sys
  
  # -- ARGPARSE --
  parser = argparse.ArgumentParser(description="extract clusters")
  parser.add_argument("model", type=str, help="k-means model pickle")
  parser.add_argument("features", type=str, help="features")

  parser.add_argument("--list", type=str, default=None, help="list file")

  parser.add_argument("--outfile", type=str, default=None, help="output file std")
  
  args = vars(parser.parse_args())
  MODEL = args["model"]
  FEATURES = args["features"]
  LST = args["list"]
  OUTFILE = args["outfile"]
  
  if OUTFILE == None:
      OUTFILE = sys.stdout
  else:
      OUTFILE = open(OUTFILE, "w")
  
  # -- READ FILE --
  features = read_file(FEATURES)
  feat_ind = index_by_id(features)

  if LST is not None:  
      lst = read_file(LST)
  else:
      lst = features

  
  kmeans = pickle.load(open(MODEL, "rb"))

  # -- CONVERT TO NUMPY --
  X = np.asarray([feat_ind[x[0][0]][x[0][3]][1] for x in lst])

  predictions = kmeans.predict(X)
  
  for i, line in enumerate(lst):
      meta = line[0]
      meta[1] = str(predictions[i])
      write_line(
          meta,
          feat_ind[meta[0]][meta[3]][1],
          OUTFILE
      )
  
  # -- CLOSE OUT FILE IF NECESSARY --
  if not OUTFILE == sys.stdout:
      OUTFILE.close()