#  ----------------------------------------------------------------------------
#
#    Train Multi-Layer Perceptron on inputs and outptus data provided in
#    two zipped CSV files and output weights, biases, metrics to GDX container
#    Up to four hidden layers defined from user input.
#
#    The training process minimizes the MSE and uses a step-wise decrease
#    in the learning rate and an increase in batch sizes, potentially running
#    multiple training processes in parallel, restarting from the best one
#    in the next step. There is inter-thread communication: if a step stalls,
#    it loads the best parameters found so far

#    Wolfgang Britz, May 2024
#
#  ----------------------------------------------------------------------------

import argparse
import gams.transfer as gt
import multiprocessing as mp
import numpy  as np
import math
import os.path
import pandas as pd
import torch
import torch.nn as nn
import time
import sys
from copy import deepcopy
from datetime import datetime
from matplotlib import pyplot
from numpy import floor
from pandas import read_csv,DataFrame,HDFStore
from pickle import dump,load,HIGHEST_PROTOCOL
from sklearn import model_selection
from sys import stdout
from tabulate import tabulate
from torch.utils.data import DataLoader, TensorDataset
from torchsummary import summary
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

#  ----------------------------------------------------------------------------
#
#    (a)  command line options
#
#  ----------------------------------------------------------------------------

parser = argparse.ArgumentParser()
parser.add_argument("-anew","--anew",          action="store_true",               help="Switch on to update data base")
parser.add_argument("-as",  "--alt_solvers",   action="store_true",               help="Switch to use ADAM,ADAMW,AdaGrad and SGD each on 1/4 of threads, default is ADAM for all")
parser.add_argument("-ai",  "--alt_init",      action="store_true",               help="Switch to use xavier, orthogonal and default on 1/3 of threads for init, otherwise default only")
parser.add_argument("-bs",  "--batch_size",    type=int,           default=2048,  help="Batchsize in last fitting step, default is 2048")
parser.add_argument("-cpu", "--cpu",           action="store_true",               help="Use CPU even if GPU is present")
parser.add_argument("-do",  "--dropout_rate",  type=float,         default=0.0,   help="Drop out rate to remove weigts, default is zero")
parser.add_argument("-ep",  "--epochs",        type=int,           default=1000,  help="Number of epochs used in each fitting step, default is 1000")
parser.add_argument("-epl", "--epochs_last",   type=int,           default=2000,  help="Number of epochs used in last fitting step, default is 2000")
parser.add_argument("-epr", "--epochs_restart",type=int,           default=50,    help="Number of epochs for restarts, default is 50")
parser.add_argument("-h1",  "--hidden1",       type=int,           default=40,    help="Number of neurons in first hidden layer, default is 40")
parser.add_argument("-h2",  "--hidden2",       type=int,           default=40,    help="Number of neurons in second hidden layer, default is 40")
parser.add_argument("-h3",  "--hidden3",       type=int,           default=40,    help="Number of neurons in third hidden layer, default is  40")
parser.add_argument("-h4",  "--hidden4",       type=int,           default=0,     help="Number of neurons in fourth hidden layer, default is 0")
parser.add_argument("-lr",  "--learning_rate", type=float,         default=0.0001,help="Learning rate in last fitting step, default is 0.0001")
parser.add_argument("-lrm", "--lr_mult",       type=float,         default=1,     help="Multiplier for learning rates, default is 1")
parser.add_argument("-nr",  "--nRestart",      type=int,           default=10,    help="Number of restarts to define starting point for first step, default is 10")
parser.add_argument("-nt",  "--nProc",         type=int,           default=-1,    help="Number of parallel starts over epochs from best try in previous step, default is 8")
parser.add_argument("-pf",  "--postFix",       type=str,           default="",    help="Additional post fix apened to name of output files")
parser.add_argument("-r",   "--reload",        action="store_true",               help="Take an existing MLP without a prefix as a start point if existing")
parser.add_argument("-s",   "--scen",          type=str,           default="ccdr",help="Scenario name, determines inputs files to read, default is 'ccdr'")
parser.add_argument("-sit", "--stall_iter",    type=int,           default=100,   help="Number of iteration after which current best parameters are loaded, default is 100")
parser.add_argument("-ts",  "--test_size",     type=float,         default=0.001, help="Share of observations used to define the test set, default is 0.001")
args = parser.parse_args()

#  ----------------------------------------------------------------------------
#
#    (0)  Internal helper routines
#
#  ----------------------------------------------------------------------------

#
# --- calculate and print R2 for two vector
#
def RMSE(experimental, simulated):
    return math.sqrt(np.square(np.subtract(experimental, simulated)).mean())

#
# --- helper function which converts a array into the format expected by GAMS
#     transfer: the first columns comprise the labels for the rows and the columns
#
def toGamsArray(new,arr,rowNames,columnNames,layerName):
    temp = pd.DataFrame( arr.reshape( (arr.shape[0],arr.shape[1] ) ) )
    for i in temp.index.values:
        for j in temp.columns:
            new.append( [rowNames[j],columnNames[i],layerName,arr[i,j]])
    return new

#
# --- helper function which converts a vector into the format expected by GAMS
#     transfer: the first columns comprise the labels for the row
#
def toGamsVector(new,arr,names,layerName):
    for i in range(len(arr)):
        new.append( [names[i],layerName,arr[i]] )
    return new

#
# ---- calculate R2 for all predicted outputs and put in to data frame,
#       including avg, min and max across outputs
#
def r2AsDataFrame(pred,outputs):
   r2Vec  = []
   maxErr = []

   yhat_df = DataFrame(pred.reshape( (outputs.shape[0],outputs.shape[1]) ) )
   yhat_df.columns = outputs.columns

   r2Avg = 0
   r2Min = 1
   r2Max = 0
   for y in outputs.columns:
       r2 = np.corrcoef(yhat_df[y],outputs[y])[0,1]**2
       r2Avg += r2
       r2Min = min(r2,r2Min)
       r2Max = max(r2,r2Max)
       r2Vec.append((y,r2))

       maxObs  = np.argmax((yhat_df[y] - outputs[y]))
       maxDiff = np.max((yhat_df[y] - outputs[y]))
       minObs  = np.argmin((yhat_df[y] - outputs[y]))
       minDiff = np.min((yhat_df[y] - outputs[y]))

       maxErr.append((y,draws[maxObs],maxDiff,draws[minObs],minDiff))

   maxErr = pd.DataFrame(maxErr,columns=["y1","maxObs","maxDiff","minObs","minDiff"])


   r2Avg = r2Avg/outputs.shape[1]
   r2Vec.append(("max",r2Max))
   r2Vec.append(("avg",r2Avg))
   r2Vec.append(("min",r2Min))

   r2VecPd = pd.DataFrame(r2Vec,columns=["y","R2"])

   return r2Vec,r2VecPd,maxErr
#
# --- set dropout rate for all drop-out layers
#
def set_dropout(model, drop_rate=0.1):
    for layer in model.children():
        if isinstance(layer, torch.nn.Dropout):
            layer.p = drop_rate
            set_dropout(layer,drop_rate)

