text: # This file is part of the Astrometry.net suite.
  # Licensed under a 3-clause BSD style license - see LICENSE
  from __future__ import print_function
  import os
  from astrometry.util.file import *
  
  '''
  
  Stages: a utility for saving and resuming computation, savings
  intermediate results as pickle files.
  
  
  '''
  
  class CallGlobal(object):
      def __init__(self, pattern, globals, *args, **kwargs):
          self.pat = pattern
          self.args = args
          self.kwargs = kwargs
          self.globals = globals
      def getfunc(self, stage):
          func = self.pat % stage
          func = eval(func, self.globals)
          return func
      def getkwargs(self, stage, **kwargs):
          kwa = self.kwargs.copy()
          kwa.update(kwargs)
          return kwa
      def __call__(self, stage, **kwargs):
          func = self.getfunc(stage)
          kwa = self.getkwargs(stage, **kwargs)
          return func(*self.args, **kwa)
  
  class CallGlobalTime(CallGlobal):
      def __call__(self, stage, **kwargs):
          from astrometry.util.ttime import Time
          from datetime import datetime
          t0 = Time()
          print('Running stage', stage, 'at', datetime.now().isoformat())
          rtn = super(CallGlobalTime, self).__call__(stage, **kwargs)
          t1 = Time()
          print('Stage', stage, ':', t1-t0)
          print('Stage', stage, 'finished:', datetime.now().isoformat())
          return rtn
  
  def runstage(stage, picklepat, stagefunc, force=[], forceall=False, prereqs={},
               update=True, write=True, initial_args={}, **kwargs):
      '''
      Run to a given *stage*.
  
      Each stage is a function.
      
      Each stage takes a dict and returns None or a dict.
  
      Each stage (except the first one!) has a prerequisite stage.
  
      Results after running each stage can be written to a pickle file,
      and later read instead of running the stage.
  
      Parameters
      ----------
      stage - int or string
          Stage name
      picklepat - string
          Filename pattern to which pickle data should be written
      stagefunc - function
          Function to call for each stage.
      force - list of stage name/number
          List of stages to run, ignoring existing pickle files
      forceall - boolean
          Force running all stages; ignore all existing pickle files
      prereqs - dict of stage->stage mappings
          Defines the dependencies between stages
      update - boolean
          Update prerequiste dict with results before writing out pickle
      write - boolean, or list of stages
          Write a pickle for this stage / all stages?
      initial_args - dict
          Arguments to pass to the first stage
      kwargs - dict
          Keyword args added to the prerequisites before running each stage
      '''
      # NOTE, if you add or change args here, be sure to update the recursive
      # "runstage" call below!!
      print('Runstage', stage)
  
      try:
          pfn = picklepat % stage
      except:
          pfn = picklepat % dict(stage=stage)
      
      if os.path.exists(pfn):
          if forceall or stage in force:
              print('Ignoring pickle', pfn, 'and forcing stage', stage)
          else:
              print('Reading pickle', pfn)
              R = unpickle_from_file(pfn)
              return R
  
      try:
          prereq = prereqs[stage]
      except KeyError:
          prereq = stage - 1
  
      if prereq is None:
          P = initial_args
      else:
          P = runstage(prereq, picklepat, stagefunc,
                       force=force, forceall=forceall, prereqs=prereqs, update=update,
                       write=write, initial_args=initial_args, **kwargs)
  
      #P.update(kwargs)
      Px = P.copy()
      Px.update(kwargs)
  
      print('Running stage', stage)
      # print('Prereq keys:', P.keys())
      # print('Adding kwargs keys:', kwargs.keys())
      # print('Combined keys:', Px.keys())
  
      R = stagefunc(stage, **Px)
      print('Stage', stage, 'finished')
      # if R is not None:
      #     print('Result keys:', R.keys())
  
      if update:
          if R is not None:
              P.update(R)
          R = P
  
      if not write:
          pass
      elif (write is True or stage in write):
          print('Saving pickle', pfn)
          #print('Pickling keys:', R.keys())
          # Create directory, if necessary.
          dirnm = os.path.dirname(pfn)
          if len(dirnm) and not os.path.exists(dirnm):
              try:
                  os.makedirs(dirnm)
              except:
                  pass
          pickle_to_file(R, pfn)
          print('Saved', pfn)
      return R
  