#!/search/odin/liyaozong/tools/python3/bin/python3
# coding: utf8

import random
import numpy as np
from collections import defaultdict, deque
from game import Board, Game
from mcts_pure import MCTSPlayer as MCTS_Pure
from mcts_alphaZero import MCTSPlayer
import multiprocessing
from multiprocessing import Manager, Pool
import time
import logging
from tensorflow.python.platform import gfile
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

model_dir='./multi_3'
log_name = model_dir + '/logs'
current_model_name = model_dir + '/current_policy_model'
temp_current_model_name = model_dir + '/temp_current_policy_model'
best_model_name = model_dir + '/best_policy_model'

class TrainPipeline():

    def __init__(self, init_model=None):
        # params of the board and the game
        self.board_width = 8
        self.board_height = 8
        self.n_in_row = 5
        # training params
        self.learn_rate = 2e-3
        # self.lr_multiplier = 1.0
        self.temp = 1.0  # the temperature param
        self.n_playout = 400  # num of simulations for each move
        self.c_puct = 5
        self.batch_size = 512  # mini-batch size for training
        self.buffer_num = self.batch_size * 10
        self.epochs = 5  # num of train_steps for each update
        self.kl_targ = 0.02
        self.check_freq = 60
        self.game_batch_num = 1000000000
        self.process_num = 12
        self.summary_record_freq = 5
        # self.global_update_step = 0

    def collect_selfplay_data_thread(self, thread_id, shared_queue, net_lock, data_lock):
        logging.info('selfpaly process start: {}'.format(thread_id))
        os.environ["CUDA_VISIBLE_DEVICES"] = str(thread_id % 6 + 2)
        from policy_value_net_tensorflow import PolicyValueNet
        # 读取模型文件，加锁
        with net_lock:
            current_policy = PolicyValueNet(self.board_width, self.board_height, model_file=current_model_name)
        local_board = Board(width=self.board_width,
                           height=self.board_height,
                           n_in_row=self.n_in_row)
        local_game = Game(local_board)
        local_mcts_player = MCTSPlayer(current_policy.policy_value_fn,
                                      c_puct=self.c_puct,
                                      n_playout=self.n_playout,
                                      is_selfplay=1)
        winner, play_data = local_game.start_self_play(local_mcts_player,
                                                       temp=self.temp)
        play_data = list(play_data)
        play_data = self.get_equi_data(play_data)
        # 添加对弈数据，加锁
        with data_lock:
            shared_queue.extend(play_data)
            while len(shared_queue) > self.buffer_num:
                shared_queue.pop(0)
        logging.info('selfpaly process finished: {}'.format(thread_id))

    def get_equi_data(self, play_data):
        """augment the data set by rotation and flipping
        play_data: [(state, mcts_prob, winner_z), ..., ...]
        """
        extend_data = []
        for state, mcts_porb, winner in play_data:
            for i in [1, 2, 3, 4]:
                # rotate counterclockwise
                equi_state = np.array([np.rot90(s, i) for s in state])
                equi_mcts_prob = np.rot90(np.flipud(
                    mcts_porb.reshape(self.board_height, self.board_width)), i)
                extend_data.append((equi_state,
                                    np.flipud(equi_mcts_prob).flatten(),
                                    winner))
                # flip horizontally
                equi_state = np.array([np.fliplr(s) for s in equi_state])
                equi_mcts_prob = np.fliplr(equi_mcts_prob)
                extend_data.append((equi_state,
                                    np.flipud(equi_mcts_prob).flatten(),
                                    winner))
        return extend_data

    def policy_update(self, current_policy_value_net, shared_queue, net_lock, global_update_step, lr_multiplier):
        """update the policy-value net"""
        random_index = list(range(len(shared_queue)))
        random.shuffle(random_index)
        mini_batch = []
        for i in range(self.batch_size):
            mini_batch.append(shared_queue[random_index[i]])
        state_batch = [data[0] for data in mini_batch]
        mcts_probs_batch = [data[1] for data in mini_batch]
        winner_batch = [data[2] for data in mini_batch]
        old_probs, old_v = current_policy_value_net.policy_value(state_batch)
        loss, entropy = current_policy_value_net.train_step(
                state_batch,
                mcts_probs_batch,
                winner_batch,
                self.learn_rate * lr_multiplier.value)
        new_probs, new_v = current_policy_value_net.policy_value(state_batch)
        kl = np.mean(np.sum(old_probs * (
                np.log(old_probs + 1e-10) - np.log(new_probs + 1e-10)),
                axis=1)
        )
        # adaptively adjust the learning rate
        if kl > self.kl_targ * 2 and lr_multiplier.value > 0.1:
            lr_multiplier.value /= 1.5
        elif kl < self.kl_targ / 2 and lr_multiplier.value < 10:
            lr_multiplier.value *= 1.5

        explained_var_old = (1 -
                             np.var(np.array(winner_batch) - old_v.flatten()) /
                             np.var(np.array(winner_batch)))
        explained_var_new = (1 -
                             np.var(np.array(winner_batch) - new_v.flatten()) /
                             np.var(np.array(winner_batch)))
        logging.info("update current model process kl:{:.5f},lr_multiplier:{:.3f},loss:{},entropy:{},explained_var_old:{:.3f},explained_var_new:{:.3f}".format(
            kl,
            lr_multiplier.value,
            loss,
            entropy,
            explained_var_old,
            explained_var_new))
        # summary for tensorboard
        if global_update_step.value % self.summary_record_freq == 0:
            current_policy_value_net.summary_record(
                    state_batch,
                    mcts_probs_batch,
                    winner_batch,
                    global_update_step.value,
                    )

    def update_net(self, shared_queue, net_lock, update_best_model, global_update_step, lr_multiplier, stop_update_process, update_or_selfplay):
        os.environ["CUDA_VISIBLE_DEVICES"] = "1"
        from policy_value_net_tensorflow import PolicyValueNet
        current_policy_value_net = PolicyValueNet(self.board_width, self.board_height, model_dir)
        current_policy_value_net.save_model(current_model_name)
        current_policy_value_net.save_model(best_model_name)
        while global_update_step.value <= self.game_batch_num:
            if update_or_selfplay.value == 0:
                if len(shared_queue) >= self.batch_size:
                    for _ in range(self.epochs):
                        global_update_step.value += 1
                        logging.info('update current model process start self train: {}'.format(global_update_step.value))
                        self.policy_update(current_policy_value_net, shared_queue, net_lock, global_update_step, lr_multiplier)
                        if (global_update_step.value) % self.check_freq == 0:
                            update_best_model.value = 1
                    # 这里更新最新模型文件
                    with net_lock:
                        logging.info('update process update current model')
                        current_policy_value_net.save_model(current_model_name)
                update_or_selfplay.value = 1
            else:
                time.sleep(1)
        stop_update_process.value = 1

    def update_best_model_thread(self, current_model_name, best_model_name, net_lock, update_best_model, stop_update_process):
        os.environ["CUDA_VISIBLE_DEVICES"] = "1"
        def update_best_model_local_process_func(name1, name2):
            from policy_value_net_tensorflow import PolicyValueNet
            PolicyValueNet(self.board_width, self.board_height, model_file=name1).save_model(name2)
        update_best_model_global_num = 0
        update_best_model_time = 0
        logging.info('update best model process start!')
        while stop_update_process.value == 0:
            time.sleep(1)
            if update_best_model.value == 1:
                update_best_model.value = 0
                update_best_model_global_num += 1
                logging.info('update best model process global_num:{}'.format(update_best_model_global_num))
                win_num = 0
                start_player = 0
                with net_lock:
                    p = multiprocessing.Process(target=update_best_model_local_process_func, args=(current_model_name, temp_current_model_name))
                    p.start()
                    p.join()
                for i in range(30):
                    board = Board(width=self.board_width, height=self.board_height, n_in_row=self.n_in_row)
                    game = Game(board)
                    win_player = game.two_net_play(temp_current_model_name, best_model_name, net_lock, start_player=start_player, is_shown=0)
                    start_player = 1 - start_player  # play first in turn
                    if win_player == 0:
                        win_num += 1
                logging.info('update best model process global_num:{} finished, current model win total {}'.format(update_best_model_global_num, win_num))
                if win_num >= 16:
                    update_best_model_time += 1
                    logging.info('update best model process get new best model:{}'.format(update_best_model_time))
                    p = multiprocessing.Process(target=update_best_model_local_process_func, args=(temp_current_model_name, best_model_name))
                    p.start()
                    p.join()

    def run(self):
        """run the training pipeline"""
        try:
            # 必须在一个线程中引入tensorflow，否则会造成其他线程由于错误阻塞。
            # ERROR: could not retrieve CUDA device count: CUDA_ERROR_NOT_INITIALIZED
            m = Manager()
            shared_queue = m.list()
            net_lock = m.Lock()
            data_lock = m.Lock()
            stop_update_process = multiprocessing.Value('i', 0)
            update_best_model = multiprocessing.Value('i', 0)
            global_update_step = multiprocessing.Value('i', 0)
            lr_multiplier = multiprocessing.Value('d', 1.0)
            update_or_selfplay = multiprocessing.Value('i', 0)
            update_best_model_process = multiprocessing.Process(target=self.update_best_model_thread,
                    args=(current_model_name, best_model_name, net_lock, update_best_model, stop_update_process))
            update_best_model_process.start()
            update_current_model_process = multiprocessing.Process(target=self.update_net,
                    args=(shared_queue, net_lock, update_best_model, global_update_step, lr_multiplier, stop_update_process, update_or_selfplay))
            update_current_model_process.start()
            while stop_update_process.value == 0:
                if update_or_selfplay.value == 1:
                    # 产生训练数据
                    pro_list = []
                    for i in range(self.process_num):
                        pro = multiprocessing.Process(target=self.collect_selfplay_data_thread, args=(i, shared_queue, net_lock, data_lock))
                        pro_list.append(pro)
                        pro.start()
                    for pro in pro_list:
                        pro.join()
                    update_or_selfplay.value = 0
                else:
                    time.sleep(0.1)
            # 等待更新最有模型的进程结束，再结束主程序
            stop_update_process.value = 1
            while update_best_model_process.is_alive():
                time.sleep(5)
        except Exception as e:
            logging.error(e)

if __name__ == '__main__':
    if gfile.Exists(model_dir):
        gfile.DeleteRecursively(model_dir)
    gfile.MakeDirs(model_dir)
    # log 配置
    logging.basicConfig(filename=log_name, level=logging.INFO, format="[%(levelname)s]\t%(asctime)s\tLINENO:%(lineno)d\t%(message)s", datefmt="%Y-%m-%d %H:%M:%S")
    training_pipeline = TrainPipeline()
    logging.info('start training')
    training_pipeline.run()
    logging.info('all finished')


#!/usr/bin/env python