#
# --- In standard mode (alt_solvers = False), set-up an ADAM
#     optimizer with default settings. Otherwise, 1/4 Adagrad,
#     1/4 SGD, 1/4 AdamW and 1/4 ADAM
#
def getOptimizer(model,jTry,fused):
   #
   # --- define ADAM optimizer
   #
   if jTry % 4 == 0 and args.alt_solvers:
      optName   = "Adagrad"
      optimizer = torch.optim.Adagrad(model.parameters())
   elif jTry % 4 == 2 and args.alt_solvers:
      optName   = "SGD    "
      optimizer = torch.optim.SGD(model.parameters(),momentum=0.9,fused=fused)
   elif jTry % 4 == 3 and args.alt_solvers:
      optName   = "AdamW  "
      optimizer = torch.optim.AdamW(model.parameters(),fused=fused)
   else:
      optName   = "Adam   "
      optimizer = torch.optim.Adam(model.parameters(),fused=fused)

   return optName,optimizer

#
# --- a number constructed from current time in milli seconds
#
def getRandomSeed():
   t = int( time.time() * 1000.0)
   r =          (((t & 0xff000000) >> 24) +
                 ((t & 0x00ff0000) >>  8) +
                 ((t & 0x0000ff00) <<  8) +
                 ((t & 0x000000ff) << 24)  )
   return r
#
#  --- set up model, with tensors, dataloader, optimizer and loss function
#
def setup(modelLayout,device,txTrain,tyTrain,dropout_rate,batch_size,fused,iProc):

      getRandomSeed()
      #
      # --- define the model and move it to the device
      #
      model = nn.Sequential(*modelLayout)
      model.to(device)
      #
      # --- Use torch.utils.data to create a DataLoader
      #     that will take care of creating batches

      dataset = TensorDataset(txTrain,tyTrain)
      dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
      #
      # --- set desired drop out rate: it looks as if the model cannot be overfitted,
      #     set it to zero by default
      #
      set_dropout(model, drop_rate=dropout_rate)
      #
      # --- MSE loss function
      #
      loss_fn = nn.MSELoss()
      #
      # --- define optimizer (default is ADAM)
      #
      optName,optimizer = getOptimizer(model,iProc,fused)

      return model,dataloader,loss_fn,optName,optimizer
#
#  --- ANSI escape sequences for coloring command prompt output
#
class bcolors:
    GREEN       = '\033[32;40m'
    BRIGHTGREEN = '\033[92;40m'
    YELLOW      = '\033[33;40m'
    NORMAL      = '\033[37;40m'

#
# --- initialization of weights and biases. If the alt_init is False
#     (default) use the standard. Othwise, six different methods
#     are used
#
def initParam(model,iProc,bestLoss):

   init = ""
   if bestLoss == 1:

      if iProc % 3 == 0 and args.alt_init:
         init = " xavier "
         if iProc % 6 == 0:
            init += ",2  "
         else:
            init += ",0.5"
      elif iProc % 3 == 1 and args.alt_init:
         init = " ortho  "
         if iProc % 6 == 1:
            init += ",2  "
         else:
            init += ",0.5"
      else:
         if iProc % 6 == 2 and args.alt_init:
            init = " kaiming,out"
         else:
            init = " kaiming,in "

      for layer in model.children():

          if isinstance(layer,nn.Linear):
             if layer.weight.size(1)>0:
                stdv = 1. / math.sqrt(layer.weight.size(1))
             else:
                stdv = 1.

             if iProc % 3 == 0 and args.alt_init:
                 if iProc % 6 == 0:
                    torch.nn.init.xavier_uniform_(layer.weight,gain=0.90)
                    layer.bias.data.uniform_(-stdv*0.5, stdv*0.5)
                 else:
                    torch.nn.init.xavier_uniform_(layer.weight,gain=1.10)
                    layer.bias.data.uniform_(-stdv*2.0, stdv*2.0)
             elif iProc % 3 == 1 and args.alt_init:
                 if iProc % 6 == 1:
                    torch.nn.init.orthogonal_(layer.weight,gain=0.90)
                    layer.bias.data.uniform_(-stdv*0.5, stdv*0.5)
                 else:
                    torch.nn.init.orthogonal_(layer.weight,gain=1.10)
                    layer.bias.data.uniform_(-stdv*2.0, stdv*2.0)
             else:
                if iProc % 6 == 2 and args.alt_init:
                   torch.nn.init.kaiming_uniform_(layer.weight,mode='fan_out')
                else:
                   layer.reset_parameters()

   return init

#
# --- load biases and weights from file
#
def load_state(model,name,device):
   try:
       state = torch.load(name,map_location=device,weights_only=True)
       model.load_state_dict(state)
   except:
      print("error in load_state")
      pass
#
# --- save biases and weigths to file
#
def save_state(model,name):
    try:
      torch.save(model.state_dict(),name)
    except:
      pass
