text: from collections import deque
  from email.parser import BytesParser
  import numpy as np
  #import sortednp as snp
  import torch
  import matplotlib
  import matplotlib.pyplot as plt
  import math
  import time
  import bisect
  from datetime import timedelta
  from training_utils import loss_function, loss_from_residuals, dynamic_loss_from_residuals
  #import training_config as config
  from pipeline import Molecule
  from resizable import H5Array as Array
  import os
  from shutil import copyfile
  import h5py
  import statistics
  import traceback
  from collections.abc import Sequence
  import time
  from filelocks import ExistLock, Timeout
  
  n_to_s = {
  1:'H',
  6:'C',
  7:'N',
  8:'O',
  16:'S',
  9:'F',
  17:'Cl'
  }
  
  
  class TrainTestHistory:
  
      @classmethod
      def get_lock(cls, filename=None, save_prefix=None, load=True, writeable=False, block=True, hard_lock=True, timeout=None, polling_interval=.05):
          if not block:
              timeout = 0.0
          start_time = time.time()
          
          if filename is None:
              filename = save_prefix + "-history.torch"
          
          if hard_lock:
              #print(f"Acquiring ExistLock('{filename + '.lock'}')... ", end='')
              with ExistLock(filename + '.lock'):
                  #print("Done.")
                  pass
  
          mode = {
              (True, True)   : 'r+',
              (True, False)  : 'r',
              (False, True)  : 'w'
             }[load, writeable]
          file = None
          #print(f"Opening hdf5 file {filename} ... ", flush=True)
          while file is None:
              try:
                  file = h5py.File(filename, mode)
              except BlockingIOError:
                  #print("Except BlockingIOError.", flush=True)
                  if timeout is None or time.time() - start_time < timeout:
                      #print(f"Timeout ({timeout:.2f} s) not reached", flush=True)
                      time.sleep(polling_interval)
                  else:
                      #print(f"Timeout reached. Raising Timeout exception.", flush=True)
                      raise Timeout(f"Failed to open {filename}")
          #print("Done opening hdf5.")
  
          if os.path.isfile(filename + ".lock"):
              os.remove(filename + ".lock")
  
          return file
    
      def __init__(self, testing_batches=[], device='cpu', examples_per_epoch=None, relevant_elements=None, 
                   run_name=None, number_to_symbol=None, writeable=False, smoothing_window=10,
                   train_dynamic=None, store_residuals=False, wandb_log=None, wandb_interval=1,
                   file=None, save_prefix=None, use_backup=True, load=None,
                   block=True, hard_lock=True, timeout=None, polling_interval=0.05):
  
          if file is not None:
              self.filename = file
          elif save_prefix is not None:
              self.filename = save_prefix + "-history.torch"
          else:
              raise ValueError("Must provide either 'file' or 'save_prefix'.")
  
          self.device = device
          self.use_backup = use_backup
              
          if not os.path.isfile(self.filename):
              if load:
                  if self.use_backup and os.path.isfile(self.filename + '.bak'):
                      print(f"Restoring {self.filename} from backup.")
                      copyfile(self.filename + ".bak", self.filename)
                  else:
                      raise ValueError(f"Can't find history file '{self.filename}'")
              elif not os.path.isfile(self.filename + '.bak'):
                  load = False
          elif load is None:
              load = True
          if self.use_backup and os.path.isfile(self.filename + '.bak'):
              print("History file failed to close properly last time.")
              if os.path.getsize(self.filename) < os.path.getsize(self.filename + '.bak'):
                  print(f"Restoring {self.filename} from backup.")
                  copyfile(self.filename + ".bak", self.filename)
                      
          writeable = writeable or not load
          self.writeable = writeable
          
          self.file = TrainTestHistory.get_lock(
              self.filename,
              load=load,
              writeable=writeable,
              block=block,
              hard_lock=hard_lock,
              timeout=timeout,
              polling_interval=polling_interval)
          
          if not load:
              # create new history file
              self.train = TrainingHistory(
                  examples_per_epoch, smoothing_window, bool(train_dynamic),
                  wandb_log=wandb_log, wandb_interval=wandb_interval,
                  file=self.file.create_group('train'), load=False, writeable=writeable)
              self.test = TestingHistory(
                  examples_per_epoch, testing_batches, relevant_elements, train_dynamic, device,
                  number_to_symbol=number_to_symbol, store_residuals=store_residuals, wandb_log=wandb_log,
                  file=self.file.create_group('test'), load=False, writeable=writeable, name=run_name)
              self.name = run_name
              self.file.attrs['name'] = run_name
              self.file.attrs['examples_per_epoch'] = examples_per_epoch
              self.file.attrs['train_dynamic'] = bool(train_dynamic)
  
          else:            
              try: # in case the file is corrupted
                  self.name = self.file.attrs['name']
                  if 'examples_per_epoch' in self.file.attrs and examples_per_epoch is None:
                      examples_per_epoch = self.file.attrs['examples_per_epoch']
                  else:
                      self.file.attrs['examples_per_epoch'] = examples_per_epoch
                  if train_dynamic is None and 'train_dynamic' in self.file.attrs:
                      train_dynamic = self.file.attrs['train_dynamic']
                  else:
                      train_dynamic = bool(train_dynamic)
                      self.file.attrs['train_dynamic'] = train_dynamic
                  self.train = TrainingHistory(
                      examples_per_epoch,
                      train_dynamic=train_dynamic,
                      wandb_log=wandb_log, wandb_interval=wandb_interval, 
                      file=self.file['train'], load=True, writeable=writeable)
                  self.test = TestingHistory(
                      examples_per_epoch, testing_batches=testing_batches,
                      relevant_elements=relevant_elements, test_dynamic=train_dynamic,
                      device=device, wandb_log=wandb_log,
                      file=self.file['test'], load=True, writeable=writeable, name=run_name)
              except: # file is corrupted
                  self.close(save=False)
                  raise
                   
          #if os.path.isfile(self.file.filename + '.bak'):
          #    os.remove(self.file.filename + '.bak')
              
          
      def save(self, verbose = True):
          self.file.flush()
          if not self.writeable:
              raise RuntimeError(f"Tried to save history, but {self.filename} in non-saving mode.")
              if verbose: print(f"{self.filename} in non-saving mode.")
              return
          if verbose: print("Saving " + ("and backing up " if self.use_backup else "") + "history...")
          if self.use_backup:
              copyfile(self.filename, self.filename + ".bak")
          
      def close(self, save=True, verbose = True):
          if self.file:
              if save and self.writeable:
                  self.save()
              if (os.path.isfile(self.filename + '.bak') and
                  os.path.getsize(self.filename) >= os.path.getsize(self.filename + '.bak')):
                  os.remove(self.filename + '.bak')
              self.file.close()
          self.file = None       
              
  
      def plot(self, figure=None, x_axis='batch_number', y_axis='smoothed_loss', show=True):
          # x_axis = 'time' is also allowed
          if figure is None:
              plt.figure(figsize=(12,8))
          self.train.plot(plt.gcf(), x_axis, y_axis)
          self.test.plot(plt.gcf(), x_axis)
          if figure is None:
              plt.legend(loc="best")
              if show:
                  plt.show()
  
      def log_batch(self, *args, **kwargs):
          self.train.log_batch(*args, **kwargs)
  
      def elapsed_time(self):
          return self.train.elapsed_time[-1]
  
      def run_test(self, model, batch_number=None, epoch=None, batch_in_epoch=None, example_in_epoch=None, example=None, elapsed_time=None, *args, **kwargs):
          if batch_number is None: batch_number = len(self.train.loss)
          if epoch is None: epoch = self.train.epoch[-1]
          if batch_in_epoch is None: batch_in_epoch = self.train.batch_in_epoch[-1]
          if example_in_epoch is None: example_in_epoch = self.train.example_in_epoch[-1]
          if example is None: example = self.train.example[-1]
          if elapsed_time is None: elapsed_time = self.train.elapsed_time[-1]
          return self.test.run_test(model, batch_number, epoch, batch_in_epoch, example_in_epoch, example, elapsed_time, *args, **kwargs)
  
  class BaseHistory:
  
      # methods for finding comparison points
  
      def max_epoch_batch(self):
          return self.epoch[-1], self.batch[-1]
  
      def max_batch(self):
          raise Exception("This should be overloaded!")
  
      def max_epoch_example(self):
          return self.epoch[-1], self.example_in_epoch[-1]
          
  
  class TrainingHistory(BaseHistory):
      def __init__(self, examples_per_epoch, smoothing_window=10, train_dynamic=False,
                   wandb_log=None, wandb_interval=1, file=None, load=True, writeable=None):
          self.examples_per_epoch = examples_per_epoch
          self.wandb_log = wandb_log
          self.wandb_interval = wandb_interval
          self.file = file
          self.writeable = writeable
          
          if not load:
              self.smoothing_window = smoothing_window
              file.attrs['smoothing_window'] = smoothing_window
              
              self.last_wandb = 0
              file.attrs['last_wandb'] = 0
  
              # initialize the lists we will be accumulating
              # these lists correspond to each other
              # one entry per batch
              # batch 0 is a dummy batch
              self.epoch=Array(file, 'epoch', [0])
  
              # batch n covers example[example_number[n-1]:example_number[n]]
              self.example_in_epoch=Array(file, 'example_in_epoch', [examples_per_epoch])
              self.example=Array(file, 'example', [0])
              
              self.examples_in_batch=Array(file, 'examples_in_batch', [0])
  
              self.batch_in_epoch=Array(file, 'batch_in_epoch', [0])
              self.elapsed_time=Array(file, 'elapsed_time', [0.0])
  
              self.atoms_in_batch=Array(file, 'atoms_in_batch', [0])
  
              self.loss=Array(file, 'loss', [float("inf")])
              self.smoothed_loss=Array(file, 'smoothed_loss', [float("inf")])
              
              if train_dynamic:
                  self.dynamic_loss=Array(file, 'dynamic_loss', [float("inf")])
                  self.smoothed_dynamic_loss=Array(file, 'smoothed_dynamic_loss', [float("inf")])
  
              # this list has one entry per epoch
              # epoch e starts at index epoch_start[e]
              self.epoch_start=Array(file, 'epoch_start', [0]) # includes a dummy epoch 0 starting at batch 0
          else:
              self.smoothing_window = file.attrs['smoothing_window']
              self.last_wandb = file.attrs['last_wandb']
              
              self.epoch=Array(file, 'epoch')
              self.example_in_epoch=Array(file, 'example_in_epoch')
              self.example=Array(file, 'example')               
              self.examples_in_batch=Array(file, 'examples_in_batch')
              self.batch_in_epoch=Array(file, 'batch_in_epoch')
              self.elapsed_time=Array(file, 'elapsed_time')
              self.atoms_in_batch=Array(file, 'atoms_in_batch')
              self.loss=Array(file, 'loss')
              self.smoothed_loss=Array(file, 'smoothed_loss')              
              if train_dynamic:
                  self.dynamic_loss=Array(file, 'dynamic_loss', self.loss.shape, dtype=float)
                  self.smoothed_dynamic_loss=Array(file, 'smoothed_dynamic_loss', self.smoothed_loss.shape, dtype=float)
              self.epoch_start=Array(file, 'epoch_start') 
          
      def log_batch(
              self, batch_time, wait_time, examples_in_batch, atoms_in_batch, elapsed_time, example_in_epoch,
              example, loss, dynamic_loss=None, epoch=None, batch_in_epoch=None, verbose=True):
          self.elapsed_time.append(elapsed_time)
  
          if example is None:
              example = (epoch-1) * self.examples_per_epoch + example_in_epoch
          elif example_in_epoch is None:
              example_in_epoch = (example - 1) % self.examples_per_epoch + 1
  
          epoch = self.epoch[-1] if epoch is None else epoch
          if example_in_epoch > self.examples_per_epoch:
              epoch += 1
          if epoch > self.epoch[-1]:
              self.epoch_start.append(len(self.example))
          self.epoch.append(epoch)
  
          self.examples_in_batch.append(examples_in_batch)
          self.example_in_epoch.append(example_in_epoch)
          self.example.append(example)
          self.atoms_in_batch.append(atoms_in_batch)
          batch_in_epoch = math.ceil(example_in_epoch / self.examples_per_batch)
          self.batch_in_epoch.append(batch_in_epoch)
          
          self.loss.append(loss)
          window = min(len(self.loss)-1, self.smoothing_window)
          self.smoothed_loss.append(sum(self.loss[-window:]) / window)
          
          if dynamic_loss is not None:
              self.dynamic_loss.append(dynamic_loss)
              window = min(len(self.dynamic_loss)-1, self.smoothing_window)
              self.smoothed_dynamic_loss.append(sum(self.dynamic_loss[-window:]) / window)            
          
          if self.wandb_log is not None and len(self.loss) - 1 - self.last_wandb >= self.wandb_interval:
              self.last_wandb = len(self.loss) - 1
              self.file.attrs['last_wandb'] = self.last_wandb
              log_dict = {
                  'elapsed_time':self.elapsed_time[-1],
                  'epoch':epoch,
                  'batch_in_epoch':batch_in_epoch,
                  'example_in_epoch':example_in_epoch,
                  'example':example,
                  'examples_in_batch':examples_in_batch,
                  'atoms_in_batch':atoms_in_batch,
                  'loss':loss,
                  'smoothed_loss':self.smoothed_loss[-1]
              }
              if dynamic_loss is not None:
                  log_dict['dynamic_loss'] = dynamic_loss
                  log_dict['smoothed_dynamic_loss'] = self.smoothed_dynamic_loss[-1]
              self.wandb_log(log_dict)
  
          if verbose:
              print(f"{self.epoch[-1]} : {self.batch_in_epoch[-1]} / {self.batch_in_epoch[-1] + self.batches_remaining_in_epoch()}  " +
                    ("loss =" if dynamic_loss is None else "abs_loss =") + f"{self.smoothed_loss[-1]:8.3f}  " +
                    ("" if dynamic_loss is None else f"dyn_loss ={self.smoothed_dynamic_loss[-1]:8.3f}") +
                    f"  t_train = {batch_time:.2f} s  " +
                    (f"t_wait = {wait_time:.2f} s  " if dynamic_loss is None else "") +
                    f"t = {str(timedelta(seconds=self.elapsed_time[-1]))[:-5]}   ",
                     end="\r", flush=True)
  
      def num_batches(self):
          return len(self.epoch) - 1
  
      # number of epochs we have seen
      def num_epochs(self):
          return len(self.epoch_start) - 1
  
      def current_epoch(self, batch=-1):
          """
          epoch of the next incoming batch
          """
          return self.epoch[batch] + self.example_in_epoch[batch] // self.examples_per_epoch
          
      def next_example_in_epoch(self, batch=-1):
          """
          example in epoch that would start next batch
          """
          return self.example_in_epoch[-1] % self.examples_per_epoch
  
      def next_batch_in_epoch(self, batch=-1):
          """
          batch number that would follow given batch
          wraps around to 1 if epoch ends
          """
          return 1 if self.example_in_epoch[batch] >= self.examples_per_epoch else self.batch_in_epoch[batch] + 1
          
      @property
      def examples_per_batch(self):
          return max(self.examples_in_batch[-2:])
  
      def batches_remaining_in_epoch(self, batch=-1):
          return math.ceil((self.examples_per_epoch - self.example_in_epoch[batch]) / self.examples_per_batch)
  
      def max_batch(self):
          return len(self.loss) - 1
          
      # x_axis = 'time' and y_axis = 'loss' are also allowed
      def plot(self, figure=None, x_axis='example', y_axis='smoothed_loss', color="r", show=True):
          x_axis = self.example[1:] if x_axis=='example' else self.elapsed_time[1:]
          y_axis = self.smoothed_loss[1:] if y_axis=='smoothed_loss' else self.loss[1:]
          
          if figure is None:
              plt.figure(figsize=(12,8))
          plt.plot(np.array(x_axis), np.array(y_axis), "o-", color=color, label="train")
          if figure is None:
              plt.legend(loc="best")
              if show:
                  plt.show()
              
      def __getstate__(self):
          d = self.__dict__.copy()
          del d['wandb_log']
          return d
          
  class TestingHistory(BaseHistory):
  
      def __init__(self, examples_per_epoch, testing_batches, relevant_elements=None, test_dynamic=False, device='cuda',
                   number_to_symbol=None, store_residuals=False, wandb_log=None,
                   file=None, load=True, writeable=None, name=None):
          self.examples_per_epoch = examples_per_epoch
          self.testing_batches = testing_batches
          self.test_dynamic = test_dynamic
          self.device = device
          self.number_to_symbol = number_to_symbol if number_to_symbol else n_to_s
          self.wandb_log = wandb_log
          self.name = name
          self.writeable=writeable
  
          if not load:
              self.relevant_elements = relevant_elements
              file.attrs['relevant_elements'] = relevant_elements
              self.store_residuals = store_residuals
              file.attrs['store_residuals'] = store_residuals
  
              # initialize the lists we will be accumulating
              self.batch_number = Array(file, 'batch_number', dtype=int)
              self.epoch = Array(file, 'epoch', dtype=int)
              self.batch_in_epoch = Array(file, 'batch_in_epoch', dtype=int)
              self.example_in_epoch = Array(file, 'example_in_epoch', dtype=int)
              self.example = Array(file, 'example', dtype=int)
              self.elapsed_time = Array(file, 'elapsed_time', dtype=float)
              self.loss = Array(file, 'loss', dtype=float)
              self.mean_error_by_element = Array(file, 'mean_error_by_element', (0,len(relevant_elements)), dtype=float, resizable_cross=True)
              self.RMSE_by_element = Array(file, 'RMSE_by_element', (0,len(relevant_elements)), dtype=float, resizable_cross=True)
              if test_dynamic:
                  self.dynamic_loss = Array(file, 'dynamic_loss', dtype=float)
                  self.dynamic_mean_error_by_element = Array(file, 'dynamic_mean_error_by_element', (0,len(relevant_elements)), dtype=float, resizable_cross=True)
                  self.dynamic_RMSE_by_element = Array(file, 'dynamic_RMSE_by_element', (0,len(relevant_elements)), dtype=float, resizable_cross=True)                
          else:
              self.relevant_elements = file.attrs['relevant_elements']
              self.store_residuals = file.attrs['store_residuals']
  
              self.batch_number = Array(file, 'batch_number')
              self.epoch = Array(file, 'epoch')
              self.batch_in_epoch = Array(file, 'batch_in_epoch')
              self.example_in_epoch = Array(file, 'example_in_epoch')
              self.example = Array(file, 'example')
              self.elapsed_time = Array(file, 'elapsed_time')
              self.loss = Array(file, 'loss')
              self.mean_error_by_element = Array(file, 'mean_error_by_element', resizable_cross=True)
              self.RMSE_by_element = Array(file, 'RMSE_by_element', resizable_cross=True)
              if test_dynamic:
                  self.dynamic_loss = Array(file, 'dynamic_loss', self.loss.shape, dtype=float)
                  self.dynamic_mean_error_by_element = Array(file, 'dynamic_mean_error_by_element', self.mean_error_by_element.shape, dtype=float, resizable_cross=True)
                  self.dynamic_RMSE_by_element = Array(file, 'dynamic_RMSE_by_element', self.RMSE_by_element.shape, dtype=float, resizable_cross=True)
  
              # legacy code, shouldn't be used in newer history files!
              if (self.mean_error_by_element.maxshape[1] is not None) or (self.RMSE_by_element.maxshape[1] is not None):
                  mean_error_by_element = np.array(self.mean_error_by_element).copy()
                  RMSE_by_element = np.array(self.RMSE_by_element).copy()
                  del file['mean_error_by_element']
                  del file['RMSE_by_element']
                  self.mean_error_by_element = Array(file, 'mean_error_by_element', mean_error_by_element.shape, data=mean_error_by_element, dtype=float, resizable_cross=True)
                  self.RMSE_by_element = Array(file, 'RMSE_by_element', RMSE_by_element.shape, data=RMSE_by_element, dtype=float, resizable_cross=True)
                  #self.mean_error_by_element[:] = mean_error_by_element
                  #self.RMSE_by_element[:] = RMSE_by_element
  
              # allows the resizing of by-element data arrays if relevant elements change
              if relevant_elements is not None and len(self.relevant_elements) != len(relevant_elements):
                  print("relevant_elements has changed!")
                  if len(self.relevant_elements) < len(relevant_elements):
                      self.mean_error_by_element.resize_cross(len(relevant_elements))
                      self.RMSE_by_element.resize_cross(len(relevant_elements))
                      if test_dynamic:
                          self.dynamic_mean_error_by_element.resize_cross(len(relevant_elements))
                          self.dynamic_RMSE_by_element.resize_cross(len(relevant_elements))
                  self.relevant_elements = relevant_elements
  
          # if we have no testing batches, we won't be running tests, so no need to prep
          if not testing_batches:
              return
  
          # for each relevant element, gives you a list of atom indices in the testing set
          atom_indices = {e:[] for e in self.relevant_elements}
          atom_index = 0
          for batch in testing_batches:
              if test_dynamic: # batches come in pairs. inspect first of pair.
                  batch = batch[0]
              atomic_numbers = Molecule.get_atomic_numbers(batch.x)
              for e in atomic_numbers:
                  e = e.item()
                  if e in self.relevant_elements:
                      atom_indices[e].append(atom_index)
                  atom_index += 1
          self.atom_indices = {e:np.array(ai) for e,ai in atom_indices.items()}
  
          # precompute weight per testing batch and total testing weight
          self.batch_weights = torch.tensor([torch.sum(batch.weights) for batch in testing_batches])
          self.total_weight = sum(self.batch_weights)
  
          
      def run_test(self, model, batch_number, epoch, batch_in_epoch, example_in_epoch, example, elapsed_time, verbose=True, log=True):
          if verbose: print("")
  
          time0 = time.time()
          losses = []
          residual_chunks = []
          if self.test_dynamic:
              dynamic_losses = []
              dynamic_residual_chunks = []
          model.eval() # don't compute running means
          with torch.no_grad(): # don't compute gradients
              for i, batch in enumerate(self.testing_batches):
                  if verbose and i % 10 == 0: print(f"Testing batches...  {i:3} / {len(self.testing_batches)}   ", end="\r", flush=True)
                  if self.test_dynamic:
                      weights = batch[0].weights.to(self.device)
                      residuals = [model(b.to(self.device).x, b.edge_index, b.edge_attr) - b.y for b in batch]
                      
                      loss = sum(loss_from_residuals(r, weights) for r in residuals) / len(residuals)
                      dynamic_residuals = residuals[1] - residuals[0]
                      dynamic_loss = loss_from_residuals(dynamic_residuals, weights)
                      
                      residuals = residuals[0]
                      
                      dynamic_losses.append(dynamic_loss.detach())
                      dynamic_residual_chunks.append(dynamic_residuals.detach())
                  else:
                      batch.to(self.device)
                      loss, residuals = loss_function(model(batch.x, batch.edge_index, batch.edge_attr), batch)
                      #if use_tensor_constraint:
                      #    chunk = chunk[...,0] # Keep only scalar part of residuals
                  losses.append(loss.detach())
                  residual_chunks.append(residuals)
  
          if verbose: print("Collating batch results...", end="\r", flush=True)
          # Using average loss over batches, rather than weighted average, to better mirror
          # running average of testing loss:
          #loss = (torch.dot(torch.tensor(losses), self.batch_weights) / self.total_weight).sqrt()
          loss = torch.tensor(losses).sqrt().mean()
          residuals = torch.cat(residual_chunks)
          residuals_by_element = {e:residuals[self.atom_indices[e]] for e in self.relevant_elements}
  
          # compute mean errors and RMSEs
          mean_error_by_element = torch.tensor([residuals_by_element[e].mean() for e in self.relevant_elements])
          RMSE_by_element = torch.tensor([residuals_by_element[e].square().mean().sqrt() for e in self.relevant_elements])
  
          if self.test_dynamic:
              dynamic_loss = torch.tensor(dynamic_losses).sqrt().mean()
              dynamic_residuals = torch.cat(dynamic_residual_chunks)
              dynamic_residuals_by_element = {e:dynamic_residuals[self.atom_indices[e]] for e in self.relevant_elements}
  
              # compute mean errors and RMSEs
              dynamic_mean_error_by_element = torch.tensor([dynamic_residuals_by_element[e].mean() for e in self.relevant_elements])
              dynamic_RMSE_by_element = torch.tensor([dynamic_residuals_by_element[e].square().mean().sqrt() for e in self.relevant_elements])
          else:
              dynamic_loss=None
              dynamic_mean_error_by_element=None
              dynamic_RMSE_by_element=None
              dynamic_residuals_by_element=None
  
          time1 = time.time()
          test_time = time1 - time0
  
          if verbose:
              if self.test_dynamic:
                  print(f"  Test loss = {loss:6.3f}   Dynamic loss = {dynamic_loss:6.3f}   Test time = {test_time:.2f}")
                  print(f"  Element   Mean Error    RMSE   Dyn.Mean   Dyn.RMSE")
                  #print(f"<4> Ee  <7>  012.345 <5> 012.345")
                  for i, e in enumerate(self.relevant_elements):
                      print(f"    {self.number_to_symbol[e].rjust(2)}       {mean_error_by_element[i]:3.3f}     {RMSE_by_element[i]:3.3f}     {dynamic_mean_error_by_element[i]:3.3f}     {dynamic_RMSE_by_element[i]:3.3f}")
              else:
                  print(f"  Test loss = {loss:6.3f}   Test time = {test_time:.2f}")
                  print(f"  Element   Mean Error    RMSE")
                  #print(f"<4> Ee  <7>  012.345 <5> 012.345")
                  for i, e in enumerate(self.relevant_elements):
                      print(f"    {self.number_to_symbol[e].rjust(2)}       {mean_error_by_element[i]:3.3f}     {RMSE_by_element[i]:3.3f}")
  
          if log:
              if example is None:
                  example = (epoch-1) * self.examples_per_epoch + example_in_epoch
              elif example_in_epoch is None:
                  example_in_epoch = example % self.examples_per_epoch
              self.log_test(batch_number, epoch, batch_in_epoch, example_in_epoch, example, elapsed_time,
                      loss, mean_error_by_element, RMSE_by_element, residuals_by_element,
                      dynamic_loss, dynamic_mean_error_by_element,
                      dynamic_RMSE_by_element, dynamic_residuals_by_element) 
          return loss
  
      def log_test(self, batch_number, epoch, batch_in_epoch, example_in_epoch, example, elapsed_time,
              loss, mean_error_by_element, RMSE_by_element, residuals_by_element=None,
              dynamic_loss=None, dynamic_mean_error_by_element=None,
              dynamic_RMSE_by_element=None, dynamic_residuals_by_element=None):
          self.batch_number.append(batch_number)
          self.epoch.append(epoch)
          self.batch_in_epoch.append(batch_in_epoch)
          self.example_in_epoch.append(example_in_epoch)
          self.example.append(example)
          self.elapsed_time.append(elapsed_time)
          self.loss.append(loss)
          self.mean_error_by_element.append(mean_error_by_element)
          self.RMSE_by_element.append(RMSE_by_element)
          if self.test_dynamic:
              self.dynamic_loss.append(dynamic_loss)
              self.dynamic_mean_error_by_element.append(dynamic_mean_error_by_element)
              self.dynamic_RMSE_by_element.append(dynamic_RMSE_by_element)            
          
          if self.store_residuals:
              self.residuals_by_element = residuals_by_element
              if self.test_dynamic:
                  self.dynamic_residuals_by_element = dynamic_residuals_by_element
              
          if self.wandb_log is not None:
              wandb_dict = {
                  'batch_number':batch_number,
                  'epoch':epoch,
                  'batch_in_epoch':batch_in_epoch,
                  'example_in_epoch':example_in_epoch,
                  'example':example,
                  'elapsed_time':elapsed_time,
                  'test_loss':loss,
                  'mean_error_by_element':{
                      int(e):mean_error_by_element[i].item() for i, e in enumerate(self.relevant_elements)},
                  'RMSE_by_element':{
                      int(e):RMSE_by_element[i].item() for i, e in enumerate(self.relevant_elements)}}
              if self.test_dynamic:
                  wandb_dict = {**wandb_dict,
                  'dynamic_test_loss':dynamic_loss,
                  'dynamic_mean_error_by_element':{
                      int(e):dynamic_mean_error_by_element[i].item() for i, e in enumerate(self.relevant_elements)},
                  'dynamic_RMSE_by_element':{
                      int(e):dynamic_RMSE_by_element[i].item() for i, e in enumerate(self.relevant_elements)},                    
                  }
              self.wandb_log(wandb_dict)
                  
      def smoothed_loss(self, i = -1, window=5):
          if i < 0:
              j = len(self.loss) - i + 1
              assert j > 0, "index out of range"
          else:
              j = i + 1
              assert j <= len(self.loss), "index out of range"
          window = min(window, j)
          return sum(self.loss[j-window:j]) / window
          
      def coord(self, coord):
          if coord == 'time':
              return self.elapsed_time
          elif coord == 'example':
              return self.example
          else:
              raise ValueError("Argument 'coord' should be 'time' or 'example'")
      
      # WARNING: don't share between processes if you use this        
      x_coord = 'example'
      
      @property
      def x(self):
          return getattr(self, TestingHistory.x_coord)  
                  
      def get_spectral(self, smooth=.1, coord='example', log_loss=True, decay=0):
          return Spectral(self, log_loss=True, smooth=smooth, coord=coord, decay=decay)
                  
      def __getstate__(self):
          d = self.__dict__.copy()
          del d['testing_batches']
          del d['wandb_log']
          return d
  
  class Spectral:
      
      def __init__(self, h, log_loss=True, smooth=.1, coord='example', decay=0):
          self.h = h
          self.log_loss = log_loss
          self.smooth = smooth
          self.coord = coord
          self.decay = decay
          self.compute()
      
      def compute(self):
          set_x(self.coord)
          x0 = self.h.x[0]
          x1 = self.h.x[1]
          chunk_size = round(len(self.h.x) * self.smooth)
          chunk_num = len(self.h.x) + 1 - chunk_size
          loss = self.h.loss if not self.log_loss else np.log(self.h.loss)
          cumsum = np.cumsum(loss)
          #print(f"cumsum = {cumsum}")
          #print(f"chunk_size = {chunk_size}")
          #print(f"cumsum[chunk_size:] = {cumsum[chunk_size:]}")
          #print(f"cumsum[:chunk_size] = {cumsum[:chunk_size]}")
          self.chunk_val = (cumsum[chunk_size:] - cumsum[:-chunk_size]) / chunk_num
          #self.chunk_edges = np.arange(chunk_size-1, chunk_size)
          #assert len(self.chunk_edges) == len(self.chunks)
          #self.x = self.h.x[self.chunk_edges]
          self.spectrum = np.fft.fft(self.chunk_val)
          
                     
      def recompute(self):
          self.compute()
          
          
  
  
          
  # Static methods for comparing runs of different lengths,
  # computing uncertainty in long-runs predicions from short data
  # and other things?
  def set_x(coord):
      TestingHistory.x_coord = 'example' if coord=='example' else 'elapsed_time'
  
  def sort_by_last_coord(histories, coord=None):
      set_x(coord)
      indices, h_sorted = zip(*sorted(enumerate(histories), key = lambda th : th[1].x[-1], reverse=True))
      return np.array(indices), h_sorted
  
  # Handy unsorting function      
  
  def unsort(indices, data):
      indices = np.array(indices)
      if isinstance(data, np.ndarray):
          fixed = np.empty_like(data)
          np.put(fixed, indices, data)
      elif isinstance(data, tuple) and len(data) != len(indices):
          fixed = tuple(unsort(indices, d) for d in data)
      elif isinstance(data, Sequence) and len(data) == len(indices):
          fixed = [None] * len(indices)
          for t, s in enumerate(indices):
              fixed[t] = data[s]
      else:
          fixed = data
      return fixed
      
  class SList(list):
      pass
      
  class seg_pass:
      def __init__(self, seg):
          self.seg = seg
  class start_seg:
      def __init__(self, seg):
          self.seg = seg
          
          
      
  def needs_sorted_histories(f, assume=False, h_arg=0):
      """
      Decorator for functions that need histories sorted from longest to shortest
      """
      def new_f(*args, sorted=assume, sort_indices=None, **kwargs):
          if sorted or isinstance(args[h_arg], SList):
              return f(*args, **kwargs)
          else:
              if sort_indices is None:
                  sort_indices = np.empty(len(args[h_arg]), dtype=int)
              args=list(args)
              sort_indices[:], args[h_arg] = sort_by_last_coord(args[h_arg], kwargs['coord'])
              args[h_arg] = SList(args[h_arg])
              value = f(*args, **kwargs)
              return unsort(sort_indices, value)
      #new_f.__name__ = f.__name__
      return new_f
  
              
  def decay_weight(a0, a1, R, end):
      """
      Tells you the integrated weight of [a0,a1], given that
      infinitestimal weights are w(x) = exp(R * (end - x))
  
      a0 < a1 are the endpoints of the segment.
      end is the coord at which the decay is 1.0.
      R is the rate of (leftwards) decay.
      """
      return (np.exp(R * (end - a0)) - np.exp(R * (end - a1))) / R
      
  def predictions_with_errors(histories, coord, errors=True, e_lives=4, verbose=0):
      """
      Doing pairwise comparisons.
      Weighted by shared length, and distance from lead time
      """    
      set_x(coord)
      histories = list(histories)
      N = len(histories)
  
      log_loss = [np.log(h.loss) for h in histories]
  
      pred = np.zeros((N,), dtype=float)
      rel_avg = np.zeros((N,N), dtype=float)
      if errors:
          err = np.zeros((N,), dtype=float)
          rel_var = np.zeros((N,N), dtype=float)
      rel_weight = np.zeros((N,N), dtype=float)
      
      for k, h in enumerate(histories):
          end = h.x[-1]
          d = end - h.x[0]
          #R = e_lives / d
          
          for j, g in enumerate(histories):
              
              if j == k:
                  continue
                  #if len(log_loss[k]) < 4 or not errors:
                  #    continue
                  #
                  #diff = log_loss[k][:-2] - log_loss[k][2:]
                  #weights = np.exp(-R * h.x[1:-1])
                  #avg_diff = np.average(diff, weights=weights)
                  #rel_var[k,k] = np.average((diff - avg_diff)**2, weights=weights)
                  #rel_weight[k,k] = decay_weight(h.x[1], h.x[-2], R, end)
              
              #else:
              b0 = max(h.x[0], g.x[0])
              b1 = min(h.x[-1], g.x[-1])
              if b1 <= b0:
                  continue
              
              X, Y = interleave((h.x, g.x), (log_loss[k], log_loss[j]), (b0, b1))
              
              weights = np.exp(-e_lives * X / b1)
              diff = Y[0] - Y[1]
              rel_avg[k,j] = np.average(diff, weights=weights)
              
              if rel_avg[k,j] < -1000:
                  rel_avg[k,j] = 0
                  continue
              
              if errors:
                  rel_var[k,j] = np.average((diff-rel_avg[k,j])**2, weights=weights)
              rel_weight[k,j] = decay_weight(b0, b1, e_lives / d, end)
              
              if rel_weight[k,j] == 0:
                  print(f"   rel_weight[{k},{j}] == 0")
          
          weight = np.sum(rel_weight[k])
          if weight == 0:
              print(f" weight[{k}] == 0")
              pred[k] = np.nan
              if errors:
                  err[k] = np.inf
          else:
              pred[k] = np.average(rel_avg[k]) #, weights=rel_weight[k])
              if pred[k] < -1000:
                  print(f" pred[{k}] == {pred[k]}")
              if errors:
                  err[k] = np.sqrt(np.average(rel_var[k])) #, weights=rel_weight[k]))
      if errors:
          return np.stack((pred, err), axis=1)
      else:
          return pred
  
  def log_inferred_loss(histories, x, coord='example', window=.1):
      """
      calculates a score for each history, at coordinate x,
      predictive of long-term relative performance.
      """
      raise NotImplementedError("not implemented yet!")
          
          
  def interleave(X, Y=None, bounds=None):
      """
      Each X[i] is an increasing series.
      Interleave them into one series, truncate to bounds=[b0,b1].
      This is W.
      
      If Y is not None, it should have a value for each X
      Z is a collection of interpolations of each Y to the values in W.
      
      Return W, Z  (or just W if Y is None)
      """
      if bounds=='inner':
          b0 = max(x[0] for x in X)
          b1 = min(x[-1] for x in X)
      elif bounds=='outer':
          # same as if no bounds specified
          bounds = None  
      elif bounds:
          b0, b1 = bounds
          
      W = np.array(X[0])
      for x in X[1:]:
          b_indices = np.searchsorted(W, x)
          np.insert(W, b_indices, x)
      if bounds:
          i0 = np.searchsorted(W, b0)
          i1 = np.searchsorted(W, b1)
          W = np.insert(W, (i0,i1), (b0,b1))[i0:i1+2]      
      W = np.unique(W)
      
      if Y:
          return W, [np.interp(W, x, y) for (x,y) in zip(X,Y)]
      else:
          return W 
  
  
  def savitzky_golay(y, window_size, order, deriv=0, rate=1):
      r"""Smooth (and optionally differentiate) data with a Savitzky-Golay filter.
      The Savitzky-Golay filter removes high frequency noise from data.
      It has the advantage of preserving the original shape and
      features of the signal better than other types of filtering
      approaches, such as moving averages techniques.
       Parameters
       ----------
       y : array_like, shape (N,)
             the values of the time history of the signal.
         window_size : int
             the length of the window. Must be an odd integer number.
         order : int
             the order of the polynomial used in the filtering.
             Must be less then `window_size` - 1.
         deriv: int
             the order of the derivative to compute (default = 0 means only smoothing)
         Returns
         -------
         ys : ndarray, shape (N)
             the smoothed signal (or it's n-th derivative).
         Notes
         -----
         The Savitzky-Golay is a type of low-pass filter, particularly
         suited for smoothing noisy data. The main idea behind this
         approach is to make for each point a least-square fit with a
         polynomial of high order over a odd-sized window centered at
         the point.
         Examples
         --------
         t = np.linspace(-4, 4, 500)
         y = np.exp( -t**2 ) + np.random.normal(0, 0.05, t.shape)
         ysg = savitzky_golay(y, window_size=31, order=4)
         import matplotlib.pyplot as plt
         plt.plot(t, y, label='Noisy signal')
         plt.plot(t, np.exp(-t**2), 'k', lw=1.5, label='Original signal')
         plt.plot(t, ysg, 'r', label='Filtered signal')
         plt.legend()
         plt.show()
         References
         ----------
         .. [1] A. Savitzky, M. J. E. Golay, Smoothing and Differentiation of
            Data by Simplified Least Squares Procedures. Analytical
            Chemistry, 1964, 36 (8), pp 1627-1639.
         .. [2] Numerical Recipes 3rd Edition: The Art of Scientific Computing
            W.H. Press, S.A. Teukolsky, W.T. Vetterling, B.P. Flannery
            Cambridge University Press ISBN-13: 9780521880688
      """
      from math import factorial
      
      try:
          window_size = np.abs(np.int(window_size))
          order = np.abs(np.int(order))
      except ValueError as msg:
          raise ValueError("window_size and order have to be of type int")
      if window_size % 2 != 1 or window_size < 1:
          raise TypeError("window_size size must be a positive odd number")
      if window_size < order + 2:
          raise TypeError("window_size is too small for the polynomials order")
      order_range = range(order+1)
      half_window = (window_size -1) // 2
      # precompute coefficients
      b = np.mat([[k**i for i in order_range] for k in range(-half_window, half_window+1)])
      m = np.linalg.pinv(b).A[deriv] * rate**deriv * factorial(deriv)
      # pad the signal at the extremes with
      # values taken from the signal itself
      firstvals = y[0] - np.abs( y[1:half_window+1][::-1] - y[0] )
      lastvals = y[-1] + np.abs(y[-half_window-1:-1][::-1] - y[-1])
      y = np.concatenate((firstvals, y, lastvals))
      return np.convolve( m[::-1], y, mode='valid')
    
                      
  TH = TestingHistory    
      
      
          
  
  #def exp_fn(x, a, speed, limit):
  #    return a * np.exp(speed * x) + limit
  def exp_fn(x, a, speed, limit):
      return np.exp(a * np.exp(speed * x) + limit)
  #def exp_fn(x, a, speed, limit, k, w):
  #    return a * np.exp(speed * x) + limit + k / (x + w)
  #def exp_fn(x, k, w, limit):
  #    return limit + k / (x + w)
          
  
  # if function is called
  if __name__ == '__main__':
      import sys
      if len(sys.argv) < 2:
          exit()
      input("Exploring the evolution of exp fit over time...")
      history = TrainTestHistory(file=sys.argv[1])
      history.test.show_fit_evolution()
      history.close()
      input("Press Enter to continue...")
  
  
  
  