import os
import time
import numpy as np
import socket
import logging
import tempfile
import shutil
from glob import glob
from datetime import datetime
from dlnpyutils import utils as dln, coords
import photred
from photred import io,daomatch
from photred import allframe as alf
import dill as pickle
from astropy.io import fits
from astropy.wcs import WCS
from astropy.table import Table
from astropy.coordinates import SkyCoord
from dustmaps.sfd import SFDQuery
import subprocess

def make_brick_wcs(brickstr):
    """
    This creates a brick WCS given the brick structure

    Each brick covers 0.25 x 0.25 deg at a pixel scale of 0.262"/pix
    with 3600x3600 pixels that gives us an overlap of ~82 pixels on
    each side.  The *UNIQUE* brick area (0.25x0.25 deg) is defined by
    the BRAMIN/MAX and BDECMIN/MAX keywords.
    """

    # Make the tiling file
    #---------------------
    nx = 3600
    ny = 3600
    step = 0.262 / 3600    # 0.262" per pixel, DECam pixel scale
    xref = nx/2
    yref = ny/2

    #  Make the header as well
    tilehead = fits.Header()
    tilehead['NAXIS1'] = nx
    tilehead['CDELT1'] = step
    tilehead['CRPIX1'] = xref+1
    tilehead['CRVAL1'] = brickstr['ra']
    tilehead['CTYPE1'] = 'RA---TAN'
    tilehead['NAXIS2'] = ny
    tilehead['CDELT2'] = step
    tilehead['CRPIX2'] = yref+1
    tilehead['CRVAL2'] = brickstr['dec']
    tilehead['CTYPE2'] = 'DEC--TAN'
    tilehead['BRAMIN'] = brickstr['ra1'],'RA min of unique brick area'
    tilehead['BRAMAX'] = brickstr['ra2'],'RA max of unique brick area'
    tilehead['BDECMIN'] = brickstr['dec1'],'DEC min of unique brick area'
    tilehead['BDECMAX'] = brickstr['dec2'],'DEC max of unique brick area'
    wcs = WCS(tilehead)

    # Create the TILE structure
    tile = {'type':'WCS','naxis':np.array([nx,ny]),'cdelt':np.array([step,step]).astype(float),
            'crpix':np.array([xref+1,yref+1]).astype(float),
            'crval':np.array([brickstr['ra'],brickstr['dec']]).astype(float),'ctype':['RA--TAN','DEC--TAN'],
            'head':tilehead,'wcs':wcs,'xrange':[0,nx],'yrange':[0,ny],'nx':nx,'ny':ny}

    return tile