#
#  --- core function for parallel training
#
def train(modelLayout,init,device,tyTest,txTest,tyTrain,txTrain,dropout_rate,iRun,nRuns,
          batch_size,learning_rate,learning_rate_red,num_epochs,
          NNName,fused,bestLossIn,deviceType,startTime,iProc,nProc,q,overAllBest,overAllBestT):

      #
      # --- define overall settings, such as profiling
      #
      torchSettings(nProc)
      #
      # --- time at which task was started, last improvement, last report to screen
      #
      taskTime     = datetime.now()
      lastTime     = taskTime
      lastRepTime  = taskTime
      #
      # -- allow for colored output to command prompt
      #
      os.system("color")
      #
      # --- pin tensors to memory if CUDA is available
      #
      if deviceType == 'cuda':
         tyTest,txTest,tyTrain,txTrain = tyTest.pin_memory(),txTest.pin_memory(),tyTrain.pin_memory(),txTrain.pin_memory()
      #
      #  --- move tensores with training and test data to device (cpu or gpu)
      #
      tyTest,txTest,tyTrain,txTrain = tyTest.to(device),txTest.to(device),tyTrain.to(device),txTrain.to(device)
      #
      #  --- define model, data loader, loss function calcultor, optimizer
      #
      torch.manual_seed(getRandomSeed())
      myModel,dataloader,loss_fn,optName,optimizer = setup(modelLayout,device,txTrain,tyTrain,dropout_rate,batch_size,fused,iProc)
      #
      # -- set up model, tensors, data loader, optimzer and loss function on device
      #
      init = optName + " "+init
      #
      # --- set startlearning rate
      #
      optimizer.param_groups[0]['lr'] = learning_rate
      #
      # --- use second best previous start point
      #     for first half of parallel processes
      #
      if iProc < nProc/2+1 and nProc > 1:
         torchFileName = "pt\\"+NNName+args.postFix+"_snd.pt"
      else:
         torchFileName = "pt\\"+NNName+args.postFix+".pt"

      #  --- start loss and number of reloads
      startLoss  = bestLossIn
      bestLoss   = bestLossIn + 2.E-10
      nReload    = 0
      #
      # --- load start point if existing from file defined above
      #
      if os.path.isfile(torchFileName):
         load_state(myModel,torchFileName,device)
         optName,optimizer = getOptimizer(myModel,iProc,fused)
      elif args.reload:
         if iProc < nProc/2+1 and nProc > 1:
            torchFileName = "pt\\"+NNName+"_snd.pt"
         else:
            torchFileName = "pt\\"+NNName+".pt"
         if os.path.isfile(torchFileName):
            load_state(myModel,torchFileName,device)
            optName,optimizer = getOptimizer(myModel,iProc,fused)

      torchFileName     = "tmp\\"+NNName+args.postFix+"_"+str(iProc)+".pt"
      best_state  = deepcopy(myModel.state_dict())
      # --- store start point
      overAllBest[0] = iProc
      overAllBest[1] = startLoss
      save_state(myModel,torchFileName)
      #
      # --- epoch where last best value was found and last reload happened
      #
      lastBest   = 0
      lastLoad   = 0
      #
      #  --- these lists will collect the history of improvements to plot for this run
      #
      losses     = []
      lossesTest = []
      indices    = []
      #
      # --- training iterations (epochs)
      #
      for n in range(num_epochs):
          id_batch = 0
          #  skip first epoch to report on start point
          if n > 0:
             #
             # --- set model in training mode (important if drop out rate is not zero)
             #
             myModel.train()
             # --- Loop over batches in an epoch using DataLoader
             for id_batch, (x_batch, y_batch) in enumerate(dataloader):
                #
                # --- make a prediction and calculate loss
                #     and let optimizer make a step, report to screen
                #
                #
                optimizer.zero_grad(set_to_none=True)
                if fused:
                  with torch.autocast(device_type=deviceType):
                      y_pred = myModel(x_batch)
                      loss   = loss_fn(y_pred, y_batch)
                else:
                      y_pred = myModel(x_batch)
                      loss   = loss_fn(y_pred, y_batch)

                #
                # --- gradient calculation and next optimization step
                #
                loss.backward()
                optimizer.step()

                #  --- update log every 1 secs (as processes run in parallel, there are more frequent update on screen)
                curTime = datetime.now()
                timeDiff = curTime - lastRepTime
                if timeDiff.total_seconds() > 1:
                   lastRepTime  = curTime
                   timeDiffT    = curTime - taskTime
                   timeDiff     = curTime - startTime
                   print(bcolors.NORMAL    +'\r  %2d\t%3d(%3d) %3d\t%4d %4d\t%4.4f' % ((iProc+1),n,num_epochs,id_batch,
                                                         timeDiffT.total_seconds(),timeDiff.total_seconds(),learning_rate),end="")
                   stdout.flush()
          else:
             time.sleep(2)

          #
          # --- predict on full training set after all batches are run:
          #     set model in evaluation mode to avoid that drop out becomes active
          #
          myModel.eval()
          with torch.inference_mode():
             y_pred = myModel(txTrain)
             loss   = loss_fn(y_pred,tyTrain)
          curLoss = loss.item()

          # --- get current best loss for group of parallel processes
          overAllBestCur = overAllBest[0]
          # --- get current winner metrics across all parallel threads
          overAllBestTCur = overAllBestT[0]
          #
          # --- check if we should reload (depenbding on # of stalled iterations
          #     and the difference to the current best loss for group of processes)
          #
          if (n - lastLoad > args.stall_iter-1) or \
               ((bestLoss > overAllBestCur*1.01) and (n - lastLoad > max(50,args.stall_iter/2))) or \
               ((bestLoss > overAllBestCur*1.05) and (n - lastLoad > max(20,args.stall_iter/3))) or \
               ((bestLoss > overAllBestCur*1.10) and (n - lastLoad > max(10,args.stall_iter/4))) or \
               ((bestLoss > overAllBestCur*1.20) and (n - lastLoad > max( 5,args.stall_iter/5))):

             curBestProc = int(overAllBest[1])
             if curBestProc > -1:
                 torchFileNameBest = "tmp\\"+NNName+args.postFix+"_"+str(curBestProc)+".pt"
                 load_state(myModel,torchFileNameBest,device)
                 best_state = deepcopy(myModel.state_dict())
             else:
                 load_state_dict(best_state)

             optName,optimizer = getOptimizer(myModel,iProc,fused)
             optimizer.param_groups[0]['lr'] = learning_rate

             if  n - lastLoad > args.stall_iter-1:
                 limit = overAllBestCur*1.005
             elif  n - lastLoad > max(50,args.stall_iter/2):
                 limit = overAllBestCur*1.01
             elif  n - lastLoad > max(20,args.stall_iter/3):
                 limit = overAllBestCur*1.05
             elif  n - lastLoad > max(10,args.stall_iter/4):
                 limit = overAllBestCur*1.10
             else:
                 limit = overAllBestCur*1.20

             nTry = 0
             curLossTry  = 1
             bestCurLoss = 1
             while( (curLossTry > limit) and (nTry < 100)):

                nTry+=1
                # --- set model in evaluation mode, predict on training set with loaded parameters and store resulting loss
                #
                myModel.eval()
                with torch.inference_mode():
                   y_pred   = myModel(txTrain)
                   loss     = loss_fn(y_pred,tyTrain)
                curLossTry  = loss.item()

                if curLossTry < bestCurLoss:
                   bestCurLoss = curLossTry
                   best_state = deepcopy(myModel.state_dict())

             #
             # --- no improvement for a larger number of epochs: reload current best
             #     or current thread made no progress and is 10% worse than best in group
             #
             if iProc < nProc/2+1 and nProc > 1:
                flag = ' (2)'
             else:
                flag = ' (1)'
             reloadStr = "Rel "+str(curBestProc+1)+" "+str(nTry)
             curTime = datetime.now()
             timeDiff  = curTime - startTime
             timeDiffT = curTime - taskTime
             print(bcolors.NORMAL+'\r  %2d\t%3d(%3d) %3d\t%4d %4d\t%4.4f' % ((iProc+1),n,num_epochs,id_batch,
                                               timeDiffT.total_seconds(),timeDiff.total_seconds(),learning_rate),end="")
             timeDiff = curTime - lastTime
             print(bcolors.YELLOW+'     | %10.10f%s\t| %s\t| %3d\t%2d\t%+8.4f%% %+8.4f%% %+8.4f%% %3d %3d %1d \t%s' %
                                     (bestCurLoss,flag,reloadStr,(n-lastBest),
                                      timeDiff.total_seconds(),((bestCurLoss/bestLoss-1)*100),
                                      ((bestCurLoss/startLoss-1)*100),((bestCurLoss/overAllBestT[0] -1)*100),
                                      (len(losses)+1),nReload,(iProc+1),init))
             stdout.flush()

             myModel.load_state_dict(best_state)
             lastLoad  = n
             nReload  += 1
             bestLoss  = bestCurLoss

          if ((bestLoss - curLoss > 1.E-10) or (n == 0)):
             # --- store best loss so far in memoty
             best_state = deepcopy(myModel.state_dict())

             curTime = datetime.now()
             timeDiff        = curTime - startTime
             timeDiffT       = curTime - taskTime
             #
             # --- calculate loss for test data set (collected for reporting)
             #
             myModel.eval()
             with torch.inference_mode():
                y_pred          = myModel(txTest)
                lossTest        = loss_fn(y_pred,tyTest)

             # --- repeat first part of status information
             print(bcolors.NORMAL+'\r  %2d\t%3d(%3d) %3d\t%4d %4d\t%4.4f' % ((iProc+1),n,num_epochs,id_batch,
                                               timeDiffT.total_seconds(),timeDiff.total_seconds(),learning_rate),end="")
             timeDiff  = curTime - lastTime
             if iProc < nProc/2+1 and nProc > 1:
                # --- indicates an improvement in the second best branch of trials
                flag = ' (2)'
             else:
                # --- indicates an improvement in the first best branch of trials
                flag = ' (1)'
             #
             #
             # --- add to lists used to plot history of improvements
             #
             losses.append(bestLoss)
             lossesTest.append(lossTest.item())
             indices.append(n)
             if curLoss < overAllBestCur*0.9999:
                # --- if there is some minimal improvement, store to temp file (to be loaded by others)
                #     this is not synchronized for higher speed but can provoke quite seldom a crash
                #     of a parallel process if it tries to load in parallel the file just being written
                save_state(myModel,torchFileName)
                overAllBest[1] = iProc
                overAllBest[0] = curLoss
                # --- indicates an improvement in this group of parallel processes (first or second half)
                if curLoss < overAllBestTCur*0.9999:
                   # --- store best overall loss over the two groups
                   overAllBestT[0] = curLoss

             # --- indicates an improvement for the current or all parallel processes
             if curLoss <= overAllBestT[0]:
                bColor = bcolors.BRIGHTGREEN
                flag += '**'
             elif curLoss <= overAllBest[0]:
                bColor = bcolors.GREEN
                flag += '*'
             else:
                bColor = bcolors.NORMAL

             # --- show improvement in log and store time and epoch of this improvement
             print(bColor+'     | %10.10f%s\t| %10.10f\t| %3d\t%2d\t%+8.4f%% %+8.4f%% %+8.4f%% %3d %3d %1d \t%s' % (curLoss,flag,lossTest.item(),(n-lastBest),
                                   timeDiff.total_seconds(),((curLoss/bestLoss-1)*100),((curLoss/startLoss-1)*100),((curLoss/overAllBestTCur -1)*100),
                                  (len(losses)+1),nReload,(iProc+1),init))
             stdout.flush()
             #
             # --- update iteration control variables
             #
             lastTime = curTime
             lastBest = n
             lastLoad = n
             bestLoss = curLoss

          #
          # --- reduce learning rate linearly and set in optimizer after each epoch
          #
          learning_rate -=learning_rate_red
          optimizer.param_groups[0]['lr'] = learning_rate
      #
      # --- all epochs done, store best weights and biases for current process
      #
      myModel.load_state_dict(best_state)
      save_state(myModel,torchFileName)
      #
      # --- store history in temp. file
      #
      outputs = [losses,lossesTest,indices,loss]
      with open('tmp\\'+str(iProc)+'.pkl', 'wb') as handle:
        dump(outputs, handle, protocol=HIGHEST_PROTOCOL)
      #
      # --- put results in queue (collected by main)
      #
      q.put((iProc,bestLoss,init,optimizer.param_groups[0]['lr']))

