import tensorflow as tf

def contra_traj_lossV9(pred, tgt, horizon=12, margin=1):
    horizon_pred, horizon_tgt = horizon_sumV1(pred, horizon), horizon_sumV1(tgt, horizon)
    # horizon_pred, horizon_tgt = horizon_sumV2(pred, tgt, horizon)

    pred_flat1, pred_flat2 = tf.reshape(horizon_pred, [-1, 1]), tf.reshape(horizon_pred, [1, -1])
    tgt_flat1, tgt_flat2 = tf.reshape(horizon_tgt, [-1, 1]), tf.reshape(horizon_tgt, [1, -1])
    tgt_dif = tgt_flat1 - tgt_flat2
    pred_dif = pred_flat1 - pred_flat2
    geq = tf.cast(tgt_dif > 0, tf.bool)
    # tgt_posi_dif = tf.where(geq, tgt_dif, -tgt_dif)
    pred_posi_dif = tf.where(geq, pred_dif, -pred_dif)
    loss = tf.maximum(0.,  margin-pred_posi_dif)

    cstr_pct = tf.math.count_nonzero(loss, dtype=tf.float32) / tf.cast(tf.reduce_prod(tf.shape(loss)), tf.float32)
    final_loss = tf.reduce_mean(loss)
    return final_loss, cstr_pct

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()
