text: import numpy as np
  import pickle as pickle
  from .reader import sessionizer
  from .featurizer import extract_features
  
  from sklearn.model_selection import train_test_split
  from sklearn.neural_network import MLPClassifier
  from sklearn.metrics import f1_score
  
  from .training_utils import read_data
  from .training_utils import select_features
  
  
  class OneLayerModel:
      def __init__(self, duration, hidden_size=None, labels=None):
          '''
          Initializes a model with a single hidden layer.  Features are
          aggregated over the time specified by the duration and the hidden
          layer size is a hyperparameter set at initialization.
  
          Args:
              duration: Time duration to aggregate features for
          '''
  
          self.duration = duration
          self.hidden_size = hidden_size
  
          self.means = None
          self.stds = None
          self.feature_list = None
          self.model = None
          self.labels = labels
  
          self.sessions = None
  
      def _augment_data(self, X, y):
          '''
          Augments the data with randomly permuted samples. Added samples are
          labeled Unknown.
  
          Args:
              X: Numpy 2D array containing data to augment
              y: Numpy 1D array of labels to augment
  
          Returns:
              X_aug: Numpy 2D array containing augmented data
              y_aug: Numpy 1D array containing labels of augmented data
          '''
  
          # Randomly permute the inputs, and label the permutations Unknown
          X_permed = np.copy(X)
          for i in range(X_permed.shape[1]):
              np.random.shuffle(X_permed[:,i])
  
          y_permed = [self.labels.index("Unknown")]*X_permed.shape[0]
          y_permed = np.stack(y_permed)
  
          X_aug = np.concatenate((X, X_permed), axis=0)
          y_aug = np.concatenate((y, y_permed), axis=0)
  
          return X_aug, y_aug
  
      def get_features(self, filepath, source_ip=None):
          '''
          Reads a pcap specified by the file path and returns an array of the
          computed model inputs
  
          Args:
              filepath: Path to pcap to compute features for
  
          Returns:
              features: Numpy 2D array containing features for each time bin
              timestamp: datetime of the last observed packet
          '''
  
          # Read the capture into a feature array
          X = []
          timestamps = []
          binned_sessions = sessionizer(filepath, duration=self.duration)
          self.sessions = binned_sessions
  
          if len(binned_sessions) is 0:
              return None, None, None, None
  
          for session_dict in binned_sessions:
              if len(session_dict) > 0:
                  if source_ip is None:
                      feature_list, source_ip, other_ips = extract_features(
                                                                      session_dict
                                                                           )
                  else:
                      feature_list, _, other_ips = extract_features(
                                                        session_dict,
                                                        capture_source=source_ip
                                                                   )
                  X.append(feature_list)
                  last_packet = list(session_dict.items())[-1]
                  timestamps.append(last_packet[1][0][0])
  
          if len(X) == 0:
              return None, None, None, None
  
          full_features = np.stack(X)
  
          # Mean normalize the features
          full_features -= np.expand_dims(self.means, 0)
          full_features /= np.expand_dims(self.stds, 0)
          features = full_features[:, self.feature_list]
  
          '''
          last_packet = list(binned_sessions[-1].items())[-1]
          timestamp = last_packet[1][0][0]
          '''
  
          return features, source_ip, timestamps, other_ips
  
  
      def train(self, data_dir):
          '''
          Trains a single layer model on the data contained in the specified
          directory.  Labels found in the directory are augmented with an
          unknown label.
  
          Args:
              data_dir: Directory containing the training data
          '''
  
          print("Reading data")
          # First read the data directory for the features and labels
          X_all, y_all, new_labels = read_data(
                                                data_dir,
                                                duration=self.duration,
                                                labels=self.labels
                                              )
          self.labels = new_labels
  
          print("Making data splits")
          # Split the data into training, validation, and testing sets
          X_train, X_test, y_train, y_test = train_test_split(
                                                              X_all,
                                                              y_all,
                                                              test_size=0.2,
                                                              random_state=0
                                                             )
  
          print("Normalizing features")
          # Mean normalize the features, saving the means and variances
          self.means = X_train.mean(axis=0)
          self.stds = X_train.std(axis=0)
          # Set the zero standard deviations to 1
          zero_stds = self.stds <= 1
          self.stds[zero_stds] = 1
          # Apply the mean normalization transformation to the training dataj
          X_normed = X_train - np.expand_dims(self.means, 0)
          X_normed /= np.expand_dims(self.stds, 0)
  
          print("Doing feature selection")
          # Select the relevant features from the training set
          self.feature_list = select_features(X_normed, y_train)
          print(self.feature_list)
  
          # If hidden size wasn't specified, default to the mean of the number
          # of features and the size of the label space
          if self.hidden_size is None:
              self.hidden_size = int(1/2*(
                                          len(self.labels) + \
                                          len(self.feature_list)
                                         )
                                    )
  
          # Augment the data with randomly permuted samples
          X_aug, y_aug = self._augment_data(X_normed, y_train)
  
          # Fit the one layer model to the augmented training data
          X_input = X_aug[:, self.feature_list]
          self.model = MLPClassifier(
                                      (self.hidden_size),
                                      alpha=0.1,
                                      activation='relu',
                                      max_iter=1000
                                    )
  
          self.model.fit(X_input, y_aug)
  
          # Evaulate the model on the augmented test data
          X_test_input = X_test - np.expand_dims(self.means, 0)
          X_test_input /= np.expand_dims(self.stds, 0)
          X_test_aug, y_test_aug = self._augment_data(X_test_input, y_test)
          predictions = self.model.predict(X_test_aug[:, self.feature_list])
          print("F1 score:",
                  f1_score(y_test_aug, predictions, average='weighted'))
  
      def predict(self, filepath, source_ip=None):
          '''
          Read a capture file from the specified path and make a prediction
          of the source.
  
          Args:
              filepath: Path of capture file to read
  
          Returns:
              prediction: list of tuples formatted as (source, probability)
          '''
  
          features, _, _, _ = self.get_features(filepath, source_ip=source_ip)
  
          if features is None:
              return None
          predictions = self.model.predict_proba(features)
          mean_predictions = np.mean(predictions, axis=0)
  
          prediction = [
                        (self.labels[i], prob)
                        for i, prob in enumerate(mean_predictions)
                       ]
          prediction = sorted(prediction, key=lambda x: x[1], reverse=True)
          return prediction
  
      def get_representation(self, filepath, mean=True, source_ip=None):
          '''
          Computes the mean hidden representation of the input file.
  
          Args:
              filepath: Path of capture file to represent
              mean: If true(default), averages all the representations into one
  
          Returns:
              representation:  representation vector of the input file
          '''
  
          features, source_ip, timestamp, other_ips = self.get_features(
                                                             filepath,
                                                             source_ip=source_ip,
                                                                       )
          if features is None:
              return None, None, None, None, None
  
          L1_weights = self.model.coefs_[0]
          L1_biases = self.model.intercepts_[0]
          representation = np.maximum(
                                       np.matmul(features, L1_weights)+L1_biases,
                                       0
                                     )
  
          mean_rep = np.mean(representation, axis=0)
  
          L2_weights = self.model.coefs_[1]
          L2_biases = self.model.intercepts_[1]
          probabilities = np.matmul(representation, L2_weights) + L2_biases
          probabilities = np.exp(probabilities)
          probabilities /= np.expand_dims(np.sum(probabilities, axis=1), axis=1)
          probabilities = np.mean(probabilities, axis=0)
          prediction = [
                          (self.labels[i], prob)
                          for i, prob in enumerate(probabilities)
                       ]
          prediction = sorted(prediction, key=lambda x: x[1], reverse=True)
  
          if mean:
              representation = mean_rep
              timestamp = timestamp[-1]
  
          return representation, source_ip, timestamp, prediction, other_ips
  
      def classify_representation(self, representation):
          '''
          Takes in a representation and produces a classification
          '''
          L2_weights = self.model.coefs_[1]
          L2_biases = self.model.intercepts_[1]
          probabilities = np.matmul(representation, L2_weights) + L2_biases
          probabilities = np.exp(probabilities)
          probabilities /= np.sum(probabilities)
          prediction = [
                          (self.labels[i], prob)
                          for i, prob in enumerate(probabilities)
                       ]
          prediction = sorted(prediction, key=lambda x: x[1], reverse=True)
  
          return prediction
  
      def save(self, save_path):
          '''
          Saves the model to the specified file path
  
          Args:
              save_path: Path to store the saved model at.
          '''
  
          model_attributes = {
                              'duration': self.duration,
                              'hidden_size': self.hidden_size,
                              'means': self.means,
                              'stds': self.stds,
                              'feature_list': self.feature_list,
                              'model': self.model,
                              'labels': self.labels
                             }
  
          with open(save_path, 'wb') as handle:
              pickle.dump(model_attributes, handle)
  
      def load(self, load_path):
          '''
          Load the model parameters from the specified path.
  
          Args:
              load_path: Path to load the model parameters from
          '''
  
          with open(load_path, 'rb') as handle:
              model_attributes = pickle.load(handle)
  
          self.duration = model_attributes['duration']
          self.hidden_size = model_attributes['hidden_size']
          self.means = model_attributes['means']
          self.stds = model_attributes['stds']
          self.feature_list = model_attributes['feature_list']
          self.model = model_attributes['model']
          self.labels = model_attributes['labels']
  