# -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow as tf
import random
import os
import time
from modelbin2 import VAE

from sklearn.metrics import accuracy_score
import numpy.matlib

tf.app.flags.DEFINE_integer("epoch_size", 40000000, "epoch size")
tf.app.flags.DEFINE_integer("batch_size", 2048, "batch size")
tf.app.flags.DEFINE_integer("Lbatch_size", 2, "Labled batch size")
tf.app.flags.DEFINE_float("gamma", 1, "gamma param for latent loss")
tf.app.flags.DEFINE_float("capacity_limit", 25.0,
                          "encoding capacity limit paramd for latent loss")
tf.app.flags.DEFINE_integer("capacity_change_duration", 100000,
                            "encoding capacity change duration")
tf.app.flags.DEFINE_float("learning_rate", 2.5e-5, "learning rate")
tf.app.flags.DEFINE_string("checkpoint_dir", "checkpoints-VAE-CFR", "checkpoint directory")
tf.app.flags.DEFINE_string("log_file", "./log", "log file directory")
tf.app.flags.DEFINE_boolean("training", True, "training or not")

flags = tf.app.flags.FLAGS

def train(sess,
          model,saveba):

  step = 0
  rl=[]
  prr=[]
  ll=[]
  lera=4e-5

  print(lera)
  minees=[]
  minee2s=[]
  MIIs=[]
  MIIsk=[]
  MIIsr=[]
  randombers2=[]
  n=16
  k=5
  r=11
  randomMSGs=np.array(np.zeros((2**r,r)))
  for i in range(2**r):
      c='{0:011b}'.format(i)
      randomMSGs[i]=[int(j) for j in list(c)]
  print(randomMSGs)
  randomMSGs=np.matlib.repmat(randomMSGs, 20000, 1)
  secretMSGs=np.array(np.zeros((2**k,k)))
  for i in range(2**k):
      c='{0:05b}'.format(i)
      secretMSGs[i]=[int(j) for j in list(c)]
  secretMSGs=np.matlib.repmat(secretMSGs, 200000, 1)

  for epoch in range(flags.epoch_size):

    print(epoch)
    total_batch = 8*10

    if (epoch > 0 and epoch%400==0  and epoch<400*2+2):
      lera=lera/2
      print('rate just got halved!')

    minee2=0
    minee=0
    mll1=0
    mll2=0
    randombers=0
    cc=[]

    MII=0
    MIIk=0
    MIIr=0

    if epoch<999999999999:
      for i in range(total_batch):
        message_bits=np.random.choice([0, 1], size=(flags.batch_size,k), p=[1./2, 1./2])

        random_bits=randomMSGs[i*flags.batch_size:(i+1)*flags.batch_size]#np.matlib.repmat(randomMSGs, 1, 1)#randomMSGs

        np.random.shuffle(random_bits)

        n1=np.random.normal(0,1,[flags.batch_size,n])
        n2=np.random.normal(0,1,[flags.batch_size,n])

        if epoch<-20:

        else:
          BobBer,EveBer,cc,ml1,ml2,CM,CR,randomber,wg  = model.partial_fit2(sess,message_bits,random_bits,lera,n1,n2)
          MI,MIk,MIr  = model.partial_fit_MINE222(sess,message_bits,random_bits,lera,n1,n2)

          mll1=mll1+ml1
          mll2=mll2+ml2
          minee=minee+BobBer
          minee2=minee2+EveBer
          randombers=randombers+randomber

        MII=MII+MI
        MIIk=MIIk+MIk
        MIIr=MIIr+MIr

    step += 1

    if epoch>0:
      randombers2.append(randombers/total_batch)
      minees.append(minee/total_batch)
      minee2s.append(minee2/total_batch)
      MIIs.append(MII/total_batch)
      MIIsk.append(MIIk/total_batch)
      MIIsr.append(MIIr/total_batch)
      #print('scorse : '+str(scorse))
      #print(wg)
      print(cc)
      print("ML1:",mll1/total_batch)
      print("ML2:",mll2/total_batch)
      print("MAIN CHANNEL BER:",minee/total_batch)
      print("Random BER:",randombers/total_batch)
      print("EVE's MI:",MII/total_batch)
      print("K's MI:",MIIk/total_batch)
      print("R's MI:",MIIr/total_batch)

    # Save checkpoint
    if epoch>0 and epoch%100 ==0:
      saveba.save(sess, flags.checkpoint_dir + '/' + 'checkpoint', global_step = step)

def load_checkpoints(sess):
   saver = tf.train.Saver()
   checkpoint = tf.train.get_checkpoint_state(flags.checkpoint_dir)
   if checkpoint and checkpoint.model_checkpoint_path:
     saver.restore(sess, checkpoint.model_checkpoint_path)
     print("loaded checkpoint: {0}".format(checkpoint.model_checkpoint_path))
   else:
     print("Could not find old checkpoint")
     if not os.path.exists(flags.checkpoint_dir):
       os.mkdir(flags.checkpoint_dir)
   return saver

def main(argv):

  sess = tf.Session()

  model = VAE(gamma=flags.gamma,
              capacity_limit=flags.capacity_limit,
              capacity_change_duration=flags.capacity_change_duration,
              learning_rate=flags.learning_rate,n3=16,k3=5,r3=11,EbN0dB=-4,EbN0dBb=10)

  sess.run(tf.global_variables_initializer())

  saver = load_checkpoints(sess)

  if flags.training:
    # Train
    train(sess, model, saver)

if __name__ == '__main__':
  tf.app.run()


import matplotlib as mpl
mpl.use('Agg')      # training mode, no screen should be open. (It will block training loop)

import argparse
import logging
import os
import time

import cv2
import numpy as np
import tensorflow as tf
from tqdm import tqdm

from pose_dataset import get_dataflow_batch, DataFlowToQueue, CocoPose
from pose_augment import set_network_input_wh, set_network_scale
from common import get_sample_images
from networks import get_network