def modelLayout(NNName):

      modelLayout = [nn.Linear(nX,nInputLayer)]
      if args.dropout_rate > 0:
         modelLayout.append(nn.Dropout(args.dropout_rate))
      modelLayout.append(nn.Sigmoid())

      if nHiddenLayer1 > 0:
         NNName += 'x' + str(nHiddenLayer1)
         modelLayout.append(nn.Linear(nInputLayer,nHiddenLayer1))
         if args.dropout_rate > 0:
            modelLayout.append(nn.Dropout(args.dropout_rate))
         modelLayout.append(nn.Sigmoid())
      else:
          modelLayout.append(nn.Linear(nInputLayer,nY))

      if nHiddenLayer2 > 0:
         NNName += 'x' + str(nHiddenLayer2)
         modelLayout.append(nn.Linear(nHiddenLayer1,nHiddenLayer2))
         if args.dropout_rate > 0:
            modelLayout.append(nn.Dropout(args.dropout_rate))
         modelLayout.append(nn.Sigmoid())
      elif nHiddenLayer1 > 0:
         modelLayout.append(nn.Linear(nHiddenLayer1,nY))

      if nHiddenLayer3 > 0:
         NNName += 'x' + str(nHiddenLayer3)
         modelLayout.append(nn.Linear(nHiddenLayer2,nHiddenLayer3))
         if args.dropout_rate > 0:
            modelLayout.append(nn.Dropout(args.dropout_rate))
         modelLayout.append(nn.Sigmoid())
         modelLayout.append(nn.Linear(nHiddenLayer3,nY))
      elif nHiddenLayer2 > 0:
         modelLayout.append(nn.Linear(nHiddenLayer2,nY))

      return NNName,modelLayout


def delTempFiles(nProc):
   #
   # --- delete temporary files from trials
   #
   torchFileName = "tmp\\"+NNName+args.postFix+"_restart.pt"
   if os.path.exists(torchFileName): os.remove(torchFileName)
   torchFileName = "tmp\\"+NNName+args.postFix+"_restart_snd.pt"
   if os.path.exists(torchFileName): os.remove(torchFileName)
   for iProc in range(nProc):
       torchFileName = "tmp\\"+NNName+args.postFix+"_"+str(iProc)+".pt"
       if os.path.exists(torchFileName): os.remove(torchFileName)
       pickleFileName = 'tmp\\'+str(iProc)+'.pkl'
       if os.path.exists(pickleFileName): os.remove(pickleFileName)

def saveBest(bestTry,secondBestTry,nProc):
   #
   # --- store best weigths and biases for potential restart
   #
   try:
     if bestTry > -1:
        os.system("copy tmp\\"+NNName+args.postFix+"_"+str(bestTry)+".pt pt\\"+NNName+args.postFix+".pt")

     if nProc > 1:
        if secondBestTry > -1:
           os.system("copy tmp\\"+NNName+args.postFix+"_"+str(secondBestTry)+".pt pt\\"+NNName+args.postFix+"_snd.pt")
   except:
     print("Error in saveBest")
     pass

