text: """
  Last change: 29.07.2018
  """
  
  import torch
  import torch.optim as optim
  from torch.autograd import Variable
  from torch.autograd import gradcheck
  import torch.nn as nn
  import numpy as np
  import matplotlib.pyplot as plt
  import matplotlib
  import time
  
  
  class Solver:
  
      """
      Functions:
  
       - __init__() : Initializer
       - reset_histories() : Resets loss and accuracy histories
       - train_manual() : Takes one whole dataset and trains it by going through it with many iterations
       - train_loader() : Takes one, into many mini-batches divided dataset as an input and trains it...
       - check_accuracy() : Checks accuracy of provided model on inserted dataset
       - test(): for testing
      """
  
      # Set global font for all matplotlib plots
      matplotlib.rcParams.update({'font.size': 24})
  
      # Attributes
      default_adam_args = {"lr": 1e-3,
                           "betas": (0.9, 0.999),
                           "eps": 1e-8,
                           "weight_decay": 0.0}
  
      default_rmsprop_args = {"lr": 1e-3,
                              "alpha": 0.99,
                              "eps": 1e-8,
                              "weight_decay": 0.0,
                              "momentum": 0.9}
  
      default_sgd_args = {"lr": 1e-3,
                          "momentum": 0.9}
  
      def __init__(self, optimizer=optim.Adam, optimizer_name='adam', loss_func=nn.MSELoss(), optim_args={}):
  
          """
          Initializes solver
  
          - optimizer: optimization algorithm
          - optimizer_name: name of the optimization algorithm for easy distinguishing in the if-block
          - loss_func: type of the loss function
          - optim_args: parameters for a particular optimization algorithm
          """
  
          self.loss_func = loss_func
  
          # Distinguish between different optimization algorithms
          if optimizer_name == 'adam':
              optim_args_merged = self.default_adam_args.copy()
          elif optimizer_name == 'sgd':
              optim_args_merged = self.default_sgd_args
          elif optimizer_name == 'rmsprop':
              optim_args_merged = self.default_rmsprop_args
  
          optim_args_merged.update(optim_args)
          self.optim_args = optim_args_merged
          self.optimizer = optimizer
          #self.scheduler = optim.lr_scheduler.MultiStepLR(optim.Adam, gamma=0.1, milestones=[0, 500])
  
      def reset_histories(self):
  
          """
          Resets all histories
          """
  
          self.train_loss_history = []
          self.val_loss_history = []
          self.train_av_loss_history = []     # only for train_loader()
          self.val_av_loss_history = []       # only for train_loader()
          self.train_acc_history = []
          self.val_acc_history = []
  
      def train_manual(self,
                       model,
                       X_train, y_train,
                       X_val, y_val,
                       save_name, save_name_best,
                       num_iterations=100, log_nth=1, plot=False, validate=False):
  
          """
          Takes one whole dataset and trains it by going through it with many iterations...
  
          Inputs:
          - model: model to train
          - X_train: training set
          - y_train: training labels
          - X_val: training set
          - y_val: training labels
          - save_name: defines how to name the model by saving it after the training
          - save_name_best: defines how to name the best model by saving while training
          - num_iterations: number of iterations for one set
          - log_nth: every log_nth iterations some statistics will be logged
          - plot: plot something (loss function). Default: False
          - validate: if True, the function will validate on val_set every log_nth-step while training. Default: False
  
          Returns:
          - self.train_loss_history: loss history
          - self.train_loss_history: accuracy history
          """
  
          # Set best_acc to zero
          best_acc = 0
  
          # Set a timer
          t = time.time()
          print('Set timer...')
  
          self.reset_histories()
          optimizer = self.optimizer(model.parameters(), **self.optim_args)
          #scheduler = self.scheduler
  
          i = 0
          for iter in range(num_iterations):
  
              # Zero the gradient buffers of all parameters
              optimizer.zero_grad()
  
              # Output scores
              output = model(X_train)
              if validate == True:
                  output_val = model(X_val)
  
              # Loss computation
              loss = self.loss_func(output, y_train)
              self.train_loss_history.append(loss)
              if validate == True:
                  loss_val = self.loss_func(output_val, y_val)
                  self.val_loss_history.append(loss_val)
  
              # Every log_nth-step will be evaluated
              if iter % log_nth == 0:
  
                  # Compute accuracy
                  # pred_np and y will not be needed in this function
                  train_acc, pred_np_t, y_t = self.check_accuracy(model=model, X=X_train, y=y_train, internal_print=False)
                  self.train_acc_history.append(train_acc)
                  val_acc = 0.0
                  if validate == True:
                      val_acc, pred_np_v, y_v = self.check_accuracy(model=model, X=X_val, y=y_val, internal_print=False)
                      self.val_acc_history.append(val_acc)
  
                      # Save best model
                      if val_acc > best_acc:
                          best_acc = val_acc
                          best_model = model
                          torch.save(best_model, save_name_best)
  
                  # Print results
                  print('Iteration: %d/%d; Loss: %.4f; train_acc; %.4f; val_acc: %.4f' % \
                        (iter+log_nth, num_iterations, loss, train_acc, val_acc))
  
              # Backpropagation
              loss.backward()
  
              # Optimization Step
              optimizer.step()
  
              # Learning rate decay
              #scheduler.step()
  
              # Free delete tensors to avoid memory run out
              #del output, output_val
  
          # Evaluate elapsed time
          t_elapsed = (time.time() - t) / 60
          print('Elapsed time: %.3f min.' % t_elapsed)
  
          # # Save trained model
          torch.save(model, save_name)
          print('Last model saved.')
  
          # Plot results
          if plot == True:
              fig = plt.figure()
              fig.set_figheight(12)
              fig.set_figwidth(15)
  
              plt.subplot(2, 1, 1)
              x = np.linspace(1, num_iterations, num_iterations)                  # setting x axis range
              plt.plot(x, self.train_loss_history, linewidth=1.8)
              if validate == True:
                  plt.plot(x, self.val_loss_history, linewidth=1.8)
              plt.xlabel('Iteration')
              plt.autoscale(enable=True, axis='x', tight=True)                    # x-axis tight
              plt.ylabel('Loss')
              plt.grid()
              if validate == True:
                  plt.legend(labels=['train_loss', 'val_loss'], loc='best')
              else:
                  plt.legend(labels=['train_loss'], loc='best')
  
              plt.subplot(2, 1, 2)
              """
              fig.text(0.82, 0.92,
                      '#Iterations: %d\n log: %d' % (num_iterations, log_nth),
                      size=20, ha='center', va='center',
                      bbox={'facecolor': 'moccasin'}) 
              """
              x = np.linspace(1, num_iterations/log_nth, num_iterations/log_nth)  # setting x axis range
              plt.plot(x, self.train_acc_history, linewidth=1.8)
              if validate == True:
                  plt.plot(x, self.val_acc_history, linewidth=1.8)
              plt.xlabel('Iteration')                                             # correct = 'Iteration / log'
              plt.autoscale(enable=True, axis='x', tight=True)  # x-axis tight
              plt.ylabel('Accuracy')
              plt.grid()
              if validate == True:
                  plt.legend(labels=['train_acc', 'val_acc'], loc='best')
              else:
                  plt.legend(labels=['train_acc'], loc='best')
              plt.savefig('train_manual_plt.pdf')
              plt.show()
  
          return self.train_loss_history, self.train_acc_history, self.val_loss_history, self.val_acc_history
  
      # TODO: implement model saving like in train_manual()
      def train_loader(self,
                       model, loader,
                       X_train, y_train,
                       X_val, y_val,
                       save_name,
                       num_epochs=1, log_nth_ep=1, plot=False, validate=False):
  
          """
          Takes one, into many mini-batches divided dataset as an input and trains it...
  
          Inputs:
          - model: model to train
          - loader: dataloader with X, y
          - X_train: will be needed for accuracy check (not for the training - batches will be used for the training)
          - y_train: will be needed for accuracy check (not for the training - batches will be used for the training)
          - X_val: training set
          - y_val: training labels
          - save_name: defines how to name the model by saving it after the training
          - num_epochs: number of epochs
          - log_nth_ep: every log_nth_ep epochs some statistics will be logged
          - plot: plot something (loss function). Default: False
          - validate: if True, the function will validate on val_set every log_nth-step while training. Default: False
  
          Returns:
          - self.train_loss_history: loss history
          - self.train_loss_history: accuracy history
          """
  
          # Set best_acc to zero
          best_acc = 0
  
          # Set a timer
          t = time.time()
          print('Set timer...')
  
          self.reset_histories()
          optim = self.optim(model.parameters(), **self.optim_args)
  
          for ep in range(num_epochs):
  
              accum_loss = []
              accum_loss_val = []
  
              for batchIdx, (X, y) in enumerate(loader):
  
                  # Zero the gradient buffers of all parameters
                  optim.zero_grad()
  
                  # Output scores
                  output = model(X)
                  if validate == True:
                      output_val = model(X_val)
  
                  # Loss computation
                  loss = self.loss_func(output, y)
                  self.train_loss_history.append(loss)
                  if validate == True:
                      loss_val = self.loss_func(output_val, y_val)
                      self.val_loss_history.append(loss_val)
  
                  # Accumulative loss for the mean value of the epoch loss over all batches
                  accum_loss.append(loss)
                  if validate == True:
                      accum_loss_val.append(loss_val)
  
                      if val_acc > best_acc:
                          best_acc = val_acc
                          best_model = model
  
                  # Backpropagation
                  loss.backward()
  
                  # Optimization Step
                  optim.step()
  
                  # Free delete tensors to avoid memory run out
                  if validate == True:
                      del output, output_val
  
              # Every log_nth_ep-step will be evaluated
              if ep % log_nth_ep == 0:
  
                  # Compute average loss over all batches in one epoch
                  accum_loss_len  = len(accum_loss)
                  average_loss = sum(accum_loss) / accum_loss_len
                  self.train_av_loss_history.append(average_loss)
                  if validate == True:
                      accum_loss_val_len = len(accum_loss_val)
                      average_loss_val = sum(accum_loss_val) / accum_loss_val_len
                      self.val_av_loss_history.append(average_loss_val)
  
                  # Compute accuracy
                  # pred_np and y will not be needed in this function
                  train_acc, pred_np, y = self.check_accuracy(model=model, X=X_train, y=y_train, internal_print=False)
                  self.train_acc_history.append(train_acc)
                  val_acc = 0.0
                  if validate == True:
                      val_acc, pred_np_v, y_v = self.check_accuracy(model=model, X=X_val, y=y_val, internal_print=False)
                      self.val_acc_history.append(val_acc)
  
                  # Print results
                  print('Epoch: %d/%d; Loss: %.4f; train_acc; %.4f; val_acc: %.4f' % \
                        (ep+log_nth_ep, num_epochs, loss, train_acc, val_acc))
  
          # Evaluate elapsed time
          t_elapsed = (time.time() - t) / 60
          print('Elapsed time: %.3f min.' % t_elapsed)
  
          # # Save trained model
          torch.save(best_model, save_name)
          print('Model saved.')
  
          # Plot results
          if plot == True:
              fig = plt.figure()
              fig.set_figheight(12)
              fig.set_figwidth(15)
  
              plt.subplot(2, 1, 1)
              """
              fig.text(0.84, 0.92,
                      '#Epochs: %d\n log: %d' % (num_epochs, log_nth_ep),
                      size=20, ha='center', va='center',
                      bbox={'facecolor': 'moccasin'})         # 0.033
              """
              x = np.linspace(1, num_epochs/log_nth_ep, num_epochs/log_nth_ep)  # setting x axis range
              plt.plot(x, self.train_av_loss_history, linewidth=1.8)
              if validate == True:
                  plt.plot(x, self.val_av_loss_history, linewidth=1.8)
              plt.xlabel('Epoch')
              plt.autoscale(enable=True, axis='x', tight=True)  # x-axis tight
              plt.ylabel('Loss')
              plt.grid()
              if validate == True:
                  plt.legend(labels=['train_loss', 'val_loss'], loc='best')
              else:
                  plt.legend(labels=['train_loss'], loc='best')
  
              plt.subplot(2, 1, 2)
              x = np.linspace(1, num_epochs/log_nth_ep, num_epochs/log_nth_ep)  # setting x axis range
              plt.plot(x, self.train_acc_history, linewidth=1.8)
              if validate == True:
                  plt.plot(x, self.val_acc_history, linewidth=1.8)
              plt.xlabel('Epoch')
              plt.autoscale(enable=True, axis='x', tight=True)  # x-axis tight
              plt.ylabel('Accuracy')
              plt.grid()
              if validate == True:
                  plt.legend(labels=['train_acc', 'val_acc'], loc='best')
              else:
                  plt.legend(labels=['train_acc'], loc='best')
              plt.savefig('train_loader_plt.pdf')
              plt.show()
  
          return self.train_av_loss_history, self.train_acc_history, self.val_av_loss_history, self.val_acc_history
  
      def check_accuracy(self, model, X, y, internal_print=False):
  
          """
          Checks accuracy of provided model on inserted dataset
  
          Inputs:
          - model: model on which accuracy should be tested
          - X: input
          - y: corresponding labels
          - internal_print: prints something out from the function call
  
          Returns:
          - accuracy: accuracy of the model
          """
  
          # Make prediction on X
          pred = model(X)
  
          # Convert pred from Tensor to numpy.ndarray
          pred_np = pred.detach().numpy()
  
          # Build a vector that compares pred_np with labels
          logical_arr = np.logical_or(np.logical_and(pred_np >= 0.5, y == 1),
                                      np.logical_and(pred_np < 0.5, y == 0))
  
          # Evaluate predictions ans compute accuracy
          correct = 0
          for i in logical_arr:
              if i == True:
                  correct = correct + 1
          #print('Number of elements in logical_arr (should be equal to the number of samples): ', logical_arr.shape[0])
          accuracy = correct / logical_arr.shape[0] * 100
  
          # Print accuracies
          if internal_print == True:
              print('Accuracy: %.2f' % accuracy)
  
          return accuracy, pred_np, y
  
      def test(self, input_filename, solver, model, X, y, plot=False):
  
          """
          Checks accuracy on a provided set and plots results.
  
          Inputs
          - input_filename: name of the dataset generated by MATLAB, where the testing subset is from
          - solver: solver from main
          - model: model from main
          - X: test set
          - y: test labels
          - plot: when True, results will be plotted
  
          Returns:
          - test_accuracy: accuracy achieved on the test set
          """
  
          # Call check_accuracy
          test_accuracy, pred_np, y = solver.check_accuracy(model=model.double(), X=X, y=y, internal_print=False)
  
          print('TEST: Computed accuracy on %s test set: %.4f' % (input_filename, test_accuracy))
  
          # Evaluate and plot predictions of the trained model
          if plot:
              pred_orig = model(X)
  
              pred_orig_np = pred_orig.detach().numpy()
  
              fig1 = plt.figure()
              fig1.set_figheight(10)
              fig1.set_figwidth(15)
  
              # Subplot 1
              plt.subplot(2, 1, 1)
              plt.plot(y, 'b', linewidth=3)
              plt.ylabel('Target labels')
              # plt.legend(labels=['Original labels'], loc='center right')
              plt.autoscale(enable=True, axis='x', tight=True)
              plt.grid()
              fig1.text(0.81, 0.92,
                        'Accuracy: %.3f %%' % test_accuracy,
                        size=20, ha='center', va='center',
                        bbox={'facecolor': 'moccasin'})
  
              # Subplot 2
              plt.subplot(2, 1, 2)
              plt.plot(pred_orig_np, 'r')
              plt.ylabel('Prediction\n(probability)')
              # plt.legend(labels=['prediction (probability)'], loc='center right')
              plt.autoscale(enable=True, axis='x', tight=True)
              plt.grid()
  
              # Subplot 3
              """
              pred_orig_np[pred_orig_np >= 0.5] = 1
              pred_orig_np[pred_orig_np < 0.5] = 0
  
              plt.subplot(3, 1, 3)
              plt.plot(pred_orig_np, 'r')
              plt.plot(y, 'b')
              plt.ylabel('Predicted (red) and \noriginal (blue) labels')
              # plt.legend(labels=['original labels', 'predicted labels'], loc='center right')
              plt.xlabel('Samples')
              plt.autoscale(enable=True, axis='x', tight=True)
              plt.grid()
              """
              plt.xlabel('Samples')
              plt.savefig('test_plt.pdf')
              # Show plots
              plt.show()
  
          return test_accuracy
  