def forcebrick(brick,scriptsdir=None,irafdir=None,workdir=None,redo=False,
               update=False,logfile=None,delvedir=None,delvereddir=None):
    """
    Process a single DELVE brick and perform ALLFRAME FORCED photometry

    Parameters
    ----------
    brick : str
      The DELVE brick name, e.g. 1234m045

    By D. Nidever  August 2019
    Translated to Python by D. Nidever,  April 2022
    """

    # Limit the number of threads
    #CPU,tpool_nthreads=4

    # This bricks pre-processing script gets DELVE and community MC data ready
    # to run PHOTRED ALLFRAME on it.

    t0 = time.time()
    curdir = os.path.abspath(os.curdir)

    # Defaults
    if delvedir is not None:
        if delvedir.endswith('/')==False:
            delvedir += '/'
    else:
        delvedir = '/net/dl2/dnidever/delve/'
    if os.path.exists(delvedir)==False:
        os.makedirs(delvedir)
    if delvereddir is not None:
        if delvereddir.endswith('/')==False:
            delvereddir += '/'
    else:
        delvereddir = '/'.join(os.path.abspath(__file__).split('/')[:-3])+'/'
    if scriptsdir is not None:
        if scriptsdir.endswith('/')==False:
            scriptsdir += '/'
    else:
        scriptsdir = '/'.join(photred.__file__.split('/')[:-3])+'/scripts/'
    if irafdir is not None:
        if irafdir.endswith('/')==False:
            irafdir += '/'
    else:
        irafdir = os.path.expanduser('~/iraf/')
    tempdir = '/tmp/'
    if workdir is None:
        host = socket.gethostname()
        hostname = host.split('.')[0]
        workdir = '/data0/dnidever/delve/'
        if os.path.exists(workdir)==False:
            os.makedirs(workdir)
    # Exposures directory
    expdir = delvedir+'exposures/'
    # Bricks directory
    brickdir = delvedir+'bricks/'
    if os.path.exists(brickdir)==False:
        os.makedirs(brickdir)
    logsdir = brickdir+'logs/'
    if os.path.exists(logsdir)==False:
        os.makedirs(logsdir)

    # Start the logfile
    #------------------
    # format is delvered_forcebrick.brick.host.DATETIME.log
    host = socket.gethostname()
    hostname = host.split('.')[0]
    logtime = datetime.now().strftime("%Y%m%d%H%M%S")
    # Set up logging to screen and logfile
    logFormatter = logging.Formatter("%(asctime)s [%(levelname)-5.5s]  %(message)s")
    logger = logging.getLogger()
    while logger.hasHandlers(): # some existing loggers, remove them
        logger.removeHandler(logger.handlers[0])
    logger = logging.getLogger()
    logtime = datetime.now().strftime("%Y%m%d%H%M%S")
    if logfile is None:
        logfile = logsdir+'delvered_brick.'+brick+'.'+hostname+'.'+logtime+'.log'
    if os.path.exists(logfile): os.remove(logfile)
    fileHandler = logging.FileHandler(logfile)
    fileHandler.setFormatter(logFormatter)
    logger.addHandler(fileHandler)
    consoleHandler = logging.StreamHandler()
    consoleHandler.setFormatter(logFormatter)
    logger.addHandler(consoleHandler)
    logger.setLevel(logging.NOTSET)

    # Chip name and CCDNUM relations
    decam = Table.read(delvereddir+'data/decam.txt',format='ascii')
    for n in decam.colnames: decam[n].name = n.lower()  # lowercase column names

    # Load the brick information
    brickstr = Table.read(delvereddir+'data/delvemc_bricks_0.25deg.fits.gz')
    for n in brickstr.colnames: brickstr[n].name = n.lower()  # lowercase column names

    # Get the brick information
    bind, = np.where(brickstr['brickname'] == brick)
    if len(bind) == 0:
        logger.info(brick+' not in DELVE-MC brick list')
        return
    brickstr1 = brickstr[bind[0]]

    # Subdirectory is the first 4 digits of the brickname, e.g., 0952 of 0952m462, the RA portion of the name
    subdir = brickdir+brick[0:4]+'/'
    if os.path.exists(subdir)==False:
        os.makedirs(subdir)
    # Create brick directory if it doesn't exist yet
    bdir = subdir+brick+'/'
    if os.path.exists(bdir)==False:
        os.makedirs(bdir)
    logfile = bdir+brick+'.'+logtime+'.log'

    # Check the output file
    photfile = bdir+brick+'.fits'
    if os.path.exists(photfile+'.gz') and redo==False and update==False:
        logger.info(photfile+'.gz EXISTS and redo or update NOT set')
        return

    # DECam imager
    thisimager = {'telescope':'BLANCO','instrument':'DECam','namps':62,'separator':'_'}

    # Print information
    logger.info('--------------------------------------------------------------------')
    logger.info(' Run ALLFRAME forced photometry on DELVE Brick = '+brick)
    logger.info('--------------------------------------------------------------------')
    logger.info('DELVEDIR = '+delvedir)
    logger.info('SCRIPTSDIR = '+scriptsdir)
    logger.info('IRAFDIR = '+irafdir)
    logger.info('EXPOSUREDIR = '+expdir)
    logger.info('WORKINGDIR = '+workdir)
    logger.info('HOST = '+host)
    logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))

    # The photed.setup file
    setup = ['##### REQUIRED #####',
             'scriptsdir  '+scriptsdir,
             'irafdir     '+irafdir,
             'telescope   Blanco',
             'instrument  DECAM',
             'observatory CTIO',
             'nmulti      10',
             'nmulti_wcs       40',
             'nmulti_daophot   30',
             'nmulti_allframe  10',
             'filtref     g,i,r,z,u',
             'trans       delve.trans',
             '##### OPTIONAL #####',
             'sepfielddir  1',
             'sepchipdir   1',
             'keepmef      0',
             'catformat    FITS',
             'workdir      '+workdir,
             'clean        1',
             'skipcheck    1',
             'redo         0',
             'wcsrefname   GAIADR2',
             'searchdist   20',
             '#wcsrmslim   1.0',
             'hyperthread   1',
             'daopsfva      1',
             'daofitradfwhm 1.0',
             'psfcomsrc     0',
             'psfcomglobal  0',
             'psfcomgauss   0',
             '#mchmaxshift  50.0',
             'finditer      1',       # 2->1 on 7/22/20
             'alfdetprog  sextractor',
             '#alfnocmbimscale 0',
             'alftrimcomb   0',
             '#ddo51radoffset  1',
             'cmbforce      1',
             'keepinstr     1',
             'avgmag        1',
             'avgonlymag    0',
             'todered       u,g,r,i,z,g-i',
             '##### STAGES #####',
             '#rename',
             '#split',
             '#wcs',
             '#daophot',
             '#zeropoint',
             '#match',
             ' allframe',
             '#apcor',
             '#astrom',
             '#calib',
             '#combine',
             '#deredden',
             '#save',
             '#html']
    dln.writelines(bdir+'photred.setup',setup)
    setup = io.readsetup(setupdir=bdir)

    # Print out some information about the brick
    logger.info('RA = %.5f' % brickstr1['ra'])
    logger.info('DEC = %.3f' % brickstr1['dec'])
    logger.info('RA range  = [ %.5f,%.5f ]' % (brickstr1['ra1'],brickstr1['ra2']))
    logger.info('DEC range = [ %.5f,%.5f ]' % (brickstr1['dec1'],brickstr1['dec2']))

    # Get the Brick WCS information
    tile = make_brick_wcs(brickstr1)

    # Step 1: Get the list of exposures/chips that overlap this brick
    #----------------------------------------------------------------
    logger.info('Step 1: Get list of chip files that overlap this brick')
    #logger.info('Getting chips that overlap this brick'
    cenra = brickstr1['ra']
    cendec = brickstr1['dec']
    tid,tmpfile = tempfile.mkstemp(prefix="tmp",suffix='.fits',dir=tempdir)
    dln.touch(tmpfile)
    for f in [tmpfile]:
        if os.path.exists(f): os.remove(f)
    # /noshell causes problems on gp09 because it gives python2 instead of python3
    out = subprocess.check_output(delvereddir+'bin/query_delvered_summary_table '+str(cenra)+' '+str(cendec)+' '+tmpfile+' --lim 0.5',shell=True)
    exists = os.path.exists(tmpfile)
    size = os.path.getsize(tmpfile)
    if size == 0:
        logger.info('No overlapping chips found')
        if os.path.exists(tmpfile): os.remove(tmpfile)
        return
    chstr = Table.read(tmpfile)
    if os.path.exists(tmpfile): os.remove(tmpfile)
    nchstr = len(chstr)
    logger.info('Found '+str(nchstr)+' overlapping chips within 0.5 deg of brick center')
    logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))

    # Make sure the chips are unique,  some were duplicated on SMASH nights
    chid = np.char.array(chstr['expnum'].astype(str))+'-'+np.char.array(chstr['chip'].astype(str))
    uchstr,ui = np.unique(chid,return_index=True)
    chstr = chstr[ui]
    nchstr = len(chstr)

    # Do more rigorous overlap checking
    #  the brick region with overlap
    print('Performing more rigorous overlap checking')

    bcoo = tile['wcs'].pixel_to_world([0,tile['nx']-1,tile['nx']-1,0],[0,0,tile['ny']-1,tile['ny']-1])
    bvra = bcoo.ra.deg
    bvdec = bcoo.dec.deg
    #HEAD_XYAD,tile.head,[0,tile.nx-1,tile.nx-1,0],[0,0,tile.ny-1,tile.ny-1],bvra,bvdec,/deg
    olap = np.zeros(nchstr,bool)
    vxarr = np.zeros((nchstr,4),float)
    vyarr = np.zeros((nchstr,4),float)
    for i in range(nchstr):
        if (i % 100 == 0) and (i > 0):
            print(i)
        hd1 = io.readfile(chstr['file'][i],header=True)
        nx = hd1['naxis1']
        ny = hd1['naxis2']
        w1 = WCS(hd1)
        vcoo = w1.pixel_to_world([0,nx-1,nx-1,0],[0,0,ny-1,ny-1])
        vra = vcoo.ra.deg
        vdec = vcoo.dec.deg
        olap[i] = coords.doPolygonsOverlap(bvra,bvdec,vra,vdec)
        vx,vy = tile['wcs'].world_to_pixel(vcoo)
        vxarr[i,:] = vx
        vyarr[i,:] = vy
    # Require at least a 2 pixel overlap in X and Y
    g, = np.where((olap==True) & (np.max(vxarr,axis=1) >= 2) & (np.max(vyarr,axis=1) >= 2) &
                  (np.min(vxarr,axis=1) <= tile['nx']-3) & (np.min(vyarr,axis=1) <= tile['ny']-3))
    if len(g) == 0:
        logger.info('No chips overlap this brick')
        return
    logger.info(str(len(g))+' chips overlap this brick')
    chstr = chstr[g]
    nchstr = len(chstr)

    # APPLY cuts on the exposures
    #-----------------------------
    logger.info('Applying quality, zero-point, filter and exptime cuts')

    # Zero-point structure, from NSC
    zpstr = np.zeros(7,dtype=np.dtype([('instrument',np.str,10),('filter',np.str,10),('amcoef',(float,2)),('thresh',float)]))
    #zpstr = replicate({instrument:'',filter:'',amcoef:fltarr(2),thresh:0.5},7)
    zpstr['instrument'] = 'c4d'
    zpstr['filter'] = ['u','g','r','i','z','Y','VR']
    zpstr['amcoef'][0] = [-1.60273, -0.375253]  # c4d-u
    zpstr['amcoef'][1] = [0.277124, -0.198037]  # c4d-g
    zpstr['amcoef'][2] = [0.516382, -0.115443]  # c4d-r
    zpstr['amcoef'][3] = [0.380338, -0.067439]  # c4d-i
    zpstr['amcoef'][4] = [0.074517, -0.067031]  # c4d-z
    zpstr['amcoef'][5] = [-1.07800, -0.060014]  # c4d-Y
    zpstr['amcoef'][6] = [1.111859, -0.083630]  # c4d-VR

    # Convert to additive zero-point as used in NSC
    zpterm = -chstr['calib_zpterm']
    # Fix early DES exposures that used different units/gain
    gdes, = np.where(chstr['gain'] < 2)
    if len(gdes) > 0:
        zpterm[gdes] -= 1.55
    # global zpterm and airmass correction
    for i in range(len(zpstr)):
        ind, = np.where(chstr['filter']==zpstr['filter'][i])
        if len(ind)>0:
            zpterm[ind] -= np.polyval(np.flip(zpstr['amcoef'][i]),chstr['airmass'][ind])

    fwhmthresh = 2.0  # seeing 2.0" threshold
    filt = np.char.array(chstr['filter'].astype(str)).ljust(1)  # first characater
    gdch, = np.where((chstr['fwhm']*chstr['pixscale'] <= fwhmthresh) & (chstr['exptime'] >= 90.) & (zpterm >= -0.5) &
                     np.isfinite(zpterm) & np.isfinite(chstr['apcor']) &
                     ((filt== 'u') | (filt=='g') | (filt=='r') | (filt=='i') | (filt=='z') | (filt=='Y')))
    if len(gdch) == 0:
        logger.info('No chips passed the cuts')
        return
    logger.info(str(len(gdch))+' chips passed the cuts')
    chstr = chstr[gdch]
    nchstr = len(chstr)
    chstr['file'] = np.char.array(chstr['file'].astype(str)).strip()
    chstr['base'] = np.char.array(chstr['base'].astype(str)).strip()

    # Check if we need to update
    if update and redo==False:
        # Load previous meta
        metafile = bdir+brick+'_meta.fits'
        if os.path.exists(metafile):
            meta0 = Table.read(metafile,1)
            meta0['base'] = str(meta0['base'])
            ind1,ind2 = dln.match(meta0['base'],chstr['base'])
            if nchstr == len(meta0) and nchstr == nmatch:
                logger.info('Nothing to UPDATE')
                return
            # Moving previous catalogs to a backup
            oldfiles = glob(bdir+brick+'*fits*')
            oldfiles += glob(bdir+brick+'allframe.opt')
            oldfiles += glob(bdir+brick+'default.*')
            if len(oldfiles)>0:
                bakdir = bdir+'bak'+logtime
                logger.info('Backing up old files to '+bakdir)
                os.makedirs(bakdir)
                for f in oldfiles:
                    shutil.move(f,bakdir)
            else:
                logger.info('No old files to backup')
            logger.info('There are new exposures to include.  UPDATING')

    # Create temporary local directory to perform the work/processing
    #  copy everything to it
    if os.path.exists(workdir)==False:
        os.makedirs(workdir)

    procdir = tempfile.mkdtemp(prefix="brk",dir=workdir)
    procdir += '/'
    os.chmod(procdir,0o755)
    logger.info('Working in temporary directory '+procdir)
    logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))

    # Copy the files
    #----------------
    logger.info('Copying over the necessary files to '+procdir)
    # Chip loop
    for i in range(nchstr):
        logger.info('Copying over files for '+chstr['file'][i])
        # New filename
        #chstr[i].newbase = chstr[i].base
        odir = os.path.dirname(chstr['file'][i])+'/'
        obase = photred.utils.fitsext(chstr['file'][i],basename=True)
        # Need symlinks to .psf, .als
        if os.path.exists(odir+obase+'.psf') == False:
            logger.info(odir+obase+'.psf NOT FOUND.  Skipping this chip')
            continue
        if os.path.exists(odir+obase+'.als') == False:
            logger.info(odir+obase+'.als NOT FOUND.  Skipping this chip')
            continue
        for e in ['.psf','.als','.ap','.opt','.als.opt','.log']:
            dln.remove(procdir+str(chstr['base'][i])+e,allow=True)
        for e in ['.psf','.als','.ap','.opt','.als.opt','.log']:
            dln.remove(procdir+'/'+obase+e,allow=True)
            shutil.copyfile(odir+obase+e,procdir+obase+e)
            os.chmod(procdir+chstr['base'][i]+e,0o755)  # make sure they are writable
        # Copy the fits, fits resource file and header files locally
        if os.path.exists(odir+obase+'.fits') == False:
            logger.info(odir+obase+'.fits NOT FOUND.  Skipping this chip')
            continue
        dln.remove(procdir+obase+'.fits',allow=True)
        shutil.copyfile(odir+obase+'.fits',procdir+obase+'.fits')
        if os.path.exists(odir+'.'+obase+'.fits'):
            dln.remove(procdir+'.'+obase+'.fits',allow=True)
            shutil.copyfile(odir+'.'+obase+'.fits',procdir+'.'+obase+'.fits')
        if os.path.exists(odir+obase+'.fits.head'):
            dln.remove(procdir+obase+'.fits.head',allow=True)
            shutil.copyfile(odir+obase+'.fits.head',procdir+obase+'.fits.head')
        os.chmod(procdir+chstr['base'][i]+'.fits',0o755)

    # Step 2: Run DAOMATCH_TILE.PRO on the files
    #--------------------------------------------
    logger.info('Step 2: Matching up objects with DAOMATCH_TILE')
    logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))
    os.chdir(procdir)
    group = {'x0':0,'y0':0}
    daomatch.daomatch_tile(np.char.array(chstr['base'])+'.als',tile,group)

    # Step 3: Run ALLFRAME
    #----------------------
    logger.info('Step 3: Run ALLFRAME')
    logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))
    # DO I NEED TO HAVE IT TRIM THE COMBINED IMAGE???
    mchbase = procdir+chstr['base'][0]
    mchfile = mchbase+'.mch'# allframe needs absolute path
    import pdb; pdb.set_trace()
    alf.allframe(mchfile,tile=tile,setupdir=bdir,scriptsdir=scriptsdir,irafdir=irafdir,
                 logfile=logfile,catformat='FITS',imager=thisimager,geocoef=0)
    magfile = chstr['base'][0]+'.mag'
    if os.path.exists(magfile) == False:
        logger.info(magfile+' NOT FOUND')
        return

    # Step 4: Calculate coordinates
    #-------------------------------
    logger.info('Step 4: Adding coordinates')
    logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))
    # Load the MCH file
    alsfiles,trans,magoff = io.readmch(mchfile)
    nalsfiles = len(alsfiles)
    # Load the photometry file
    instphot = io.readfile(magfile)
    ninstphot = len(instphot)
    logger.info('Nstars = '+str(ninstphot))
    # Converting to IDL X/Y convention, starting at (0,0)
    # DAOPHOT has X/Y start at (1,1)
    ra,dec = tile['wcs'].pixel_to_world(instphot.x-1.0,instphot.y-1.0)
    #HEAD_XYAD,tile.head,instphot.x-1.0,instphot.y-1.0,ra,dec,/degree

    # Step 5: Calibrating photometry with zero-points
    #-------------------------------------------------
    logger.info('Step 5: Calibrating photometry with zero-points')
    logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))
    cmag = np.zeros((ninstphot,nchstr),float)+99.99
    cerr = np.zeros((ninstphot,nchstr),float)+9.99
    # Chip loop
    for i in range(nchstr):
        # id, x, y, ra, dec, unsolved magnitudes, chi, sharp
        imag = instphot['mag'+str(i+1)]
        ierr = instphot['err'+str(i+1)]
        gdmag, = np.where(imag < 50)
        if len(gdmag) > 0:
            # exptime, aperture correction, zero-point
            # aperture correction is SUBTRACTIVE, makes it brighter
            # ZPTERM is a SUBTRACTIVE constant offset
            cmag[gdmag,i] = imag[gdmag] + 2.5*np.log10(chstr['exptime'][i]) - chstr['apcor'][i] - chstr['calib_zpterm'][i]
            # Add zero-point error in quadrature
            cerr[gdmag,i] = np.sqrt(ierr[gdmag]**2+chstr['calib_zptermsig'][i]**2)
    # Calculate average photometry per filter
    ufilt = np.unique(chstr['filter'])
    nufilt = len(ufilt)
    avgmag = np.zeros((ninstphot,nufilt),float)
    avgerr = np.zeros((ninstphot,nufilt),float)
    ndet = np.zeros((ninstphot,nufilt),int)
    for i in range(nufilt):
        gdf, = np.where(chstr['filter'] == ufilt[i])
        # Single exposure
        if len(gdf) == 1:
            avgmag[:,i] = cmag[:,gdf[0]]
            avgerr[:,i] = cerr[:,gdf[0]]
            ndet[:,i] = np.sum(cmag[:,gdf[0]] < 50)
            # Multiple exposures
        else:
            # Loop through all of the exposures and add up the flux, totalwt, etc.
            totalwt = np.zeros(ninstphot,float)
            totalfluxwt = np.zeros(ninstphot,float)
            for k in range(ngdf):
                gdmag, = np.where(cmag[:,gdf[k]] < 50)
                if len(gdmag) > 0:
                    totalwt[gdmag] += 1.0/cerr[gdmag,gdf[k]]**2
                    totalfluxwt[gdmag] += 2.5118864**cmag[gdmag,gdf[k]] * (1.0/cerr[gdmag,gdf[k]]**2)
                    ndet[gdmag,i] += 1
            newflux = totalfluxwt/totalwt
            newmag = 2.5*np.log10(newflux)
            newerr = np.sqrt(1.0/totalwt)
            bdmag, = np.where(~np.isfinite(newmag))
            if len(bdmag) > 0:
                newmag[bdmag] = 99.99
                newerr[bdmag] = 9.99
            avgmag[:,i] = newmag
            avgerr[:,i] = newerr
    # Measure scatter
    scatter = np.zeros((ninstphot,nufilt),float)+99.99
    for i in range(nufilt):
        gdf, = np.where(chstr['filter'] == ufilt[i])
        if len(gdf) > 1:
            totaldiff = np.zeros(ninstphot,float)
            for k in range(ngdf):
                gdmag, = np.where(cmag[:,gdf[k]] < 50)
                if len(gdmag) > 0:
                    totaldiff[gdmag] += (avgmag[gdmag,i]-cmag[gdmag,gdf[k]])**2
            scatter[:,i] = np.sqrt( totaldiff/np.maximum(ndet[:,i],1) )
            bd, = np.where(ndet[:,i] <= 1)
            if len(bd) > 0 :
                scatter[bd,i]=99.99
    # Create final catalog schema
    photdt = [('objid',(np.str,100)),('x',float),('y',float),('ra',float),('dec',float)]

    # Add columns for calibrated single-epoch photometry columns
    cmagnames = np.zeros(nchstr,(np.str,100))
    cerrnames = np.zeros(nchstr,(np.str,100))
    for i in range(nufilt):
        ind, = np.where(chstr['filter'] == ufilt[i])
        cmagnames[ind] = [ufilt[i].upper()+str(j+1)+'MAG' for j in np.arange(nind)]
        cerrnames[ind] = [ufilt[i].upper()+str(j+1)+'ERR' for j in np.arange(nind)]
    for i in range(nchstr):
        photdt += [(cmagnames[i],np.float32),(cerrnames[i],np.float32)]
    # Add columns for average photometry per filter
    for i in range(nufilt):
        photdt += [(ufilt[i]+'MAG',np.float32),(ufilt[i]+'ERR',np.float32),(ufilt[i]+'SCATTER',np.float32),('NDET'+ufilt[i],int)]
    # Extra columns
    photdt += [('chi',np.float32),('sharp',np.float32),('prob',np.float32),('ebv',np.float32)]
    # other SE columns
    photdt += [('mag_auto',np.float32),('magerr_auto',np.float32),('asemi',np.float32),('bsemi',np.float32),
               ('theta',np.float32),('ellipticity',np.float32),('fwhm',np.float32)]
    # in unique brick area
    dt += [('brickuniq',boolean)]
    # Create final catalog
    phot = np.zeros(ninstphot,dtype=np.dtype(photdt))
    for n in instphot:
        phot[n] = instphot[n]
    phtags = phot.colnames
    # object IDs
    phot['objid'] = brick+'.'+np.char.array(instphot['id'])
    # Stuff in the coordinates calculated above
    phot['ra'] = ra
    phot['dec'] = dec
    # Stuff in the calibrated single-epoch photometry columns
    for i in range(nchstr):
        phot[cmagnames[i]] = cmag[:,i]
        phot[cerrnames[i]] = cerr[:,i]
    # Stuff in the average photometry per filter
    for i in range(nufilt):
        phot[ufilt[i].upper()+'MAG'] = avgmag[:,i]
        phot[ufilt[i].upper()+'ERR'] = avgerr[:,i]
        phot[ufilt[i].upper()+'SCATTER'] = scatter[:,i]
        phot['NDET'+ufilt[i].upper()] = ndet[:,i]

    # Calculate SFD E(B-V)
    #GLACTC,phot.ra,phot.dec,2000.0,glon,glat,1,/deg
    coo = SkyCoord(phot['ra'], phot['dec'], unit='deg', frame='icrs')
    coo_gal = coo.transform_to('galactic')
    sfd = SFDQuery()
    phot['ebv'] = sfd(coords_gal)

    # THIS IS NOW BEING DONE IN DELVERED_FINALCAT.PRO THAT COMBINES ALL CATALOGS
    # Only include objects that are INSIDE the UNIQUE brick area
    #   INCLUSIVE at the lower RA and DEC limit
    # Getting objects that are in the UNIQUE brick area
    if brickstr1.dec == -90:
        # the brick right at the pole does not have any RA limits
        ginside, = np.where(phot['dec'] < brickstr1['dec2'])
    else:
        ginside, = np.where((phot['ra'] >= brickstr1['ra1']) & (phot['ra'] < brickstr1['ra2']) &
                            (phot['dec'] >= brickstr1['dec1']) & (phot['dec'] < brickstr1['dec2']))
    if len(ginside)>0:
        phot['brickuniq'][ginside] = True

    # Get some meta-data
    for i in range(nchstr):
        alffile = chstr['base'][i]+'.alf'
        if os.path.exists(alffile):
            chstr['alf_nsources'][i] = dln.numlines(alffile)-3

    # Make the exposure-level forced photometry catalog
    #--------------------------------------------------
    # Load the individual ALF files to get chi, sharp
    # Load TFR file to conver ALF IDs to final object ID
    tfrfile = mchbase+'_comb.tfr'
    alffiles,tfrstr = io.readtfr(tfrfile)
    expdt = [('id',(np.str,100)),('objid',(np.str,100)),('exposure',(np.str,100)),('ccdnum',int),('filter',(np.str,10)),
             ('mjd',float),('x',float),('y',float),('ra',float),('dec',float),('imag',float),('ierr',float),
             ('mag',float),('err',float),('sky',float),('chi',float),('sharp',float)]
    expcat = np.zeros(int(np.sum(cmag < 50))+10000,dtype=np.dtype(expdt))
    cnt = 0
    for i in range(nchstr):
        base1 = chstr['base'][i]
        fitsfile = base1+'.fits'
        if os.path.exists(alffiles[i]) and os.path.exists(fitsfile):
            alf = io.readals(alffiles[i])
            if len(alf)==0:
                continue
            # Sometimes the rows are duplicated in the ALF file
            ui = np.uniq(alf.id,np.argsort(alf.id))
            if len(ui) < nalf:
                alf = alf[ui]
                nalf = len(alf)
            head = io.readfile(fitsfile,header=True)

            # Calibrate the photometry
            # exptime, aperture correction, zero-point
            # aperture correction is SUBTRACTIVE, makes it brighter
            # ZPTERM is a SUBTRACTIVE constant offset
            cmag1 = alf['mag'] + 2.5*np.log10(chstr['exptime'][i]) - chstr['apcor'][i] - chstr['calib_zpterm'][i]
            # Add zero-point error in quadrature
            cerr1 = np.sqrt(alf['err']**2+chstr['calib_zptermsig'][i]**2)

            # Coordinates
            ra1,dec1 = head.pixel_to_world(alf['x']-1,alf['y']-1)
            #HEAD_XYAD,head,alf.x-1,alf.y-1,ra1,dec1,/deg

            # MATCH up ALF IDs to TFR INDEX
            # One row per unique object
            # the INDEX values are 1-based indices into the ALF files
            ind1,ind2 = dln.match(tfrstr.index[i],lindgen(nalf)+1)
            #MATCH,tfrstr.index[i],lindgen(nalf)+1,ind1,ind2,/sort,count=nmatch
            objid = brick+'.'+np.char.array(tfrstr['id'][ind1])

            # Create the new catalog
            newcat = np.zeros(nalf,dtype=np.dtype(dt))
            newcat['objid'] = objid
            newcat['id'] = chstr['expnum'][i]+'_'+chstr['chip'][i]+'.'+np.char.array(alf['id'])
            newcat['exposure'] = chstr['base'][i]
            newcat['ccdnum'] = chstr['chip'][i]
            newcat['filter'] = chstr['filter'][i]
            dateobs = chstr['utdate'][i]+'T'+chstr['uttime'][i]
            t = Time(dateobs)
            newcat['mjd'] = t.mjd
            newcat['x'] = alf['x']
            newcat['y'] = alf['y']
            newcat['ra'] = ra1
            newcat['dec'] = dec1
            newcat['imag'] = alf['mag']
            newcat['ierr'] = alf['err']
            newcat['mag'] = cmag1
            newcat['err'] = cerr1
            newcat['sky'] = alf['sky']
            newcat['chi'] = alf['chi']
            newcat['sharp'] = alf['sharp']

            # Add more elements
            if cnt+nalf > len(expcat):
                expcat = add_elements(expcat,100000>nalf)

            # Add to global catalog
            expcat[cnt:cnt+nalf-1] = newcat
            cnt += nalf

        else:
            logger.info(alffile+' NOT FOUND')
        expcat = expcat[0:cnt]   # this should not be needed

        # Object catalog
        #---------------
        # do not want all of the individual epoch photometry
        #  these are between the DEC and the first unique mean magnitude (e.g. GMAG, and GERR)
        lo = where(phtags=='DEC')[0][0]
        hi = where(np.char.array(phtags).upper() == ufilt[0].upper()+'MAG')[0][0]
        objdt = photdt[0:lo]
        objdt += photdt[hi:]
        obj = np.zeros(ninstphot,dtype=np.dtype(objdt))
        for n in obj.colnames:
            obj[n] = phot[h]

        # Saving final catalog
        logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))
        photfile = bdir+brick+'.fits'
        if len(phot.colnames) <= 999:
            logger.info('Writing photometry to '+photfile+'.gz')
            phot.writeto(photfile,overwrite=True)
            out = subprocess.check_output(['gzip','-f',photfile],shell=False)
        else:
            logger.info('Too many columns for FITS.  Saving as pickle file instead. '+bdir+brick+'.pkl')
            with open(bdir+brick+'.pkl','wb') as f:
                pickle.dump(phot,f)

        # Saving object photometry catalog
        objfile = bdir+brick+'_object.fits'
        obj.writeto(objfile,overwrite=True)
        out = subprocess.check_output(['gzip','-f',objfile],shell=False)

        # Saving exposure-level forced photometry catalog
        expfile = bdir+brick+'_expforced.fits'
        expcat.writeto(expfile,overwrite=True)
        out = subprocess.check_output(['gzip','-f',expfile],shell=False)

        # Save metadata
        metafile = bdir+brick+'_meta.fits'
        logger.info('Writing meta-data to '+metafile)
        chstr.writeto(metafile,overwrite=True)

        # Delete files we don't want to keep
        #-----------------------------------
        # Individual fits files
        alsbase = os.path.basename(alsfiles,'.als')
        # if fits files have resources files then replace the fits file by a
        #  dummy file
        for i in range(len(alsbase)):
            exists = os.path.exists(alsbase[i]+'.fits')
            size = os.path.getsize(alsbase[i]+'.fits')
            if exists and size > 1:
                os.remove(alsbase[i]+'.fits')
                if os.path.exists('.'+alsbase[i]+'.fits'):
                    dln.writelines(alsbase[i]+'.fits','')
        # Combined files
        #   _comb  lst, lst1, lst2, lst1.chi, grp, nst, lst2.chi, plst.chi, psfini.ap
        #   nei, als.inp, a.fits, cmn.log, cmn.coo, cmn.ap, cmn.lst,
        #   _sub.fits, _sub.cat, _sub.als, _all.coo, makemag
        base = os.path.basename(mchfile,'.mch')
        ext = ['.lst','.lst1','.lst2','.lst1.chi','.lst2.chi','.grp','.nst','.plst.chi','.nei',
               '.als.inp','.cmn.log','.cmn.coo','.cmn.ap','.cmn.lst','a.fits','a.fits.fz',
               '_sub.fits','_sub.cat','_sub.als','_all.coo','.makemag']
        for e in ext:
            if os.path.exists(base+e): os.remove(base+e)
        if os.path.exists('check.fits'): os.remove('check.fits')

        # fpack _comb.fits and _combs.fits
        for f in [base+'_comb.fits.fz',base+'_combs.fits.fz']:
            if os.path.exists(f): os.remove(f)
        out = subprocess.check_output(['fpack','-D','-Y',base+'_comb.fits'],shell=False)
        out = subprocess.check_output(['fpack','-D','-Y',base+'_combs.fits'],shell=False)
        # gzip _comb.mask.fits and _comb.bpm.fits
        out = subprocess.check_output(['gzip','-f',base+'_comb.mask.fits'],shell=False)
        out = subprocess.check_output(['gzip','-f',base+'_comb.bpm.fits'],shell=False)

        # Copy everything left back to the original directory
        logger.info('Copying files back to '+bdir)
        allfiles = glob(procdir+'*')
        allfiles += glob(procdir+'.*.fits')
        nallfiles = len(allfiles)
        for f in allfiles:
            base = os.path.basename(f)
            if os.path.exists(bdir+'/'+base): os.remove(bdir+'/'+base)
            shutil.move(f,bdir)
        os.rmdir(procdir)  # delete temporary processing directory

        logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))
        logger.info('DELVERED_FORCEBRICK:ne after '+str(time.time()-t0,2)+' sec.')

        os.chdir(curdir)  # back to original directory

        # Create JOINT catalogs
        logger.info('')
        logger.info('CREATE JOINT CATALOGS')
        logger.info('')
        logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))

        if redo or update:
            jntredo = True
        else:
            jntredo = False
        delvered_jointbrickcats(brick,logfile=logfile,redo=jntredo)