def loadData(anew):
   scen = args.scen
   if aNew == 1:

      # --- load the base dataset with the inputs
      print("--- Read input data from '"+scen+"_inputs.zip'")
      inputs  = pd.read_csv('obs/'+scen+'_inputs.zip', delimiter=';')
      draws   = inputs['draws']
      inputs  = inputs.drop('draws',axis=1)

      # --- add new observations (if existing)
      #     this is also introduced to keep the memory load from the GAMS
      #     manageable
      if os.path.isfile('obs/'+scen+'_inputs_add.zip'):
         print("--- Add input data from '"+scen+"_inputs_add.zip'")
         inputsAdd  = pd.read_csv('obs/'+scen+'_inputs_add.zip', delimiter=';')
         drawsAdd   = inputsAdd['draws']
         draws = pd.concat([draws,drawsAdd])
         inputsAdd  = inputsAdd.drop('draws',axis=1)
         inputsAdd.index+=inputs.shape[0]
         inputs = pd.concat([inputs,inputsAdd])

      # -- load the dataset with the outputs
      print("--- Read output data from '"+scen+"_outputs.zip'")
      outputs  = read_csv('obs/'+scen+'_outputs.zip', delimiter=';')
      outputs  = outputs.drop('draws',axis=1)

      # --- add new observations (if existing)
      if os.path.isfile('obs/'+scen+'_outputs_add.zip'):
         print("--- Read additional output data from '"+scen+"_outputs_add.zip'")
         outputsAdd  = read_csv('obs/'+scen+'_outputs_add.zip', delimiter=';')
         outputsAdd  = outputsAdd.drop('draws',axis=1)
         outputsAdd.index+=outputs.shape[0]
         outputs = pd.concat([outputs,outputsAdd])

      print("--- Calculate stats of raw data and export to 'stats_raw.xlsx'")
      outputs_stats = outputs.describe(percentiles=[0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9])
      inputs_stats  = inputs.describe(percentiles=[0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9])
      outputs_stats = outputs_stats.join(inputs_stats)
      outputs_stats.to_excel("obs/stats_raw.xlsx")
      #
      # --- min/max normalization for inputs, store min and max for later de-normalization in GAMS
      #
      print("--- Min/max normalization of inputs")
      inputs_min = inputs.min()
      inputs_max = inputs.max()
      inputs     = (inputs-inputs.min())/(inputs.max()-inputs.min())
      #
      # --- min/max normalization for outputs, store min and max for later de-normlization
      #
      print("--- Min/max normalization of outputs")
      outputs_min = outputs.min()
      outputs_max = outputs.max()
      outputs     = (outputs-outputs.min())/(outputs.max()-outputs.min())

      print("--- Calculate stats of normalized data and export to 'stats_norm.xlsx'")
      outputs_stats = outputs.describe(percentiles=[0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9])
      inputs_stats  = inputs.describe(percentiles=[0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9])

      # define outliers as 3 std devs from mean

      inputs_stats.loc["low"] = inputs_stats.loc["50%"] - 3*inputs_stats.loc["std"]
      inputs_stats.loc["upp"] = inputs_stats.loc["50%"] + 3*inputs_stats.loc["std"]

      # filter inputs accordingly

      limits = inputs_stats.loc[["low","upp"]]
      inputs = inputs[
       (inputs[list(limits.keys())]
        .apply(lambda col: col.between(*limits[col.name])))
        .all(axis=1)
      ]

      outputs_stats = outputs_stats.join(inputs_stats)
      outputs_stats.to_excel("obs/stats_norm.xlsx")

      # remove outliers from output and input set
      inputs   = inputs.loc[inputs.index.intersection(outputs.index)]
      outputs  = outputs.loc[outputs.index.intersection(inputs.index)]

      #
      # --- output the data frames to  h5 file for later reload
      #
      print("--- Store data to h5 formatted file 'tmp\\obsSample.h5'")
      inputs.to_hdf('tmp\\obsSample.h5',key='inputs',mode="w",complevel=9)
      draws.to_hdf('tmp\\obsSample.h5',key='draws',mode="a",complevel=9)
      inputs_min.to_hdf('tmp\\obsSample.h5',key='inputs_min',mode="a",complevel=9)
      inputs_max.to_hdf('tmp\\obsSample.h5',key='inputs_max',mode="a",complevel=9)

      outputs.to_hdf('tmp\\obsSample.h5',key='outputs',mode="a",complevel=9)
      outputs_min.to_hdf('tmp\\obsSample.h5',key='outputs_min',mode="a",complevel=9)
      outputs_max.to_hdf('tmp\\obsSample.h5',key='outputs_max',mode="a",complevel=9)

   else:
      #
      # --- read data from compressed h5 file (faster than zipped CSV
      #     inputs, spares the normalization step)
      #

      outputs     = pd.read_hdf('tmp\\obsSample.h5',key='outputs')
      inputs      = pd.read_hdf('tmp\\obsSample.h5',key='inputs')
      draws       = pd.read_hdf('tmp\\obsSample.h5',key='draws')
      inputs_min  = pd.read_hdf('tmp\\obsSample.h5',key='inputs_min')
      inputs_max  = pd.read_hdf('tmp\\obsSample.h5',key='inputs_max')

      outputs_min = pd.read_hdf('tmp\\obsSample.h5',key='outputs_min')
      outputs_max = pd.read_hdf('tmp\\obsSample.h5',key='outputs_max')

   return draws,outputs,outputs_min,outputs_max,inputs,inputs_min,inputs_max

#
#   --- define step specific training parameters
#
def trainingPars(reload,iRun,deviceType):
   learning_rate_mult = args.lr_mult
   nRestart = 1

   if reload:
      #
      # --- after a reload, use settings for step 5, do not restart
      #
      num_epochs         =    args.epochs_last
      batch_size         =    args.batch_size
      learning_rate      =    args.learning_rate*learning_rate_mult
      learning_rate_red  =    0
   elif iRun == 0:
      num_epochs         =    args.epochs
      nRestart           =    args.nRestart
      batch_size         =    64
      if deviceType == 'cpu':
         learning_rate      =    0.004*learning_rate_mult
         learning_rate_red  =    (learning_rate-0.002*learning_rate_mult)/num_epochs
      else:
         learning_rate      =    0.040*learning_rate_mult
         learning_rate_red  =    (learning_rate-0.015*learning_rate_mult)/num_epochs
   elif iRun == 1:
      num_epochs         =    args.epochs
      batch_size         =    128
      if deviceType == 'cpu':
         learning_rate      =    0.002*learning_rate_mult
         learning_rate_red  =    (learning_rate-0.001*learning_rate_mult)/num_epochs
      else:
         learning_rate      =    0.015*learning_rate_mult
         learning_rate_red  =    (learning_rate-0.010*learning_rate_mult)/num_epochs
   elif iRun == 2:
      num_epochs         =    args.epochs
      batch_size         =    256
      if deviceType == 'cpu':
         learning_rate      =    0.001*learning_rate_mult
         learning_rate_red  =    (learning_rate-0.0005*learning_rate_mult)/num_epochs
      else:
         learning_rate      =    0.010
         learning_rate_red  =    (learning_rate-0.0050*learning_rate_mult)/num_epochs
   elif iRun == 3:
      num_epochs         =    args.epochs
      batch_size         =    512
      if deviceType == 'cpu':
         learning_rate      =    0.0005*learning_rate_mult
         learning_rate_red  =    (learning_rate-0.00025*learning_rate_mult)/num_epochs
      else:
         learning_rate      =    0.0050*learning_rate_mult
         learning_rate_red  =    (learning_rate-0.0025*learning_rate_mult)/num_epochs
   elif iRun == 4:
      num_epochs         =    args.epochs
      batch_size         =    1024
      if deviceType == 'cpu':
         learning_rate      =    0.00025*learning_rate_mult
         learning_rate_red  =    (learning_rate-0.00015*learning_rate_mult)/num_epochs
      else:
         learning_rate      =    0.0025*learning_rate_mult
         learning_rate_red  =    (learning_rate-max(args.learning_rate,0.0015*learning_rate_mult))/num_epochs
   elif iRun == 5:
      num_epochs         =    args.epochs_last
      batch_size         =    args.batch_size
      if deviceType == 'cpu':
         learning_rate      =    max(args.learning_rate*learning_rate_mult,0.00015*learning_rate_mult)
         learning_rate_red  =    (learning_rate-0.00015*learning_rate_mult)/num_epochs
      else:
         learning_rate      =    max(args.learning_rate*learning_rate_mult,0.0015*learning_rate_mult)
         learning_rate_red  =    (max(args.learning_rate*learning_rate_mult,0.0015*learning_rate_mult)-args.learning_rate)/num_epochs

   return nRestart,num_epochs,batch_size,learning_rate,learning_rate_red

