import tensorflow as tf


def contra_traj_lossV4(pred, tgt, horizon=12, resample=1, hard_ratio=1.0):
    horizon_pred = horizon_sumV1(pred, horizon)
    horizon_tgt = horizon_sumV1(tgt, horizon)

    pred_flat = tf.reshape(horizon_pred, [-1])
    tgt_flat = tf.reshape(horizon_tgt, [-1])
    batch = tf.stack([pred_flat, tgt_flat], 1)
    sample_func = sample_pair(batch)

    def sample_compute(_):
        pairs = sample_func()
        loss = compute_contra_loss(*pairs, hard_ratio=hard_ratio)

        pct = tf.math.count_nonzero(loss, dtype=tf.float32) / tf.size(loss, out_type=tf.float32)
        p = tf.cond(tf.random_uniform((), dtype=tf.float32) < 1e-4,
                    lambda: tf.print('csrt acc ', [pct]),
                    lambda: tf.no_op())
        with tf.control_dependencies([p]):
            return tf.reduce_mean(loss)

    loss = tf.map_fn(fn=lambda inp: sample_compute(inp), elems=tf.range(resample), dtype=tf.float32,
                     parallel_iterations=32)
    final_loss = tf.reduce_mean(loss)

    return final_loss


