import sys
import csv
from ch2017_train.ch2017_mimo_trainer import Ch2017MimoTrainer 
from data_io.mimo_wfdb_adaptor import MimoWfdbAdaptor 
from utils import my_numpy_utils as my_np
import gc
import keras.backend.tensorflow_backend as K
from mimo_sequence.mimo_sequence import MimoSequence

#block_seq_base_length=6000
#out_dec_factor=18
class PredictClassWfdb():
  def __init__(self):
    self.ch2017_mimo_trainer = Ch2017MimoTrainer(train_batch_size = 1, block_seq_base_length=sys.maxsize)
    self.mimo_wfdb_adaptor = MimoWfdbAdaptor()
    self.weights_filename = "best_weights.hdf5"
        
  def read_file(self, filename):
#     self.mimo_wfdb_adaptor.mimo = mimo_sequence(tag=filename, n_input_dim=1, 
#                         n_output_dim = 1, max_length=None, labels=self.mimo_wfdb_adaptor.labels, 
#                         out_dec_factor=self.Ch2017MimoTrainer.out_dec_factor)
    self.mimo_wfdb_adaptor.read_single_input(filename, 
                      out_dec_factor=self.ch2017_mimo_trainer.out_dec_factor)    
    #self.mimo_wfdb_adaptor.mimo.generate_block_input()
    
  def prepare_data(self, mimo, batch_size):
    # data is ready; nothing to do
    return self.mimo_wfdb_adaptor.mimo
  
  def do_test(self, weights_filename=None, do_plot=False):
    batch_size = 1
    test_model = self.ch2017_mimo_trainer.create_model(
      self.mimo_wfdb_adaptor.mimo, batch_size, device_id = '/cpu:0');
    self.ch2017_mimo_trainer.do_test(test_model, 
      self.mimo_wfdb_adaptor.mimo, batch_size, weights_filename=weights_filename, 
      csvOutputFile='answer.txt', do_evaluate=False)

    if do_plot:    
      print(self.mimo_wfdb_adaptor.mimo.block_output)
      self.mimo_wfdb_adaptor.mimo.plot_input()
      self.mimo_wfdb_adaptor.mimo.plot_output()

    # use last prediction as classification
    predictions = my_np.prob_to_class(self.mimo_wfdb_adaptor.mimo.block_output[0, :, :])

    # -1 to remove the 'no-signal class'
    print('Classification:' + self.mimo_wfdb_adaptor.labels[predictions[-1]-1])
    
    #K.clear_session()
    self.ch2017_mimo_trainer.shut_down()
    
    
    #gc.collect()
  

    

if __name__ == "__main__":
  pc = PredictClassWfdb()
  pc.read_file(sys.argv[1])
  if len(sys.argv) > 2:
    weights_filename=sys.argv[2]
  else:
    weights_filename=None  
  pc.do_test(weights_filename)
  #print('Usage: predict_class [filename]')
   