def outputToGDX(model,NNName,stepReport,outputs,inputs,outputs_min,outputs_max,inputs_min,inputs_max,r2Vec):

   #
   #  --- output to GAMS in GDX container
   #
   nLayer  = 0
   for layer in model.children():
       if isinstance(layer, nn.Linear):
          nLayer+=1

   m = gt.Container()
   weigths = []
   biases  = []
   layers  = []
   hiddenLayers  = []
   outputLayer   = []
   inputLayer    = []
   iLayer  = 0
   firstNamesS   = []
   firstNames_L  = []
   secondNamesS  = []
   secondNames_L = []
   iFirst  = 0
   iSecond = 0
   for layer in model.children():
      if isinstance(layer, nn.Linear):

         inputs_w  = layer.weight.cpu().detach().numpy()
         inputs_b  = layer.bias.cpu().detach().numpy()

         if iLayer == 0:
         #
         #  --- input layer
         #
            firstNames  = inputs.columns
            layerName   = "InputLayer"
            inputLayer.append(layerName)
            secondNames = []
            for i in range(inputs_w.shape[0]):
               secondNames.append(str(iSecond))
               iSecond+=1
         elif iLayer == nLayer-1:
         #
         #  --- output layer
         #
            firstNames  = secondNames
            secondNames = outputs.columns
            layerName   = "OutputLayer"
            outputLayer.append(layerName)
         else:
            firstNames  = secondNames
            layerName   = "HiddenLayer_"+str(iLayer)
            hiddenLayers.append(layerName)
            secondNames = []
            for i in range(inputs_w.shape[0]):
               secondNames.append(str(iSecond))
               iSecond+=1

         for firstName in firstNames:
             firstNames_L.append((firstName,layerName))
             firstNamesS.append(firstName)
         for secondName in secondNames:
             secondNames_L.append((secondName,layerName))
             secondNamesS.append(secondName)

         weights  = toGamsArray(weigths,inputs_w,firstNames,secondNames,layerName)
         biases   = toGamsVector(biases,inputs_b,secondNames,layerName)
         layers.append(layerName)
         iLayer+=1

   l         = gt.Set(m,"l",                         records=layers,        description="Layers")
   rows      = gt.Set(m,"rows",                      records=firstNamesS,   description="Input rows")
   colummns  = gt.Set(m,"columns",                   records=secondNamesS,  description="Output columns")
   rows_L    = gt.Set(m,"rows_L",    ["rows","l"],   records=firstNames_L,  description="Input rows to layers")
   columns_L =  gt.Set(m,"columns_L",["columns","l"],records=secondNames_L, description="Output columns to layers")
   il        = gt.Set(m,"il",        ["l"],          records=inputLayer,    description="Input layer")
   hl        = gt.Set(m,"hl",        ["l"],          records=hiddenLayers,  description="Hidden layers")
   ol        = gt.Set(m,"ol",        ["l"],          records=outputLayer,   description="Output layer")
   ii        = gt.Set(m,"ii",        ["rows"],       records=inputs.columns,description="Rows of input layer")
   jo        = gt.Set(m,"jo",        ["columns"],    records=outputs.columns,description="Neurons of outputs layer")

   gt.Parameter( m, "weights",["rows","columns",l],records=pd.DataFrame(weigths),description="weights")
   gt.Parameter( m, "biases",[ "colulmns",l],      records=pd.DataFrame(biases), description="biases")
   #
   # --- normalization information
   #
   maxO = gt.Parameter( m, "maxO",[jo],records=outputs_max,description="Max of output layer")
   minO = gt.Parameter( m, "minO",[jo],records=outputs_min,description="Max of output layer")
   maxI = gt.Parameter( m, "maxI",[ii],records=inputs_max,description="Max of output layer")
   minI = gt.Parameter( m, "minI",[ii],records=inputs_min,description="Max of output layer")
   #
   #  --- generate GAMS parameter with R2s, write GDX to disk
   #
   r2Gams     = gt.Parameter( m, "r2",["*"],records=r2Vec,description="R squared over full sample")
   stepReport = stepReport.melt(id_vars=['step'],value_vars=['loss','delta loss','avg R2','min R2','learning_rate','seconds']).values.tolist()
   stepReport = gt.Parameter( m, "stepReport",["*","*"],records=stepReport,description="Report with metrics for training steps")

   m.write("gdx/"+NNName+args.postFix+".gdx")
   print("    "+NNName+" done, output saved to GDX")
   return

def torchSettings(nProc):
   #
   # --- some settings to increase precision and speed up processing
   #
   torch.backends.cudnn.allow_tf32 = False
   torch.backends.cudnn.benchmark  = True
   torch.autograd.set_detect_anomaly(False, check_nan=False)
   torch.autograd.profiler.profile(enabled=False)

   #
   # --- assign each process a share of the CPU cores
   #
   nCores = os.cpu_count()
   torch.set_num_threads( int( max(1,floor(nCores/nProc)) ))

