'''
Created on Apr 18, 2017

@author: philip
'''
import os
import sys
import numpy as np

from sklearn.model_selection import StratifiedKFold
from sklearn.model_selection import train_test_split

from train.mimo_trainer import MimoTrainer
from mimo_sequence.mimo_sequence import MimoSequence 

class MimoExperiment:
    '''
    Perform experiments on MimoSequence classes
    '''


    def __init__(self):
        '''
        Constructor
        '''
            
    def stratified_cross_validation(self, trainer, mimo, n_folds):
        
      skf = StratifiedKFold(n_splits=n_folds)
      skf.get_n_splits()
      print(skf)  
      
      n_seq = mimo.n_sequences()
      targets = np.zeros((n_seq,))
      for i in range(0, n_seq):
        targets[i] = mimo.target[i][0]
      
      dummy = np.zeros(n_seq)
      i=0
      for train_index, test_index in skf.split(dummy, targets):
        i += 1
        if i < 1:
          continue 
        if i > 10:
          break      
        
        print("*************Fold " + str(i) + '**************')
            
        fold_dir = 'Fold' + str(i); 
        if not os.path.exists(fold_dir):
          os.makedirs(fold_dir)       
        os.chdir(fold_dir)
        
        train_index2, validation_index = train_test_split(train_index, test_size=1/(n_folds), random_state=42, stratify=targets[train_index])
  
        print('Fold:' + str(i))
        print("TRAIN: " + str(len(train_index2)) + " examples", train_index2)
        print(targets[train_index2])       
        for j in range(0, mimo.n_labels):
          n_matches = len(np.where(targets[train_index2] == j)[0])
          print('# ' + mimo.labels[j] +  ': ' + str(n_matches))
        
        print("VALIDATION: " + str(len(validation_index)) + " examples", validation_index)
        print(targets[validation_index])       
        for j in range(0, mimo.n_labels):
          n_matches = len(np.where(targets[validation_index] == j)[0])
          print('# ' + mimo.labels[j] +  ': ' + str(n_matches))
        
        print("TEST: " + str(len(test_index)) + " examples", test_index)
        print(targets[test_index])
        for j in range(0, mimo.n_labels):
          n_matches = len(np.where(targets[test_index] == j)[0])
          print('# ' + mimo.labels[j] +  ': ' + str(n_matches))
         
         
        train_mimo, dummy = mimo.filter_by_index(train_index2, 'train')
        validation_mimo, dummy = mimo.filter_by_index(validation_index, 'validation')
        test_mimo, dummy = mimo.filter_by_index(test_index, 'test')
        print('Initial # training sequences: ' + str(train_mimo.n_sequences())) 
        print('Initial # validation sequences: ' + str(validation_mimo.n_sequences()))
        print('Initial # test sequences: ' + str(test_mimo.n_sequences()))
        
        train_mimo = train_mimo.split_long(trainer.block_seq_base_length)
        validation_mimo = validation_mimo.split_long(trainer.block_seq_base_length)
        
        train_mimo.generate_block_sequences()
        validation_mimo.generate_block_sequences()
        
        train_mimo.plot_input()
        validation_mimo.plot_input()
       
        train_mimo.plot_target() 
        validation_mimo.plot_target()    
        
        print('Batch size: ' + str(trainer.train_batch_size))
        
        train_mimo.pad_number_of_sequences(trainer.train_batch_size)
        validation_mimo.pad_number_of_sequences(trainer.train_batch_size)
         
        print('Padded # training sequences: ' + str(train_mimo.n_sequences())) 
        print('Padded # validation sequences: ' + str(validation_mimo.n_sequences()))
        
        train_model = trainer.create_model(train_mimo, trainer.train_batch_size, train_mimo.max_length)
        cycleTag = '';
        for j in range(trainer.n_trainTestCycles):
          if trainer.n_trainTestCycles>1:
            cycleTag = '_Cycle' + str(j);
          trainer.weights_filename = trainer.weights_filebase + cycleTag + trainer.weights_file_ext
     
          trainer.do_train(train_model, train_mimo, validation_mimo, trainer.train_batch_size, tag=cycleTag)
       
          trainer.do_test(train_model, train_mimo, trainer.train_batch_size, tag=cycleTag, csvOutputFile='.'+'/trainOutput.csv')
          trainer.do_test(train_model, validation_mimo, trainer.train_batch_size, tag=cycleTag, csvOutputFile='.'+'/validationOutput.csv')
      
        trainer.shut_down()
  
        batch_test_size = 1
        test_mimo = trainer.prepare_data(test_mimo, batch_test_size)
        
        test_mimo = test_mimo.split_long(sys.maxsize)
        test_mimo.generate_block_sequences()
  
        testModel = trainer.create_model(test_mimo, batch_test_size, test_mimo.max_length)
        trainer.do_test(testModel, test_mimo, batch_test_size, tag='', csvOutputFile='.'+'/testOutput.csv')
    
        trainer.shut_down()
        os.chdir('../')
         
      print('Cross validation complete')