import warnings
import numpy as np
import copy

from astropy.time import Time
from astropy.utils.iers.iers import IERSRangeError

__all__ = ["ModifiedJulianDate", "MJDWarning", "UTCtoUT1Warning"]

# Filter out ERFA's complaints that we are simulating dates which
# are in the future
warnings.filterwarnings("ignore",
                        message='.*taiutc.*dubious.year.*')

class MJDWarning(Warning):
    """
    A sub-class of Warning.  All of the warnings raised by ModifiedJulianDate
    will be of this class (or its sub-classes), so that users can filter them
    out by creating a simple filter targeted at category=MJDWarning.
    """
    pass

class UTCtoUT1Warning(MJDWarning):
    """
    A sub-class of MJDWarning meant for use when astropy.Time cannot interpolate
    UT1-UTC as a function of UTC because UTC is out of bounds of the data.
    This class exists so that users can filter these warnings out by creating
    a simple filter targeted at category=UTCtoUT1Warning.
    """
    pass

class ModifiedJulianDate(object):

    @classmethod
    def _get_ut1_from_utc(cls, UTC):
        """
        Take a numpy array of UTC values and return a numpy array of UT1 and dut1 values
        """

        time_list = Time(UTC, scale='utc', format='mjd')

        try:
            dut1_out = time_list.delta_ut1_utc
            ut1_out = time_list.ut1.mjd
        except IERSRangeError:
            ut1_out = np.copy(UTC)
            dut1_out = np.zeros(len(UTC))
            warnings.warn("ModifiedJulianData.get_list() was given date values that are outside "
                          "astropy's range of interpolation for converting from UTC to UT1. "
                          "We will treat UT1=UTC for those dates, lacking a better alternative.",
                          category=UTCtoUT1Warning)
            from astropy.utils.iers import TIME_BEFORE_IERS_RANGE, TIME_BEYOND_IERS_RANGE
            dut1_test, status = time_list.get_delta_ut1_utc(return_status=True)
            good_dexes = np.where(np.logical_and(status != TIME_BEFORE_IERS_RANGE,
                                                 status != TIME_BEYOND_IERS_RANGE))

            if len(good_dexes[0]) > 0:
                time_good = Time(UTC[good_dexes], scale='utc', format='mjd')
                dut1_good = time_good.delta_ut1_utc
                ut1_good = time_good.ut1.mjd

                ut1_out[good_dexes] = ut1_good
                dut1_out[good_dexes] = dut1_good

        return ut1_out, dut1_out

    @classmethod
    def get_list(cls, TAI=None, UTC=None):
        """
        Instantiate a list of ModifiedJulianDates from a numpy array of either TAI
        or UTC values.

        @param[in] TAI (optional) a numpy array of MJD' in TAI

        @param[in] UTC (optional) a numpy array of MJDs in UTC

        @param[out] a list of ModifiedJulianDate instantiations with all of their
        properties already set (so the code does not waste time converting from TAI
        to TT, TDB, etc. when those time scales are called for).
        """

        if TAI is None and UTC is None:
            return None

        if TAI is not None and UTC is not None:
            raise RuntimeError("You should not specify both TAI and UTC in ModifiedJulianDate.get_list()")

        if TAI is not None:
            time_list = Time(TAI, scale='tai', format='mjd')
            tai_list = TAI
            utc_list = time_list.utc.mjd
        elif UTC is not None:
            time_list = Time(UTC, scale='utc', format='mjd')
            utc_list = UTC
            tai_list = time_list.tai.mjd

        tt_list = time_list.tt.mjd
        tdb_list = time_list.tdb.mjd

        ut1_list, dut1_list = cls._get_ut1_from_utc(utc_list)

        values = np.array([tai_list, utc_list, tt_list, tdb_list,
                           ut1_list, dut1_list]).transpose()

        output = []
        for vv in values:
            mjd = ModifiedJulianDate(TAI=40000.0)
            mjd._force_values(vv)
            output.append(mjd)

        return output

    def __init__(self, TAI=None, UTC=None):
        """
        Must specify either:

        @param [in] TAI = the International Atomic Time as an MJD

        or

        @param [in] UTC = Universal Coordinate Time as an MJD
        """

        if TAI is None and UTC is None:
            raise RuntimeError("You must specify either TAI or UTC to "
                               "instantiate ModifiedJulianDate")

        if TAI is not None:
            self._time = Time(TAI, scale='tai', format='mjd')
            self._tai = TAI
            self._utc = None
            self._initialized_with = 'TAI'
        else:
            self._time = Time(UTC, scale='utc', format='mjd')
            self._utc = UTC
            self._tai = None
            self._initialized_with = 'UTC'

        self._tt = None
        self._tdb = None
        self._ut1 = None
        self._dut1 = None

    def _force_values(self, values):
        """
        Force the properties of this ModifiedJulianDate to have specific values.

        values is a list of [TAI, UTC, TT, TDB, UT1, UT1-UTC] values.

        This method exists so that, when instantiating lists of ModifiedJulianDates,
        we can use astropy.time.Time's vectorized methods to quickly perform many
        conversions at once.  Users should not try to use this method by hand.
        """
        self._tai = values[0]
        self._utc = values[1]
        self._tt = values[2]
        self._tdb = values[3]
        self._ut1 = values[4]
        self._dut1 = values[5]

    def __eq__(self, other):
        return self._time == other._time

    def __ne__(self, other):
        return not self.__eq__(other)

    def __deepcopy__(self, memo):
        if self._initialized_with == 'TAI':
            new_mjd = ModifiedJulianDate(TAI=self.TAI)
        else:
            new_mjd = ModifiedJulianDate(UTC=self.UTC)

        new_mjd._tai = copy.deepcopy(self._tai, memo)
        new_mjd._utc = copy.deepcopy(self._utc, memo)
        new_mjd._tt = copy.deepcopy(self._tt, memo)
        new_mjd._tdb = copy.deepcopy(self._tdb, memo)
        new_mjd._ut1 = copy.deepcopy(self._ut1, memo)
        new_mjd._dut1 = copy.deepcopy(self._dut1, memo)

        return new_mjd

    def _warn_utc_out_of_bounds(self, method_name):
        """
        Raise a standard warning if UTC is outside of the range that can
        be interpolated on the IERS tables.

        method_name is the name of the method that caused this warning.
        """
        warnings.warn("UTC is outside of IERS table for UT1-UTC.\n"
                      "Returning UT1 = UTC for lack of a better idea\n"
                      "This warning was caused by calling ModifiedJulianDate.%s\n" % method_name,
                      category=UTCtoUT1Warning)

    @property
    def TAI(self):
        """
        International Atomic Time as an MJD
        """
        if self._tai is None:
            self._tai = self._time.tai.mjd

        return self._tai

    @property
    def UTC(self):
        """
        Universal Coordinate Time as an MJD
        """
        if self._utc is None:
            self._utc = self._time.utc.mjd

        return self._utc

    @property
    def UT1(self):
        """
        Universal Time as an MJD
        """
        if self._ut1 is None:
            try:
                self._ut1 = self._time.ut1.mjd
            except IERSRangeError:
                self._warn_utc_out_of_bounds('UT1')
                self._ut1 = self.UTC

        return self._ut1

    @property
    def dut1(self):
        """
        UT1-UTC in seconds
        """

        if self._dut1 is None:
            try:
                self._dut1 = self._time.delta_ut1_utc
            except IERSRangeError:
                self._warn_utc_out_of_bounds('dut1')
                self._dut1 = 0.0

        return self._dut1

    @property
    def TT(self):
        """
        Terrestrial Time (aka Terrestrial Dynamical Time) as an MJD
        """
        if self._tt is None:
            self._tt = self._time.tt.mjd

        return self._tt

    @property
    def TDB(self):
        """
        Barycentric Dynamical Time as an MJD
        """
        if self._tdb is None:
            self._tdb = self._time.tdb.mjd

        return self._tdb