logger = logging.getLogger('train')
logger.handlers.clear()
logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
formatter = logging.Formatter('[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s')
ch.setFormatter(formatter)
logger.addHandler(ch)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Training codes for Openpose using Tensorflow')
    parser.add_argument('--model', default='mobilenet_v2_1.4', help='model name')
    parser.add_argument('--datapath', type=str, default='/data/public/rw/coco/annotations')
    parser.add_argument('--imgpath', type=str, default='/data/public/rw/coco/')
    parser.add_argument('--batchsize', type=int, default=64)
    parser.add_argument('--gpus', type=int, default=4)
    parser.add_argument('--max-epoch', type=int, default=600)
    parser.add_argument('--lr', type=str, default='0.001')
    parser.add_argument('--tag', type=str, default='test')
    parser.add_argument('--checkpoint', type=str, default='')

    parser.add_argument('--input-width', type=int, default=432)
    parser.add_argument('--input-height', type=int, default=368)
    parser.add_argument('--quant-delay', type=int, default=-1)
    args = parser.parse_args()

    modelpath = logpath = './models/train/'

    if args.gpus <= 0:
        raise Exception('gpus <= 0')

    # define input placeholder
    set_network_input_wh(args.input_width, args.input_height)
    scale = 4

    if args.model in ['cmu', 'vgg'] or 'mobilenet' in args.model:
        scale = 8

    set_network_scale(scale)
    output_w, output_h = args.input_width // scale, args.input_height // scale

    logger.info('define model+')
    with tf.device(tf.DeviceSpec(device_type="CPU")):
        input_node = tf.placeholder(tf.float32, shape=(args.batchsize, args.input_height, args.input_width, 3), name='image')
        vectmap_node = tf.placeholder(tf.float32, shape=(args.batchsize, output_h, output_w, 38), name='vectmap')
        heatmap_node = tf.placeholder(tf.float32, shape=(args.batchsize, output_h, output_w, 19), name='heatmap')

        # prepare data
        df = get_dataflow_batch(args.datapath, True, args.batchsize, img_path=args.imgpath)
        enqueuer = DataFlowToQueue(df, [input_node, heatmap_node, vectmap_node], queue_size=100)
        q_inp, q_heat, q_vect = enqueuer.dequeue()

    df_valid = get_dataflow_batch(args.datapath, False, args.batchsize, img_path=args.imgpath)
    df_valid.reset_state()
    validation_cache = []

    val_image = get_sample_images(args.input_width, args.input_height)
    logger.debug('tensorboard val image: %d' % len(val_image))
    logger.debug(q_inp)
    logger.debug(q_heat)
    logger.debug(q_vect)

    # define model for multi-gpu
    q_inp_split, q_heat_split, q_vect_split = tf.split(q_inp, args.gpus), tf.split(q_heat, args.gpus), tf.split(q_vect, args.gpus)

    output_vectmap = []
    output_heatmap = []
    losses = []
    last_losses_l1 = []
    last_losses_l2 = []
    outputs = []
    for gpu_id in range(args.gpus):
        with tf.device(tf.DeviceSpec(device_type="GPU", device_index=gpu_id)):
            with tf.variable_scope(tf.get_variable_scope(), reuse=(gpu_id > 0)):
                net, pretrain_path, last_layer = get_network(args.model, q_inp_split[gpu_id],sess_for_load=None, trainable=True)
                if args.checkpoint:
                    pretrain_path = args.checkpoint
                vect, heat = net.loss_last()
                output_vectmap.append(vect)
                output_heatmap.append(heat)
                outputs.append(net.get_output())

                l1s, l2s = net.loss_l1_l2()
                for idx, (l1, l2) in enumerate(zip(l1s, l2s)):
                    loss_l1 = tf.nn.l2_loss(tf.concat(l1, axis=0) - q_vect_split[gpu_id], name='loss_l1_stage%d_tower%d' % (idx, gpu_id))
                    loss_l2 = tf.nn.l2_loss(tf.concat(l2, axis=0) - q_heat_split[gpu_id], name='loss_l2_stage%d_tower%d' % (idx, gpu_id))
                    losses.append(tf.reduce_mean([loss_l1, loss_l2]))

                last_losses_l1.append(loss_l1)
                last_losses_l2.append(loss_l2)

    outputs = tf.concat(outputs, axis=0)

    with tf.device(tf.DeviceSpec(device_type="GPU")):
        # define loss
        total_loss = tf.reduce_sum(losses) / args.batchsize
        total_loss_ll_paf = tf.reduce_sum(last_losses_l1) / args.batchsize
        total_loss_ll_heat = tf.reduce_sum(last_losses_l2) / args.batchsize
        total_loss_ll = tf.reduce_sum([total_loss_ll_paf, total_loss_ll_heat])

        # define optimizer
        step_per_epoch = 121745 // args.batchsize
        global_step = tf.Variable(0, trainable=False)
        if ',' not in args.lr:
            starter_learning_rate = float(args.lr)
            # learning_rate = tf.train.exponential_decay(starter_learning_rate, global_step,
            #                                            decay_steps=10000, decay_rate=0.33, staircase=True)
            learning_rate = tf.train.cosine_decay(starter_learning_rate, global_step, args.max_epoch * step_per_epoch, alpha=0.0)
        else:
            lrs = [float(x) for x in args.lr.split(',')]
            boundaries = [step_per_epoch * 5 * i for i, _ in range(len(lrs)) if i > 0]
            learning_rate = tf.train.piecewise_constant(global_step, boundaries, lrs)

    if args.quant_delay >= 0:
        logger.info('train using quantized mode, delay=%d' % args.quant_delay)
        g = tf.get_default_graph()
        tf.contrib.quantize.create_training_graph(input_graph=g, quant_delay=args.quant_delay)

    # optimizer = tf.train.RMSPropOptimizer(learning_rate, decay=0.0005, momentum=0.9, epsilon=1e-10)
    optimizer = tf.train.AdamOptimizer(learning_rate, epsilon=1e-8)
    # optimizer = tf.train.MomentumOptimizer(learning_rate, momentum=0.8, use_locking=True, use_nesterov=True)
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
        train_op = optimizer.minimize(total_loss, global_step, colocate_gradients_with_ops=True)
    logger.info('define model-')

    # define summary
    tf.summary.scalar("loss", total_loss)
    tf.summary.scalar("loss_lastlayer", total_loss_ll)
    tf.summary.scalar("loss_lastlayer_paf", total_loss_ll_paf)
    tf.summary.scalar("loss_lastlayer_heat", total_loss_ll_heat)
    tf.summary.scalar("queue_size", enqueuer.size())
    tf.summary.scalar("lr", learning_rate)
    merged_summary_op = tf.summary.merge_all()

    valid_loss = tf.placeholder(tf.float32, shape=[])
    valid_loss_ll = tf.placeholder(tf.float32, shape=[])
    valid_loss_ll_paf = tf.placeholder(tf.float32, shape=[])
    valid_loss_ll_heat = tf.placeholder(tf.float32, shape=[])
    sample_train = tf.placeholder(tf.float32, shape=(4, 640, 640, 3))
    sample_valid = tf.placeholder(tf.float32, shape=(12, 640, 640, 3))
    train_img = tf.summary.image('training sample', sample_train, 4)
    valid_img = tf.summary.image('validation sample', sample_valid, 12)
    valid_loss_t = tf.summary.scalar("loss_valid", valid_loss)
    valid_loss_ll_t = tf.summary.scalar("loss_valid_lastlayer", valid_loss_ll)
    merged_validate_op = tf.summary.merge([train_img, valid_img, valid_loss_t, valid_loss_ll_t])

    saver = tf.train.Saver(max_to_keep=1000)
    config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False)
    config.gpu_options.allow_growth = True
    with tf.Session(config=config) as sess:
        logger.info('model weights initialization')
        sess.run(tf.global_variables_initializer())

        if args.checkpoint and os.path.isdir(args.checkpoint):
            logger.info('Restore from checkpoint...')
            # loader = tf.train.Saver(net.restorable_variables())
            # loader.restore(sess, tf.train.latest_checkpoint(args.checkpoint))
            saver.restore(sess, tf.train.latest_checkpoint(args.checkpoint))
            logger.info('Restore from checkpoint...Done')
        elif pretrain_path:
            logger.info('Restore pretrained weights... %s' % pretrain_path)
            if '.npy' in pretrain_path:
                net.load(pretrain_path, sess, False)
            else:
                try:
                    loader = tf.train.Saver(net.restorable_variables(only_backbone=False))
                    loader.restore(sess, pretrain_path)
                except:
                    logger.info('Restore only weights in backbone layers.')
                    loader = tf.train.Saver(net.restorable_variables())
                    loader.restore(sess, pretrain_path)
            logger.info('Restore pretrained weights...Done')

        logger.info('prepare file writer')
        file_writer = tf.summary.FileWriter(os.path.join(logpath, args.tag), sess.graph)

        logger.info('prepare coordinator')
        coord = tf.train.Coordinator()
        enqueuer.set_coordinator(coord)
        enqueuer.start()

        logger.info('Training Started.')
        time_started = time.time()
        last_gs_num = last_gs_num2 = 0
        initial_gs_num = sess.run(global_step)

        last_log_epoch1 = last_log_epoch2 = -1
        while True:
            _, gs_num = sess.run([train_op, global_step])
            curr_epoch = float(gs_num) / step_per_epoch

            if gs_num > step_per_epoch * args.max_epoch:
                break

            if gs_num - last_gs_num >= 500:
                train_loss, train_loss_ll, train_loss_ll_paf, train_loss_ll_heat, lr_val, summary = sess.run([total_loss, total_loss_ll, total_loss_ll_paf, total_loss_ll_heat, learning_rate, merged_summary_op])

                # log of training loss / accuracy
                batch_per_sec = (gs_num - initial_gs_num) / (time.time() - time_started)
                logger.info('epoch=%.2f step=%d, %0.4f examples/sec lr=%f, loss=%g, loss_ll=%g, loss_ll_paf=%g, loss_ll_heat=%g' % (gs_num / step_per_epoch, gs_num, batch_per_sec * args.batchsize, lr_val, train_loss, train_loss_ll, train_loss_ll_paf, train_loss_ll_heat))
                last_gs_num = gs_num

                if last_log_epoch1 < curr_epoch:
                    file_writer.add_summary(summary, curr_epoch)
                    last_log_epoch1 = curr_epoch

            if gs_num - last_gs_num2 >= 2000:
                # save weights
                saver.save(sess, os.path.join(modelpath, args.tag, 'model_latest'), global_step=global_step)

                average_loss = average_loss_ll = average_loss_ll_paf = average_loss_ll_heat = 0
                total_cnt = 0

                if len(validation_cache) == 0:
                    for images_test, heatmaps, vectmaps in tqdm(df_valid.get_data()):
                        validation_cache.append((images_test, heatmaps, vectmaps))
                    df_valid.reset_state()
                    del df_valid
                    df_valid = None

                # log of test accuracy
                for images_test, heatmaps, vectmaps in validation_cache:
                    lss, lss_ll, lss_ll_paf, lss_ll_heat, vectmap_sample, heatmap_sample = sess.run(
                        [total_loss, total_loss_ll, total_loss_ll_paf, total_loss_ll_heat, output_vectmap, output_heatmap],
                        feed_dict={q_inp: images_test, q_vect: vectmaps, q_heat: heatmaps}
                    )
                    average_loss += lss * len(images_test)
                    average_loss_ll += lss_ll * len(images_test)
                    average_loss_ll_paf += lss_ll_paf * len(images_test)
                    average_loss_ll_heat += lss_ll_heat * len(images_test)
                    total_cnt += len(images_test)

                logger.info('validation(%d) %s loss=%f, loss_ll=%f, loss_ll_paf=%f, loss_ll_heat=%f' % (total_cnt, args.tag, average_loss / total_cnt, average_loss_ll / total_cnt, average_loss_ll_paf / total_cnt, average_loss_ll_heat / total_cnt))
                last_gs_num2 = gs_num

                sample_image = [enqueuer.last_dp[0][i] for i in range(4)]
                outputMat = sess.run(
                    outputs,
                    feed_dict={q_inp: np.array((sample_image + val_image) * max(1, (args.batchsize // 16)))}
                )
                pafMat, heatMat = outputMat[:, :, :, 19:], outputMat[:, :, :, :19]

                sample_results = []
                for i in range(len(sample_image)):
                    test_result = CocoPose.display_image(sample_image[i], heatMat[i], pafMat[i], as_numpy=True)
                    test_result = cv2.resize(test_result, (640, 640))
                    test_result = test_result.reshape([640, 640, 3]).astype(float)
                    sample_results.append(test_result)

                test_results = []
                for i in range(len(val_image)):
                    test_result = CocoPose.display_image(val_image[i], heatMat[len(sample_image) + i], pafMat[len(sample_image) + i], as_numpy=True)
                    test_result = cv2.resize(test_result, (640, 640))
                    test_result = test_result.reshape([640, 640, 3]).astype(float)
                    test_results.append(test_result)

                # save summary
                summary = sess.run(merged_validate_op, feed_dict={
                    valid_loss: average_loss / total_cnt,
                    valid_loss_ll: average_loss_ll / total_cnt,
                    valid_loss_ll_paf: average_loss_ll_paf / total_cnt,
                    valid_loss_ll_heat: average_loss_ll_heat / total_cnt,
                    sample_valid: test_results,
                    sample_train: sample_results
                })
                if last_log_epoch2 < curr_epoch:
                    file_writer.add_summary(summary, curr_epoch)
                    last_log_epoch2 = curr_epoch

        saver.save(sess, os.path.join(modelpath, args.tag, 'model'), global_step=global_step)
    logger.info('optimization finished. %f' % (time.time() - time_started))


# -*- coding: utf-8 -*-
"""
Spyder Editor

This is a temporary script file.
"""

# Need to put the functions in the path
# Probably not necessary if I understood Python/Git/modules better
import os
import sys
parent_directory = os.getcwd().rsplit(sep='\\',maxsplit=1)[0]
if parent_directory not in sys.path:
    sys.path.insert(1, parent_directory)

import initElements_P3
ed = initElements_P3.initElements()

import numpy as np

def In_doped_GaN():
    #                           N      Ga      In  Da
    # Define possible peaks
    pk_data =   np.array(    [  (1,     0,      0,  ed['N'].isotopes[14][0]/2),
                                (1,     0,      0,  ed['N'].isotopes[14][0]/1),
                                (1,     0,      0,  ed['N'].isotopes[15][0]/1),
                                (1,     0,      0,  ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]/3),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]/3),
                                (2,     0,      0,  ed['N'].isotopes[14][0]*2),
                                (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]),
                                (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]/2),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]/2),
                                (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[69][0])/2),
                                (3,     0,      0,  ed['N'].isotopes[14][0]*3),
                                (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[71][0])/2),
                                (3,     1,      0,  (ed['Ga'].isotopes[69][0]+3*ed['N'].isotopes[14][0])/2),
                                (3,     1,      0,  (ed['Ga'].isotopes[71][0]+3*ed['N'].isotopes[14][0])/2),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]+ed['H'].isotopes[1][0]),
                                (0,     0,      1,  ed['In'].isotopes[113][0]/2),
                                (0,     0,      1,  ed['In'].isotopes[115][0]/2),
                                (0,     0,      1,  ed['In'].isotopes[113][0]),
                                (0,     0,      1,  ed['In'].isotopes[115][0]),
                                ],
                                dtype=[('N','i4'),('Ga','i4'),('In','i4'),('m2q','f4')] )
    return pk_data

