//  [python]
def mcmc_tracker(flow: spt.DataFlow, baseline, model, input_x, input_u, mask, max_iter, total_window_num,
                 window_length, x_dim, mask_last=False, pos_mask=False, use_rand_mask=False, n_mc_chain=1):
    # the baseline is the avg total score in a window on training set.
    session = spt.utils.get_default_session_or_error()
    last_x = tf.placeholder(dtype=tf.float32, shape=[None, window_length, x_dim], name='last_x')

    x_r = masked_reconstruct(model.reconstruct, last_x, input_u, mask)
    score, recons_mean, recons_std = model.get_score(x_embed=x_r, x_eval=input_x, u=input_u)
    tot_score = tf.reduce_sum(tf.multiply(score, tf.cast((1-mask), score.dtype)))

    def avg_multi_chain(x, n_chain):
        shape = (-1,) + (n_chain,) + x.shape[1:]
        return np.mean(x.reshape(shape), axis=1)

    res = {}
    with flow.threaded(5) as flow:
        for batch_x, batch_u, batch_score, batch_ori_recons, batch_ori_std, batch_idx \
                in tqdm(flow, unit='step', total=total_window_num, ascii=True):
            batch_idx = batch_idx[0]
            res[batch_idx] = {'x': [batch_x], 'recons': [batch_ori_recons], 'std': [batch_ori_std], 'score': [batch_score],
                              'K': [0], 'iter': [-1], 'mask': [np.zeros(shape=batch_x.shape)],
                              'total_score': [np.mean(np.sum(batch_score, axis=-1))]}
            best_score = batch_score
            best_total_score = np.mean(np.sum(batch_score, axis=-1))
            best_K = 0
            if pos_mask:
                pos_scores = np.mean(batch_score, axis=0)   # (window, x_dim)
                sorted_pos_idx = np.argsort(pos_scores, axis=None)
                potential_dim_num = np.sum((pos_scores < (baseline/(x_dim*window_length))).astype(np.int32))
            else:
                dim_scores = np.mean(batch_score, axis=(-2,-3))     # (x_dim, )
                sorted_dim_idx = np.argsort(dim_scores)
                potential_dim_num = np.sum((dim_scores < (baseline/(x_dim*window_length))).astype(np.int32))    # num of dims whose avg score < baseline

            if potential_dim_num > 0:
                K_init = max(potential_dim_num//5, 1)
                K_inc = max(potential_dim_num//10, 1)
            else:
                res[batch_idx]['best_score'] = best_score
                res[batch_idx]['best_total_score'] = best_total_score
                res[batch_idx]['best_K'] = best_K
                continue
            if use_rand_mask:
                rand_x = np.random.random(size=batch_x.shape)
            if pos_mask:
                max_K = x_dim * window_length
            else:
                max_K = x_dim
            for K in range(K_init, min(potential_dim_num+1, max_K), K_inc):
                if pos_mask:
                    mask_idx = sorted_pos_idx[:K]
                    batch_mask = np.zeros(shape=batch_x.shape)
                    batch_mask = batch_mask.reshape([batch_x.shape[0], -1])
                    batch_mask[:, mask_idx] = 1
                    batch_mask = batch_mask.reshape(batch_x.shape)
                else:
                    mask_idx = sorted_dim_idx[:K]
                    batch_mask = np.zeros(shape=batch_x.shape)
                    batch_mask[:, :, mask_idx] = 1
                if mask_last:
                    batch_mask[:, -1, :] = 1

                batch_last_x = batch_x
                if use_rand_mask:
                    batch_last_x = np.where(batch_mask.astype(np.bool), rand_x, batch_last_x)
                if n_mc_chain > 1:
                    init_x = np.repeat(batch_x, n_mc_chain, axis=0)
                    init_u = np.repeat(batch_u, n_mc_chain, axis=0)
                    init_mask = np.repeat(batch_mask, n_mc_chain, axis=0)
                    init_last_x = np.repeat(batch_last_x, n_mc_chain, axis=0)
                for i in range(max_iter):
                    if n_mc_chain > 1:
                        x_mc, x_recons, x_std, x_score, x_tot_score = \
                            session.run([x_r, recons_mean, recons_std, score, tot_score],
                                        feed_dict={input_x: init_x, input_u: init_u, mask: init_mask,
                                                   last_x: init_last_x})
                        init_last_x = x_mc
                        x_mc = avg_multi_chain(x_mc, n_mc_chain)
                        x_recons = avg_multi_chain(x_recons, n_mc_chain)
                        x_std = avg_multi_chain(x_std, n_mc_chain)
                        x_score = avg_multi_chain(x_score, n_mc_chain)
                        x_tot_score = float(x_tot_score) / float(n_mc_chain)
                    else:
                        x_mc, x_recons, x_std, x_score, x_tot_score = \
                            session.run([x_r, recons_mean, recons_std, score, tot_score],
                                        feed_dict={input_x: batch_x, input_u: batch_u, mask: batch_mask, last_x: batch_last_x})
                        batch_last_x = x_mc
                    total_score = float(x_tot_score) / (window_length * x_dim - np.sum(batch_mask)) / batch_x.shape[0] * x_dim
                    res[batch_idx]['x'].append(x_mc)
                    res[batch_idx]['recons'].append(x_recons)
                    res[batch_idx]['std'].append(x_std)
                    res[batch_idx]['score'].append(x_score)
                    res[batch_idx]['K'].append(K)
                    res[batch_idx]['iter'].append(i)
                    res[batch_idx]['mask'].append(batch_mask)
                    res[batch_idx]['total_score'].append(total_score)

                last_score = res[batch_idx]['total_score'][-1]
                if last_score >= best_total_score:
                    best_total_score = last_score
                    best_score = res[batch_idx]['score'][-1]
                    best_K = res[batch_idx]['K'][-1]

                if best_total_score >= (baseline/window_length):
                    break
            res[batch_idx]['best_score'] = best_score
            res[batch_idx]['best_total_score'] = best_total_score
            res[batch_idx]['best_K'] = best_K
    return res