import numpy as np
import pymc3 as pm
import theano as th
import theano.tensor as tt
import inspect
import io

import logging
import time
import tqdm
import argparse
from .. models.fancy_model import GeneralizedContinuousModel
from typing import List, Callable, Optional, Set, Tuple, Any, Dict
from abc import abstractmethod
from ..inference.convergence_tracker import NoisyELBOConvergenceTracker
from ..inference.param_tracker import VariationalParameterTrackerConfig, VariationalParameterTracker
from ..inference import fancy_optimizers
from ..inference.deterministic_annealing import ADVIDeterministicAnnealing
from .. import types

_logger = logging.getLogger(__name__)

class Sampler:
    """Base class for log emission posterior probability samplers to be used in the hybrid ADVI scheme."""
    def __init__(self, hybrid_inference_params: 'HybridInferenceParameters'):
        self.hybrid_inference_params = hybrid_inference_params

    @abstractmethod
    def update_approximation(self, approx: pm.approximations.MeanField) -> None:
        """Take a new mean-field approximation and update the sampler routine accordingly.

        Args:
            approx: an instance of PyMC3 mean-field posterior approximation

        Returns:
            None
        """
        raise NotImplementedError

    @abstractmethod
    def draw(self) -> np.ndarray:
        """Draw one sample (or average of several samples) from log emission posterior probability."""
        raise NotImplementedError

    @abstractmethod
    def reset(self):
        """Reset the internal state of the sampler (e.g. previously accumulated samples)."""
        raise NotImplementedError

    @abstractmethod
    def increment(self, update: np.ndarray):
        """Add an incremental update to the current estimate of the log emission posterior mean."""
        raise NotImplementedError

    @abstractmethod
    def get_latest_log_emission_posterior_mean_estimator(self) -> np.ndarray:
        """Returns the latest estimate of the log emission posterior mean."""
        raise NotImplementedError