def Mg_doped_GaN():
    #                            N      Ga      Mg  Da
    # Define possible peaks
    pk_data =   np.array(    [  (1,     0,      0,  ed['N'].isotopes[14][0]/2),
                                (1,     0,      0,  ed['N'].isotopes[14][0]/1),
                                (1,     0,      0,  ed['N'].isotopes[15][0]/1),
                                (1,     0,      0,  ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]/3),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]/3),
                                (2,     0,      0,  ed['N'].isotopes[14][0]*2),
                                (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]),
                                (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]/2),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]/2),
                                (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[69][0])/2),
                                (3,     0,      0,  ed['N'].isotopes[14][0]*3),
                                (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[71][0])/2),
                                (3,     1,      0,  (ed['Ga'].isotopes[69][0]+3*ed['N'].isotopes[14][0])/2),
                                (3,     1,      0,  (ed['Ga'].isotopes[71][0]+3*ed['N'].isotopes[14][0])/2),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]+ed['H'].isotopes[1][0]),
                                (0,     0,      1,  ed['Mg'].isotopes[24][0]/2),
                                (0,     0,      1,  ed['Mg'].isotopes[24][0])
                                ],
                                dtype=[('N','i4'),('Ga','i4'),('Mg','i4'),('m2q','f4')] )
    return pk_data

def AlGaN():
    #                           N      Ga      Al  Da
    # Define possible peaks
    pk_data =   np.array(    [  (1,     0,      0,  ed['N'].isotopes[14][0]/2),
                                (1,     0,      0,  ed['N'].isotopes[14][0]/1),
                                (1,     0,      0,  ed['N'].isotopes[15][0]/1),
                                (1,     0,      0,  ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]/3),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]/3),
                                (2,     0,      0,  ed['N'].isotopes[14][0]*2),
                                (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]),
                                (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]/2),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]/2),
                                (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[69][0])/2),
                                (3,     0,      0,  ed['N'].isotopes[14][0]*3),
                                (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[71][0])/2),
                                (3,     1,      0,  (ed['Ga'].isotopes[69][0]+3*ed['N'].isotopes[14][0])/2),
                                (3,     1,      0,  (ed['Ga'].isotopes[71][0]+3*ed['N'].isotopes[14][0])/2),
                                (0,     1,      0,  ed['Ga'].isotopes[69][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]),
                                (0,     1,      0,  ed['Ga'].isotopes[71][0]+ed['H'].isotopes[1][0]),
                                (0,     0,      1,  ed['Al'].isotopes[27][0]/3),
                                (0,     0,      1,  ed['Al'].isotopes[27][0]/2),
                                (1,     0,      1,  (ed['Al'].isotopes[27][0]+ed['N'].isotopes[14][0])/2),
                                (0,     0,      1,  ed['Al'].isotopes[27][0]/1)
                                ],
                                dtype=[('N','i4'),('Ga','i4'),('Al','i4'),('m2q','f4')] )
    return pk_data

def GaN():
    #                           N      Ga  Da
    # Define possible peaks
    pk_data =   np.array(    [  (1,     0,  ed['N'].isotopes[14][0]/2),
                                (1,     0,  ed['N'].isotopes[14][0]/1),
                                (1,     0,  ed['N'].isotopes[15][0]/1),
                                (1,     0,  ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,  ed['Ga'].isotopes[69][0]/3),
                                (0,     1,  ed['Ga'].isotopes[71][0]/3),
                                (2,     0,  ed['N'].isotopes[14][0]*2),
                                (2,     0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]),
                                (2,     0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,  ed['Ga'].isotopes[69][0]/2),
                                (0,     1,  ed['Ga'].isotopes[71][0]/2),
                                (1,     1,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[69][0])/2),
                                (3,     0,  ed['N'].isotopes[14][0]*3),
                                (1,     1,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[71][0])/2),
                                (3,     1,  (ed['Ga'].isotopes[69][0]+3*ed['N'].isotopes[14][0])/2),
                                (3,     1,  (ed['Ga'].isotopes[71][0]+3*ed['N'].isotopes[14][0])/2),
                                (0,     1,  ed['Ga'].isotopes[69][0]),
                                (0,     1,  ed['Ga'].isotopes[71][0]),
                                (0,     1,  ed['Ga'].isotopes[71][0]+ed['H'].isotopes[1][0])
                                ],
                                dtype=[('N','i4'),('Ga','i4'),('m2q','f4')] )
    return pk_data

def GaN_with_H():
    #                           N      Ga  Da
    # Define possible peaks
    pk_data =   np.array(    [  (0,     0,  ed['H'].isotopes[1][0]),
                                (0,     0,  ed['H'].isotopes[1][0]*2),
                                (0,     0,  ed['H'].isotopes[1][0]*3),
                                (1,     0,  ed['N'].isotopes[14][0]/2),
                                (1,     0,  ed['N'].isotopes[14][0]/1),
                                (1,     0,  ed['N'].isotopes[15][0]/1),
                                (1,     0,  ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,  ed['Ga'].isotopes[69][0]/3),
                                (0,     1,  ed['Ga'].isotopes[71][0]/3),
                                (2,     0,  ed['N'].isotopes[14][0]*2),
                                (2,     0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]),
                                (2,     0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                                (0,     1,  ed['Ga'].isotopes[69][0]/2),
                                (0,     1,  ed['Ga'].isotopes[71][0]/2),
                                (1,     1,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[69][0])/2),
                                (3,     0,  ed['N'].isotopes[14][0]*3),
                                (1,     1,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[71][0])/2),
                                (3,     1,  (ed['Ga'].isotopes[69][0]+3*ed['N'].isotopes[14][0])/2),
                                (3,     1,  (ed['Ga'].isotopes[71][0]+3*ed['N'].isotopes[14][0])/2),
                                (0,     1,  ed['Ga'].isotopes[69][0]),
                                (0,     1,  ed['Ga'].isotopes[71][0]),
                                (0,     1,  ed['Ga'].isotopes[71][0]+ed['H'].isotopes[1][0])
                                ],
                                dtype=[('N','i4'),('Ga','i4'),('m2q','f4')] )
    return pk_data


import base64
import random
import sys
from io import BytesIO

import jax
import numpy as np
from PIL import Image
from dalle_mini.model import CustomFlaxBartForConditionalGeneration
from flask import Flask, request, jsonify
from flask_cors import CORS, cross_origin
from tqdm.notebook import tqdm
from transformers import BartTokenizer
from vqgan_jax.modeling_flax_vqgan import VQModel

app = Flask(__name__)
CORS(app)

# make sure we use compatible versions
DALLE_REPO = 'flax-community/dalle-mini'
DALLE_COMMIT_ID = '4d34126d0df8bc4a692ae933e3b902a1fa8b6114'

