import numpy as np

from libensemble.message_numbers import STOP_TAG, PERSIS_STOP, FINISHED_PERSISTENT_GEN_TAG, EVAL_GEN_TAG
from libensemble.tools.persistent_support import PersistentSupport


def persistent_updater_after_likelihood(H, persis_info, gen_specs, libE_info):
    """
    """
    ub = gen_specs['user']['ub']
    lb = gen_specs['user']['lb']
    n = len(lb)
    subbatch_size = gen_specs['user']['subbatch_size']
    num_subbatches = gen_specs['user']['num_subbatches']
    ps = PersistentSupport(libE_info, EVAL_GEN_TAG)

    # Receive information from the manager (or a STOP_TAG)
    batch = -1
    tag = None
    w = np.nan
    while tag not in [STOP_TAG, PERSIS_STOP]:
        batch += 1
        H_o = np.zeros(subbatch_size*num_subbatches, dtype=gen_specs['out'])
        if np.all(~np.isnan(w)):
            H_o['weight'] = w
        for j in range(num_subbatches):
            for i in range(subbatch_size):
                row = subbatch_size*j + i
                H_o['x'][row] = persis_info['rand_stream'].uniform(lb, ub, (1, n))
                H_o['subbatch'][row] = j
                H_o['batch'][row] = batch
                H_o['prior'][row] = np.random.randn()
                H_o['prop'][row] = np.random.randn()

        # Send data and get next assignment
        tag, Work, calc_in = ps.send_recv(H_o)
        if calc_in is not None:
            w = H_o['prior'] + calc_in['like'] - H_o['prop']

    return H_o, persis_info, FINISHED_PERSISTENT_GEN_TAG