class Caller:
    """Base class for callers, i.e. routines that update the posterior of discrete RVs, to be used in the
    hybrid ADVI scheme."""
    @abstractmethod
    def call(self) -> 'CallerUpdateSummary':
        """Update the posterior of discrete RVs and return a summary"""
        raise NotImplementedError

    @abstractmethod
    def update_auxiliary_vars(self) -> None:
        """Update auxiliary variables in workspaces, if any. This routine is called after one
        (or more) round(s) of invoking `Caller.call`."""
        raise NotImplemented

class CallerUpdateSummary:
    """Represents a summary of updates made to discrete RV posteriors by a `Caller`."""
    @abstractmethod
    def reduce_to_scalar(self) -> float:
        """A function that reduces arrays to scalars. It is used to summarize tensor-valued updates."""
        raise NotImplementedError

    @abstractmethod
    def __repr__(self):
        """Represents the caller update summary in a human readable format (used in logging)."""
        raise NotImplementedError

class LoggerTQDMAdapter(io.StringIO):
    """Output stream for `tqdm` which will output to logger module instead of `stderr`."""
    logger = None
    level = None
    buf = ''

    def __init__(self, logger, level=logging.INFO):
        super().__init__()
        self.logger = logger
        self.level = level

    def write(self, buf):
        self.buf = buf.strip('\r\n\t ')

    def flush(self):
        self.logger.log(self.level, self.buf)

class InferenceTask:
    """Base class of all inference tasks."""

    # Lay in a course for starbase one two warp nine point five...
    @abstractmethod
    def engage(self):
        """Initiate inference algorithm."""
        raise NotImplementedError("Core breach imminent!")

    @abstractmethod
    def disengage(self):
        """Wrap up the inference algorithm (clean up workspaces, etc.)"""
        raise NotImplementedError