# set up tokenizer and model
tokenizer = BartTokenizer.from_pretrained(DALLE_REPO, revision=DALLE_COMMIT_ID)
model = CustomFlaxBartForConditionalGeneration.from_pretrained(DALLE_REPO, revision=DALLE_COMMIT_ID)

# make sure we use compatible versions
VQGAN_REPO = 'flax-community/vqgan_f16_16384'
VQGAN_COMMIT_ID = '90cc46addd2dd8f5be21586a9a23e1b95aa506a9'

# set up VQGAN
vqgan = VQModel.from_pretrained(VQGAN_REPO, revision=VQGAN_COMMIT_ID)

def decode_images(encoded_images):
    decoded_images = [vqgan.decode_code(encoded_image) for encoded_image in tqdm(encoded_images)]
    # normalize images
    clipped_images = [img.squeeze().clip(0., 1.) for img in decoded_images]
    # convert to image
    images = [Image.fromarray(np.asarray(img * 255, dtype=np.uint8)) for img in clipped_images]

    return images

def generate_images_new(prompt, num_predictions=4):
    tokenized_prompt = tokenizer(prompt, return_tensors='jax', padding='max_length', truncation=True, max_length=128)

    # create random keys
    seed = random.randint(0, 2 ** 32 - 1)
    key = jax.random.PRNGKey(seed)
    subkeys = jax.random.split(key, num=num_predictions)
    # generate sample predictions
    encoded_images = [model.generate(**tokenized_prompt, do_sample=True, num_beams=1, prng_key=subkey) for subkey in tqdm(subkeys)]

    # remove first token (BOS)
    encoded_images = [img.sequences[..., 1:] for img in encoded_images]
    return decode_images(encoded_images)

@app.route('/dalle', methods=['POST'])
@cross_origin()
def generate_images():
    json_data = request.get_json(force=True)
    text_prompt = json_data["text"]
    num_images = json_data["num_images"]
    generated_imgs = generate_images_new(text_prompt, num_images)

    generated_images = []
    for img in generated_imgs:
        buffered = BytesIO()
        img.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
        generated_images.append(img_str)

    print(f'Created {num_images} images from text prompt [{text_prompt}]')
    return jsonify(generated_images)

@app.route('/', methods=['GET'])
@cross_origin()
def health_check():
    return jsonify(success=True)

if __name__ == '__main__':
    app.run(host="0.0.0.0", port=int(sys.argv[1]), debug=False)


#!/usr/bin/env python
description = "> process lsc data  "
usage = "%prog  -e epoch [-s stage -n name -f filter -d idnumber]\n available stages [wcs,psf,psfmag,zcat,abscat,mag,local,getmag]\n"

import string
import re
import os
import sys
from numpy import take, argsort, asarray, array
from optparse import OptionParser
import datetime
import agnkey
from multiprocessing import Pool

def multi_run_makestamp(args):
    return agnkey.agnloopdef.makestamp(*args)

def multi_run_apmag(args):
    return agnkey.agnloopdef.run_apmag(*args)

def multi_run_cosmic(args):
    return agnkey.agnloopdef.run_cosmic(*args)

def multi_run_diff(args):
    return agnkey.agnloopdef.run_diff(*args)
# ####################################################################################