if __name__ == '__main__':

   nCores = os.cpu_count()
   nProc = args.nProc
   if args.nProc == -1:
      nProc = nCores
   torchSettings(nProc)
   #
   # --- generate directory used by the code
   #
   os.makedirs('tmp', exist_ok=True)
   os.makedirs('pt', exist_ok=True)
   os.makedirs('png', exist_ok=True)
   os.makedirs('gdx', exist_ok=True)
   os.makedirs('obs', exist_ok=True)

   #  ----------------------------------------------------------------------------
   #
   #    (1)  Input data, normalization, split in training and test set
   #
   #  ----------------------------------------------------------------------------

   #
   # --- set anew to 1 to read the data from the zipped CSV file and store them in h5
   #     format (for instance, if new data are added to CSV file)
   #
   aNew = 0
   if not os.path.isfile('tmp\\obsSample.h5'):
     aNew = 1
   if args.anew:
     aNew = 1

   scen = args.scen
   draws,outputs,outputs_min,outputs_max,inputs,inputs_min,inputs_max = loadData(aNew)
   #
   # -- use 10% of sample for validation, define # of obs used for training
   #
   nTrain     = int(inputs.shape[0]*(1-args.test_size))
   nTest      = int(inputs.shape[0]*args.test_size)

   X          = inputs.values
   Y          = outputs.values

   #
   # --- split randomly the observations into a training and a test set
   #
   XTest, XTrain, YTest, YTrain = model_selection.train_test_split(X, Y, test_size=(1-args.test_size))
   nX         = inputs.shape[1]
   nY         = outputs.shape[1]
   #
   # --- build tensors from numpy arrays
   #

   tx = torch.from_numpy(X).float()
   ty = torch.from_numpy(Y).float()

   txTrain = torch.from_numpy(XTrain).float()
   tyTrain = torch.from_numpy(YTrain).float()

   txTest = torch.from_numpy(XTest).float()
   tyTest = torch.from_numpy(YTest).float()

   # --- Define the model and name for it, used to store files (pt,gdx). The sigmoid function works well in GAMS
   #     The dropout layers will be deactivated in the example later by settings the rate zo zero
   #
   # --- Define layout of the NN
   #
   nInputLayer   = args.hidden1
   nHiddenLayer1 = args.hidden2
   nHiddenLayer2 = args.hidden3
   nHiddenLayer3 = args.hidden4

   NNName = str(nInputLayer)
   NNName,modelLayout = modelLayout(NNName)
   #
   # --- define re-load name
   #
   torchFileName = "pt\\"+NNName+args.postFix+".pt"
   reload = False
   #
   # --- number of default training steps
   #
   nSteps = 6
   if os.path.isfile(torchFileName):
       #
       # --- start point found: repeat last training step
       #
       reload = True
       nSteps  = 1

   if not reload and args.reload:
      #
      # --- try to reload from base case (no post fix)
      #
      torchFileName = "pt/"+NNName+".pt"
      if os.path.isfile(torchFileName):
         reload = True
         nSteps  = 1
   #
   # --- Check if GPU is available
   #
   if args.cpu:
      device     = torch.device("cpu")
      deviceType = 'cpu'
      fused      = False

   else:
      device     = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
      deviceType = 'cuda'               if torch.cuda.is_available() else 'cpu'
      fused      = torch.cuda.is_available()

   #
   #  --- start time for reporting on screen
   #
   startTime = datetime.now()
   lastTime  = startTime

   dropout_rate  = args.dropout_rate
   #
   # --- list of statistics to report by step
   #
   minR2List         = []
   avgR2List         = []
   stepList          = []
   lossList          = []
   deltaLossList     = []
   epochsList        = []
   lrList            = []
   secList           = []
   secTotalList      = []

   #
   # --- generate a model and loss function for reporting in the main
   #
   model = nn.Sequential(*modelLayout)
   set_dropout(model, drop_rate=dropout_rate)
   loss_fn = nn.MSELoss()
   #
   #  --- count trainable paramters
   #
   model_parameters = filter(lambda p: p.requires_grad, model.parameters())
   params = sum([np.prod(p.size()) for p in model_parameters])

   bestTry = 0
   secondBestLoss = 1
   if reload:
      #
      # --- previous results exist: load old best state
      #     define start fit
      #
      model = nn.Sequential(*modelLayout)
      load_state(model,torchFileName,device)
      model.eval()
      y_pred     = model(txTrain)
      loss       = loss_fn(y_pred,tyTrain)
      bestLoss   = loss.item()
      y_pred     = model(tx)
      pred       = y_pred.cpu().detach().numpy()
      r2Vec,r2Full,maxErr = r2AsDataFrame(pred,outputs)
      minR2      = r2Vec[-1][1]
   else:
      bestLoss   = 1
      minR2      = 0
   #
   # --- synchronized array which keep track of best loss and related parallel thread
   #     for reloading if a thread stalls. The first array is for the first 4 pararell
   #     threads, the second for the remainders
   #
   mp.set_start_method('spawn')
   overAllBest      = mp.Array('d',[bestLoss])
   overAllBestFst   = mp.Array('d',[bestLoss,bestTry])
   overAllBestSnd   = mp.Array('d',[bestLoss,bestTry])
   #
   # --- remove potential left overs from previous runs
   #
   delTempFiles(nProc)
   #
   # --- loop over training steps
   #
   bestTry        = -1
   secondBestTry  = -1
   resultsAll = []
   for iStep in range(nSteps):

      stepStartTime  = datetime.now()
      startLoss      = bestLoss
      nRestart,cur_num_epochs,batch_size,learning_rate,learning_rate_red = trainingPars(reload,iStep,deviceType)
      print("\n")

      bestOverTries    = 1
      sndBestOverTries = 1
      for iRestart in range(nRestart):
         if iRestart < nRestart -1:
            num_epochs = args.epochs_restart
         else:
            num_epochs = cur_num_epochs
         #
         # -- report model set-up, stats and chosen hyper-parameters to screen
         #
         print(bcolors.NORMAL+"\n    Model setup: nX="+str(nX)+" => "+NNName+"  => ny="+str(nY))
         print(model)
         print("    Run # %d(%d) on %s, restart # %d(%d),     batchSize:  %5d, epochs    : %6d, start learning rate: %4.4f, dropout rate: %4.4f, train size: %d, test size: %d"
                   % (iStep,(nSteps-1),str(device),iRestart,(nRestart-1),batch_size,num_epochs,learning_rate,dropout_rate,YTrain.shape[0],YTest.shape[0]))

         print("    Number of trainable parameters: %5d, nY*nTrain : %6d, d.f.: %6d, parallel tries %d, loss at start: %10.10f,%10.10f,%10.10f"
                  % (params,(nY*nTrain),(nY*nTrain-params),nProc,overAllBest[0],overAllBestFst[0],overAllBestSnd[0]))
         print('\n  Proc \tEpoch \tBatch\tSecs\t\tLearn rate | Train loss\t\t| Test loss\t| Epochs Secs\tDelta loss Total to best #improved #reloads Proc Solver')
         print("\n")

         q = mp.Queue()
         processes = []
         for iProc in range(nProc):
             if iRestart < nRestart - 1:
                # -- last restart: reset
                initFlag = 1
             else:
                initFlag = bestLoss
             init = initParam(model,iProc,initFlag)

             if iProc < nProc/2+1 and nProc > 1:
                # --- set group 1 to second half of processes
                passBest = overAllBestSnd
             else:
                # --- first half assigned to group 0
                passBest = overAllBestFst

             if iProc % 4 == 0 and nProc > 1:
                #
                # --- let a quarter of the processes uses a quite low learning rate
                #     this might help to make progress in case the rest does not work
                #
                cur_learning_rate     = max(learning_rate*0.01,0.0001)
                if learning_rate < 0.0005:
                   if deviceType == 'cuda':
                      cur_learning_rate_red = (cur_learning_rate-0.001)/cur_num_epochs
                   else:
                      cur_learning_rate_red = (cur_learning_rate-0.00005)/cur_num_epochs
                else:
                   cur_learning_rate_red = 0
             else:
                cur_learning_rate      = learning_rate
                cur_learning_rate_red  = learning_rate_red

             #
             # --- start process on parallel queue
             #
             p = mp.Process(target=train, args=(modelLayout,init,device,tyTest,txTest,tyTrain,txTrain,dropout_rate,iStep,nSteps,
                                                batch_size,cur_learning_rate,cur_learning_rate_red,num_epochs,NNName,
                                                fused,passBest[0],deviceType,startTime,iProc,nProc,q,passBest,overAllBest))
             p.start()
             processes.append(p)
         #
         # --- wait for all processes to finish
         #
         for p in processes:
             p.join()

         os.system("color")
         #
         #  --- get best losses of different parallel trials, print to screen and find best and second bestone
         #
         results    = []
         while not q.empty():
            res = q.get()
            results.append(res)
            into = list(res)
            into.insert(0,iRestart)
            into.insert(0,iStep)
            resultsAll.append(into)
         #
         #  --- find best and second best loss and report to screen
         #
         print(bcolors.NORMAL+"\n")
         curBest = 1

         for res in results:
             #
             # --- store best and second best loss and related # of try
             #
             if res[1] < bestLoss:
                bestTry        = res[0]
                bestLoss       = res[1]
                overAllBestFst[0] = bestLoss
                overAllBestFst[1] = bestTry

         print("\n\n   Outcome of parallel proceses for "+NNName+" : ")
         dfRes = pd.DataFrame(resultsAll,columns=['Step','Restart','Proc','Loss','Solver','Learning rate'])
         dfRes = dfRes.sort_values(by='Loss',axis=0,ascending=False)
         print(tabulate(dfRes, headers='keys', tablefmt='psql',showindex="never",floatfmt=("2d","2d","s",".10f","s",".5f")))

         if iProc > 1:
            for res in results:
                if res[1] > bestLoss and res[1] < secondBestLoss:
                   secondBestTry  = res[0]
                   secondBestLoss = res[1]
                   overAllBestSnd[0] = secondBestLoss
                   overAllBestSnd[1] = secondBestTry

         if iRestart < nRestart - 1:
