from __future__ import print_function

from pylib import util

try:
    import tensorflow as tf
    from tensorflow.python.ops import nn

    relu = nn.relu
    slim = tf.contrib.slim
    sigmoid = nn.sigmoid
    softmax = nn.softmax
except:
    print("tensorflow is not installed, util.tf can not be used.")


def is_gpu_available(cuda_only=True):
    """
    code from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/platform/test.py
    Returns whether TensorFlow can access a GPU.
    Args:
      cuda_only: limit the search to CUDA gpus.
    Returns:
      True iff a gpu device of the requested kind is available.
    """
    from tensorflow.python.client import device_lib as _device_lib

    if cuda_only:
        return any((x.device_type == 'GPU')
                   for x in _device_lib.list_local_devices())
    else:
        return any((x.device_type == 'GPU' or x.device_type == 'SYCL')
                   for x in _device_lib.list_local_devices())


def get_available_gpus(num_gpus=None):
    """
    Modified on http://stackoverflow.com/questions/38559755/how-to-get-current-available-gpus-in-tensorflow
    However, the original code will occupy all available gpu memory.
    The modified code need a parameter: num_gpus. It does nothing but return the device handler name
    It will work well on single-maching-training, but I don't know whether it will work well on a cluster.
    """
    if num_gpus == None:
        from tensorflow.python.client import device_lib as _device_lib
        local_device_protos = _device_lib.list_local_devices()
        return [x.name for x in local_device_protos if x.device_type == 'GPU']
    else:
        return ['/gpu:%d' % (idx) for idx in range(num_gpus)]


def get_latest_ckpt(path):
    # tf.train.latest_checkpoint
    from pylib import util

    path = util.io.get_absolute_path(path)
    if util.io.is_dir(path):
        ckpt = tf.train.get_checkpoint_state(path)
        if ckpt is not None:
            ckpt_path = ckpt.model_checkpoint_path
        else:
            ckpt_path = None
    else:
        ckpt_path = path
    return ckpt_path


def get_all_ckpts(path):
    ckpt = tf.train.get_checkpoint_state(path)
    all_ckpts = ckpt.all_model_checkpoint_paths
    ckpts = [str(c) for c in all_ckpts]
    return ckpts


def get_iter(ckpt):
    from pylib import util
    iter_ = int(util.str.find_all(ckpt, '.ckpt-\d+')[0].split('-')[-1])
    return iter_


def get_init_fn(checkpoint_path, train_dir, ignore_missing_vars=False,
                checkpoint_exclude_scopes=None, model_name=None, checkpoint_model_scope=None):
    """
    code from github/SSD-tensorflow/tf_utils.py
    Returns a function run by the chief worker to warm-start the training.
    Note that the init_fn is only run when initializing the model during the very
    first global step.

    checkpoint_path: the checkpoint to be restored
    train_dir: the directory where checkpoints are stored during training.
    ignore_missing_vars: if False and there are variables in the model but not in the checkpoint, an error will be raised.
    checkpoint_model_scope and model_name: if the root scope of checkpoints and the model in session is different,
            (but the sub-scopes are all the same), specify them clearly
    checkpoint_exclude_scopes: variables to be excluded when restoring from checkpoint_path.
    Returns:
      An init function run by the supervisor.
    """
    from pylib import util
    if util.str.is_none_or_empty(checkpoint_path):
        return None
    # Warn the user if a checkpoint exists in the train_dir. Then ignore.
    if tf.train.latest_checkpoint(train_dir):
        tf.logging.info(
            'Ignoring --checkpoint_path because a checkpoint already exists in %s'
            % train_dir)
        return None

    exclusions = []
    if checkpoint_exclude_scopes:
        exclusions = [scope.strip()
                      for scope in checkpoint_exclude_scopes.split(',')]

    # TODO(sguada) variables.filter_variables()
    variables_to_restore = []
    for var in slim.get_model_variables():
        excluded = False
        for exclusion in exclusions:
            if var.op.name.startswith(exclusion):
                excluded = True
                break
        if not excluded:
            variables_to_restore.append(var)
    # Change model scope if necessary.
    if checkpoint_model_scope is not None:
        variables_to_restore = {checkpoint_model_scope + '/' + var.op.name: var for var in variables_to_restore}
        tf.logging.info('variables_to_restore: %r' % (variables_to_restore))
    checkpoint_path = get_latest_ckpt(checkpoint_path)
    tf.logging.info('Fine-tuning from %s. Ignoring missing vars: %s' % (checkpoint_path, ignore_missing_vars))
    return slim.assign_from_checkpoint_fn(
        checkpoint_path,
        variables_to_restore,
        ignore_missing_vars=ignore_missing_vars)