if __name__ == "__main__":
    parser = OptionParser(usage=usage, description=description, version="%prog 1.0")
    parser.add_option("-e", "--epoch", dest="epoch", default='20121212', type="str",
                      help='epoch to reduce  \t [%default]')
    parser.add_option("-T", "--telescope", dest="telescope", default='all', type="str",
                      help='-T telescope ' + ', '.join(agnkey.util.telescope0['all']) + ', '.join(agnkey.util.site0) + \
                           ', fts, ftn, 1m0, kb, fl \t [%default]')
    parser.add_option("-R", "--RA", dest="ra", default='', type="str",
                      help='-R  ra    \t [%default]')
    parser.add_option("-D", "--DEC", dest="dec", default='', type="str",
                      help='-D dec   \t [%default]')
    parser.add_option("-n", "--name", dest="name", default='', type="str",
                      help="-n image name   \t [%default]")
    parser.add_option("-d", "--id", dest="id", default='', type="str",
                      help="-d identification id   \t [%default]")
    parser.add_option("-f", "--filter", dest="filter", default='', type="str",
                      help="-f filter [sloan,landolt,apass,u,g,r,i,z,U,B,V,R,I] \t [%default]")
    parser.add_option("-F", "--force", dest="force", action="store_true")
    parser.add_option("-b", "--bad", dest="bad", default='', type="str",
                      help="-b bad stage [wcs,psf,psfmag,zcat,abscat,mag,goodcat,getmag," +
                           'merge,diff,template,apmag,update] \t [%default]')
    parser.add_option("-s", "--stage", dest="stage", default='', type="str",
                      help='-s stage [wcs,psf,psf2,psfmag,zcat,abscat,mag,getmag,merge,diff,mergeall' +
                           'makestamp,template,apmag,cosmic,idlstart,copy,remove] \t [%default]')
    parser.add_option("--RAS", dest="ras", default='', type="str",
                      help='-RAS  ra    \t [%default]')
    parser.add_option("--DECS", dest="decs", default='', type="str",
                      help='-DECS dec   \t [%default]')
    parser.add_option("-x", "--xord", dest="xord", default=3, type=int,
                      help='-x order for bg fit   \t [%default]')
    parser.add_option("-y", "--yord", dest="yord", default=3, type=int,
                      help='-y order for bg fit \t [%default]')
    parser.add_option("--bkg", dest="bkg", default=4, type=float,
                      help=' bkg radius for the fit  \t [%default]')
    parser.add_option("--size", dest="size", default=7, type=float,
                      help='size of the stamp for the fit \t [%default]')
    parser.add_option("-t", "--threshold", dest="threshold", default=5.,
                      type='float', help='Source detection threshold \t\t\t %default')
    parser.add_option("-i", "--interactive", action="store_true",
                      dest='interactive', default=False, help='Interactive \t\t\t [%default]')
    parser.add_option("--show", action="store_true",
                      dest='show', default=False, help='show psf fit \t\t\t [%default]')
    parser.add_option("-c", "--center", action="store_false",
                      dest='recenter', default=True, help='recenter \t\t\t [%default]')
    parser.add_option("--fix", action="store_false",
                      dest='fix', default=True, help='fix color \t\t\t [%default]')
    parser.add_option("--cutmag", dest="cutmag", default=99., type="float",
                      help='--cutmag  [magnitude instrumental cut for zeropoint ]  \t [%default]')
    parser.add_option("--field", dest="field", default='', type="str",
                      help='--field  [landolt,sloan,apass]  \t [%default]')
    parser.add_option("--ref", dest="ref", default='', type="str",
                      help='--ref  sn22_20130303_0111.sn2.fits get sn position from this file \t [%default]')
    parser.add_option("--catalog", dest="catalogue", default='', type="str",
                      help='--catalogue  sn09ip.cat    \t [%default]')
    parser.add_option("--calib", dest="calib", default='', type="str",
                      help='--calib  (sloan,natural,sloanprime)   \t [%default]')
    parser.add_option("--type", dest="type", default='fit', type="str",
                      help='--type mag for zero point   [fit,ph,mag,appmagap1,appmagap2,appmagap3,flux1]   '
                           ' \t [%default]')
    parser.add_option("--standard", dest="standard", default='', type="str",
                      help='--standard namestd  \t use the zeropoint from this standard    \t [%default]')
    parser.add_option("--xshift", dest="xshift", default=0, type="int",
                      help='x shift in the guess astrometry \t [%default]')
    parser.add_option("--fwhm", dest="fwhm", default='', type="str",
                      help='fwhm (in pixel)  \t [%default]')
    parser.add_option("--mode", dest="mode", default='sv', type="str",
                      help='mode for wcs (sv,astrometry)  \t [%default]')
    parser.add_option("--combine", dest="combine", default=1e-10, type="float",
                      help='range to combine (in days)  \t [%default]')
    parser.add_option("--datamax", dest="dmax", default = 0, type="float",
                      help='data max for saturation (counts)  \t [%default]')
    parser.add_option("--yshift", dest="yshift", default=0, type="int",
                      help='y shift in the guess astrometry \t [%default]')
    parser.add_option("--filetype", dest="filetype", default=1, type="int",
                      help='filetype  1 [single], 2 [merge], 3 differences \t [%default]')
    parser.add_option("-o", "--output", dest="output", default='', type="str",
                      help='--output  write magnitude in the output file \t [%default]')
    parser.add_option("--tempdate", dest="tempdate", default='', type="str",
                      help='--tempdate  tamplate date \t [%default]')
    parser.add_option("--normalize", dest="normalize", default='i', type="str",
                      help='--normalize image [i], template [t] \t hotpants parameter  \t [%default]')
    parser.add_option("-X", "--xwindow", action="store_true", \
                      dest='xwindow', default=False, help='xwindow \t\t\t [%default]')
    parser.add_option("--z1", dest="z1", default=None, type="int",
                      help='z1 \t [%default]')
    parser.add_option("--z2", dest="z2", default=None, type="int",
                      help='z2 \t [%default]')
    parser.add_option("--bgo", dest="bgo", default=3, type=float,
                      help=' bgo parameter for hotpants  \t [%default]')
    parser.add_option("-p", "--psf", dest="psf", default='', type=str,
                      help='psf image for template \t\t\t %default')
    parser.add_option("--mag", dest="mag", type=float, default=0,
                      help='mag to subtract from template image \t\t [%default]')
    parser.add_option("--uncleaned", dest="clean", action='store_false', default=True,
                      help='do not use cosmic ray cleaned image as template \t\t [%default]')
    parser.add_option("--subtract-mag-from-header", action='store_true',\
                      help='automatically subtract mag from header of template image \t\t [%default]')
    parser.add_option("--fixpix", dest="fixpix", action="store_true", default=False,
                      help='Run fixpix on the images before doing image subtraction')
    parser.add_option("--multicore", dest="multicore", default=4, type=int,
                      help='--multicore numbers of cores   \t [%default]')
    parser.add_option("--header", dest="header", default='', type=str,
                      help='header image \t\t\t %default')
    parser.add_option("--column", dest="column", default='', type=str,
                      help='column database \t\t\t %default')
    parser.add_option("--table", dest="table", default='dataredulco', type=str,
                      help='table database \t\t\t %default')

    option, args = parser.parse_args()
    # _instrument=option.instrument
    _telescope = option.telescope
    _normalize = option.normalize
    _type = option.type
    _mode = option.mode
    _column = option.column
    _header = option.header
    _table = option.table
    _multicore = option.multicore
    _stage = option.stage
    _bad = option.bad
    _clean = option.clean
    _subtract_mag_from_header = option.subtract_mag_from_header
    _fixpix = option.fixpix
    _mag = option.mag
    _psf = option.psf
    _bgo = option.bgo

    if _normalize not in ['i', 't']:
        sys.argv.append('--help')
    if _telescope not in agnkey.util.telescope0['all'] + agnkey.util.site0 + ['all', 'ftn', 'fts', '1m0', 'kb', 'fl','fs']:
        sys.argv.append('--help')
    if option.force == None:
        _redo = False
    else:
        _redo = True
    if option.recenter == False:
        _recenter = True
    else:
        _recenter = False
    if _type not in ['fit', 'ph', 'mag', 'appmagap1', 'appmagap2', 'appmagap3', 'flux1','flux2','flux3']:
        sys.argv.append('--help')
    if _stage:
        if _stage not in ['wcs', 'psf', 'psf2', 'psfmag', 'zcat', 'abscat', 'mag', 'local', 'getmag',
                          'merge', 'mergeall','diff', 'template', 'apmag', 'makestamp', 'cosmic', 'idlstart','update','copy','remove']:
            sys.argv.append('--help')
    if _bad:
        if _bad not in ['wcs', 'psf', 'psfmag', 'zcat', 'abscat', 'mag', 'goodcat', 'quality', 'apmag','diff']:
            sys.argv.append('--help')
    option, args = parser.parse_args()
    _id = option.id
    _filter = option.filter
    _ra = option.ra
    _dec = option.dec
    _ras = option.ras
    _output = option.output
    _decs = option.decs
    _name = option.name
    _fwhm = option.fwhm
    _xord = option.xord
    _yord = option.yord
    _bkg = option.bkg
    _size = option.size
    _standard = option.standard
    _threshold = option.threshold
    _interactive = option.interactive
    _xwindow = option.xwindow
    _show = option.show
    _fix = option.fix
    _catalogue = option.catalogue
    _calib = option.calib
    _ref = option.ref
    _field = option.field
    _cutmag = option.cutmag
    _xshift = option.xshift
    _yshift = option.yshift
    _bin = option.combine
    _dmax = option.dmax
    _filetype = option.filetype
    _tempdate = option.tempdate
    _z1 = option.z1
    _z2 = option.z2

    if _xwindow:
        from stsci.tools import capable

        capable.OF_GRAPHICS = False
        import matplotlib

        matplotlib.use('Agg')
        XX = ' -X '
    else:
        XX = ''

    if _filter:
        if _filter not in ['landolt', 'sloan', 'apass', 'u', 'g', 'r', 'i', 'z', 'U', 'B', 'V', 'R', 'I',
                           'SDSS-I', 'SDSS-G', 'SDSS-R', 'Pan-Starrs-Z', 'Bessell-B', 'Bessell-V', 'Bessell-R',
                           'Bessell-I', 'ug', 'gr', 'gri', 'griz', 'riz', 'BVR', 'BV', 'BVRI', 'VRI']:
            sys.argv.append('--help')
        else:
            try:
                _filter = agnkey.sites.filterst(_telescope)[_filter]
            except:
                pass

    if _filter and not _field:
        if _filter == 'landolt':
            _field = 'landolt'
        elif _filter == 'sloan':
            _field = 'sloan'
        elif _filter == 'apass':
            _field = 'apass'
    if _field and not _filter:
        if _field == 'landolt':
            _filter = 'landolt'
        elif _field == 'sloan':
            _filter = 'sloan'
        elif _field == 'apass':
            _filter = 'apass'

    option, args = parser.parse_args()
    epoch = option.epoch
    if '-' not in str(epoch):
        epoch0 = datetime.date(int(epoch[0:4]), int(epoch[4:6]), int(epoch[6:8]))
        listepoch = [epoch0]
    else:
        epoch1, epoch2 = string.split(epoch, '-')
        start = datetime.date(int(epoch1[0:4]), int(epoch1[4:6]), int(epoch1[6:8]))
        stop = datetime.date(int(epoch2[0:4]), int(epoch2[4:6]), int(epoch2[6:8]))
        listepoch = [re.sub('-', '', str(i)) for i in
                     [start + datetime.timedelta(days=x) for x in range(0, 1 + (stop - start).days)]]

    if not _stage or _stage in ['local', 'getmag', 'wcs', 'psf', 'psf2', 'psfmag', 'makestamp', 'mergeall','apmag',\
                                'cosmic', 'idlstart','diff','update','copy','remove']:
        if len(listepoch) == 1:
            lista = agnkey.agnsqldef.getlistfromraw(agnkey.agnsqldef.conn, _table, 'dateobs', str(listepoch[0]),
                                                    '', '*', _telescope)
        else:
            lista = agnkey.agnsqldef.getlistfromraw(agnkey.agnsqldef.conn, _table, 'dateobs', str(listepoch[0]),
                                                    str(listepoch[-1]), '*', _telescope)
        if lista:
            ll0 = {}
            for jj in lista[0].keys():
                ll0[jj] = []
            for i in range(0, len(lista)):
                for jj in lista[0].keys():
                    ll0[jj].append(lista[i][jj])
            inds = argsort(ll0['mjd'])  # sort by jd
            for i in ll0.keys():
                ll0[i] = take(ll0[i], inds)
            ll0['ra'] = ll0['ra0'][:]
            ll0['dec'] = ll0['dec0'][:]
            ll = agnkey.agnloopdef.filtralist(ll0, _filter, _id, _name, _ra, _dec, _bad, _filetype)
            print '##' * 50
            print "# IMAGE                                  OBJECT           FILTER           WCS           PSF   " + \
                  "        PSFMAG    APMAG       ZCAT          MAG      ABSCAT"
            for i in range(0, len(ll['namefile'])):
                try:
                    print '%s\t%12s\t%9s\t%9s\t%9s\t%9s\t%9s\t%9s\t%9s\t%9s' % \
                          (str(re.sub('.fits', '', ll['namefile'][i])), str(ll['objname'][i]), str(ll['filter'][i]),
                           str(ll['wcs'][i]), str(re.sub('.fits', '', ll['psf'][i])),
                           str(round(ll['psfmag'][i], 4)), str(ll['apmag'][i]), str(re.sub('.cat', '', ll['zcat'][i])),
                           str(round(ll['mag'][i], 4)), str(re.sub('.cat', '', ll['abscat'][i])))
                except:
                    print '%s\t%12s\t%9s\t%9s\t%9s\t%9s\t%9s\t%9s\t%9s\t%9s' % \
                          (str(ll['namefile'][i]), str(ll['objname'][i]), str(ll['filter'][i]), str(ll['wcs'][i]),
                           str(ll['psf'][i]), str(ll['psfmag'][i]), str(ll['apmag'][i]), str(ll['zcat'][i]),
                           str(ll['mag'][i]), str(ll['abscat'][i]))
            print '\n###  total number = ' + str(len(ll['namefile']))
            # ####################################
            if _stage == 'local':  # calibrate local sequence from .cat files
                agnkey.agnloopdef.run_local(ll['namefile'], _field, _interactive)

            elif _stage == 'getmag':  # get final magnitude from mysql
                if not _field:
                    sys.exit('use option --field landolt or sloan')
                else:
                    fields = [_field]
                for ff in fields:
                    setup = agnkey.agnloopdef.run_getmag(ll, _field, _output, _interactive, _show, _bin, _type,
                                                         _table, _ra, _dec)

            elif _stage == 'psf':
                agnkey.agnloopdef.run_psf(ll['namefile'], _threshold, _interactive, _fwhm, _show, _redo,
                                          XX, _fix, _catalogue, _table, _dmax)

            elif _stage == 'psf2':
                agnkey.agnloopdef.run_psf2(ll['namefile'], _threshold, _interactive, _fwhm, _show, _redo,
                                           XX, _fix, _catalogue, _table, _dmax)

            elif _stage == 'psfmag':
                agnkey.agnloopdef.run_fit(ll['namefile'], _ras, _decs, _xord, _yord, _bkg, _size, _recenter, _ref,
                                          _interactive, _show, _redo, _dmax)

            elif _stage == 'wcs':
                listfile = [k + v for k, v in zip(ll['wdirectory'], ll['namefile'])]
                agnkey.agnloopdef.run_wcs(listfile, _interactive, _redo, _xshift, _yshift, _catalogue,_table,_mode)

            elif _stage == 'copy':
                listfile = [k + v for k, v in zip(ll['wdirectory'], ll['namefile'])]
                for imgname in listfile:
                    if _clean:
                        if os.path.isfile(re.sub('.fits','.clean.fits',imgname)):
                            print('cp '+re.sub('.fits','.clean.fits',imgname)+' ./')
                            os.system('cp '+re.sub('.fits','.clean.fits',imgname)+' ./')
                        else:
                            print('warning: cosmic rejection not done: '+imgname)
                    else:
                            os.system('cp ' + imgname + ' ./')
                            print('cp ' + imgname + ' ./')
            elif _stage == 'remove':
                listfile = [k + v for k, v in zip(ll['wdirectory'], ll['namefile'])]
                agnkey.agnloopdef.run_remove(listfile, _filetype, _redo)
            elif _stage == 'makestamp':
                listfile = [k + v for k, v in zip(ll['wdirectory'], ll['namefile'])]
                if _show:
                    p = Pool(1)
                else:
                    p = Pool(_multicore)
                inp = [([i],_table,_z1, _z2, _interactive,_redo,_output) for i in listfile]
                p.imap_unordered(multi_run_makestamp, inp)
                p.close()
                p.join()
