import tensorflow as tf
import os
import re
import sys
import time
import numpy as np
from model.common import shape_list
from six.moves import range

from model.multitask_v1.common import make_phone_masks
from model.loss import softmax, asoftmax, additive_margin_softmax, additive_angular_margin_softmax
from misc.utils import substring_in_list, activation_summaries, remove_params_prefix, add_dict_prefix
from collections import OrderedDict
from dataset.data_loader import DataOutOfRange
from dataset.multitask.data_loader_v2 import KaldiDataRandomQueueV2, KaldiDataSeqQueueV2
from model.multitask_v1.tdnn import build_speaker_encoder, build_phone_encoder

loss_network = {"softmax": softmax,
                "asoftmax": asoftmax,
                "additive_margin_softmax": additive_margin_softmax,
                "additive_angular_margin_softmax": additive_angular_margin_softmax}

class BaseMT(object):
    """Handle the training, validation and prediction

    Trainer is a simple class that deals with examples having feature-label structure.
    """

    def __init__(self, params, model_dir, dim, num_speakers=None, num_phones=None, single_cpu=False):
        """
        Args:
            params: Parameters loaded from JSON.
            model_dir: The model directory.
            dim: The dimension of the feature.
            num_speakers: The total number of speakers. Used in softmax-like network. Can be None.
            num_phones: The number of phones. Can be None.
            single_cpu: Run Tensorflow on one cpu. (default = False)
        """
        self.params = params

        if single_cpu:
            self.sess_config = tf.ConfigProto(intra_op_parallelism_threads=1,
                                              inter_op_parallelism_threads=1,
                                              device_count={'CPU': 1},
                                              allow_soft_placement=True)
        else:
            self.sess_config = tf.ConfigProto(allow_soft_placement=True)
        self.sess = tf.Session(config=self.sess_config)
        self.model = os.path.join(model_dir, "nnet")

        self.train_summary = None
        self.valid_summary = None

        self.spk_embeddings = None
        self.spk_logits = None
        self.phn_embeddings = None
        self.phn_posteriors = None
        self.phn_logits_subset = None

        self.endpoints = OrderedDict()

        self.optimizer = None
        self.loss = {}
        self.total_loss = None

        self.train_op = None
        self.train_ops = {}
        self.valid_ops = {}

        self.saver = None
        self.summary_writer = None
        self.valid_summary_writer = None

        self.is_built = False
        self.is_loaded = False

        # Required to build the train and valid network
        self.dim = dim
        self.num_speakers = num_speakers
        self.num_phones = num_phones

        self.features = tf.placeholder(tf.float32, shape=[None, None, dim], name="train_features")
        # self.recon = tf.placeholder(tf.float32, shape=[None, None, None, dim], name="recon_features")
        self.spk_labels = tf.placeholder(tf.int32, shape=[None, ], name="spk_labels")
        self.phn_labels = tf.placeholder(tf.int32, shape=[None, None], name="phn_labels")
        # self.vad = tf.placeholder(tf.float32, shape=[None, None], name="vad")  # Not used at this moment
        self.feat_length = tf.placeholder(tf.int32, shape=[None, ], name="lengths")
        # phn_masks is used to choose the valid phone examples
        # and follows the format of gather_nd. [[index1], [index2], ...]
        self.phn_masks = tf.placeholder(tf.int32, shape=[None, 2], name="phn_masks")
        self.global_step = tf.placeholder(tf.int64, name="global_step")
        self.params.dict["global_step"] = self.global_step

        # TODO: We have multiple networks. How many learning rates do we need?
        self.learning_rate = tf.placeholder(tf.float32, name="learning_rate")

        # The discriminant loss network
        self.spk_loss_network = None
        self.phn_loss_network = None

        # The training statistics
        self.spk_training_count = np.zeros((num_speakers), dtype=np.int64)
        self.phn_training_count = np.zeros((num_phones), dtype=np.int64)
        self.phn_expansion_count = np.zeros((2), dtype=np.int64)

        # TODO: The context of the speaker and phone networks can be different.
        # TODO: We make a hypothesis that the context of the phone network will be larger than the speaker network.
        # TODO: If this is not true, the feature slicing should be differnet in the network building.
        # We need to expand the features to make #posteriors == #alignments
        # And convert to [b, 1, l, d]
        assert(self.params.phone_left_context > self.params.speaker_left_context and
               self.params.phone_right_context > self.params.speaker_right_context), \
            "The speaker context is expected to be smaller than the phone context (which may be not true). " \
            "If larger speaker context is used, change the feature expansion code."

        # The loaded feautre is already expanded. Nothing to do here.
        self.expand_features = tf.expand_dims(self.features, axis=1)
        return

    def reset(self):
        """Reset the graph so we can create new input pipeline or graph. (Or for other purposes)"""
        try:
            self.sess.close()
        except tf.errors.OpError:
            # Maybe the session is closed before
            pass
        tf.reset_default_graph()
        # The session should be created again after the graph is reset.
        self.sess = tf.Session(config=self.sess_config)
        # After the graph is reset, the flag should be set
        self.is_built = False
        self.is_loaded = False
        # After reset the graph, it is important to reset the seed.
        tf.set_random_seed(self.params.seed)

        # Reset some variables. The previous ones have become invalid due to the graph reset.
        self.saver = None
        self.summary_writer = None
        self.valid_summary_writer = None

    def close(self):
        """Close the session we opened."""
        try:
            self.sess.close()
        except tf.errors.OpError:
            pass

    def load(self):
        """Load the saved variables.

        If the variables have values, the current values will be changed to the saved ones
        :return The step of the saved model.
        """
        tf.logging.info("Reading checkpoints...")
        ckpt = tf.train.get_checkpoint_state(self.model)
        if ckpt and ckpt.model_checkpoint_path:
            ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
            step = int(next(re.finditer("(\d+)(?!.*\d)", ckpt_name)).group(0))
            self.saver.restore(self.sess, os.path.join(self.model, ckpt_name))
            tf.logging.info("Succeed to load checkpoint {}".format(ckpt_name))
        else:
            sys.exit("Failed to find a checkpoint in {}".format(self.model))
        self.is_loaded = True
        return step

    def save(self, step):
        """Save the model.

        Args:
            step: The global step.
        """
        self.saver.save(self.sess, os.path.join(self.model, "model"), global_step=step)

    def build(self, mode, noupdate_var_list=None):
        """ Build a network.

        Currently, I use placeholder in the graph and feed data during sess.run. So no need to parse
        features and labels.

        Args:
            mode: `train`, `valid` or `predict`.
            noupdate_var_list: In the fine-tuning, some variables are fixed. The list contains their names (or part of their names).
                               We use `noupdate` rather than `notrain` because some variables are not trainable, e.g.
                               the mean and var in the batchnorm layers.
        """
        assert(mode == "train" or mode == "valid" or mode == "predict")
        reuse_variables = True if self.is_built else None

        # Create a new path for prediction, since the training may build a tower the support multi-GPUs
        if mode == "predict":
            with tf.name_scope("predict") as scope:
                self.endpoints.clear()
                _, _, _ = build_speaker_encoder(self.expand_features,
                                                self.phn_labels,
                                                self.feat_length,
                                                self.params,
                                                self.endpoints,
                                                reuse_variables,
                                                is_training=False)
                _, mu_zp, _ = build_phone_encoder(self.expand_features,
                                                  self.spk_labels,
                                                  self.feat_length,
                                                  self.params,
                                                  self.endpoints,
                                                  reuse_variables,
                                                  is_training=False)

                # Sometimes we need the posteriors of the phones, so we have the build the softmax layer
                # If marginal angular softmax is used, correct the margin to compute the logits
                if self.params.phn_loss_type == "softmax":
                    pass
                elif self.params.phn_loss_type == "asoftmax":
                    self.params.phn_asoftmax_m = 1
                elif self.params.phn_loss_type == "additive_margin_softmax":
                    self.params.phn_amsoftmax_m = 0
                elif self.params.phn_loss_type == "additive_angular_margin_softmax":
                    self.params.phn_arcsoftmax_m = 0
                params = remove_params_prefix(self.params, "phn")

                dummy_labels = tf.zeros(shape_list(mu_zp)[:2], dtype=tf.int32)
                _, endpoints_phn_loss = loss_network[self.params.phn_loss_type](mu_zp, dummy_labels,
                                                                                self.num_phones, params,
                                                                                is_training=False,
                                                                                reuse_variables=reuse_variables,
                                                                                name="phn_softmax")
                self.endpoints.update(add_dict_prefix(endpoints_phn_loss, "phn"))

                # Try to use double float in the decoding process
                self.endpoints["phn_logits"] = tf.cast(self.endpoints["phn_logits"], dtype=tf.float64)

                # The posteriors node.
                self.endpoints["phn_post"] = tf.nn.softmax(self.endpoints["phn_logits"], axis=-1)
                # Compute a special log-posteriors node.
                self.endpoints["log-output"] = tf.nn.log_softmax(self.endpoints["phn_logits"], axis=-1)

                tf.logging.info("The parameters have changed. Do not train the network!")
                if self.saver is None:
                    self.saver = tf.train.Saver()
            return

        if mode == "valid":
            tf.logging.info("Building valid network...")
            with tf.name_scope("valid") as scope:
                self.endpoints.clear()
                _, mu_zs, _ = build_speaker_encoder(self.expand_features,
                                                    self.phn_labels,
                                                    self.feat_length,
                                                    self.params,
                                                    self.endpoints,
                                                    reuse_variables,
                                                    is_training=False)
                _, mu_zp, _ = build_phone_encoder(self.expand_features,
                                                  self.spk_labels,
                                                  self.feat_length,
                                                  self.params,
                                                  self.endpoints,
                                                  reuse_variables,
                                                  is_training=False)

                with tf.control_dependencies([tf.assert_equal(shape_list(mu_zp)[1], shape_list(self.phn_labels)[1])]):
                    mu_zp_subset = tf.gather_nd(mu_zp, self.phn_masks)

                phn_labels_subset = tf.gather_nd(self.phn_labels, self.phn_masks)

                # The speaker discriminator
                if self.params.spk_loss_type == "softmax":
                    pass
                elif self.params.spk_loss_type == "asoftmax":
                    train_spk_margin = self.params.spk_asoftmax_m
                    self.params.spk_asoftmax_m = 1
                elif self.params.spk_loss_type == "additive_margin_softmax":
                    train_spk_margin = self.params.spk_amsoftmax_m
                    self.params.spk_amsoftmax_m = 0
                elif self.params.spk_loss_type == "additive_angular_margin_softmax":
                    train_spk_margin = self.params.spk_arcsoftmax_m
                    self.params.spk_arcsoftmax_m = 0
                else:
                    pass
                params = remove_params_prefix(self.params, "spk")
                valid_spk_loss, endpoints_spk_loss = loss_network[self.params.spk_loss_type](mu_zs, self.spk_labels,
                                                                                self.num_speakers, params,
                                                                                is_training=False,
                                                                                reuse_variables=reuse_variables,
                                                                                name="spk_softmax")
                self.endpoints.update(add_dict_prefix(endpoints_spk_loss, "spk"))

                # Restore the original parameters
                if self.params.spk_loss_type == "softmax":
                    pass
                elif self.params.spk_loss_type == "asoftmax":
                    self.params.spk_asoftmax_m = train_spk_margin
                elif self.params.spk_loss_type == "additive_margin_softmax":
                    self.params.spk_amsoftmax_m = train_spk_margin
                elif self.params.spk_loss_type == "additive_angular_margin_softmax":
                    self.params.spk_arcsoftmax_m = train_spk_margin
                else:
                    pass

                # Then the phone discriminator
                if self.params.phn_loss_type == "softmax":
                    pass
                elif self.params.phn_loss_type == "asoftmax":
                    train_phn_margin = self.params.phn_asoftmax_m
                    self.params.phn_asoftmax_m = 1
                elif self.params.phn_loss_type == "additive_margin_softmax":
                    train_phn_margin = self.params.phn_amsoftmax_m
                    self.params.phn_amsoftmax_m = 0
                elif self.params.phn_loss_type == "additive_angular_margin_softmax":
                    train_phn_margin = self.params.phn_arcsoftmax_m
                    self.params.phn_arcsoftmax_m = 0
                else:
                    pass
                params = remove_params_prefix(self.params, "phn")
                valid_phn_loss, endpoints_phn_loss = loss_network[self.params.phn_loss_type](mu_zp_subset, phn_labels_subset,
                                                                                self.num_phones, params,
                                                                                is_training=False,
                                                                                reuse_variables=reuse_variables,
                                                                                name="phn_softmax")
                self.endpoints.update(add_dict_prefix(endpoints_phn_loss, "phn"))

                # Restore the original parameters
                if self.params.phn_loss_type == "softmax":
                    pass
                elif self.params.phn_loss_type == "asoftmax":
                    self.params.phn_asoftmax_m = train_phn_margin
                elif self.params.phn_loss_type == "additive_margin_softmax":
                    self.params.phn_amsoftmax_m = train_phn_margin
                elif self.params.phn_loss_type == "additive_angular_margin_softmax":
                    self.params.phn_arcsoftmax_m = train_phn_margin
                else:
                    pass

                # During validation, the embedding is the input of the loss layers
                self.spk_embeddings = mu_zs
                self.phn_embeddings = mu_zp
                self.spk_logits = self.endpoints["spk_logits"]
                self.phn_logits_subset = self.endpoints["phn_logits"]

                # TODO: define the loss
                valid_loss = self.params.spk_loss_weight * valid_spk_loss + self.params.phn_loss_weight * valid_phn_loss

                mean_valid_loss, mean_valid_loss_op = tf.metrics.mean(valid_loss)
                mean_valid_spk_loss, mean_valid_spk_loss_op = tf.metrics.mean(valid_spk_loss)
                mean_valid_spk_acc, mean_valid_spk_acc_op = tf.metrics.accuracy(labels=self.spk_labels,
                                                                                predictions=tf.argmax(self.spk_logits, axis=-1))
                mean_valid_phn_loss, mean_valid_phn_loss_op = tf.metrics.mean(valid_phn_loss)
                mean_valid_phn_acc, mean_valid_phn_acc_op = tf.metrics.accuracy(labels=phn_labels_subset,
                                                                                predictions=tf.argmax(self.phn_logits_subset, axis=-1))
                valid_ops = {"valid_loss": mean_valid_loss,
                             "valid_loss_op": mean_valid_loss_op,
                             "valid_spk_loss": mean_valid_spk_loss,
                             "valid_spk_loss_op": mean_valid_spk_loss_op,
                             "valid_spk_acc": mean_valid_spk_acc,
                             "valid_spk_acc_op": mean_valid_spk_acc_op,
                             "valid_phn_loss": mean_valid_phn_loss,
                             "valid_phn_loss_op": mean_valid_phn_loss_op,
                             "valid_phn_acc": mean_valid_phn_acc,
                             "valid_phn_acc_op": mean_valid_phn_acc_op}
                self.valid_ops.update(valid_ops)
                valid_loss_summary = tf.summary.scalar("loss", mean_valid_loss)
                valid_spk_loss_summary = tf.summary.scalar("speaker_loss", mean_valid_spk_loss)
                valid_spk_acc_summary = tf.summary.scalar("speaker_acc", mean_valid_spk_acc)
                valid_phn_loss_summary = tf.summary.scalar("phone_loss", mean_valid_phn_loss)
                valid_phn_acc_summary = tf.summary.scalar("phone_acc", mean_valid_phn_acc)
                self.valid_summary = tf.summary.merge([valid_loss_summary,
                                                       valid_spk_loss_summary,
                                                       valid_spk_acc_summary,
                                                       valid_phn_loss_summary,
                                                       valid_phn_acc_summary])
                if self.saver is None:
                    self.saver = tf.train.Saver(max_to_keep=self.params.keep_checkpoint_max)
                if self.valid_summary_writer is None:
                    self.valid_summary_writer = tf.summary.FileWriter(os.path.join(self.model, "eval"), self.sess.graph)
            return

        tf.logging.info("Building training network...")
        if "optimizer" not in self.params.dict:
            # The default optimizer is sgd
            self.params.dict["optimizer"] = "sgd"

        if self.params.optimizer == "sgd":
            if "momentum" in self.params.dict:
                sys.exit("Using sgd as the optimizer and you should not specify the momentum.")
            tf.logging.info("***** Using SGD as the optimizer.")
            opt = tf.train.GradientDescentOptimizer(self.learning_rate, name="optimizer")
        elif self.params.optimizer == "momentum":
            # SGD with momentum
            # It is also possible to use other optimizers, e.g. Adam.
            tf.logging.info("***** Using Momentum as the optimizer.")
            opt = tf.train.MomentumOptimizer(self.learning_rate, self.params.momentum, use_nesterov=self.params.use_nesterov, name="optimizer")
        elif self.params.optimizer == "adam":
            tf.logging.info("***** Using Adam as the optimizer.")
            opt = tf.train.AdamOptimizer(self.learning_rate, name="optimizer")
        else:
            sys.exit("Optimizer %s is not supported." % self.params.optimizer)
        self.optimizer = opt

        with tf.name_scope("train") as scope:
            self.endpoints.clear()
            _, mu_zs, _ = build_speaker_encoder(self.expand_features,
                                                self.phn_labels,
                                                self.feat_length,
                                                self.params,
                                                self.endpoints,
                                                reuse_variables,
                                                is_training=True)
            _, mu_zp, _ = build_phone_encoder(self.expand_features,
                                              self.spk_labels,
                                              self.feat_length,
                                              self.params,
                                              self.endpoints,
                                              reuse_variables,
                                              is_training=True)

            with tf.control_dependencies([tf.assert_equal(shape_list(mu_zp)[1], shape_list(self.phn_labels)[1])]):
                # The length of the outputs should be the same with the length of the alignments
                mu_zp_subset = tf.gather_nd(mu_zp, self.phn_masks)

            phn_labels_subset = tf.gather_nd(self.phn_labels, self.phn_masks)
            self.endpoints["mu_zp_subset"] = mu_zp_subset
            self.endpoints["phn_labels_subset"] = phn_labels_subset

            # TODO: we should use the sampled data or the mean?
            tf.logging.info("Speaker loss")
            params = remove_params_prefix(self.params, "spk")
            spk_disc_loss, endpoints_spk_loss = loss_network[self.params.spk_loss_type](mu_zs, self.spk_labels,
                                                                                self.num_speakers, params,
                                                                                is_training=True,
                                                                                reuse_variables=reuse_variables,
                                                                                name="spk_softmax")
            self.endpoints.update(add_dict_prefix(endpoints_spk_loss, "spk"))

            tf.logging.info("Phone loss")
            params = remove_params_prefix(self.params, "phn")
            phn_disc_loss, endpoints_phn_loss = loss_network[self.params.phn_loss_type](mu_zp_subset, phn_labels_subset,
                                                                                self.num_phones, params,
                                                                                is_training=True,
                                                                                reuse_variables=reuse_variables,
                                                                                name="phn_softmax")
            self.endpoints.update(add_dict_prefix(endpoints_phn_loss, "phn"))

            # TODO: Define the loss
            loss = self.params.spk_loss_weight * spk_disc_loss + self.params.phn_loss_weight * phn_disc_loss

            regularization_loss = tf.losses.get_regularization_loss()
            total_loss = loss + regularization_loss
            penalty_loss = tf.get_collection("PENALTY")
            if len(penalty_loss) != 0:
                penalty_loss = tf.reduce_sum(penalty_loss)
                total_loss += penalty_loss
                self.train_summary.append(tf.summary.scalar("penalty_term", penalty_loss))

            self.train_summary = tf.get_collection(tf.GraphKeys.SUMMARIES, scope)
            self.train_summary.append(tf.summary.scalar("loss", loss))
            self.train_summary.append(tf.summary.scalar("regularization_loss", regularization_loss))

            self.total_loss = total_loss
            self.train_summary.append(tf.summary.scalar("total_loss", total_loss))
            self.train_summary.append(tf.summary.scalar("learning_rate", self.learning_rate))

            # The gradient ops is inside the scope to support multi-gpus
            if noupdate_var_list is not None:
                old_batchnorm_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope)
                batchnorm_update_ops = []
                for op in old_batchnorm_update_ops:
                    if not substring_in_list(op.name, noupdate_var_list):
                        batchnorm_update_ops.append(op)
                        tf.logging.info("[Info] Update %s" % op.name)
                    else:
                        tf.logging.info("[Info] Op %s will not be executed" % op.name)
            else:
                batchnorm_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope)

            if noupdate_var_list is not None:
                variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
                train_var_list = []

                for v in variables:
                    if not substring_in_list(v.name, noupdate_var_list):
                        train_var_list.append(v)
                        tf.logging.info("[Info] Train %s" % v.name)
                    else:
                        tf.logging.info("[Info] Var %s will not be updated" % v.name)
                grads = opt.compute_gradients(total_loss, var_list=train_var_list)
            else:
                grads = opt.compute_gradients(total_loss)

            # Once the model has been built (even for a tower), we set the flag
            self.is_built = True

        # There are some things we can do to the gradients, i.e. learning rate scaling.
        # Here we only apply gradient clipping
        if self.params.clip_gradient:
            grads, vars = zip(*grads)  # compute gradients of variables with respect to loss
            grads_clip, _ = tf.clip_by_global_norm(grads, self.params.clip_gradient_norm)  # l2 norm clipping

            # # we follow the instruction in ge2e paper to scale the learning rate for w and b
            # # Actually, I wonder that we can just simply set a large value for w (e.g. 20) and fix it.
            # if self.loss_type == "ge2e":
            #     # The parameters w and b must be the last variables in the gradients
            #     grads_clip = grads_clip[:-2] + [0.01 * grad for grad in grads_clip[-2:]]
            #     # Simply check the position of w and b
            #     for var in vars[-2:]:
            #         assert("w" in var.name or "b" in var.name)
            grads = zip(grads_clip, vars)

        self.train_summary.append(activation_summaries(self.endpoints))
        for var in tf.trainable_variables():
            self.train_summary.append(tf.summary.histogram(var.op.name, var))
        self.train_summary = tf.summary.merge(self.train_summary)

        with tf.control_dependencies(batchnorm_update_ops):
            self.train_op = opt.apply_gradients(grads)

        # We want to inspect other values during training?
        # self.train_ops["vae_loss"] = vae_loss
        self.train_ops["spk_disc_loss"] = spk_disc_loss
        self.train_ops["phn_disc_loss"] = phn_disc_loss
        self.train_ops["loss"] = loss
        self.train_ops["total_loss"] = total_loss
        # self.train_ops.update(vae_all_losses)

        # The model saver
        if self.saver is None:
            self.saver = tf.train.Saver(max_to_keep=self.params.keep_checkpoint_max)

        # The training summary writer
        if self.summary_writer is None:
            self.summary_writer = tf.summary.FileWriter(self.model, self.sess.graph)
        return

    # def _build_decoder(self, zs, zp, recon):
    #     """Build the decoder to reconstruct features
    #
    #     Args:
    #         zs: The sampled speaker variable
    #         zp: The sampled phone variable
    #         recon:
    #     :return: The sampled output, mean and logvar
    #     """
    #     return sampled_recon, mu, logvar
    #
    # def vae_loss(self, mu, logvar, recon):
    #     return loss, all_losses

    def train(self, data_dir, ali_dir, spklist, learning_rate, aux_data=None):
        """Train the model.

        Args:
            data_dir: The training data directory.
            ali_dir: The ali directory.
            spklist: The spklist is a file map speaker name to the index.
            learning_rate: The learning rate is passed by the main program. The main program can easily tune the
                           learning rate according to the validation accuracy or anything else.
            aux_data: The auxiliary data (maybe useful in child class.)
        """
        # initialize all variables
        self.sess.run(tf.global_variables_initializer())

        # curr_step is the real step the training at.
        curr_step = 0

        # Load the model if we have
        if os.path.isfile(os.path.join(self.model, "checkpoint")):
            curr_step = self.load()

        left_context = max(self.params.phone_left_context, self.params.speaker_left_context)
        right_context = max(self.params.phone_right_context, self.params.speaker_right_context)

        # The data loader
        data_loader = KaldiDataRandomQueueV2(data_dir,
                                             ali_dir,
                                             spklist,
                                             num_parallel=self.params.num_parallel_datasets,
                                             max_qsize=self.params.max_queue_size,
                                             left_context=left_context,
                                             right_context=right_context,
                                             num_speakers=self.params.num_speakers_per_batch,
                                             num_segments=self.params.num_segments_per_speaker,
                                             min_len=self.params.min_segment_len,
                                             max_len=self.params.max_segment_len,
                                             shuffle=True)
        data_loader.start()

        epoch = int(curr_step / self.params.num_steps_per_epoch)
        for step in range(curr_step % self.params.num_steps_per_epoch, self.params.num_steps_per_epoch):
            try:
                # Load the data and subsample for phone loss
                features, vad, ali, length, labels, resample, valid_pos = data_loader.fetch()
                assert(features.shape[1] == ali.shape[1] + left_context + right_context)
                phn_masks = make_phone_masks(length, resample, self.params.num_frames_per_utt)
                self._training_egs_stat(labels, ali, valid_pos, phn_masks)

                if step % self.params.save_summary_steps == 0 or step % self.params.show_training_progress == 0:
                    train_ops = [self.train_ops, self.train_op]
                    if step % self.params.save_summary_steps == 0:
                        train_ops.append(self.train_summary)
                    start_time = time.time()
                    train_val = self.sess.run(train_ops, feed_dict={self.features: features,
                                                                    self.spk_labels: labels,
                                                                    self.phn_labels: ali,
                                                                    self.feat_length: length,
                                                                    self.phn_masks: phn_masks,
                                                                    self.global_step: curr_step,
                                                                    self.learning_rate: learning_rate})
                    end_time = time.time()
                    tf.logging.info(
                        "Epoch: [%2d] step: [%2d/%2d] time: %.4f s/step, spk loss: %f, phn loss: %f, loss: %f, total loss: %f" %
                          (epoch, step, self.params.num_steps_per_epoch, end_time - start_time,
                           train_val[0]["spk_disc_loss"], train_val[0]["phn_disc_loss"],
                           train_val[0]["loss"], train_val[0]["total_loss"]))
                    if step % self.params.save_summary_steps == 0:
                        self.summary_writer.add_summary(train_val[-1], curr_step)
                else:
                    # Only compute optimizer.
                    _ = self.sess.run(self.train_op, feed_dict={self.features: features,
                                                                self.spk_labels: labels,
                                                                self.phn_labels: ali,
                                                                self.feat_length: length,
                                                                self.phn_masks: phn_masks,
                                                                self.global_step: curr_step,
                                                                self.learning_rate: learning_rate})
                if step % self.params.save_checkpoints_steps == 0 and curr_step != 0:
                    self.save(curr_step)
                curr_step += 1
            except DataOutOfRange:
                tf.logging.info("Finished reading features.")
                break

        self._print_egs_stat()
        self._save_egs_stat()
        data_loader.stop()
        self.save(curr_step)
        return

    def train_tune_lr(self, data_dir, ali_dir, spklist, tune_period=100, aux_data=None):
        """Tune the learning rate.

        According to: https://www.kdnuggets.com/2017/11/estimating-optimal-learning-rate-deep-neural-network.html

        Args:
            data_dir: The data directory.
            ali_dir: The ali directory.
            spklist: The spklist is a file map speaker name to the index.
            tune_period: How many steps per learning rate.
            aux_data: The auxiliary data directory.
        """
        # initialize all variables
        self.sess.run(tf.global_variables_initializer())

        # We need to load the model sometimes, since we may try to find the learning rate for fine-tuning.
        if os.path.isfile(os.path.join(self.model, "checkpoint")):
            self.load()

        data_loader = KaldiDataRandomQueueV2(data_dir,
                                             ali_dir,
                                             spklist,
                                             num_parallel=self.params.num_parallel_datasets,
                                             max_qsize=self.params.max_queue_size,
                                             left_context=max(self.params.phone_left_context, self.params.speaker_left_context),
                                             right_context=max(self.params.phone_right_context, self.params.speaker_right_context),
                                             num_speakers=self.params.num_speakers_per_batch,
                                             num_segments=self.params.num_segments_per_speaker,
                                             min_len=self.params.min_segment_len,
                                             max_len=self.params.max_segment_len,
                                             shuffle=True)
        data_loader.start()

        # The learning rate normally varies from 1e-5 to 1
        # Some common values:
        # 1. factor = 1.15
        #    tune_period = 200
        #    tune_times = 100
        init_learning_rate = 1e-5
        factor = 1.15
        tune_times = 100

        fp_lr = open(os.path.join(self.model, "learning_rate_tuning"), "w")
        for step in range(tune_period * tune_times):
            lr = init_learning_rate * (factor ** (step // tune_period))
            try:
                features, vad, ali, length, labels, resample, valid_pos = data_loader.fetch()
                phn_masks = make_phone_masks(length, resample, self.params.num_frames_per_utt)

                if step % tune_period == 0:
                    train_ops = [self.train_ops, self.train_op, self.train_summary]
                    # train_ops = [self.train_ops, self.train_op]
                    start_time = time.time()
                    train_val = self.sess.run(train_ops, feed_dict={self.features: features,
                                                                    self.spk_labels: labels,
                                                                    self.phn_labels: ali,
                                                                    self.feat_length: length,
                                                                    self.phn_masks: phn_masks,
                                                                    self.global_step: 0,
                                                                    self.learning_rate: lr})
                    end_time = time.time()
                    tf.logging.info(
                        "Step: %2d time: %.4f s/step, spk loss: %f, phn loss: %f, loss: %f, total loss: %f" %
                        (step, end_time - start_time,
                         train_val[0]["spk_disc_loss"], train_val[0]["phn_disc_loss"],
                         train_val[0]["loss"], train_val[0]["total_loss"]))
                    fp_lr.write("%d %f %f\n" % (step, lr, train_val[0]["loss"]))
                    self.summary_writer.add_summary(train_val[-1], step)
                else:
                    _ = self.sess.run(self.train_op, feed_dict={self.features: features,
                                                                self.spk_labels: labels,
                                                                self.phn_labels: ali,
                                                                self.feat_length: length,
                                                                self.phn_masks: phn_masks,
                                                                self.global_step: 0,
                                                                self.learning_rate: lr})
            except DataOutOfRange:
                tf.logging.info("Finished reading features.")
                break
        data_loader.stop()
        fp_lr.close()
        return

    def valid(self, data_dir, ali_dir, spklist, batch_type="softmax", output_embeddings=False, aux_data=None):
        """Evaluate on the validation set

        Args:
            data_dir: The data directory.
            ali_dir: The ali directory.
            spklist: The spklist is a file map speaker name to the index.
            batch_type: `softmax` or `end2end`. The batch is `softmax-like` or `end2end-like`.
                        If the batch is `softmax-like`, each sample are from different speakers;
                        if the batch is `end2end-like`, the samples are from N speakers with M segments per speaker.
            output_embeddings: Set True to output the corresponding embeddings and labels of the valid set.
                               If output_embeddings, an additional valid metric (e.g. EER) should be computed outside
                               the function.
            aux_data: The auxiliary data directory.

        :return: valid_loss, embeddings and labels (None if output_embeddings is False).
        """
        # Initialization will reset all the variables in the graph.
        # The local variables are also need to be initialized for metrics function.
        self.sess.run(tf.global_variables_initializer())
        self.sess.run(tf.local_variables_initializer())
        assert batch_type == "softmax" or batch_type == "end2end", "The batch_type can only be softmax or end2end"

        curr_step = 0
        # Load the model. The valid function can only be called after training (of course...)
        if os.path.isfile(os.path.join(self.model, "checkpoint")):
            curr_step = self.load()
        else:
            tf.logging.info("[Warning] Cannot find model in %s. Random initialization is used in validation." % self.model)

        embeddings_val = None
        labels_val = None
        num_batches = 0

        left_context = max(self.params.phone_left_context, self.params.speaker_left_context)
        right_context = max(self.params.phone_right_context, self.params.speaker_right_context)

        if output_embeddings:
            # If we want to output embeddings, the features should be loaded in order
            data_loader = KaldiDataSeqQueueV2(data_dir,
                                              ali_dir,
                                              spklist,
                                              num_parallel=2,
                                              max_qsize=10,
                                              left_context=left_context,
                                              right_context=right_context,
                                              batch_size=self.params.num_speakers_per_batch * self.params.num_segments_per_speaker,
                                              min_len=self.params.min_segment_len,
                                              max_len=self.params.max_segment_len,
                                              shuffle=False)
            data_loader.start()

            tf.logging.info("Generate valid embeddings.")
            # In this mode, the embeddings and labels will be saved and output. It needs more memory and takes longer
            # to process these values.
            while True:
                try:
                    if num_batches % 100 == 0:
                        tf.logging.info("valid step: %d" % num_batches)
                    features, vad, ali, length, labels, resample, valid_pos = data_loader.fetch()
                    valid_emb_val = self.sess.run(self.spk_embeddings, feed_dict={self.features: features,
                                                                                self.feat_length: length,
                                                                                self.phn_labels: ali})
                    # Save the embeddings and labels
                    if embeddings_val is None:
                        embeddings_val = valid_emb_val
                        labels_val = labels
                    else:
                        embeddings_val = np.concatenate((embeddings_val, valid_emb_val), axis=0)
                        labels_val = np.concatenate((labels_val, labels), axis=0)
                    num_batches += 1
                except DataOutOfRange:
                    break
            data_loader.stop()

        if batch_type == "softmax":
            data_loader = KaldiDataSeqQueueV2(data_dir,
                                              ali_dir,
                                              spklist,
                                              num_parallel=2,
                                              max_qsize=10,
                                              left_context=left_context,
                                              right_context=right_context,
                                              batch_size=self.params.num_speakers_per_batch * self.params.num_segments_per_speaker,
                                              min_len=self.params.min_segment_len,
                                              max_len=self.params.max_segment_len,
                                              shuffle=True)
        elif batch_type == "end2end":
            # The num_valid_speakers_per_batch and num_valid_segments_per_speaker are only required when
            # End2End loss is used. Since we switch the loss function to softmax generalized e2e loss
            # when the e2e loss is used.
            assert "num_valid_speakers_per_batch" in self.params.dict and "num_valid_segments_per_speaker" in self.params.dict, \
                "Valid parameters should be set if E2E loss is selected"
            data_loader = KaldiDataRandomQueueV2(data_dir,
                                                 ali_dir,
                                                 spklist,
                                                 num_parallel=2,
                                                 max_qsize=10,
                                                 left_context=left_context,
                                                 right_context=right_context,
                                                 num_speakers=self.params.num_valid_speakers_per_batch,
                                                 num_segments=self.params.num_valid_segments_per_speaker,
                                                 min_len=self.params.min_segment_len,
                                                 max_len=self.params.max_segment_len,
                                                 shuffle=True)
        else:
            raise ValueError

        data_loader.start()
        num_batches = 0
        for _ in range(self.params.valid_max_iterations):
            try:
                if num_batches % 100 == 0:
                    tf.logging.info("valid step: %d" % num_batches)
                features, vad, ali, length, labels, resample, valid_pos = data_loader.fetch()
                # phn_masks = make_phone_masks(length, resample, self.params.num_frames_per_utt)
                phn_masks = make_phone_masks(length, resample, -1)
                valid_ops = [self.valid_ops["valid_loss_op"],
                             self.valid_ops["valid_spk_loss_op"],
                             self.valid_ops["valid_phn_loss_op"],
                             self.valid_ops["valid_spk_acc_op"],
                             self.valid_ops["valid_phn_acc_op"]]
                _ = self.sess.run(valid_ops, feed_dict={self.features: features,
                                                        self.spk_labels: labels,
                                                        self.phn_labels: ali,
                                                        self.feat_length: length,
                                                        self.phn_masks: phn_masks,
                                                        self.global_step: curr_step})
                num_batches += 1
            except DataOutOfRange:
                break
        data_loader.stop()

        loss, spk_loss, phn_loss, spk_acc, phn_acc, summary = self.sess.run([self.valid_ops["valid_loss"],
                                                                             self.valid_ops["valid_spk_loss"],
                                                                             self.valid_ops["valid_phn_loss"],
                                                                             self.valid_ops["valid_spk_acc"],
                                                                             self.valid_ops["valid_phn_acc"],
                                                                             self.valid_summary])
        # We only save the summary for the last batch.
        self.valid_summary_writer.add_summary(summary, curr_step)
        # The valid loss is averaged over all the batches.
        tf.logging.info("[Validation %d batches] valid loss: %f, spk loss: %f, phone loss: %f, spk acc: %.4f, phone acc: %.4f" %
                        (num_batches, loss, spk_loss, phn_loss, spk_acc, phn_acc))

        # The output embeddings and labels can be used to compute EER or other metrics
        return loss, embeddings_val, labels_val

    def predict_speaker(self, node, features, ali, length):
        """Output the speaker-related embedding of the specified node

        We may use ali when inferring the speaker embeddings.

        Args:
            node: The node of the network.
            features: A matrix which could be [l, d] or [b, l, d]
            ali: The alignment. [l], [b, l]
            length: The length of each segment. [1] or [b]
        :return: A numpy array which is the embeddings. [d] or [b, d]
        """
        curr_step = 0
        if not self.is_loaded:
            if os.path.isfile(os.path.join(self.model, "checkpoint")):
                curr_step = self.load()
            else:
                sys.exit("Cannot find model in %s" % self.model)
        rank = len(features.shape)
        assert(rank == 2 or rank == 3)
        assert(rank == len(ali.shape)+1)
        # Expand the feature if the rank is 2
        if rank == 2:
            features = np.expand_dims(features, axis=0)
            ali = np.expand_dims(ali, axis=0)

        # Feature expansion
        left_context = max(self.params.phone_left_context, self.params.speaker_left_context)
        right_context = max(self.params.phone_right_context, self.params.speaker_right_context)
        features = np.concatenate([np.tile(features[:, 0, :], [1, left_context, 1]), features,
                                   np.tile(features[:, -1, :], [1, right_context, 1])], axis=1)

        embeddings = self.sess.run(self.endpoints[node], feed_dict={self.features: features,
                                                                    self.phn_labels: ali,
                                                                    self.feat_length: length,
                                                                    self.global_step: curr_step})

        # If embeddings are speaker embedding, its shape should be [b, d];
        # or the shape could be [b, l, d].
        # If b == 1, squeeze the axis to [d] or [l, d]
        if rank == 2:
            embeddings = np.squeeze(embeddings, axis=0)
        return embeddings

    def predict_phone(self, node, features, length):
        """Output the phone-related embedding of the specified node.
        Also used for log-output (posterior)

        Args:
            node: The node of the network.
            features: A matrix which could be [l, d] or [b, l, d]
            length: The length of each segment. [1] or [b]
        :return: A numpy array. [l, d] or [b, l, d]
        """
        curr_step = 0
        if not self.is_loaded:
            if os.path.isfile(os.path.join(self.model, "checkpoint")):
                curr_step = self.load()
            else:
                sys.exit("Cannot find model in %s" % self.model)
        rank = len(features.shape)
        assert (rank == 2 or rank == 3)
        # Expand the feature if the rank is 2
        if rank == 2:
            features = np.expand_dims(features, axis=0)

        # Feature expansion
        left_context = max(self.params.phone_left_context, self.params.speaker_left_context)
        right_context = max(self.params.phone_right_context, self.params.speaker_right_context)
        features = np.concatenate([np.tile(features[:, 0, :], [1, left_context, 1]), features,
                                   np.tile(features[:, -1, :], [1, right_context, 1])], axis=1)

        embeddings = self.sess.run(self.endpoints[node], feed_dict={self.features: features,
                                                                    self.feat_length: length,
                                                                    self.global_step: curr_step})
        if rank == 2:
            embeddings = np.squeeze(embeddings, axis=0)
        return embeddings

    def _training_egs_stat(self, labels, ali, pos, phn_masks):
        """Compute the statistics of the training, i.e. the count of speakers and phones selected.

        Args:
            labels: The batch speaker labels
            ali: The batch phone labels (alignments)
            pos: [start, end]. If the start <= index < end, it can be inferred w/o feature expansion
            phn_masks: The masks to choose the alignments.
        :return: Save the statistics in self.spk_training_count, self.phn_training_count
        """
        # Note that we cannot use batch add since some classes may be sampled more than once in a batch.
        for k in labels:
            self.spk_training_count[k] += 1
        for k in phn_masks:
            self.phn_training_count[ali[k[0], k[1]]] += 1
            if pos[k[0], 0] <= k[1] < pos[k[0], 1]:
                self.phn_expansion_count[0] += 1
            else:
                self.phn_expansion_count[1] += 1

    def _print_egs_stat(self):
        total_spk_egs = np.sum(self.spk_training_count)
        per_spk = self.spk_training_count * 100.0 / total_spk_egs
        tf.logging.info("The min speaker %d has %.4f%% training egs, and the max speaker %d has %.4f%%" % (
            np.argmin(per_spk), np.min(per_spk), np.argmax(per_spk), np.max(per_spk)))

        total_phn_egs = np.sum(self.phn_training_count)
        per_phone = self.phn_training_count * 100.0 / total_phn_egs
        tf.logging.info("The min phone %d has %.4f%% training egs, and the max phone %d has %.4f%%" % (
            np.argmin(per_phone), np.min(per_phone), np.argmax(per_phone), np.max(per_phone)))

        total_frames = np.sum(self.phn_expansion_count)
        assert(total_frames == total_phn_egs)
        tf.logging.info("There are %.4f%% frames can be inferred w/o feature expansion, %.4f%% with feature expansion" %
                        (self.phn_expansion_count[0] * 100.0 / total_phn_egs,
                         self.phn_expansion_count[1] * 100.0 / total_phn_egs))

    def _save_egs_stat(self):
        fp_spk = open(os.path.join(self.model, "speaker_egs"), "w")
        for i in range(self.spk_training_count.shape[0]):
            fp_spk.write("%d\n" % self.spk_training_count[i])
        fp_spk.close()
        fp_phn = open(os.path.join(self.model, "phone_egs"), "w")
        for i in range(self.phn_training_count.shape[0]):
            fp_phn.write("%d\n" % self.phn_training_count[i])
        fp_phn.close()


"""" Takes input data from benign samples
     and adversarial examples to train model
     [not currently in use]"""

import os
import random
import keras
import torch
import numpy as np
from sklearn.model_selection import train_test_split

from models.athena import Ensemble
from models.keraswrapper import WeakDefense
from models.mnist_cnn import cnn
from utils.file import load_from_json
from utils.model import load_lenet, load_pool
from utils.data import subsampling
from utils.data import get_dataloader
# from keras.engine.saving import save_model
from keras.callbacks import ModelCheckpoint
from keras.models import save_model
from models.image_processor import transform

def shuffle(data, labels):
    # convert to labels
    if len(labels.shape) > 1:
        labels = [np.argmax(p) for p in labels]
    # shuffle the selected ids
    id_s = [i for i in range(data.shape[0])]
    random.shuffle(id_s)
    # get sampled data and labels
    samples = np.asarray([data[i] for i in id_s])
    labels = np.asarray([labels[i] for i in id_s])
    return samples, labels

# modified from toy model on stackoverflow
def add_checkpoint(filepath_p, callback_list):
    checkpoint = ModelCheckpoint(filepath_p+'_checkpoint', monitor='loss', verbose=1, save_best_only=False, mode='min')
    # callback_list.extend([checkpoint])
    callback_list = [checkpoint]

def make_ensemble(wd_config, model_config):
    pool, _ = load_pool(trans_configs=wd_config, model_configs=model_config, active_list=True)
    # turns WD into WeakDefense objects for Ensemble
    wd_models = [WeakDefense(pool[wd], wd_config.get('configs'+str(wd))) for wd in pool]
    wds = Ensemble(wd_models, strategy=None)
    return wds

def train_model(wd, model, x, y, batch_size_p=1, epochs=1,
                load_from_checkpoint=False, save_checkpoint=False, filepath_p='./temp', callback_list_p=None):
    prediction = wd.predict(x, raw=True, batch_size=batch_size_p)
    print(prediction.shape)
    prediction = np.transpose(prediction, (1, 0, 2))
    print(y.shape)
    if load_from_checkpoint and callback_list_p:
        model.fit(prediction, y, callbacks=callback_list_p, epochs=epochs, batch_size=batch_size_p)
    else:
        model.fit(prediction, y, epochs=epochs, batch_size=batch_size_p)
    if save_checkpoint and filepath_p and callback_list_p:
        add_checkpoint(callback_list_p)

def train_ensemble_model(model_p, data_config_p, wd_p, batch_size_p=1, test_size_p=10, save=True, filepath_p=None):
    # Load data
    data_file = os.path.join(data_config_p.get('dir'), data_config_p.get('bs_file'))
    data_bs = np.load(data_file)
    print(data_bs.shape)
    # data_bs = np.transpose(data_bs, (1, 0, 2))
    # exit()
    # load the corresponding true labels
    label_file = os.path.join(data_config_p.get('dir'), data_config_p.get('label_file'))
    labels = np.load(label_file)
    print('Shape of labels: ', labels.shape)
    # split train test data to use for model training
    # subsampling didn't allow for ratios that were too high, so shuffle data instead
    bs_data, bs_labels = shuffle(data_bs, labels)
    train_data, test_data, train_labels, test_labels = train_test_split(bs_data, labels, test_size=0.1)
    # train_data, test_data = bs_data[:-test_size_p], bs_data[-test_size_p:]
    # train_labels, test_labels = bs_labels[:-test_size_p], bs_labels[-test_size_p:]
    # make model trained on WD ouputs, or make model that uses WD outputs as input layer
    # target.fit(train_data, train_labels) [no more]
    call_backs = []
    print(train_labels.shape)
    train_model(wd_p, model_p, train_data, train_labels, batch_size_p=batch_size_p,
                epochs=5, save_checkpoint=True, filepath_p=filepath_p, callback_list_p=call_backs)

    # save_model(target, filepath=filepath, overwrite=False, include_optimizer=True)
    # exit()
    add_checkpoint(filepath_p, call_backs)
    ae_files = [os.path.join(data_config_p.get('dir'), AE_name) for AE_name in data_config_p.get('ae_files')]

    for AE_file in ae_files:
        ae_data = np.load(AE_file)
        # ae_data = np.transpose(ae_data, (1, 0, 2))
        ae_data, ae_labels = subsampling(ae_data, labels, 10, 0.2)
        ae_train_data, ae_test_data = ae_data[:-test_size_p], ae_data[-test_size_p:]
        ae_train_labels, ae_test_labels = ae_labels[:-test_size_p], ae_labels[-test_size_p:]
        train_model(wd_p, model_p, ae_train_data, ae_train_labels, batch_size_p=batch_size_p, epochs=2,
                    load_from_checkpoint=True, save_checkpoint=True, filepath_p=filepath_p, callback_list_p=call_backs)
        # for id, model in pool.items():
        #     if id == 0: # skip the undefended model, which does not transform the image
        #         continue
        #
        #     key = 'configs{}'.format(id)
        #     trans_args = WD_config.get(key)
        #     print('TRANS CONFIG:', trans_args)
        #     # transform a small subset
        #     data_trans = transform(data_bs[:50], trans_args)
        #     fit the transformed images by the corresponding model (weak defense)
        #     target.fit(AE_train_data, AE_train_labels, callbacks=call_backs)
        #     add_checkpoint(filepath, call_backs)

    # print('finished')
    # optional save model
    if save and filepath_p:
        model_p.save(filepath_p)
    return test_data, test_labels

if __name__ == '__main__':
    # load config files: model-config, data-config
    # Change these next 4 lines
    model_config = load_from_json('../../configs/task2/zhymir_configs/model-config.json')
    data_config = load_from_json('../../configs/task2/zhymir_configs/data-config.json')
    WD_config = load_from_json('../../configs/task2/zhymir_configs/task2-athena-mnist.json')
    filepath = os.path.join('../../../Task2/models', 'zhymir_model_batch_size_10.h5')
    # change these 2
    batch_size = 10
    test_size = 10
    # use load pool to collect WDs and UM
    wds = make_ensemble(WD_config, model_config=model_config)
    # Create target model for task
    # Change the layers of this model for needs
    # exit()
    target = keras.models.Sequential([
        keras.layers.InputLayer(input_shape=(wds._nb_classifiers, 10), name='WD_layer'),
        keras.layers.Flatten(),
        keras.layers.Dense(units=100, activation='relu', name='D1'),
        keras.layers.Dense(10, name='output_layer', activation='softmax')
    ])

    # define loss and optimizer for model
    target.compile('adam', 'categorical_crossentropy')
    train_ensemble_model(target, batch_size_p=batch_size, test_size_p=test_size,
                         data_config_p=data_config, filepath_p=filepath, wd_p=wds, save=False)


# -*- coding: utf-8 -*-

from ctypes import POINTER, c_int, c_double, c_float, c_void_p

import numpy as np

from pyfr.backends.base import ComputeKernel
from pyfr.ctypesutil import LibWrapper

# Possible CUBLAS exception types
CUBLASError = type('CUBLASError', (Exception,), {})
CUBLASNotInitialized = type('CUBLASNotInitialized', (CUBLASError,), {})
CUBLASAllocFailed = type('CUBLASAllocFailed', (CUBLASError,), {})
CUBLASInvalidValue = type('CUBLASInvalidValue', (CUBLASError,), {})
CUBLASArchMismatch = type('CUBLASArchMismatch', (CUBLASError,), {})
CUBLASMappingError = type('CUBLASMappingError', (CUBLASError,), {})
CUBLASExecutionFailed = type('CUBLASExecutionFailed', (CUBLASError,), {})
CUBLASInternalError = type('CUBLASInternalError', (CUBLASError,), {})

class CUBLASWrappers(LibWrapper):
    _libname = 'cublas'

    # Error codes
    _statuses = {
        0x1: CUBLASNotInitialized,
        0x3: CUBLASAllocFailed,
        0x7: CUBLASInvalidValue,
        0x8: CUBLASArchMismatch,
        0xb: CUBLASMappingError,
        0xd: CUBLASExecutionFailed,
        0xe: CUBLASInternalError,
        '*': CUBLASError
    }

    # Constants
    OP_N = 0
    OP_T = 1
    OP_C = 2

    # Functions
    _functions = [
        (c_int, 'cublasCreate_v2', POINTER(c_void_p)),
        (c_int, 'cublasDestroy_v2', c_void_p),
        (c_int, 'cublasSetStream_v2', c_void_p, c_void_p),
        (c_int, 'cublasDgemm_v2', c_void_p, c_int, c_int, c_int, c_int, c_int,
         POINTER(c_double), c_void_p, c_int, c_void_p, c_int,
         POINTER(c_double), c_void_p, c_int),
        (c_int, 'cublasSgemm_v2', c_void_p, c_int, c_int, c_int, c_int, c_int,
         POINTER(c_float), c_void_p, c_int, c_void_p, c_int,
         POINTER(c_float), c_void_p, c_int)
    ]

    def _transname(self, name):
        return name[:-3]

class CUDACUBLASKernels(object):
    def __init__(self, backend):
        # Load and wrap CUBLAS
        self.lib = CUBLASWrappers()

        # Init
        self._as_parameter_ = handle = c_void_p()
        self.lib.cublasCreate(handle)

    def __del__(self):
        if self._as_parameter_:
            self.lib.cublasDestroy(self)

    def mul(self, a, b, out, alpha=1.0, beta=0.0):
        lib = self.lib

        # Ensure the matrices are compatible
        if a.nrow != out.nrow or a.ncol != b.nrow or b.ncol != out.ncol:
            raise ValueError('Incompatible matrices for out = a*b')

        # CUBLAS expects inputs to be column-major (or Fortran order in
        # numpy parlance).  However as C = A*B => C^T = (A*B)^T
        # = (B^T)*(A^T) with a little trickery we can multiply our
        # row-major matrices directly.
        m, n, k = b.ncol, a.nrow, a.ncol
        A, B, C = b, a, out

        # α and β factors for C = α*(A*op(B)) + β*C
        if a.dtype == np.float64:
            cublasgemm = lib.cublasDgemm
            alpha_ct, beta_ct = c_double(alpha), c_double(beta)
        else:
            cublasgemm = lib.cublasSgemm
            alpha_ct, beta_ct = c_float(alpha), c_float(beta)

        class MulKernel(ComputeKernel):
            def run(iself, queue):
                lib.cublasSetStream(self, queue.cuda_stream_comp)
                cublasgemm(self, lib.OP_N, lib.OP_N, m, n, k,
                           alpha_ct, A, A.leaddim, B, B.leaddim,
                           beta_ct, C, C.leaddim)

        return MulKernel()


import io
import os
import datetime
import warnings

import pandas as pd
import requests
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from scipy.ndimage.filters import gaussian_filter

from astropy import units as u
from astropy.coordinates import SkyCoord
from astroquery.gaia import Gaia
from astropy.time import Time
from astropy.utils.exceptions import AstropyWarning

from astropy.wcs import WCS
from astropy.wcs.utils import pixel_to_skycoord
from astropy.io import fits
from astropy.visualization import ImageNormalize, ZScaleInterval
from reproject import reproject_adaptive

from .cache import Cache

from baselayer.log import make_log

log = make_log('finder-chart')

warnings.simplefilter('ignore', category=AstropyWarning)

facility_parameters = {
    'Keck': {
        "radius_degrees": 2.0 / 60,
        "mag_limit": 18.5,
        "mag_min": 11.0,
        "min_sep_arcsec": 4.0,
    },
    'P200': {
        "radius_degrees": 2.0 / 60,
        "mag_limit": 18.0,
        "mag_min": 10.0,
        "min_sep_arcsec": 5.0,
    },
    'Shane': {
        "radius_degrees": 2.5 / 60,
        "mag_limit": 17.0,
        "mag_min": 10.0,
        "min_sep_arcsec": 5.0,
    },
}

# ZTF ref grabber URLs. See get_ztfref_url() below
irsa = {
    "url_data": "https://irsa.ipac.caltech.edu/ibe/data/ztf/products/",
    "url_search": "https://irsa.ipac.caltech.edu/ibe/search/ztf/products/",
}

def get_url(*args, **kwargs):
    # Connect and read timeouts
    kwargs['timeout'] = (6.05, 20)
    try:
        return requests.get(*args, **kwargs)
    except requests.exceptions.RequestException:
        return None

def get_ztfref_url(ra, dec, imsize, *args, **kwargs):
    """
    From:
    https://gist.github.com/dmitryduev/634bd2b21a77e2b1de89e0bfd39d14b9

    Returns the URL that points to the ZTF reference image for the
    requested position

    Parameters
    ----------
    source_ra : float
        Right ascension (J2000) of the source
    source_dec : float
        Declination (J2000) of the source
    imsize : float, optional
        Requested image size (on a size) in arcmin
    *args : optional
        Extra args (not needed here)
    **kwargs : optional
        Extra kwargs (not needed here)

    Returns
    -------
    str
        the URL to download the ZTF image

    """
    imsize_deg = imsize / 60

    url_ref_meta = os.path.join(
        irsa['url_search'], f"ref?POS={ra:f},{dec:f}&SIZE={imsize_deg:f}&ct=csv"
    )
    r = get_url(url_ref_meta)
    if r is None:
        return ''

    s = r.content
    c = pd.read_csv(io.StringIO(s.decode('utf-8')))

    field = f"{c.loc[0, 'field']:06d}"
    filt = c.loc[0, 'filtercode']
    quad = f"{c.loc[0, 'qid']}"
    ccd = f"{c.loc[0, 'ccdid']:02d}"

    path_ursa_ref = os.path.join(
        irsa['url_data'],
        'ref',
        field[:3],
        f'field{field}',
        filt,
        f'ccd{ccd}',
        f'q{quad}',
        f'ztf_{field}_{filt}_c{ccd}_q{quad}_refimg.fits',
    )
    return path_ursa_ref

# helper dict for seaching for FITS images from various surveys
source_image_parameters = {
    'desi': {
        'url': (
            'http://legacysurvey.org/viewer/fits-cutout/'
            '?ra={ra}&dec={dec}&layer=dr8&pixscale={pixscale}&bands=r'
        ),
        'npixels': 256,
        'smooth': None,
        'str': 'DESI DR8 R-band',
    },
    'dss': {
        'url': (
            'http://archive.stsci.edu/cgi-bin/dss_search'
            '?v=poss2ukstu_red&r={ra}&dec={dec}&h={imsize}&w={imsize}&e=J2000'
        ),
        'smooth': None,
        'reproject': True,
        'npixels': 500,
        'str': 'DSS-2 Red',
    },
    'ztfref': {
        'url': get_ztfref_url,
        'reproject': True,
        'npixels': 500,
        'smooth': None,
        'str': 'ZTF Ref',
    },
}

def get_nearby_offset_stars(
    source_ra,
    source_dec,
    source_name,
    how_many=3,
    radius_degrees=2 / 60.0,
    mag_limit=18.0,
    mag_min=10.0,
    min_sep_arcsec=5,
    starlist_type='Keck',
    obstime=None,
    use_source_pos_in_starlist=True,
    allowed_queries=2,
    queries_issued=0,
):
    """Finds good list of nearby offset stars for spectroscopy
       and returns info about those stars, including their
       offsets calculated to the source of interest

    Parameters
    ----------
    source_ra : float
        Right ascension (J2000) of the source
    source_dec : float
        Declination (J2000) of the source
    source_name : str
        Name of the source
    how_many : int, optional
        How many offset stars to try to find
    radius_degrees : float, optional
        Search radius from the source position in arcmin
    mag_limit : float, optional
        How faint should we search for offset stars?
    mag_min : float, optional
        What is the brightest offset star we will allow?
    min_sep_arcsec : float, optional
        What is the closest offset star allowed to the source?
    starlist_type : str, optional
        What starlist format should we use?
    obstime : str, optional
        What datetime (in isoformat) should we assume for the observation
        (to calculate proper motions)?
    use_source_pos_in_starlist : bool, optional
        Return the source itself for in starlist?
    allowed_queries : int, optional
        How many times should we query (with looser and looser criteria)
        before giving up on getting the number of offset stars we desire?
    queries_issued : int, optional
        How many times have we issued a query? Bookkeeping parameter.

    Returns
    -------
    (list, str, int, int)
        Return a tuple which contains: a list of dictionaries for each object
        in the star list, the query issued, the number of queries issues, and
        the length of the star list (not including the source itself)
    """

    if queries_issued >= allowed_queries:
        raise Exception('Number of offsets queries needed exceeds what is allowed')

    if not obstime:
        source_obstime = Time(datetime.datetime.utcnow().isoformat())
    else:
        # TODO: check the obstime format
        source_obstime = Time(obstime)

    gaia_obstime = "J2015.5"

    center = SkyCoord(
        source_ra,
        source_dec,
        unit=(u.degree, u.degree),
        frame='icrs',
        obstime=source_obstime,
    )
    # get three times as many stars as requested for now
    # and go fainter as well
    fainter_diff = 2.0  # mag
    search_multipler = 10
    query_string = f"""
                  SELECT TOP {how_many*search_multipler} DISTANCE(
                    POINT('ICRS', ra, dec),
                    POINT('ICRS', {source_ra}, {source_dec})) AS
                    dist, source_id, ra, dec, ref_epoch,
                    phot_rp_mean_mag, pmra, pmdec, parallax
                  FROM gaiadr2.gaia_source
                  WHERE 1=CONTAINS(
                    POINT('ICRS', ra, dec),
                    CIRCLE('ICRS', {source_ra}, {source_dec},
                           {radius_degrees}))
                  AND phot_rp_mean_mag < {mag_limit + fainter_diff}
                  AND phot_rp_mean_mag > {mag_min}
                  AND parallax < 250
                  ORDER BY phot_rp_mean_mag ASC
                """
    # TODO possibly: save the offset data (cache)
    job = Gaia.launch_job(query_string)
    r = job.get_results()
    queries_issued += 1

    catalog = SkyCoord.guess_from_table(r)

    # star needs to be this far away
    # from another star
    min_sep = min_sep_arcsec * u.arcsec
    good_list = []
    for source in r:
        c = SkyCoord(
            ra=source["ra"],
            dec=source["dec"],
            unit=(u.degree, u.degree),
            pm_ra_cosdec=(
                np.cos(source["dec"] * np.pi / 180.0) * source['pmra'] * u.mas / u.yr
            ),
            pm_dec=source["pmdec"] * u.mas / u.yr,
            frame='icrs',
            distance=min(abs(1 / source["parallax"]), 10) * u.kpc,
            obstime=gaia_obstime,
        )

        d2d = c.separation(catalog)  # match it to the catalog
        if sum(d2d < min_sep) == 1 and source["phot_rp_mean_mag"] <= mag_limit:
            # this star is not near another star and is bright enough
            # precess it's position forward to the source obstime and
            # get offsets suitable for spectroscopy
            # TODO: put this in geocentric coords to account for parallax
            cprime = c.apply_space_motion(new_obstime=source_obstime)
            dra, ddec = cprime.spherical_offsets_to(center)
            good_list.append(
                (source["dist"], c, source, dra.to(u.arcsec), ddec.to(u.arcsec))
            )

    good_list.sort()

    # if we got less than we asked for, relax the criteria
    if (len(good_list) < how_many) and (queries_issued < allowed_queries):
        return get_nearby_offset_stars(
            source_ra,
            source_dec,
            source_name,
            how_many=how_many,
            radius_degrees=radius_degrees * 1.3,
            mag_limit=mag_limit + 1.0,
            mag_min=mag_min - 1.0,
            min_sep_arcsec=min_sep_arcsec / 2.0,
            starlist_type=starlist_type,
            obstime=obstime,
            use_source_pos_in_starlist=use_source_pos_in_starlist,
            queries_issued=queries_issued,
            allowed_queries=allowed_queries,
        )

    # default to keck star list
    sep = ' '  # 'fromunit'
    commentstr = "#"
    giveoffsets = True
    maxname_size = 16
    # truncate the source_name if we need to
    if len(source_name) > 10:
        basename = source_name[0:3] + ".." + source_name[-6:]
    else:
        basename = source_name

    if starlist_type == 'Keck':
        pass

    elif starlist_type == 'P200':
        sep = ':'  # 'fromunit'
        commentstr = "!"
        giveoffsets = False
        maxname_size = 20

        # truncate the source_name if we need to
        if len(source_name) > 15:
            basename = source_name[0:3] + ".." + source_name[-11:]
        else:
            basename = source_name

    else:
        log("Warning: Do not recognize this starlist format. Using Keck.")

    basename = basename.strip().replace(" ", "")
    space = " "
    star_list_format = (
        f"{basename:{space}<{maxname_size}} "
        + f"{center.to_string('hmsdms', sep=sep, decimal=False, precision=2, alwayssign=True)[1:]}"
        + f" 2000.0 {commentstr}"
    )

    star_list = []
    if use_source_pos_in_starlist:
        star_list.append(
            {
                "str": star_list_format,
                "ra": float(source_ra),
                "dec": float(source_dec),
                "name": basename,
            }
        )

    for i, (dist, c, source, dra, ddec) in enumerate(good_list[:how_many]):

        dras = f"{dra.value:<0.03f}\" E" if dra > 0 else f"{abs(dra.value):<0.03f}\" W"
        ddecs = (
            f"{ddec.value:<0.03f}\" N" if ddec > 0 else f"{abs(ddec.value):<0.03f}\" S"
        )

        if giveoffsets:
            offsets = f"raoffset={dra.value:<0.03f} decoffset={ddec.value:<0.03f}"
        else:
            offsets = ""

        name = f"{basename}_off{i+1}"

        star_list_format = (
            f"{name:{space}<{maxname_size}} "
            + f"{c.to_string('hmsdms', sep=sep, decimal=False, precision=2, alwayssign=True)[1:]}"
            + f" 2000.0 {offsets} "
            + f" {commentstr} dist={3600*dist:<0.02f}\"; {source['phot_rp_mean_mag']:<0.02f} mag"
            + f"; {dras}, {ddecs} "
            + f" ID={source['source_id']}"
        )

        star_list.append(
            {
                "str": star_list_format,
                "ra": float(source["ra"]),
                "dec": float(source["dec"]),
                "name": name,
                "dras": dras,
                "ddecs": ddecs,
                "mag": float(source["phot_rp_mean_mag"]),
            }
        )

    # send back the starlist in
    return (
        star_list,
        query_string.replace("\n", " "),
        queries_issued,
        len(star_list) - 1,
    )

def fits_image(
    center_ra,
    center_dec,
    imsize=4.0,
    image_source="desi",
    cache_dir="./cache/finder/",
    cache_max_items=5,
):

    """Returns an opened FITS image centered on the source
       of the requested size.

    Parameters
    ----------
    source_ra : float
        Right ascension (J2000) of the source
    source_dec : float
        Declination (J2000) of the source
    imsize : float, optional
        Requested image size (on a size) in arcmin
    image_source : str, optional
        Survey where the image comes from "desi" or "dss" (more to be added)
    cache_dir : str, optional
        Where should the cache live?
    cache_max_items : int, optional
        How many older files in the cache should we keep?  Set to zero
        to disable cache.

    Returns
    -------
    object
        Either a pyfits HDU object or None. If no suitable image is found
        then None is returned. The caller of `fits_image` will need to
        handle this case.
    """

    if image_source not in source_image_parameters:
        raise Exception("do not know how to grab image source")

    pixscale = 60 * imsize / source_image_parameters[image_source].get("npixels", 256)

    if isinstance(source_image_parameters[image_source]["url"], str):
        url = source_image_parameters[image_source]["url"].format(
            ra=center_ra, dec=center_dec, pixscale=pixscale, imsize=imsize
        )
    else:
        # use the URL field as a function
        url = source_image_parameters[image_source]["url"](
            ra=center_ra, dec=center_dec, imsize=imsize
        )

    cache = Cache(cache_dir=cache_dir, max_items=cache_max_items)

    def get_hdu(url):
        """Try to get HDU from cache, otherwise fetch.
        """
        hash_name = f'{center_ra}{center_dec}{imsize}{image_source}'
        hdu_fn = cache[hash_name]

        # Found entry in cache, return that
        if hdu_fn is not None:
            return fits.open(hdu_fn)[0]

        response = get_url(url, stream=True, allow_redirects=True)
        if response is None or response.status_code != 200:
            return None

        # Check if HDU is a valid FITS file
        hdu = fits.open(io.BytesIO(response.content))[0]

        # Ensure it is not empty
        if np.count_nonzero(hdu.data) == 0:
            return None

        # Save a copy in cache and return
        buf = io.BytesIO()
        hdu.writeto(buf)
        buf.seek(0)
        cache[hash_name] = buf.read()
        return hdu

    return get_hdu(url)

def get_finding_chart(
    source_ra,
    source_dec,
    source_name,
    image_source='desi',
    output_format='pdf',
    imsize=3.0,
    tick_offset=0.02,
    tick_length=0.03,
    fallback_image_source='dss',
    **offset_star_kwargs,
):

    """Create a finder chart suitable for spectroscopic observations of
       the source

    Parameters
    ----------
    source_ra : float
        Right ascension (J2000) of the source
    source_dec : float
        Declination (J2000) of the source
    source_name : str
        Name of the source
    image_source : {'desi', 'dss', 'ztfref'}, optional
        Survey where the image comes from "desi", "dss", "ztfref"
        (more to be added)
    output_format : str, optional
        "pdf" of "png" -- determines the format of the returned finder
    imsize : float, optional
        Requested image size (on a size) in arcmin. Should be between 2-15.
    tick_offset : float, optional
        How far off the each source should the tick mark be made? (in arcsec)
    tick_length : float, optional
        How long should the tick mark be made? (in arcsec)
    fallback_image_source : str, optional
        Where what `image_source` should we fall back to if the
        one requested fails
    **offset_star_kwargs : dict, optional
        Other parameters passed to `get_nearby_offset_stars`

    Returns
    -------
    dict
        success : bool
            Whether the request was successful or not, returning
            a sensible error in 'reason'
        name : str
            suggested filename based on `source_name` and `output_format`
        data : str
            binary encoded data for the image (to be streamed)
        reason : str
            If not successful, a reason is returned.
    """
    if (imsize < 2.0) or (imsize > 15):
        return {
            'success': False,
            'reason': 'Requested `imsize` out of range',
            'data': '',
            'name': '',
        }

    if image_source not in source_image_parameters:
        return {
            'success': False,
            'reason': f'image source {image_source} not in list',
            'data': '',
            'name': '',
        }

    matplotlib.use("Agg")
    fig = plt.figure(figsize=(11, 8.5), constrained_layout=False)
    widths = [2.6, 1]
    heights = [2.6, 1]
    spec = fig.add_gridspec(
        ncols=2,
        nrows=2,
        width_ratios=widths,
        height_ratios=heights,
        left=0.05,
        right=0.95,
    )

    # how wide on the side will the image be? 256 as default
    npixels = source_image_parameters[image_source].get("npixels", 256)
    # set the pixelscale in arcsec (typically about 1 arcsec/pixel)
    pixscale = 60 * imsize / npixels

    hdu = fits_image(source_ra, source_dec, imsize=imsize, image_source=image_source)

    # skeleton WCS - this is the field that the user requested
    wcs = WCS(naxis=2)

    # set the headers of the WCS.
    # The center of the image is the reference point (source_ra, source_dec):
    wcs.wcs.crpix = [npixels / 2, npixels / 2]
    wcs.wcs.crval = [source_ra, source_dec]

    # create the pixel scale and orientation North up, East left
    # pixelscale is in degrees, established in the tangent plane
    # to the reference point
    wcs.wcs.cd = np.array([[-pixscale / 3600, 0], [0, pixscale / 3600]])
    wcs.wcs.ctype = ["RA---TAN", "DEC--TAN"]

    if hdu is not None:
        im = hdu.data

        # replace the nans with medians
        im[np.isnan(im)] = np.nanmedian(im)
        if source_image_parameters[image_source].get("reproject", False):
            # project image to the skeleton WCS solution
            log("Reprojecting image to requested position and orientation")
            im, _ = reproject_adaptive(hdu, wcs, shape_out=(npixels, npixels))
        else:
            wcs = WCS(hdu.header)

        if source_image_parameters[image_source].get("smooth", False):
            im = gaussian_filter(
                hdu.data, source_image_parameters[image_source]["smooth"] / pixscale
            )

        norm = ImageNormalize(im, interval=ZScaleInterval())
        watermark = source_image_parameters[image_source]["str"]

    else:
        # if we got back a blank image, try to fallback on another survey
        # and return the results from that call
        if fallback_image_source is not None:
            if fallback_image_source != image_source:
                log(f"Falling back on image source {fallback_image_source}")
                return get_finding_chart(
                    source_ra,
                    source_dec,
                    source_name,
                    image_source=fallback_image_source,
                    output_format=output_format,
                    imsize=imsize,
                    tick_offset=tick_offset,
                    tick_length=tick_length,
                    fallback_image_source=None,
                    **offset_star_kwargs,
                )

        # we dont have an image here, so let's create a dummy one
        # so we can still plot
        im = np.zeros((npixels, npixels))
        norm = None
        watermark = None

    # add the images in the top left corner
    ax = fig.add_subplot(spec[0, 0], projection=wcs)
    ax_text = fig.add_subplot(spec[0, 1])
    ax_text.axis('off')
    ax_starlist = fig.add_subplot(spec[1, 0:])
    ax_starlist.axis('off')

    ax.imshow(im, origin='lower', norm=norm, cmap='gray_r')
    ax.set_autoscale_on(False)
    ax.grid(color='white', ls='dotted')
    ax.set_xlabel(r'$\alpha$ (J2000)', fontsize='large')
    ax.set_ylabel(r'$\delta$ (J2000)', fontsize='large')
    obstime = offset_star_kwargs.get("obstime", datetime.datetime.utcnow().isoformat())
    ax.set_title(
        f'{source_name} Finder ({obstime})', fontsize='large', fontweight='bold'
    )

    star_list, _, _, _ = get_nearby_offset_stars(
        source_ra, source_dec, source_name, **offset_star_kwargs
    )

    if not isinstance(star_list, list) or len(star_list) == 0:
        return {
            'success': False,
            'reason': 'failure to get star list',
            'data': '',
            'name': '',
        }

    ncolors = len(star_list)
    colors = sns.color_palette("colorblind", ncolors)

    start_text = [-0.35, 0.99]
    starlist_str = (
        "# Note: spacing in starlist many not copy/paste correctly in PDF\n"
        + f"#       you can get starlist directly from"
        + f" /api/sources/{source_name}/offsets?"
        + f"facility={offset_star_kwargs.get('facility', 'Keck')}\n"
        + "\n".join([x["str"] for x in star_list])
    )

    # add the starlist
    ax_starlist.text(
        0,
        0.50,
        starlist_str,
        fontsize="x-small",
        family='monospace',
        transform=ax_starlist.transAxes,
    )

    # add the watermark for the survey
    props = dict(boxstyle='round', facecolor='gray', alpha=0.5)

    if watermark is not None:
        ax.text(
            0.035,
            0.035,
            watermark,
            horizontalalignment='left',
            verticalalignment='center',
            transform=ax.transAxes,
            fontsize='medium',
            fontweight='bold',
            color="yellow",
            alpha=0.5,
            bbox=props,
        )

    ax.text(
        0.95,
        0.035,
        f"{imsize}\u2032 \u00D7 {imsize}\u2032",  # size'x size'
        horizontalalignment='right',
        verticalalignment='center',
        transform=ax.transAxes,
        fontsize='medium',
        fontweight='bold',
        color="yellow",
        alpha=0.5,
        bbox=props,
    )

    # compass rose
    # rose_center_pixel = ax.transAxes.transform((0.04, 0.95))
    rose_center = pixel_to_skycoord(int(npixels * 0.1), int(npixels * 0.9), wcs)
    props = dict(boxstyle='round', facecolor='gray', alpha=0.5)

    for ang, label, off in [(0, "N", 0.01), (90, "E", 0.03)]:
        position_angle = ang * u.deg
        separation = (0.05 * imsize * 60) * u.arcsec  # 5%
        p2 = rose_center.directional_offset_by(position_angle, separation)
        ax.plot(
            [rose_center.ra.value, p2.ra.value],
            [rose_center.dec.value, p2.dec.value],
            transform=ax.get_transform('world'),
            color="gold",
            linewidth=2,
        )

        # label N and E
        position_angle = (ang + 15) * u.deg
        separation = ((0.05 + off) * imsize * 60) * u.arcsec
        p2 = rose_center.directional_offset_by(position_angle, separation)
        ax.text(
            p2.ra.value,
            p2.dec.value,
            label,
            color="gold",
            transform=ax.get_transform('world'),
            fontsize='large',
            fontweight='bold',
        )

    for i, star in enumerate(star_list):

        c1 = SkyCoord(star["ra"] * u.deg, star["dec"] * u.deg, frame='icrs')

        # mark up the right side of the page with position and offset info
        name_title = star["name"]
        if star.get("mag") is not None:
            name_title += f", mag={star.get('mag'):.2f}"
        ax_text.text(
            start_text[0],
            start_text[1] - i / ncolors,
            name_title,
            ha='left',
            va='top',
            fontsize='large',
            fontweight='bold',
            transform=ax_text.transAxes,
            color=colors[i],
        )
        source_text = f"  {star['ra']:.5f} {star['dec']:.5f}\n"
        source_text += f"  {c1.to_string('hmsdms')}\n"
        if (star.get("dras") is not None) and (star.get("ddecs") is not None):
            source_text += f'  {star.get("dras")} {star.get("ddecs")} to {source_name}'
        ax_text.text(
            start_text[0],
            start_text[1] - i / ncolors - 0.06,
            source_text,
            ha='left',
            va='top',
            fontsize='large',
            transform=ax_text.transAxes,
            color=colors[i],
        )

        # work on making marks where the stars are
        for ang in [0, 90]:
            position_angle = ang * u.deg
            separation = (tick_offset * imsize * 60) * u.arcsec
            p1 = c1.directional_offset_by(position_angle, separation)
            separation = (tick_offset + tick_length) * imsize * 60 * u.arcsec
            p2 = c1.directional_offset_by(position_angle, separation)
            ax.plot(
                [p1.ra.value, p2.ra.value],
                [p1.dec.value, p2.dec.value],
                transform=ax.get_transform('world'),
                color=colors[i],
                linewidth=3 if imsize <= 4 else 2,
            )
        if star["name"].find("_off") != -1:
            # this is an offset star
            text = star["name"].split("_off")[-1]
            position_angle = 14 * u.deg
            separation = (tick_offset + tick_length * 1.6) * imsize * 60 * u.arcsec
            p1 = c1.directional_offset_by(position_angle, separation)
            ax.text(
                p1.ra.value,
                p1.dec.value,
                text,
                color=colors[i],
                transform=ax.get_transform('world'),
                fontsize='large',
                fontweight='bold',
            )

    buf = io.BytesIO()
    fig.savefig(buf, format=output_format)
    buf.seek(0)

    return {
        "success": True,
        "name": f"finder_{source_name}.{output_format}",
        "data": buf.read(),
        "reason": "",
    }