def get_variables_to_train(flags=None):
    """code from github/SSD-tensorflow/tf_utils.py
    Returns a list of variables to train.

    Returns:
      A list of variables to train by the optimizer.
    """
    if flags is None or flags.trainable_scopes is None:
        return tf.trainable_variables()
    else:
        scopes = [scope.strip() for scope in flags.trainable_scopes.split(',')]

    variables_to_train = []
    for scope in scopes:
        variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope)
        variables_to_train.extend(variables)
    return variables_to_train


def Print(tensor, data, msg='', file=None, mode='w'):
    from tensorflow.python.ops import control_flow_ops

    def np_print(*args):
        if util.str.contains(msg, '%'):
            message = msg % tuple(args)
        else:
            message = msg + ' %' * len(args) % tuple(args)
        if file is not None:
            file_path = util.io.get_absolute_path(file)
            print('writting message to file(%s):' % (file_path), message)
            with open(file_path, mode) as f:
                print(message, file=f)
        else:
            print(message)

    return control_flow_ops.with_dependencies([tf.py_func(np_print, data, [])], tensor)


def get_variable_names_in_checkpoint(path, return_shapes=False, return_reader=False):
    """
    Args:
        path: the path to training directory containing checkpoints,
            or path to checkpoint
    Return:
        a list of variable names in the checkpoint
    """
    ckpt = get_latest_ckpt(path)
    ckpt_reader = tf.train.NewCheckpointReader(ckpt)
    ckpt_vars = ckpt_reader.get_variable_to_shape_map()
    names = [var for var in ckpt_vars]
    if return_shapes:
        return names, ckpt_vars

    def get(name):
        return ckpt_reader.get_tensor(name)

    if return_reader:
        return names, get
    return names


def min_area_rect(xs, ys):
    rects = tf.py_func(util.img.min_area_rect, [xs, ys], xs.dtype)
    rects.set_shape([None, 5])
    return rects


def gpu_config(config=None, allow_growth=None, gpu_memory_fraction=None):
    if config is None:
        config = tf.ConfigProto()

    if allow_growth is not None:
        config.gpu_options.allow_growth = allow_growth

    if gpu_memory_fraction is not None:
        config.gpu_options.per_process_gpu_memory_fraction = gpu_memory_fraction

    return config


def wait_for_checkpoint(path):
    from tensorflow.contrib.training.python.training import evaluation
    return evaluation.checkpoints_iterator(path)


def focal_loss(labels, logits, gamma=2.0, alpha=0.75, normalize=True):
    labels = tf.where(labels > 0, tf.ones_like(labels), tf.zeros_like(labels))
    labels = tf.cast(labels, tf.float32)
    probs = tf.sigmoid(logits)
    CE = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits)

    alpha_t = tf.ones_like(logits) * alpha
    alpha_t = tf.where(labels > 0, alpha_t, 1.0 - alpha_t)
    probs_t = tf.where(labels > 0, probs, 1.0 - probs)

    focal_matrix = alpha_t * tf.pow((1.0 - probs_t), gamma)
    fl = focal_matrix * CE

    fl = tf.reduce_sum(fl)
    if normalize:
        # n_pos = tf.reduce_sum(labels)
        # fl = fl / tf.cast(n_pos, tf.float32)
        total_weights = tf.stop_gradient(tf.reduce_sum(focal_matrix))
        fl = fl / total_weights
    return fl


def focal_loss_layer_initializer(sigma=0.01, pi=0.01):
    import numpy as np
    b0 = - np.log((1 - pi) / pi)
    return tf.random_normal_initializer(stddev=sigma), tf.constant_initializer(b0)


def sum_gradients(clone_grads, do_summary=False):
    averaged_grads = []
    for grad_and_vars in zip(*clone_grads):
        grads = []
        var = grad_and_vars[0][1]
        try:
            for g, v in grad_and_vars:
                assert v == var
                grads.append(g)
            grad = tf.add_n(grads, name=v.op.name + '_summed_gradients')
        except:
            import pdb
            pdb.set_trace()

        averaged_grads.append((grad, v))

        if do_summary:
            tf.summary.histogram("variables_and_gradients_" + grad.op.name, grad)
            tf.summary.histogram("variables_and_gradients_" + v.op.name, v)
            tf.summary.scalar("variables_and_gradients_" + grad.op.name + \
                              '_mean/var_mean', tf.reduce_mean(grad) / tf.reduce_mean(var))
            tf.summary.scalar("variables_and_gradients_" + v.op.name + '_mean', tf.reduce_mean(var))
    return averaged_grads


def get_update_op():
    """
    Extremely important for BatchNorm
    """
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    if update_ops:
        return tf.group(*update_ops)
    return None


def rand_val(maxval=1, minval=0, dtype=tf.float32, shape=(1,)):
    return tf.random_uniform(shape=shape, minval=minval,
                             maxval=maxval, dtype=dtype)[0]


prob = rand_val


def do_nothing(*args):
    def _do_nothing():
        return args

    return _do_nothing


def prob_do(probability, fn, args):
    p = prob()

    def do():
        return fn(*args)

    return tf.cond(p < probability, do, do_nothing(*args))