#                agnkey.agnloopdef.makestamp(listfile, _table, _z1, _z2, _interactive, _redo, _output)

            elif _stage == 'apmag':
                listfile = [k + v for k, v in zip(ll['wdirectory'], ll['namefile'])]
                listfile = [re.sub('.fits', '.sn2.fits',i) for i in listfile]
                if _show:
                    p = Pool(1)
                else:
                    p = Pool(_multicore)
                inp = [([i],_table,_ra, _dec, _catalogue,_show) for i in listfile]
                p.map(multi_run_apmag, inp)
                p.close()
                p.join()
#                agnkey.agnloopdef.run_apmag(ll['namefile'], _table,_ra,_dec,_catalogue,_show)

            elif _stage == 'cosmic':
                listfile = [k + v for k, v in zip(ll['wdirectory'], ll['namefile'])]

                p = Pool(_multicore)
                inp = [([i], _table,4.5, 0.2, 4,_redo) for i in listfile]
                p.map(multi_run_cosmic, inp)
                p.close()
                p.join()
                #agnkey.agnloopdef.run_cosmic(listfile, _table, 4.5, 0.2, 4, _redo)

            elif _stage == 'idlstart':
                listfile = [k + v for k, v in zip(ll['wdirectory'], ll['namefile'])]
                agnkey.agnloopdef.run_idlstart(listfile, _table, _redo)

            elif _stage == 'mergeall':
                listfile = [k + v for k, v in zip(ll['wdirectory'], ll['namefile'])]
                agnkey.agnloopdef.run_merge(array(listfile), _redo)

            elif _stage == 'update':
                listfile = [k + v for k, v in zip(ll['wdirectory'], ll['namefile'])]
                agnkey.agnloopdef.updatefromheader(listfile,_header, _column, _table)

            elif _stage == 'diff':  # difference images using hotpants
                        if not _name:
                            sys.exit('you need to select one object: use option -n/--name')
                        if _tempdate:
                            lista1 = agnkey.agnsqldef.getlistfromraw(agnkey.agnsqldef.conn, _table, 'dateobs',
                                                                     str(_tempdate), '', '*', _telescope)
                        else:
                            lista1 = agnkey.agnsqldef.getlistfromraw(agnkey.agnsqldef.conn, _table, 'dateobs',
                                                                     '20120101', '20150101', '*', _telescope)
                        if lista1:
                            ll00 = {}
                            for jj in lista1[0].keys():
                                ll00[jj] = []
                            for i in range(0, len(lista1)):
                                for jj in lista1[0].keys():
                                    ll00[jj].append(lista1[i][jj])
                            inds = argsort(ll00['mjd'])  #  sort by jd
                            for i in ll00.keys():
                                ll00[i] = take(ll00[i], inds)
                            lltemp = agnkey.agnloopdef.filtralist(ll00, _filter, '', _name, _ra, _dec, _bad, 4)
                        else:
                            sys.exit('template not found')

                        imglisttar = [k + v for k, v in zip(ll['wdirectory'], ll['namefile'])]
                        imglisttemp = [k + v for k, v in zip(lltemp['wdirectory'], lltemp['namefile'])]
##################################
#################    adding here the split by filter to be able to run multicore
##################################
                        listatar = {}
                        # divide targets by filter and targetid
                        for img in imglisttar:
                            hdr = agnkey.util.readhdr(img)
                            try:
                                _targetid = agnkey.agnsqldef.targimg(img)
                            except:
                                _targetid = 1

                            _filt = agnkey.util.readkey3(hdr, 'filter')
                            _filter = agnkey.sites.filterst1(agnkey.util.readkey3(hdr, 'telescop'))[_filt]

                            _obj = agnkey.util.readkey3(hdr, 'object')
                            if _filter not in listatar:
                                listatar[_filter] = {}
                            if _targetid not in listatar[_filter]:
                                listatar[_filter][_targetid] = []
                            listatar[_filter][_targetid].append(img)

                        # divide template by filter and targetid
                        listatemp = {}
                        for img in imglisttemp:
                            hdr = agnkey.util.readhdr(img)
                            try:
                                _targetid = agnkey.agnsqldef.targimg(img)
                            except:
                                _targetid = 1
                            _filt = agnkey.util.readkey3(hdr, 'filter')
                            _filter = agnkey.sites.filterst1(agnkey.util.readkey3(hdr, 'telescop'))[_filt]

                            _obj = agnkey.util.readkey3(hdr, 'object')
                            if _filter not in listatemp:
                                listatemp[_filter] = {}
                            if _targetid not in listatemp[_filter]:
                                listatemp[_filter][_targetid] = []
                            listatemp[_filter][_targetid].append(img)
                        inp = []
                        for f in listatar:
                            for o in listatar[f]:
                                if f in listatemp:
                                    if o in listatemp[f]:
                                        imglist1 = listatar[f][o]
                                        imglist2 = listatemp[f][o]
                                        inp = list(inp)+list([(array([i]),array(imglist2[0:1]),
                                                               _show, _redo,_normalize) for i in imglist1])

                        if len(inp):

                            if _show:
                                p = Pool(1)
                            else:
                                p = Pool(_multicore)
                            p.map(multi_run_diff, inp)
                            p.close()
                            p.join()
                        else:
                            sys.exit('no data selected ')