class HybridInferenceTask(InferenceTask):
    """The hybrid inference framework is applicable to PGMs with the following general structure:

        +--------------+           +----------------+
        | discrete RVs + --------► + continuous RVs |
        +--------------+           +----------------+

    Note that discrete RVs do not have any continuous parents. The inference is approximately
    performed by factorizing the true posterior into an uncorrelated product of discrete RVs (DRVs)
    and continuous RVs (CRVs):

        p(CRVs, DRVs | observed) ~ q(CRVs) q(DRVs)

    q(CRVs) is updated via deterministic annealing mean-field ADVI.
    q(DRV) is updated by the provided "caller" and is out the scope of this class.

    Usage:
    ------

    Preliminaries. Let us decompose the log joint as follows:

        -log_P(CRVs, DRVs, observed) = F_c(CRVs, observed)
                                        + F_d(DRVs, observed)
                                            + F_cd(CRVs, DRVs, observed)

    The last term in the free energy (negative log joint) is the only term with cross terms between
    the discrete and continuous sectors.

    The user must supply the following components:

        (1) a pm.Model that yields the DRV-posterior-expectation of the free energy,

            F_c^{eff}(CRVs, observed) = E_{DRVs ~ q(DRVs)} [-log_P(CRVs, DRVs, observed)]
                                        = F_c(CRVs, observed)
                                            + E_{DRVs ~ q(DRVs)} [F_cd(CRVs, DRVs, observed)]
                                                + E_{DRVs ~ q(DRVs)} [F_d(DRVs, observed)]

            Note: the last term is fully determined by q(DRVs) and can be dropped while performing
            ADVI updates in the continuous sector.

        (2) a "sampler" that provides samples from the cross term, which we call "log emission",
            defined as:

            log_emission(DRVs) = E_{CRVs ~ q(CRVs)} [-F_cd (CRVs, DRV, observed)]

        (3) a "caller" that updates q(DRVs) given log_emission(DRV), i.e.:

            q(DRVs) \propto \exp[log_emission(DRVs) - F_d(DRVs, observed)]

            In practice, one does not need the complete joint posterior of DRVs: only sufficient
            statistics, to the extent required for calculating F_c^{eff} is needed. Calculating such
            sufficient statistics could be as simple as using the Bayes rule, or more complicated if
            the DRVs are strongly correlated.

    The general implementation motif is:

        (a) to store sufficient statistics from q(DRVs) as a shared theano tensor such that the the
            model can access it,
        (b) to store log_emission(DRVs) as a shared theano tensor (or ndarray) such that the caller
            can access it, and:
        (c) let the caller directly update the shared sufficient statistics.
    """

    # if the temperature is within the following tolerance of 1.0, it is assumed that annealing
    # has finished
    temperature_tolerance = 1e-6

    def __init__(self,
                 hybrid_inference_params: 'HybridInferenceParameters',
                 continuous_model: GeneralizedContinuousModel,
                 sampler: Optional[Sampler],
                 caller: Optional[Caller],
                 **kwargs):
        """Initializer.

        Args:
            hybrid_inference_params: inference configuration
            continuous_model: a PyMC3 model representing the continuous sector of the PGM
            sampler: log emission probability sampler
            caller: discrete RV posterior updater
            **kwargs: extra keywords

        Keyword Args:
            custom_optimizer: a custom stochastic optimizer to be used in place of the default optimizer (adamax
            elbo_normalization_factor: normalization factor of the full model ELBO (for logging)
            advi_task_name: name of the ADVI step (for logging)
            sampling_task_name: name of the sampling step (for logging)
            calling_task_name: name of the calling step (for logging)
        """
        assert hybrid_inference_params is not None
        self.hybrid_inference_params = hybrid_inference_params

        assert continuous_model is not None
        self.continuous_model = continuous_model
        self.continuous_model.verify_var_registry()

        if sampler is None:
            _logger.warning("No log emission sampler given; skipping the sampling step")
        self.sampler = sampler

        if caller is None:
            _logger.warning("No caller given; skipping the calling step")
        self.caller = caller

        if self.hybrid_inference_params.track_model_params:
            _logger.info("Instantiating the parameter tracker...")
            self.param_tracker = self._create_param_tracker()
        else:
            self.param_tracker = None

        _logger.info("Instantiating the convergence tracker...")
        self.advi_convergence_tracker = NoisyELBOConvergenceTracker(
            self.hybrid_inference_params.convergence_snr_averaging_window,
            self.hybrid_inference_params.convergence_snr_trigger_threshold,
            self.hybrid_inference_params.convergence_snr_countdown_window)

        _logger.info("Setting up DA-ADVI...")
        with self.continuous_model:
            if not hasattr(self, 'temperature'):
                initial_temperature = self.hybrid_inference_params.initial_temperature
                self.temperature: types.TensorSharedVariable = th.shared(
                    np.asarray([initial_temperature], dtype=types.floatX))
            initial_temperature = self.temperature.get_value()[0]
            if (np.abs(initial_temperature - 1.0) < self.temperature_tolerance or
                    hybrid_inference_params.disable_annealing):
                # no annealing
                temperature_update = None
            else:
                # linear annealing
                max_thermal_advi_iterations = self.hybrid_inference_params.max_advi_iter_first_epoch
                max_thermal_advi_iterations += ((self.hybrid_inference_params.num_thermal_epochs - 1)
                                                * self.hybrid_inference_params.max_advi_iter_first_epoch)
                temperature_drop_per_iter = (initial_temperature - 1.0) / max_thermal_advi_iterations
                temperature_update = [(self.temperature,
                                       tt.maximum(1.0, self.temperature - temperature_drop_per_iter))]

            self.continuous_model_advi = ADVIDeterministicAnnealing(
                random_seed=self.hybrid_inference_params.random_seed,
                temperature=self.temperature)
            self.continuous_model_approx: pm.MeanField = self.continuous_model_advi.approx

            if 'custom_optimizer' in kwargs.keys():
                opt = kwargs['custom_optimizer']
                assert issubclass(type(opt), fancy_optimizers.FancyStochasticOptimizer)
                self.fancy_opt = opt
            else:
                self.fancy_opt = fancy_optimizers.FancyAdamax(
                    learning_rate=hybrid_inference_params.learning_rate,
                    beta1=hybrid_inference_params.adamax_beta1,
                    beta2=hybrid_inference_params.adamax_beta2,
                    sample_specific_only=False)

            self.continuous_model_step_func = self.continuous_model_advi.objective.step_function(
                score=True,
                obj_optimizer=self.fancy_opt.get_optimizer(self.continuous_model, self.continuous_model_approx),
                total_grad_norm_constraint=self.hybrid_inference_params.total_grad_norm_constraint,
                obj_n_mc=self.hybrid_inference_params.obj_n_mc,
                more_updates=temperature_update)

        if self.sampler is not None:
            self.sampler.update_approximation(self.continuous_model_approx)

        if 'elbo_normalization_factor' in kwargs.keys():
            self.elbo_normalization_factor = kwargs['elbo_normalization_factor']
        else:
            self.elbo_normalization_factor = 1.0

        if 'advi_task_name' in kwargs.keys():
            self.advi_task_name = kwargs['advi_task_name']
        else:
            self.advi_task_name = "ADVI"

        if 'sampling_task_name' in kwargs.keys():
            self.sampling_task_name = kwargs['sampling_task_name']
        else:
            self.sampling_task_name = "sampling"

        if 'calling_task_name' in kwargs.keys():
            self.calling_task_name = kwargs['calling_task_name']
        else:
            self.calling_task_name = "calling_task_name"

        self._t0 = None
        self._t1 = None
        self.elbo_hist: List[float] = []
        self.rls_elbo_hist: List[float] = []
        self.snr_hist: List[float] = []
        self.i_epoch = 1
        self.i_advi = 1
        self.calling_hist: List[Tuple[int, bool, bool]] = []
        self.previous_sampling_rounds = 0
        self.latest_caller_update_summary: Optional[CallerUpdateSummary] = None
        self.tqdm_out = LoggerTQDMAdapter(_logger)

    @abstractmethod
    def disengage(self):
        raise NotImplementedError

    def engage(self):
        try:
            all_converged = False
            while self.i_epoch <= self.hybrid_inference_params.max_training_epochs:
                _logger.debug("Starting epoch {0}...".format(self.i_epoch))
                converged_continuous = self._update_continuous_posteriors()
                all_converged = converged_continuous
                if self.sampler is not None:
                    converged_sampling = self._update_log_emission_posterior_expectation()
                    all_converged = all_converged and converged_sampling
                if self.caller is not None:
                    converged_discrete = self._update_discrete_posteriors()
                    all_converged = all_converged and converged_discrete
                self.i_epoch += 1
                if all_converged and not self._premature_convergence():
                    break
            if all_converged:
                _logger.info("Inference task completed successfully and convergence achieved.")
            else:
                _logger.warning("Inference task completed successfully but convergence not achieved.")
        except KeyboardInterrupt:
            pass

    def _log_start(self, task_name: str, i_epoch: int):
        self._t0 = time.time()
        _logger.debug("Starting {0} for epoch {1}...".format(task_name, i_epoch))

    def _log_stop(self, task_name: str, i_epoch: int):
        self._t1 = time.time()
        _logger.debug('The {0} for epoch {1} successfully finished in {2:.2f}s'.format(
            task_name, i_epoch, self._t1 - self._t0))

    def _log_interrupt(self, task_name: str, i_epoch: int):
        _logger.warning('The {0} for epoch {1} was interrupted'.format(task_name, i_epoch))

    def _premature_convergence(self):
        too_few_epochs = self.i_epoch < self.hybrid_inference_params.min_training_epochs
        temperature_still_high = np.abs(self.temperature.get_value()[0] - 1) > self.temperature_tolerance
        still_in_annealing = temperature_still_high and not self.hybrid_inference_params.disable_annealing
        return too_few_epochs or still_in_annealing

    def _create_param_tracker(self):
        assert all([param_name in self.continuous_model.vars or
                    param_name in self.continuous_model.deterministics
                    for param_name in self.hybrid_inference_params.param_tracker_config.var_names]),\
            "Some of the parameters chosen to be tracker are not present in the model"
        return VariationalParameterTracker(self.hybrid_inference_params.param_tracker_config)

    def _update_continuous_posteriors(self) -> bool:
        self._log_start(self.advi_task_name, self.i_epoch)
        max_advi_iters = self.hybrid_inference_params.max_advi_iter_subsequent_epochs if self.i_epoch > 1 \
            else self.hybrid_inference_params.max_advi_iter_first_epoch
        converged = False
        with tqdm.trange(max_advi_iters,
                         desc="({0}) starting...".format(self.advi_task_name),
                         file=self.tqdm_out) as progress_bar:
            try:
                for _ in progress_bar:
                    loss = self.continuous_model_step_func() / self.elbo_normalization_factor
                    self.i_advi += 1
                    try:
                        self.advi_convergence_tracker(self.continuous_model_advi.approx, loss, self.i_advi)
                    except StopIteration:
                        if not self._premature_convergence():  # suppress signal if deemed premature
                            raise StopIteration
                    snr = self.advi_convergence_tracker.snr
                    elbo_mean = self.advi_convergence_tracker.mean
                    elbo_variance = self.advi_convergence_tracker.variance
                    if snr is not None:
                        self.snr_hist.append(snr)
                    self.elbo_hist.append(-loss)
                    self.rls_elbo_hist.append(elbo_mean)
                    progress_bar.set_description("({0} epoch {1}) ELBO: {2}, SNR: {3}, T: {4:.2f}".format(
                        self.advi_task_name,
                        self.i_epoch,
                        "{0:.3f} +/- {1:.3f}".format(-elbo_mean, np.sqrt(elbo_variance))
                        if elbo_mean is not None and elbo_variance is not None else "N/A",
                        "{0:.1f}".format(snr) if snr is not None else "N/A",
                        self.temperature.get_value()[0]),
                        refresh=False)
                    if self.param_tracker is not None \
                            and self.i_advi % self.hybrid_inference_params.track_model_params_every == 0:
                        self.param_tracker(self.continuous_model_advi.approx, loss, self.i_advi)

            except StopIteration:
                converged = True
                progress_bar.refresh()
                progress_bar.close()
                self._log_stop(self.advi_task_name, self.i_epoch)

            except KeyboardInterrupt:
                progress_bar.refresh()
                progress_bar.close()
                self._log_interrupt(self.advi_task_name, self.i_epoch)
                raise KeyboardInterrupt

        return converged

    def _update_log_emission_posterior_expectation(self):
        self._log_start(self.sampling_task_name, self.i_epoch)
        if self.i_epoch == 1:
            self.sampler.reset()  # clear out log emission
        lag = min(self.previous_sampling_rounds, self.hybrid_inference_params.sampler_smoothing_window)
        converged = False
        median_rel_err = np.nan
        with tqdm.trange(self.hybrid_inference_params.log_emission_sampling_rounds,
                         desc="({0} epoch {1})".format(self.sampling_task_name, self.i_epoch),
                         file=self.tqdm_out) as progress_bar:
            try:
                for i_round in progress_bar:
                    update_to_estimator = self.sampler.draw()
                    latest_estimator = self.sampler.get_latest_log_emission_posterior_mean_estimator()
                    update_to_estimator = (update_to_estimator - latest_estimator) / (i_round + 1 + lag)
                    self.sampler.increment(update_to_estimator)
                    latest_estimator = self.sampler.get_latest_log_emission_posterior_mean_estimator()
                    median_rel_err = np.median(np.abs(update_to_estimator / latest_estimator).flatten())
                    std_rel_err = np.std(np.abs(update_to_estimator / latest_estimator).flatten())
                    del update_to_estimator
                    progress_bar.set_description("({0} epoch {1}) relative error: {2:2.4f} +/- {3:2.4f}".format(
                        self.sampling_task_name, self.i_epoch, median_rel_err, std_rel_err),
                        refresh=False)
                    if median_rel_err < self.hybrid_inference_params.log_emission_sampling_median_rel_error:
                        _logger.debug('{0} converged after {1} rounds with final '
                                      'median relative error {2:.3}.'.format(self.sampling_task_name, i_round + 1,
                                                                             median_rel_err))
                        raise StopIteration

            except StopIteration:
                converged = True
                progress_bar.refresh()
                progress_bar.close()
                self._log_stop(self.sampling_task_name, self.i_epoch)

            except KeyboardInterrupt:
                progress_bar.refresh()
                progress_bar.close()
                raise KeyboardInterrupt

            finally:
                self.previous_sampling_rounds = i_round + 1
                if not converged:
                    _logger.warning('{0} did not converge (median relative error '
                                    '= {1:.3}). Increase sampling rounds (current: {2}) if this behavior persists.'
                                    .format(self.sampling_task_name, median_rel_err,
                                            self.hybrid_inference_params.log_emission_sampling_rounds))

        return converged

    def _update_discrete_posteriors(self):
        self._log_start(self.calling_task_name, self.i_epoch)
        first_call_converged = False  # if convergence is achieved on the first call (stronger)
        iters_converged = False  # if internal loop is converged (weaker, does not imply global convergence)
        with tqdm.trange(self.hybrid_inference_params.max_calling_iters,
                         desc="({0} epoch {1})".format(self.calling_task_name, self.i_epoch),
                         file=self.tqdm_out) as progress_bar:
            try:
                for i_calling_iter in progress_bar:
                    caller_summary = self.caller.call()
                    self.latest_caller_update_summary = caller_summary
                    progress_bar.set_description("({0} epoch {1}) {2}".format(
                        self.calling_task_name, self.i_epoch, repr(caller_summary)), refresh=False)
                    caller_update_size_scalar = caller_summary.reduce_to_scalar()
                    if caller_update_size_scalar < self.hybrid_inference_params.caller_update_convergence_threshold:
                        iters_converged = True
                        if i_calling_iter == 0:
                            first_call_converged = True
                        raise StopIteration

            except StopIteration:
                progress_bar.refresh()
                progress_bar.close()
                self._log_stop(self.calling_task_name, self.i_epoch)

            except KeyboardInterrupt:
                progress_bar.refresh()
                progress_bar.close()
                self._log_interrupt(self.calling_task_name, self.i_epoch)
                raise KeyboardInterrupt

            finally:
                self.caller.update_auxiliary_vars()
                self.calling_hist.append((self.i_advi, iters_converged, first_call_converged))
                # if there is a self-consistency loop and not converged ...
                if not iters_converged and self.hybrid_inference_params.max_calling_iters > 1:
                    _logger.warning('{0} did not converge. Increase maximum calling rounds (current: {1}) '
                                    'if this behavior persists.'.format(
                        self.calling_task_name, self.hybrid_inference_params.max_calling_iters))

        return first_call_converged