#           --- multiple restarts: save best case temp. to disk
            if bestOverTries > bestLoss:
               print("\n   New best      restart "+str(bestLoss)+" from process # "+str(bestTry+1))
               bestOverTries = bestLoss
               os.system("copy tmp\\"+NNName+args.postFix+"_"+str(bestTry)+".pt tmp\\"+NNName+args.postFix+"_restart.pt")

            if sndBestOverTries > secondBestLoss:
               print("\n   New second best restart "+str(secondBestLoss)+" from process # "+str(secondBestTry+1))
               sndBestOverTries = secondBestLoss
               os.system("copy tmp\\"+NNName+args.postFix+"_"+str(secondBestTry)+".pt tmp\\"+NNName+args.postFix+"_restart_snd.pt")

            overAllBestFst[0] = 1
            overAllBestSnd[0] = 1

         if iRestart == nRestart - 2:
#
#           --- before the last start: get best
#
            os.system("copy tmp\\"+NNName+args.postFix+"_restart_snd.pt tmp\\"+NNName+args.postFix+"_0.pt")
            overAllBestSnd[0] = sndBestOverTries
            overAllBestSnd[1] = 0

            snd = int(nProc/2+1)
            overAllBestFst[0] = bestOverTries
            overAllBestFst[1] = snd
            os.system("copy tmp\\"+NNName+args.postFix+"_restart.pt tmp\\"+NNName+args.postFix+"_"+str(snd)+".pt")

            overAllBest[0]    = bestOverTries
            bestLoss          = bestOverTries
            secondBestLoss    = sndBestOverTries
            #
            # --- save to disk as start point
            #
            saveBest(int(overAllBestFst[1]),int(overAllBestSnd[1]),nProc)

         elif iRestart == nRestart - 1:
            #
            # --- all epochs and tries done, load best and second best model from disk
            #     in tmp dir and save to pt dir
            #
            saveBest(bestTry,secondBestTry,nProc)
            #
            # --- load stats from best try
            #
            with open('tmp\\'+str(iProc)+'.pkl', 'rb') as handle:
                 losses,lossesTest,indices,loss = load(handle)
            #
            # ---  plot fit over improvment steps
            #      plots are stored to disk and not openened as a Window as this stops execution
            pyplot.figure(iStep)
            pyplot.plot(indices, losses)
            pyplot.plot(indices, lossesTest)
            pyplot.ylabel('loss')
            pyplot.xlabel('epoch')
            pyplot.legend(['train', 'test'], loc='upper left')
            if reload:
               pyplot.savefig('png/'+NNName+args.postFix+'_reload.png')
            else:
               pyplot.savefig('png/'+NNName+args.postFix+"_"+str(iStep)+'.png')

            #
            # --- predict the full sample and reports R2s
            #
            load_state(model,"pt\\"+NNName+args.postFix+".pt",device)
            model.eval()
            y_pred = model(tx)
            pred = y_pred.cpu().detach().numpy()
            #
            #
            # --- MSE, added up over outputs, as loss function
            #
            r2Vec,r2Full,maxErr = r2AsDataFrame(pred,outputs)
            print("\n\n    R2,Observations with highest max and min deviations:")

            df = pd.concat([r2Full,maxErr],axis=1)
            df = df.drop(["y1"],axis=1)
            #df = df.fillna('')

            print(tabulate(df, headers='keys', tablefmt='psql',showindex="never",floatfmt=("s",".7f","s",".7f","s",".7f")))

            timeDiff      = datetime.now() - stepStartTime
            timeDiffTotal = datetime.now() - startTime
            print('\n  Min R2 %10.10f, Start Min R2 %10.10f, delta %4.2f%%'
               % (r2Vec[-1][1],minR2,((-r2Vec[-1][1]+minR2)/(1-minR2)*100)))
            minR2 = r2Vec[-1][1]
            print('\n  Best loss %10.10f, Start loss %10.10f, delta %4.2f%%, seconds used %d, secondsTotal %d'
               % (bestLoss,startLoss,((bestLoss-startLoss)/startLoss*100),timeDiff.total_seconds(),timeDiffTotal.total_seconds()))

            #
            # --- lists with loss improvements etc. over whole step-wise training process
            #
            stepList     .append(iStep)
            lossList     .append(bestLoss)
            minR2List    .append(minR2)
            avgR2List    .append(r2Vec[-2][1])
            deltaLossList.append((bestLoss-startLoss)/startLoss*100)
            epochsList   .append(num_epochs)
            lrList       .append(learning_rate)
            secList      .append(timeDiff.total_seconds())
            secTotalList .append(timeDiffTotal.total_seconds())

   #
   # --- end of steps: show progress over while training process
   #
   print("\n Summary over steps:")
   stepReport = pd.DataFrame({'step':stepList,
                              'loss':lossList,
                              'delta loss':deltaLossList,
                              'avg R2':avgR2List,
                              'min R2':minR2List,
                              'learning_rate':lrList,
                              'epochs':epochsList,
                              'seconds':secList,
                              'seconds total':secTotalList},index=stepList)
   print(tabulate(stepReport,headers='keys', tablefmt='psql',showindex="never"))

   #
   # --- save the best results for a warm start
   #
   saveBest(bestTry,secondBestTry,nProc)
   #
   # --- delete all temporary files
   #
   delTempFiles(nProc)
   #
   # --- output results to GDX
   #
   outputToGDX(model,NNName,stepReport,outputs,inputs,outputs_min,outputs_max,inputs_min,inputs_max,r2Vec)