##################################
##################################
#                        agnkey.agnloopdef.run_diff(array(listtar), array(listtemp), _show, _redo, _normalize)
#                        if len(listtemp) == 0 or len(listtar) == 0:
#                            sys.exit('no data selected ')

        else:
            print '\n### no data selected'
    # ################################################
    else:
        for epo in listepoch:
            print '\n#### ' + str(epo)
            lista = agnkey.agnsqldef.getlistfromraw(agnkey.agnsqldef.conn, _table, 'dateobs', str(epo), '', '*',
                                                    _telescope)
            if lista:
                ll0 = {}
                for jj in lista[0].keys():
                    ll0[jj] = []
                for i in range(0, len(lista)):
                    for jj in lista[0].keys():
                        ll0[jj].append(lista[i][jj])

                inds = argsort(ll0['mjd'])  # sort by jd
                for i in ll0.keys():
                    ll0[i] = take(ll0[i], inds)
                ll0['ra'] = ll0['ra0'][:]
                ll0['dec'] = ll0['dec0'][:]
                ll = agnkey.agnloopdef.filtralist(ll0, _filter, _id, _name, _ra, _dec, _bad, _filetype)
                if len(ll['namefile']) > 0:
                    for i in range(0, len(ll['namefile'])):
                        print '%s\t%12s\t%9s\t%9s\t%9s\t%9s\t%9s\t%9s\t%9s' % \
                              (str(ll['namefile'][i]), str(ll['objname'][i]), str(ll['filter'][i]), str(ll['wcs'][i]),
                               str(ll['psf'][i]), str(ll['psfmag'][i]), str(ll['zcat'][i]), str(ll['mag'][i]),
                               str(ll['abscat'][i]))
                    print '\n###  total number = ' + str(len(ll['namefile']))
                if _stage and len(ll['namefile']) > 0:
                    print '##' * 50
                    ll3 = {}
                    for ii in ll.keys():
                        ll3[ii] = ll[ii]
                    if _stage == 'zcat':
                        if not _field:
                            if _filter in ['U', 'B', 'V', 'R', 'I', 'landolt']:
                                _field = 'landolt'
                            elif _filter in ['u', 'g', 'r', 'i', 'z', 'sloan']:
                                _field = 'sloan'
                            elif _filter in ['apass']:
                                _field = 'apass'
                            else:
                                _field = 'sloan'

                        if _field == 'apass':
                            ww0 = asarray([i for i in range(len(ll3['filter'])) if (ll['filter'][i] in ['V', 'B'])])
                            ww1 = asarray(
                                [i for i in range(len(ll3['filter'])) if (ll['filter'][i] in ['gp', 'rp', 'ip'])])
                            _color = ''
                            if len(ww0) >= 1:
                                _color = 'BV'
                                agnkey.agnloopdef.run_zero(ll3['namefile'][ww0], _fix, _type, _field, _catalogue,
                                                           _color, _interactive, _redo, _show, _cutmag, _table,
                                                           _calib)

                            if len(ww1) >= 1:
                                _color = ''
                                for jj in ['gp', 'rp', 'ip']:
                                    if jj in list(set(ll3['filter'])):
                                        _color = _color + agnkey.sites.filterst1(_telescope)[jj]
                                agnkey.agnloopdef.run_zero(ll3['namefile'][ww1], _fix, _type, _field, _catalogue,
                                                           _color, _interactive, _redo, _show, _cutmag, _table,
                                                           _calib)
                        elif _field == 'landolt':
                            ww0 = asarray([i for i in range(len(ll3['filter'])) if
                                           (ll['filter'][i] in ['U', 'I', 'R', 'V', 'B'])])
                            _color = ''
                            for jj in ['U', 'I', 'R', 'V', 'B']:
                                if jj in list(set(ll3['filter'])):
                                    _color = _color + agnkey.sites.filterst1(_telescope)[jj]
                            if len(ww0) >= 1:
                                agnkey.agnloopdef.run_zero(ll3['namefile'][ww0], _fix, _type, _field, _catalogue,
                                                           _color, _interactive, _redo, _show, _cutmag, _table,
                                                           _calib)
                        elif _field == 'sloan':
                            ww0 = asarray([i for i in range(len(ll3['filter'])) if
                                           (ll['filter'][i] in ['us', 'gp', 'rp', 'ip', 'zs'])])
                            _color = ''
                            for jj in ['gp', 'us', 'rp', 'ip', 'zs']:
                                if jj in list(set(ll3['filter'])):
                                    _color = _color + agnkey.sites.filterst1(_telescope)[jj]
                            if len(ww0) >= 1:
                                agnkey.agnloopdef.run_zero(ll3['namefile'][ww0], _fix, _type, _field, _catalogue,
                                                           _color, _interactive, _redo, _show, _cutmag, _table,
                                                           _calib)
                        else:
                            print 'warning: field not defined, zeropoint not computed'

                    elif _stage == 'abscat':  # compute magnitudes for sequence stars > img.cat
                        if _standard:
                            mm = agnkey.agnloopdef.filtralist(ll0, _filter, '', _standard, '', '', '', _filetype)
                            if len(mm['namefile']) > 0:
                                for i in range(0, len(mm['namefile'])):
                                    print '%s\t%12s\t%9s\t%9s\t%9s\t%9s\t%9s\t%9s\t%9s' % \
                                          (str(mm['namefile'][i]), str(mm['objname'][i]), str(mm['filter'][i]),
                                           str(mm['wcs'][i]), str(mm['psf'][i]),
                                           str(mm['psfmag'][i]), str(mm['zcat'][i]), str(mm['mag'][i]),
                                           str(mm['abscat'][i]))
                                agnkey.agnloopdef.run_cat(ll3['namefile'], mm['namefile'], _interactive, 1, _type, _fix,
                                                          _table, _field)
                            else:
                                print '\n### warning : standard not found for this night ' + str(epo)
                        else:
                            agnkey.agnloopdef.run_cat(ll3['namefile'], '', _interactive, 1, _type, _fix, _table,
                                                      _field)
                    elif _stage == 'mag':  # compute final magnitude using:   mag1  mag2  Z1  Z2  C1  C2
                        if _standard:
                            mm = agnkey.agnloopdef.filtralist(ll0, _filter, '', _standard, '', '', '', _filetype)
                            if len(mm['namefile']) > 0:
                                for i in range(0, len(mm['namefile'])):
                                    print '%s\t%12s\t%9s\t%9s\t%9s\t%9s\t%9s\t%9s\t%9s' % \
                                          (str(mm['namefile'][i]), str(mm['objname'][i]), str(mm['filter'][i]),
                                           str(mm['wcs'][i]), str(mm['psf'][i]),
                                           str(mm['psfmag'][i]), str(mm['zcat'][i]), str(mm['mag'][i]),
                                           str(mm['abscat'][i]))
                                agnkey.agnloopdef.run_cat(ll3['namefile'], mm['namefile'], _interactive, 2, _type,
                                                          False, _table, _field)
                            else:
                                print '\n### error: standard not found for this night' + str(epo)
                        else:
                            agnkey.agnloopdef.run_cat(ll3['namefile'], '', _interactive, 2, _type, False, _table,
                                                      _field)
                    elif _stage == 'merge':  # merge images using lacos and swarp
                        listfile = [k + v for k, v in zip(ll['wdirectory'], ll['namefile'])]
                        agnkey.agnloopdef.run_merge(array(listfile), _redo)
                    elif _stage == 'template':  # merge images using lacos and swarp
                        listfile = [k + v for k, v in zip(ll['wdirectory'], ll['namefile'])]
                        agnkey.agnloopdef.run_template(array(listfile), _show, _redo, _interactive,\
                                                       _ra, _dec, _psf, _mag, _clean, _subtract_mag_from_header)
                        #agnkey.agnloopdef.run_template(array(listfile), _show, _redo)
                    else:
                        print _stage + ' not defined'
            else:
                print '\n### no data selected'


# ActivitySim
# See full license in LICENSE.txt.
import logging

import pandas as pd
import numpy as np
from orca import orca

from activitysim.core import tracing
from activitysim.core import config
from activitysim.core import inject
from activitysim.core import pipeline
from activitysim.core import simulate
from activitysim.core import logit
from activitysim.core.util import assign_in_place, reindex

from activitysim.core import los
from activitysim.core.pathbuilder import TransitVirtualPathBuilder

from activitysim.core import los

from .util.mode import run_tour_mode_choice_simulate
from .util import trip
from .util import estimation

logger = logging.getLogger(__name__)

"""
Tour mode choice is run for all tours to determine the transportation mode that
will be used for the tour
"""

def get_alts_from_segmented_nested_logit(model_settings, segment_name, trace_label):
    """Infer alts from logit spec

    Parameters
    ----------
    model_settings : dict
    segment_column_name : str
    trace_label : str

    Returns
    -------
    list
    """

    nest_spec = config.get_logit_model_settings(model_settings)
    coefficients = simulate.get_segment_coefficients(model_settings, segment_name)
    nest_spec = simulate.eval_nest_coefficients(nest_spec, coefficients, trace_label)
    tour_mode_alts = []
    for nest in logit.each_nest(nest_spec):
        if nest.is_leaf:
            tour_mode_alts.append(nest.name)

    return tour_mode_alts

def create_logsum_trips(
        tours, segment_column_name, model_settings, trace_label):
    """
    Construct table of trips from half-tours (1 inbound, 1 outbound) for each tour-mode.

    Parameters
    ----------
    tours : pandas.DataFrame
    segment_column_name : str
        column in tours table used for segmenting model spec
    model_settings : dict
    trace_label : str

    Returns
    -------
    pandas.DataFrame
        Table of trips: 2 per tour, with O/D and purpose inherited from tour
    """
    stop_frequency_alts = inject.get_injectable('stop_frequency_alts')
    stop_freq = '0out_0in'  # no intermediate stops
    tours['stop_frequency'] = stop_freq
    tours['primary_purpose'] = tours['tour_purpose']
    trips = trip.initialize_from_tours(tours, stop_frequency_alts)
    trips['stop_frequency'] = stop_freq
    outbound = trips['outbound']
    trips['depart'] = reindex(tours.start, trips.tour_id)
    trips.loc[~outbound, 'depart'] = reindex(tours.end, trips.loc[~outbound, 'tour_id'])

    # actual segment doesn't matter. just need to grab one
    # to get a set of coefficients from the spec
    segment_name = tours.iloc[0][segment_column_name]
    tour_mode_alts = get_alts_from_segmented_nested_logit(
        model_settings, segment_name, trace_label)

    # repeat rows from the trips table iterating over tour mode
    logsum_trips = pd.DataFrame()
    for tour_mode in tour_mode_alts:
        trips['tour_mode'] = tour_mode
        logsum_trips = pd.concat((logsum_trips, trips), ignore_index=True)
    assert len(logsum_trips) == len(trips) * len(tour_mode_alts)
    logsum_trips.index.name = 'trip_id'

    return logsum_trips

def append_tour_leg_trip_mode_choice_logsums(tours):
    """Creates trip mode choice logsum column in tours table for each tour mode and leg

    Parameters
    ----------
    tours : pd.DataFrame

    Returns
    -------
    tours : pd.DataFrame
        Adds two * n_modes logsum columns to each tour row, e.g. "logsum_DRIVE_outbound"
    """
    trips = inject.get_table('trips').to_frame()
    trip_dir_mode_logsums = trips.pivot(
        index='tour_id', columns=['tour_mode', 'outbound'], values='trip_mode_choice_logsum')
    new_cols = [
        '_'.join(['logsum', mode, 'outbound' if outbound else 'inbound'])
        for mode, outbound in trip_dir_mode_logsums.columns]
    trip_dir_mode_logsums.columns = new_cols
    trip_dir_mode_logsums.reindex(tours.index)
    tours = pd.merge(tours, trip_dir_mode_logsums, left_index=True, right_index=True)

    return tours