class HybridInferenceParameters:
    """Hybrid ADVI inference parameters."""
    def __init__(self,
                 learning_rate: float = 0.05,
                 adamax_beta1: float = 0.9,
                 adamax_beta2: float = 0.99,
                 obj_n_mc: int = 1,
                 random_seed: int = 1984,
                 total_grad_norm_constraint: Optional[float] = None,
                 log_emission_samples_per_round: int = 50,
                 log_emission_sampling_median_rel_error: float = 5e-3,
                 log_emission_sampling_rounds: int = 10,
                 max_advi_iter_first_epoch: int = 100,
                 max_advi_iter_subsequent_epochs: int = 100,
                 min_training_epochs: int = 5,
                 max_training_epochs: int = 50,
                 initial_temperature: float = 2.0,
                 num_thermal_epochs: int = 20,
                 track_model_params: bool = False,
                 track_model_params_every: int = 10,
                 param_tracker_config: Optional['VariationalParameterTrackerConfig'] = None,
                 convergence_snr_averaging_window: int = 500,
                 convergence_snr_trigger_threshold: float = 0.1,
                 convergence_snr_countdown_window: int = 10,
                 max_calling_iters: int = 10,
                 caller_update_convergence_threshold: float = 1e-3,
                 caller_admixing_rate: float = 0.75,
                 sampler_smoothing_window: int = 0,
                 caller_summary_statistics_reducer: Callable[[np.ndarray], float] = np.mean,
                 disable_sampler: bool = False,
                 disable_caller: bool = False,
                 disable_annealing: bool = False):
        """See `expose_args` for the description of arguments."""
        self.learning_rate = learning_rate
        self.adamax_beta1 = adamax_beta1
        self.adamax_beta2 = adamax_beta2
        self.obj_n_mc = obj_n_mc
        self.random_seed = random_seed
        self.total_grad_norm_constraint = total_grad_norm_constraint
        self.log_emission_samples_per_round = log_emission_samples_per_round
        self.log_emission_sampling_median_rel_error = log_emission_sampling_median_rel_error
        self.log_emission_sampling_rounds = log_emission_sampling_rounds
        self.max_advi_iter_first_epoch = max_advi_iter_first_epoch
        self.max_advi_iter_subsequent_epochs = max_advi_iter_subsequent_epochs
        self.min_training_epochs = min_training_epochs
        self.max_training_epochs = max_training_epochs
        self.initial_temperature = initial_temperature
        self.num_thermal_epochs = num_thermal_epochs
        self.track_model_params = track_model_params
        self.track_model_params_every = track_model_params_every
        self.param_tracker_config = param_tracker_config
        self.convergence_snr_averaging_window = convergence_snr_averaging_window
        self.convergence_snr_trigger_threshold = convergence_snr_trigger_threshold
        self.convergence_snr_countdown_window = convergence_snr_countdown_window
        self.max_calling_iters = max_calling_iters
        self.caller_update_convergence_threshold = caller_update_convergence_threshold
        self.caller_admixing_rate = caller_admixing_rate
        self.sampler_smoothing_window = sampler_smoothing_window
        self.caller_summary_statistics_reducer = caller_summary_statistics_reducer
        self.disable_sampler = disable_sampler
        self.disable_caller = disable_caller
        self.disable_annealing = disable_annealing

        self._assert_params()

    def _assert_params(self):
        assert self.learning_rate > 0
        assert self.adamax_beta1 > 0
        assert self.adamax_beta2 > 0
        assert self.obj_n_mc > 0
        assert self.log_emission_samples_per_round > 0
        assert 0.0 < self.log_emission_sampling_median_rel_error < 1.0
        assert self.log_emission_sampling_rounds > 0
        assert self.max_advi_iter_first_epoch > 0
        assert self.max_advi_iter_subsequent_epochs > 0
        assert self.min_training_epochs > 0
        assert self.max_training_epochs > 0
        assert self.max_training_epochs >= self.min_training_epochs
        assert self.num_thermal_epochs <= self.max_training_epochs
        assert self.initial_temperature >= 1.0
        assert self.disable_annealing or self.num_thermal_epochs > 0
        assert self.track_model_params_every > 0
        assert self.convergence_snr_averaging_window > 0
        assert self.convergence_snr_trigger_threshold > 0
        assert self.max_calling_iters > 0
        assert self.caller_update_convergence_threshold > 0
        assert self.caller_admixing_rate > 0
        assert self.sampler_smoothing_window >= 0

        if self.track_model_params:
            assert self.param_tracker_config is not None

        if self.disable_annealing and self.initial_temperature > 1.0:
            _logger.warning("The initial temperature ({0}) is above 1.0 and annealing is disabled. This run "
                            "makes inferences in a thermal state. Ignore this warning if "
                            "this setting is intended.".format(self.initial_temperature))

    @staticmethod
    def expose_args(args: argparse.ArgumentParser, override_default: Dict[str, Any] = None, hide: Set[str] = None):
        """Exposes arguments of `__init__` to a given instance of `ArgumentParser`.

        Args:
            args: an instance of `ArgumentParser`
            override_default: a dictionary containing arguments the default values of which
                are to be overridden before passing to the argument parser
            hide: a set of arguments not to expose

        Returns:
            None
        """
        group = args.add_argument_group(title="Inference parameters")
        if override_default is None:
            override_default = dict()
        if hide is None:
            hide = set()

        initializer_params = inspect.signature(HybridInferenceParameters.__init__).parameters
        valid_args = {"--" + arg for arg in initializer_params.keys()}
        for hidden_arg in hide:
            assert hidden_arg in valid_args, \
                "Initializer argument to be hidden {0} is not a valid initializer arguments; possible " \
                "choices are: {1}".format(hidden_arg, valid_args)
        for override_default_arg in override_default.keys():
            assert override_default_arg in valid_args, \
                "Initializer argument of which the default is to be overridden {0} is not a valid initializer " \
                "arguments; possible choices are: {1}".format(override_default_arg, valid_args)

        def process_and_maybe_add(arg, **kwargs):
            full_arg = "--" + arg
            if full_arg in hide:
                return
            if full_arg in override_default:
                kwargs['default'] = override_default[full_arg]
            else:
                kwargs['default'] = initializer_params[arg].default
            group.add_argument(full_arg, **kwargs)

        def str_to_bool(value: str):
            if value.lower() in ('yes', 'true', 't', 'y', '1'):
                return True
            elif value.lower() in ('no', 'false', 'f', 'n', '0'):
                return False
            else:
                raise argparse.ArgumentTypeError('Boolean value expected.')

        process_and_maybe_add("learning_rate",
                              type=float,
                              help="Adamax optimizer learning rate")

        process_and_maybe_add("adamax_beta1",
                              type=float,
                              help="Adamax first moment estimator forgetting factor")

        process_and_maybe_add("adamax_beta2",
                              type=float,
                              help="Adamax second moment estimator forgetting factor")

        process_and_maybe_add("log_emission_samples_per_round",
                              type=int,
                              help="Number of log emission posterior samples per sampling round")

        process_and_maybe_add("log_emission_sampling_median_rel_error",
                              type=float,
                              help="Maximum tolerated median relative error in log emission posterior sampling")

        process_and_maybe_add("log_emission_sampling_rounds",
                              type=int,
                              help="Maximum log emission posterior sampling rounds")

        process_and_maybe_add("max_advi_iter_first_epoch",
                              type=int,
                              help="Maximum ADVI iterations in the first epoch (before sampling and calling)")

        process_and_maybe_add("max_advi_iter_subsequent_epochs",
                              type=int,
                              help="Maximum ADVI iterations after the first epoch")

        process_and_maybe_add("min_training_epochs",
                              type=int,
                              help="Minimum number of training epochs before evaluating convergence")

        process_and_maybe_add("max_training_epochs",
                              type=int,
                              help="Maximum number of training epochs before stopping")

        process_and_maybe_add("initial_temperature",
                              type=float,
                              help="Initial temperature for deterministic annealing (must be >= 1.0)")

        process_and_maybe_add("num_thermal_epochs",
                              type=int,
                              help="Annealing duration (in the units of training epochs)")

        process_and_maybe_add("convergence_snr_averaging_window",
                              type=int,
                              help="Averaging window for calculating training SNR for evaluating convergence")

        process_and_maybe_add("convergence_snr_trigger_threshold",
                              type=float,
                              help="The SNR threshold to be reached for triggering convergence")

        process_and_maybe_add("convergence_snr_countdown_window",
                              type=int,
                              help="The number of ADVI iterations during which the SNR is required to stay below the "
                                   "set threshold for convergence")

        process_and_maybe_add("max_calling_iters",
                              type=int,
                              help="Maximum number of calling internal self-consistency iterations")

        process_and_maybe_add("caller_update_convergence_threshold",
                              type=float,
                              help="Maximum tolerated calling update size for convergence")

        process_and_maybe_add("caller_admixing_rate",
                              type=float,
                              help="Admixing ratio of new and old caller posteriors (between 0 and 1; higher means "
                                   "using more of the new posterior)")

        process_and_maybe_add("disable_sampler",
                              type=str_to_bool,
                              help="(advanced) Disable sampler")

        process_and_maybe_add("disable_caller",
                              type=str_to_bool,
                              help="(advanced) Disable caller")

        process_and_maybe_add("disable_annealing",
                              type=str_to_bool,
                              help="(advanced) Keep the temperature pinned at its initial value and disable annealing")

    @staticmethod
    def from_args_dict(args_dict: Dict):
        """Initialize an instance of `HybridInferenceParameters` from a dictionary of arguments.

        Args:
            args_dict: a dictionary of arguments; the keys must match argument names in
                `HybridInferenceParameters.__init__`

        Returns:
            an instance of `HybridInferenceParameters`
        """
        relevant_keys = set(inspect.getfullargspec(HybridInferenceParameters.__init__).args)
        relevant_kwargs = {k: v for k, v in args_dict.items() if k in relevant_keys}
        return HybridInferenceParameters(**relevant_kwargs)