def get_trip_mc_logsums_for_all_modes(
        tours, segment_column_name, model_settings, trace_label):
    """Creates pseudo-trips from tours and runs trip mode choice to get logsums

    Parameters
    ----------
    tours : pandas.DataFrame
    segment_column_name : str
        column in tours table used for segmenting model spec
    model_settings : dict
    trace_label : str

    Returns
    -------
    tours : pd.DataFrame
        Adds two * n_modes logsum columns to each tour row, e.g. "logsum_DRIVE_outbound"
    """

    # create pseudo-trips from tours for all tour modes
    logsum_trips = create_logsum_trips(
        tours, segment_column_name, model_settings,
        trace_label)

    # temporarily register trips in the pipeline
    pipeline.replace_table('trips', logsum_trips)
    tracing.register_traceable_table('trips', logsum_trips)
    pipeline.get_rn_generator().add_channel('trips', logsum_trips)

    # run trip mode choice on pseudo-trips. use orca instead of pipeline to
    # execute the step because pipeline can only handle one open step at a time
    orca.run(['trip_mode_choice'])

    # add trip mode choice logsums as new cols in tours
    tours = append_tour_leg_trip_mode_choice_logsums(tours)

    # de-register logsum trips table
    pipeline.get_rn_generator().drop_channel('trips')
    tracing.deregister_traceable_table('trips')

    return tours

@inject.step()
def tour_mode_choice_simulate(tours, persons_merged,
                              network_los,
                              chunk_size,
                              trace_hh_id):
    """
    Tour mode choice simulate
    """
    trace_label = 'tour_mode_choice'
    model_settings_file_name = 'tour_mode_choice.yaml'
    model_settings = config.read_model_settings(model_settings_file_name)

    logsum_column_name = model_settings.get('MODE_CHOICE_LOGSUM_COLUMN_NAME')
    mode_column_name = 'tour_mode'
    segment_column_name = 'tour_purpose'

    primary_tours = tours.to_frame()
    assert not (primary_tours.tour_category == 'atwork').any()

    logger.info("Running %s with %d tours" % (trace_label, primary_tours.shape[0]))

    tracing.print_summary('tour_types',
                          primary_tours.tour_type, value_counts=True)

    persons_merged = persons_merged.to_frame()
    primary_tours_merged = pd.merge(primary_tours, persons_merged, left_on='person_id',
                                    right_index=True, how='left', suffixes=('', '_r'))

    constants = {}
    # model_constants can appear in expressions
    constants.update(config.get_model_constants(model_settings))

    skim_dict = network_los.get_default_skim_dict()

    # setup skim keys
    orig_col_name = 'home_zone_id'
    dest_col_name = 'destination'

    out_time_col_name = 'start'
    in_time_col_name = 'end'
    odt_skim_stack_wrapper = skim_dict.wrap_3d(orig_key=orig_col_name, dest_key=dest_col_name,
                                               dim3_key='out_period')
    dot_skim_stack_wrapper = skim_dict.wrap_3d(orig_key=dest_col_name, dest_key=orig_col_name,
                                               dim3_key='in_period')
    odr_skim_stack_wrapper = skim_dict.wrap_3d(orig_key=orig_col_name, dest_key=dest_col_name,
                                               dim3_key='in_period')
    dor_skim_stack_wrapper = skim_dict.wrap_3d(orig_key=dest_col_name, dest_key=orig_col_name,
                                               dim3_key='out_period')
    od_skim_stack_wrapper = skim_dict.wrap(orig_col_name, dest_col_name)

    skims = {
        "odt_skims": odt_skim_stack_wrapper,
        "dot_skims": dot_skim_stack_wrapper,
        "odr_skims": odr_skim_stack_wrapper,  # dot return skims for e.g. TNC bridge return fare
        "dor_skims": dor_skim_stack_wrapper,  # odt return skims for e.g. TNC bridge return fare
        "od_skims": od_skim_stack_wrapper,
        'orig_col_name': orig_col_name,
        'dest_col_name': dest_col_name,
        'out_time_col_name': out_time_col_name,
        'in_time_col_name': in_time_col_name
    }

    if network_los.zone_system == los.THREE_ZONE:
        # fixme - is this a lightweight object?

        tvpb = network_los.tvpb

        tvpb_logsum_odt = tvpb.wrap_logsum(orig_key=orig_col_name, dest_key=dest_col_name,
                                           tod_key='out_period', segment_key='demographic_segment',
                                           cache_choices=True,
                                           trace_label=trace_label, tag='tvpb_logsum_odt')
        tvpb_logsum_dot = tvpb.wrap_logsum(orig_key=dest_col_name, dest_key=orig_col_name,
                                           tod_key='in_period', segment_key='demographic_segment',
                                           cache_choices=True,
                                           trace_label=trace_label, tag='tvpb_logsum_dot')

        skims.update({
            'tvpb_logsum_odt': tvpb_logsum_odt,
            'tvpb_logsum_dot': tvpb_logsum_dot
        })

        # TVPB constants can appear in expressions
        if model_settings.get('use_TVPB_constants', True):
            constants.update(network_los.setting('TVPB_SETTINGS.tour_mode_choice.CONSTANTS'))

    # don't create estimation data bundle if trip mode choice is being called
    # from another model step (i.e. tour mode choice logsum creation)
    if pipeline.get_rn_generator().step_name != 'tour_mode_choice_simulate':
        estimator = None
    else:
        estimator = estimation.manager.begin_estimation('tour_mode_choice')
    if estimator:
        estimator.write_coefficients(model_settings=model_settings)
        estimator.write_coefficients_template(model_settings=model_settings)
        estimator.write_spec(model_settings)
        estimator.write_model_settings(model_settings, model_settings_file_name)
        # (run_tour_mode_choice_simulate writes choosers post-annotation)

    # FIXME should normalize handling of tour_type and tour_purpose
    # mtctm1 school tour_type includes univ, which has different coefficients from elementary and HS
    # we should either add this column when tours created or add univ to tour_types
    not_university = (primary_tours_merged.tour_type != 'school') | \
        ~(primary_tours_merged.is_university if 'is_university' in primary_tours_merged.columns else False)
    primary_tours_merged['tour_purpose'] = \
        primary_tours_merged.tour_type.where(not_university, 'univ')

    # if trip logsums are used, run trip mode choice and append the logsums
    if model_settings.get('COMPUTE_TRIP_MODE_CHOICE_LOGSUMS', False):
        primary_tours_merged = get_trip_mc_logsums_for_all_modes(
            primary_tours_merged, segment_column_name,
            model_settings, trace_label)

    choices_list = []
    for tour_purpose, tours_segment in primary_tours_merged.groupby(segment_column_name):

        logger.info("tour_mode_choice_simulate tour_type '%s' (%s tours)" %
                    (tour_purpose, len(tours_segment.index), ))

        if network_los.zone_system == los.THREE_ZONE:
            tvpb_logsum_odt.extend_trace_label(tour_purpose)
            tvpb_logsum_dot.extend_trace_label(tour_purpose)

        # name index so tracing knows how to slice
        assert tours_segment.index.name == 'tour_id'

        choices_df = run_tour_mode_choice_simulate(
            tours_segment,
            tour_purpose, model_settings,
            mode_column_name=mode_column_name,
            logsum_column_name=logsum_column_name,
            network_los=network_los,
            skims=skims,
            constants=constants,
            estimator=estimator,
            chunk_size=chunk_size,
            trace_label=tracing.extend_trace_label(trace_label, tour_purpose),
            trace_choice_name='tour_mode_choice')

        tracing.print_summary('tour_mode_choice_simulate %s choices_df' % tour_purpose,
                              choices_df.tour_mode, value_counts=True)

        choices_list.append(choices_df)

    choices_df = pd.concat(choices_list)

    # add cached tvpb_logsum tap choices for modes specified in tvpb_mode_path_types
    if network_los.zone_system == los.THREE_ZONE:

        tvpb_mode_path_types = model_settings.get('tvpb_mode_path_types')
        if tvpb_mode_path_types is not None:
            for mode, path_types in tvpb_mode_path_types.items():

                for direction, skim in zip(['od', 'do'], [tvpb_logsum_odt, tvpb_logsum_dot]):

                    path_type = path_types[direction]
                    skim_cache = skim.cache[path_type]

                    print(f"mode {mode} direction {direction} path_type {path_type}")

                    for c in skim_cache:

                        dest_col = f'{direction}_{c}'

                        if dest_col not in choices_df:
                            choices_df[dest_col] = np.nan if pd.api.types.is_numeric_dtype(skim_cache[c]) else ''
                        choices_df[dest_col].where(choices_df.tour_mode != mode, skim_cache[c], inplace=True)

    if estimator:
        estimator.write_choices(choices_df.tour_mode)
        choices_df.tour_mode = estimator.get_survey_values(choices_df.tour_mode, 'tours', 'tour_mode')
        estimator.write_override_choices(choices_df.tour_mode)
        estimator.end_estimation()

    tracing.print_summary('tour_mode_choice_simulate all tour type choices',
                          choices_df.tour_mode, value_counts=True)

    # so we can trace with annotations
    assign_in_place(primary_tours, choices_df)

    # update tours table with mode choice (and optionally logsums)
    all_tours = tours.to_frame()
    assign_in_place(all_tours, choices_df)

    pipeline.replace_table("tours", all_tours)

    if trace_hh_id:
        tracing.trace_df(primary_tours,
                         label=tracing.extend_trace_label(trace_label, mode_column_name),
                         slicer='tour_id',
                         index_label='tour_id',
                         warn_if_empty=True)