from random import sample
from reinforces.reinforce_utils import *
from reinforces import rephraser
from reinforces.ReplayBuffer import SharedReplayBuffer, ReplayBuffer
from reinforces.trans_env import Translate_Env
from src.utils.logging import *
from src.utils.common_utils import *
from src.data.dataset import Dataset, TextLineDataset, ZipDataset
from src.data.data_iterator import DataIterator
from torch.utils.tensorboard import SummaryWriter

import nltk.translate.bleu_score as bleu
import argparse
import math
import torch
import numpy as np
import torch.multiprocessing as _mp

os.system('export PYTHONWARNINGS="ignore:semaphore_tracker:UserWarning" ')

parser = argparse.ArgumentParser()
parser.add_argument("--n", type=int, default=1,
                    help="parallel attacker process (default as 1)")
parser.add_argument("--config_path", type=str,
                    default="./configs/cwmt_zh2en_reinforce.yaml",
                    help="the path to reinforcement config file.")
parser.add_argument("--save_to", type=str, default="./reinforces/reinforce_tf_zh2en_log",
                    help="the path for model-saving and log saving.")
parser.add_argument("--action_roll_steps", type=int, default=35,
                    help="training rolling steps (default as 35)")
parser.add_argument("--max_episode_lengths", type=int, default=200,
                    help="maximum steps for attack (default as 200)")
parser.add_argument("--max_episodes", type=int, default=5000000,
                    help="maximum environment episode for learning (default as 500k)")
parser.add_argument("--use_gpu", action="store_true", default=False,
                    help="Whether to use GPU.(default as false)")
parser.add_argument("--seed", type=int, default=1,
                    help="random seed (default as 1)")


def run():
    # default actor threads as 1
    os.environ["OMP_NUM_THREADS"] = "1"
    mp = _mp.get_context('spawn')

    args = parser.parse_args()
    if not os.path.exists(args.save_to):
        os.mkdir(args.save_to)
    # load reinforce configs
    with open(args.config_path, "r") as f, \
            open(os.path.join(args.save_to, "current_reinforce.yaml"), "w") as current_configs:
        INFO("load reinforce configures")
        configs = yaml.load(f, Loader=yaml.FullLoader)
        yaml.dump(configs, current_configs)
        reinforce_configs = configs["reinforce_configs"]
    agent_configs = configs["agent_configs"]
    rephraser_model_configs = agent_configs["rephraser_model_configs"]
    rephraser_optimizer_configs = agent_configs["rephraser_optimizer_configs"]
    annunciator_configs = configs["annunciator_configs"]

    # the Global variable of  USE_GPU is mainly used for environments
    GlobalNames.SEED = reinforce_configs["seed"]
    GlobalNames.USE_GPU = args.use_gpu
    torch.manual_seed(GlobalNames.SEED)

    # build vocabulary and data iterator for env
    with open(reinforce_configs["victim_configs"], "r") as victim_f:
        victim_configs = yaml.load(victim_f, Loader=yaml.FullLoader)
    data_configs = victim_configs["data_configs"]
    src_vocab = Vocabulary(**data_configs["vocabularies"][0])
    trg_vocab = Vocabulary(**data_configs["vocabularies"][1])
    data_set = ZipDataset(
        TextLineDataset(data_path=data_configs["train_data"][0],
                        vocabulary=src_vocab, max_len=data_configs["max_len"][0]),
        TextLineDataset(data_path=data_configs["train_data"][1],
                        vocabulary=trg_vocab, max_len=data_configs["max_len"][1]),
        shuffle=reinforce_configs["shuffle"]
    )  # we build the parallel data sets and iterate inside a thread
    # collect range of action space:
    _, _, limit_dist = load_or_extract_near_vocab(
        config_path=reinforce_configs["victim_configs"],
        model_path=reinforce_configs["victim_model"],
        init_perturb_rate=reinforce_configs["init_perturb_rate"],
        save_to=os.path.join(args.save_to, "near_vocab"),
        save_to_full=os.path.join(args.save_to, "full_near_vocab"),
        top_reserve=12,
        emit_as_id=True, use_max_dist=True)
    # build global SACAgent for the final policy (on cpu)
    global_agent = rephraser.SACAgent(
        device="cpu",
        d_word_vec=victim_configs["model_configs"]["d_word_vec"],
        d_model=rephraser_model_configs["d_model"],
        limit_dist=limit_dist,
        dropout=rephraser_model_configs["dropout"],
        learnable_temperature=rephraser_model_configs["learnable_temperature"],
        init_temperature=rephraser_model_configs["init_temperature"],
        rephraser_optimizer_configs=rephraser_optimizer_configs,
        save_to=args.save_to,
        num_kept_checkpoints=reinforce_configs["num_kept_checkpoints"]
    )
    # load global ckp (only for the AC parameters) if needed
    global_step = global_agent.load_model()
    print("global_step:", global_step)
    if global_step != 0:
        INFO("restarting at step %d"%global_step)
    else:  # save the initial model
        global_agent.save_model(global_step)
    global_summary_writer = SummaryWriter(
        log_dir=os.path.join(args.save_to, "global_summary"))
    global_replay_buffer = SharedReplayBuffer(
        max_sen_len=data_configs["max_len"][0],
        state_dim=victim_configs["model_configs"]["d_word_vec"],
        action_dim=victim_configs["model_configs"]["d_word_vec"],
        capacity=reinforce_configs["replay_buffer_capacity"])
    # test_for_throughput(global_replay_buffer)
    # make global objects shared memory
    global_agent.share_memory()
    global_replay_buffer.share_memory()

    # collect available devices and distribute env on the available gpu
    if args.use_gpu:
        device = "cuda"
        devices = []
        for i in range(torch.cuda.device_count()):
            devices += ["cuda:%d" % i]
        print("available gpus:", devices)
    else:
        device = "cpu"
        devices = [device]

    # initialize global parameters for the current training trial
    global_step_lock = mp.Lock()
    global_step_counter = mp.Value("i", global_step)  # "i is the type code for c_int"

    # train_thread(0, device, args,
    #              reinforce_configs, annunciator_configs,
    #              src_vocab, trg_vocab, data_set,
    #              global_agent, global_replay_buffer,
    #              global_step_counter, global_step_lock,
    #              agent_configs)

    # valid_thread(device, args,
    #              reinforce_configs, annunciator_configs,
    #              src_vocab, trg_vocab, data_set,
    #              global_agent, global_replay_buffer,
    #              global_step_counter, global_step_lock,
    #              agent_configs)

    # build multi thread for learning and validation
    process = []
    for rank in range(args.n):
        print("initialize training thread on cuda:%d" % (rank+1))
        p=mp.Process(
            target=train_thread,
            args=(rank, "cuda:%d"%(rank+1), args,
                  reinforce_configs, annunciator_configs,
                  src_vocab, trg_vocab, data_set,
                  global_agent, global_replay_buffer,
                  global_step_counter, global_step_lock,
                  agent_configs)
        )
        p.start()
        process.append(p)
    # run the dev thread for initiation
    print("initialize dev thread on cuda:0")
    p = mp.Process(
        target=valid_thread,
        args=("cuda:0", args,
              reinforce_configs, annunciator_configs,
              src_vocab, trg_vocab, data_set,
              global_agent, global_replay_buffer,
              global_step_counter, global_step_lock,
              agent_configs)
    )
    p.start()
    process.append(p)
    for p in process:
        p.join()


def train_thread(rank, device, args,
                 reinforce_configs, annunciator_configs,
                 src_vocab, trg_vocab, data_set,
                 global_SACAgent, global_replay_buffer,
                 global_step_counter, global_step_lock,
                 local_agent_configs):
    """
    build training thread for a local SACAgent on gpu device
    provide parameter soft-updates for the global_models.
    """
    GlobalNames.USE_GPU = args.use_gpu
    GlobalNames.SEED = reinforce_configs["seed"]
    torch.manual_seed(GlobalNames.SEED + rank)
    # build local SACAgent
    rephraser_model_configs = local_agent_configs["rephraser_model_configs"]
    rephraser_optimizer_configs = local_agent_configs["rephraser_optimizer_configs"]
    with open(reinforce_configs["victim_configs"], "r") as victim_f:
        victim_configs = yaml.load(victim_f, Loader=yaml.FullLoader)

    _, _, limit_dist = load_or_extract_near_vocab(
        config_path=reinforce_configs["victim_configs"],
        model_path=reinforce_configs["victim_model"],
        init_perturb_rate=reinforce_configs["init_perturb_rate"],
        save_to=os.path.join(args.save_to, "near_vocab"),
        save_to_full=os.path.join(args.save_to, "full_near_vocab"),
        top_reserve=12,
        emit_as_id=True)
    local_agent = rephraser.SACAgent(
        device=device,
        d_word_vec=victim_configs["model_configs"]["d_word_vec"],
        d_model=rephraser_model_configs["d_model"],
        limit_dist=limit_dist,
        dropout=rephraser_model_configs["dropout"],
        learnable_temperature=rephraser_model_configs["learnable_temperature"],
        init_temperature=rephraser_model_configs["init_temperature"],
        rephraser_optimizer_configs=rephraser_optimizer_configs,
        rank=rank,
        save_to=os.path.join(args.save_to),
        num_kept_checkpoints=reinforce_configs["num_kept_checkpoints"]
    )
    local_summary_writer = SummaryWriter(
        log_dir=os.path.join(args.save_to, "train_env%d" % rank))

    # # this is the secondary buffer
    # local_replay_buffer = ReplayBuffer(
    #     max_sen_len=global_replay_buffer.max_sen_len,
    #     state_dim=victim_configs["model_configs"]["d_word_vec"],
    #     action_dim=victim_configs["model_configs"]["d_word_vec"],
    #     capacity=max(reinforce_configs["replay_buffer_capacity"]/20, 10000)
    # )

    # build environment (include annunciator update settings)
    reinforce_iterator = DataIterator(
        dataset=data_set, batch_size=reinforce_configs["batch_size"],
        use_bucket=True, buffer_size=reinforce_configs["buffer_size"],
        numbering=True)
    local_data_iterator = reinforce_iterator.build_generator()
    local_env = Translate_Env(
        reinforce_configs=reinforce_configs,
        annunciator_configs=annunciator_configs,
        src_vocab=src_vocab, trg_vocab=trg_vocab,
        data_iterator=local_data_iterator,
        save_to=args.save_to, device=device)

    episode_count = 0  # a batch of sentences as an episode & learning episodes
    local_step = global_step_counter.value  # initiate local agent update steps
    patience_t = annunciator_configs["patience"]
    trust_acc = 0.5
    while True:  # infinite loop of data set iterator (epoch)
        # we will continue with a new iterator with refreshed environments
        # whenever the last iterator breaks with "StopIteration"
        rephraser_iterator = reinforce_iterator.build_generator()
        local_env.reset_data_iter(rephraser_iterator)

        try:
            while True: # loop training episodes
                # the environment will be initiated by an actor as self learning
                local_env.reset()
                states, _, _, _ = local_env.get_state()
                # x_emb = local_env.reset(local_agent.actor)
                # x_emb = torch.from_numpy(x_emb)
                # if device != "cpu":
                #     x_emb = x_emb.to(device)
                annunciator_base_step = local_step
                episode_length = 0  # the rollout steps
                episode_rewards = np.array([0.] * states.shape[0])
                done = True  # whether the episode is finished & should reset episode with new batch

                while episode_length <= args.max_episode_lengths:  # loop TD learning rollouts
                    # check for the environment updates using current **global** agent
                    if episode_count % local_agent_configs["rephraser_update_steps"] == 0:
                        """ stop criterion:
                        until mse reaches the bound within patience.
                        otherwise the training thread stops
                        """
                        try:  # update environment
                            INFO("Update environment")
                            # sync from the global agent to avoid unnecessary locking.
                            local_agent.sync_from(global_SACAgent)
                            annunciator_base_step, trust_acc = local_env.update_annunciator(
                                local_agent,
                                annunciator_base_step,
                                min_update_steps=annunciator_configs["valid_freq"],
                                max_update_steps=annunciator_configs["annunciator_update_steps"],
                                accuracy_bound=annunciator_configs["acc_bound"],
                                overall_update_weight=1-trust_acc,
                                summary_writer=local_summary_writer)
                            # global_SACAgent.to("cpu").train()
                            # switch back to training mode
                            local_agent = local_agent.to(device).train()
                        except StopIteration:
                            INFO("finish one training epoch, reset data_iterator")
                            break
                        annunciator_base_step += 1  # a flag to label the scorer updates
                        if trust_acc < annunciator_configs["d_converged_bound"]:
                            # GAN target reached, scorer has reached its limit.
                            patience_t -= 1
                            INFO("scorer reached GAN convergence bound: %d times" % patience_t)
                        else:  # reset patience if scorer is refreshed
                            patience_t = annunciator_configs["patience"]

                    if patience_t == 0 or episode_count > args.max_episodes:
                        WARN("maximum patience & training step reached. Thread stop")
                        break

                    local_agent = local_agent.to(device).train()  # switch the agent to training mode
                    if done:  # can't create SARSA with only one step, start a new episode
                        INFO("sync from global agent")
                        local_agent.sync_from(global_SACAgent)
                        local_agent = local_agent.to(device).train()
                        local_env.reset()
                        done = False

                    states, masks, rephrase_positions, _ = local_env.get_state()
                    roll_out_rewards = np.array([0] * states.shape[0])  # n-step TD learning accumulation
                    for i in range(args.action_roll_steps):
                        # sample actions for buffer collection (start with fully random)
                        if global_replay_buffer.size() < local_agent_configs["rephraser_learning_batch"]:
                            # fast random actions in the **action_space** as initial exploration
                            actions = local_env.limit_dist * np.random.sample([states.shape[0], states.shape[2]])
                            actions = torch.tensor(actions,dtype=torch.float32, device=device)
                        else:  # explore and exploit
                            local_agent = local_agent.train(False)  # set to validation mode  .detach().cpu().numpy()
                            label_emb = rephraser.slice_by_indices(states, rephrase_positions, device=device)
                            actions, _ = local_agent.actor.sample_normal(
                                states, x_pad_indicator=1.-masks, label_emb=label_emb, reparamization=False)
                        # able to update agent with current buffer
                        if global_replay_buffer.size() >= local_agent_configs["rephraser_learning_batch"]:
                            local_agent = local_agent.train()
                            local_agent.update_local_net(
                                local_agent_configs,
                                global_replay_buffer,
                                target_critic=global_SACAgent.critic.to(device),
                                update_step=local_step,
                                discount_factor=reinforce_configs["gamma"],
                                summary_writer=local_summary_writer,
                                update_trust_region=trust_acc
                            )
                            global_SACAgent.critic.to("cpu")

                            if local_step % local_agent_configs["soft_update_freq"] == 0:
                                with global_step_lock:  # update global model
                                    INFO("update global model by local model on: %s"%device)
                                    global_step_counter.value += 1
                                    local_agent.soft_update_target_net(global_SACAgent, reinforce_configs["tau"])

                                    if global_step_counter.value % reinforce_configs["save_freq"] == 0:
                                        INFO("save global model %d by device: %s"%(global_step_counter.value, device))
                                        with global_SACAgent.soft_update_lock:
                                            global_SACAgent.save_model(global_step_counter.value)
                                            # save final global model (as object)
                                            if trust_acc < annunciator_configs["d_converged_bound"]:
                                                global_SACAgent.save_model(save_to_final=args.save_to)

                            local_step += 1
                            INFO("finish local update by %2f"%trust_acc)

                        _, rewards, _ = local_env.step(actions)
                        new_states, _, new_rephrase_positions, survival_signals = local_env.get_state()

                        roll_out_rewards = [roll_out_rewards[t]+rewards[t] for t in range(states.shape[0])]
                        episode_rewards = [episode_rewards[t]+rewards[t] for t in range(states.shape[0])]
                        survival_signals_cpu = survival_signals.cpu().numpy().tolist()
                        survival_and_no_max = np.array([0 if episode_length == local_env.sent_len[t]
                                                        else survival_signals_cpu[t][0]
                                                        for t in range(states.shape[0])])
                        INFO("step %d of episode:%d, r:%f, device:%s"%(i, episode_count, np.array(episode_rewards).mean(), actions.device))
                        # add the tensor variables to the replay buffer,
                        # learn by rollout rewards or episode reward accumulated

                        # update the overall signals for rollout
                        if 1 in survival_signals:
                            done = False
                        else:
                            done = True
                        if local_env.index == states.shape[1]-1:
                            # note that the environment index will increase beyond the shape
                            done = True
                        if done:  # break from roll-out
                            break

                            # print(states.shape, masks.shape, actions.shape, rephrase_positions.shape, torch.tensor(roll_out_rewards).shape,
                        #   survival_signals.shape, torch.tensor(survival_and_no_max).shape)
                        INFO("write buffer begin")
                        #cleanse unecessary experience to save time
                        invalid_sample = []
                        if local_env.index>min(local_env.sent_len):
                            for i in range(states.shape[0]):  # for every record
                                if local_env.index > local_env.sent_len[i]:
                                    invalid_sample.append(i)

                        global_replay_buffer.add(
                            states=states,
                            masks=masks,
                            actions=actions,
                            rephrase_positions=rephrase_positions,
                            rewards=torch.tensor(rewards, dtype=torch.float32).unsqueeze(dim=-1),
                            survival_signals=survival_signals,
                            survive_and_no_max=torch.tensor(survival_and_no_max, dtype=torch.float32).unsqueeze(dim=-1),
                            invalid_sample=invalid_sample
                        )
                        INFO("write buffer end, size=%d"%global_replay_buffer.size())
                        episode_length += 1
                        states = new_states
                        rephrase_positions = new_rephrase_positions

                    if done:  # break from current episodes
                        episode_count += 1
                        break
                if episode_count > args.max_episodes:
                    return
        except StopIteration:
            INFO("finish one training epoch, reset data_iterator")
            continue


def valid_thread(device, args,
                 reinforce_configs, annunciator_configs,
                 src_vocab, trg_vocab, data_set,
                 global_SACAgent, global_replay_buffer,
                 global_step_counter, global_step_lock,
                 local_agent_configs):
    """
    build validation thread to check global SACAgent tests and update buffer
    """
    GlobalNames.USE_GPU = args.use_gpu
    GlobalNames.SEED = reinforce_configs["seed"]
    torch.manual_seed(GlobalNames.SEED)

    # build local agent and sync from the global model to avoid deadlock
    rephraser_model_configs = local_agent_configs["rephraser_model_configs"]
    rephraser_optimizer_configs = local_agent_configs["rephraser_optimizer_configs"]
    with open(reinforce_configs["victim_configs"], "r") as victim_f:
        victim_configs = yaml.load(victim_f, Loader=yaml.FullLoader)
    _, _, limit_dist = load_or_extract_near_vocab(
        config_path=reinforce_configs["victim_configs"],
        model_path=reinforce_configs["victim_model"],
        init_perturb_rate=reinforce_configs["init_perturb_rate"],
        save_to=os.path.join(args.save_to, "near_vocab"),
        save_to_full=os.path.join(args.save_to, "full_near_vocab"),
        top_reserve=12,
        emit_as_id=True)
    local_agent = rephraser.SACAgent(
        device=device,
        d_word_vec=victim_configs["model_configs"]["d_word_vec"],
        d_model=rephraser_model_configs["d_model"],
        limit_dist=limit_dist,
        dropout=rephraser_model_configs["dropout"],
        learnable_temperature=rephraser_model_configs["learnable_temperature"],
        init_temperature=rephraser_model_configs["init_temperature"],
        rephraser_optimizer_configs=rephraser_optimizer_configs,
        rank="dev",
        save_to=os.path.join(args.save_to),
        num_kept_checkpoints=reinforce_configs["num_kept_checkpoints"]
    )
    local_summary_writer = SummaryWriter(
        log_dir=os.path.join(args.save_to, "valid_env"))

    # build an independent dev env (annunciator)
    reinforce_iterator = DataIterator(
        dataset=data_set, batch_size=reinforce_configs["batch_size"],
        use_bucket=True, buffer_size=reinforce_configs["buffer_size"],
        numbering=True)
    local_data_iterator = reinforce_iterator.build_generator()
    local_env = Translate_Env(
        reinforce_configs=reinforce_configs,
        annunciator_configs=annunciator_configs,
        src_vocab=src_vocab, trg_vocab=trg_vocab,
        data_iterator=local_data_iterator,
        save_to=args.save_to, device=device)

    local_step = 0
    episode_count = 0  # determines the global sync
    update_signal = False # whether update the environment annunciator
    while True: # inifinite loop of dataset iterator
        rephraser_iterator = reinforce_iterator.build_generator()
        local_env.reset_data_iter(rephraser_iterator)
        try:
            annunciator_base_step = local_step
            while True:  # inifinite loop of the episodes
                # check for the environment annunciator updates;
                with global_step_lock:
                    if global_step_counter.value % (local_agent_configs["rephraser_update_steps"]) == 0:
                        update_signal=True
                if update_signal:
                    try:
                        INFO("update valid env annunciator")
                        update_signal=False
                        annunciator_base_step, trust_acc = local_env.update_annunciator(
                            local_agent,
                            annunciator_base_step,
                            min_update_steps=annunciator_configs["valid_freq"],
                            max_update_steps=annunciator_configs["annunciator_update_steps"],
                            accuracy_bound=annunciator_configs["acc_bound"],
                            summary_writer=local_summary_writer)
                        # global_SACAgent.to("cpu").train()
                        local_agent = local_agent.to(device).train(False)
                    except StopIteration:
                        INFO("finish one training epoch, reset data_iterator")
                        break
                # prepare a new episode
                local_agent.sync_from(global_SACAgent)
                local_agent = local_agent.to(device).train(False)
                local_env.reset()
                states, masks, rephrase_positions, _ = local_env.get_state()
                episode_rewards = np.array([0.] * states.shape[0])
                episode_length = 0
                done = False
                while not done:
                    states, masks, rephrase_positions, _ = local_env.get_state()
                    label_embs = rephraser.slice_by_indices(states, rephrase_positions, device=device)
                    # directly use the expectation as validation
                    mu, _ = local_agent.actor.forward(
                        states, x_pad_indicator=1.-masks, label_emb=label_embs
                    )

                    actions = mu * local_agent.actor.action_range

                    _, rewards, _ = local_env.step(actions)
                    episode_rewards = [episode_rewards[t]+rewards[t] for t in range(states.shape[0])]
                    new_states, _, _, survival_signals = local_env.get_state()
                    survival_signals_cpu = survival_signals.cpu().numpy().tolist()
                    survival_and_no_max = np.array([0 if episode_length == local_env.sent_len[t]
                                                    else survival_signals_cpu[t][0]
                                                    for t in range(states.shape[0])])
                    INFO("step %d survive and no max:%d, %s"%(episode_length, survival_and_no_max.sum(), actions.device))

                    if 1 in survival_signals:
                        done = False
                    else:
                        episode_count += 1
                        done = True
                    if local_env.index == states.shape[1]-1:
                        # note that the environment index will increase beyond the shape
                        done = True
                    if done:  # skip the final step
                        break
                    INFO("write buffer begin")
                    invalid_sample = []
                    if local_env.index>min(local_env.sent_len):
                        for i in range(states.shape[0]):  # for every record
                            if local_env.index > local_env.sent_len[i]:
                                invalid_sample.append(i)
                    # print(states.shape, masks.shape, actions.shape, rephrase_positions.shape, torch.tensor(roll_out_rewards).shape,
                    #   survival_signals.shape, torch.tensor(survival_and_no_max).shape)
                    global_replay_buffer.add(
                        states=states,
                        masks=masks,
                        actions=actions,
                        rephrase_positions=rephrase_positions,
                        rewards=torch.tensor(rewards, dtype=torch.float32).unsqueeze(dim=-1),
                        survival_signals=survival_signals,
                        survive_and_no_max=torch.tensor(survival_and_no_max, dtype=torch.float32).unsqueeze(dim=-1),
                        invalid_sample=invalid_sample
                    )
                    INFO("write buffer end, size=%d"%global_replay_buffer.size())
                    local_step += 1
                    episode_length += 1
                INFO("mean episode rewards on batch:%f"%np.array(episode_rewards).mean())
                local_summary_writer.add_scalar(
                    "valid_rewards",
                    scalar_value=np.array(episode_rewards).mean(), global_step=local_step)

        except StopIteration:
            INFO("finish one validation epoch, reset data_iterator")
            continue


if __name__ == "__main__":
    run()

from random import sample
from reinforces.reinforce_utils import *
from reinforces import rephraser
from reinforces.ReplayBuffer import SharedReplayBuffer, ReplayBuffer
from reinforces.trans_env import Translate_Env
from src.utils.logging import *
from src.utils.common_utils import *
from src.data.dataset import Dataset, TextLineDataset, ZipDataset
from src.data.data_iterator import DataIterator
from torch.utils.tensorboard import SummaryWriter

import nltk.translate.bleu_score as bleu
import argparse
import math
import torch
import numpy as np
import torch.multiprocessing as _mp

os.system('export PYTHONWARNINGS="ignore:semaphore_tracker:UserWarning" ')

parser = argparse.ArgumentParser()
parser.add_argument("--n", type=int, default=1,
                    help="parallel attacker process (default as 1)")
parser.add_argument("--config_path", type=str,
                    default="./configs/cwmt_zh2en_reinforce.yaml",
                    help="the path to reinforcement config file.")
parser.add_argument("--save_to", type=str, default="./reinforces/reinforce_tf_zh2en_log",
                    help="the path for model-saving and log saving.")
parser.add_argument("--action_roll_steps", type=int, default=35,
                    help="training rolling steps (default as 35)")
parser.add_argument("--max_episode_lengths", type=int, default=200,
                    help="maximum steps for attack (default as 200)")
parser.add_argument("--max_episodes", type=int, default=5000000,
                    help="maximum environment episode for learning (default as 500k)")
parser.add_argument("--use_gpu", action="store_true", default=False,
                    help="Whether to use GPU.(default as false)")
parser.add_argument("--seed", type=int, default=1,
                    help="random seed (default as 1)")


def run():
    # default actor threads as 1
    os.environ["OMP_NUM_THREADS"] = "1"
    mp = _mp.get_context('spawn')

    args = parser.parse_args()
    if not os.path.exists(args.save_to):
        os.mkdir(args.save_to)
    # load reinforce configs
    with open(args.config_path, "r") as f, \
            open(os.path.join(args.save_to, "current_reinforce.yaml"), "w") as current_configs:
        INFO("load reinforce configures")
        configs = yaml.load(f, Loader=yaml.FullLoader)
        yaml.dump(configs, current_configs)
        reinforce_configs = configs["reinforce_configs"]
    agent_configs = configs["agent_configs"]
    rephraser_model_configs = agent_configs["rephraser_model_configs"]
    rephraser_optimizer_configs = agent_configs["rephraser_optimizer_configs"]
    annunciator_configs = configs["annunciator_configs"]

    # the Global variable of  USE_GPU is mainly used for environments
    GlobalNames.SEED = reinforce_configs["seed"]
    GlobalNames.USE_GPU = args.use_gpu
    torch.manual_seed(GlobalNames.SEED)

    # build vocabulary and data iterator for env
    with open(reinforce_configs["victim_configs"], "r") as victim_f:
        victim_configs = yaml.load(victim_f, Loader=yaml.FullLoader)
    data_configs = victim_configs["data_configs"]
    src_vocab = Vocabulary(**data_configs["vocabularies"][0])
    trg_vocab = Vocabulary(**data_configs["vocabularies"][1])
    data_set = ZipDataset(
        TextLineDataset(data_path=data_configs["train_data"][0],
                        vocabulary=src_vocab, max_len=data_configs["max_len"][0]),
        TextLineDataset(data_path=data_configs["train_data"][1],
                        vocabulary=trg_vocab, max_len=data_configs["max_len"][1]),
        shuffle=reinforce_configs["shuffle"]
    )  # we build the parallel data sets and iterate inside a thread
    # collect range of action space:
    _, _, limit_dist = load_or_extract_near_vocab(
        config_path=reinforce_configs["victim_configs"],
        model_path=reinforce_configs["victim_model"],
        init_perturb_rate=reinforce_configs["init_perturb_rate"],
        save_to=os.path.join(args.save_to, "near_vocab"),
        save_to_full=os.path.join(args.save_to, "full_near_vocab"),
        top_reserve=12,
        emit_as_id=True, use_max_dist=True)
    # build global SACAgent for the final policy (on cpu)
    global_agent = rephraser.SACAgent(
        device="cpu",
        d_word_vec=victim_configs["model_configs"]["d_word_vec"],
        d_model=rephraser_model_configs["d_model"],
        limit_dist=limit_dist,
        dropout=rephraser_model_configs["dropout"],
        learnable_temperature=rephraser_model_configs["learnable_temperature"],
        init_temperature=rephraser_model_configs["init_temperature"],
        rephraser_optimizer_configs=rephraser_optimizer_configs,
        save_to=args.save_to,
        num_kept_checkpoints=reinforce_configs["num_kept_checkpoints"]
    )
    # load global ckp (only for the AC parameters) if needed
    global_step = global_agent.load_model()
    print("global_step:", global_step)
    if global_step != 0:
        INFO("restarting at step %d"%global_step)
    else:  # save the initial model
        global_agent.save_model(global_step)
    global_summary_writer = SummaryWriter(
        log_dir=os.path.join(args.save_to, "global_summary"))
    global_replay_buffer = SharedReplayBuffer(
        max_sen_len=data_configs["max_len"][0],
        state_dim=victim_configs["model_configs"]["d_word_vec"],
        action_dim=victim_configs["model_configs"]["d_word_vec"],
        capacity=reinforce_configs["replay_buffer_capacity"])
    # test_for_throughput(global_replay_buffer)
    # make global objects shared memory
    global_agent.share_memory()
    global_replay_buffer.share_memory()

    # collect available devices and distribute env on the available gpu
    if args.use_gpu:
        device = "cuda"
        devices = []
        for i in range(torch.cuda.device_count()):
            devices += ["cuda:%d" % i]
        print("available gpus:", devices)
    else:
        device = "cpu"
        devices = [device]

    # initialize global parameters for the current training trial
    global_step_lock = mp.Lock()
    global_step_counter = mp.Value("i", global_step)  # "i is the type code for c_int"

    # train_thread(0, device, args,
    #              reinforce_configs, annunciator_configs,
    #              src_vocab, trg_vocab, data_set,
    #              global_agent, global_replay_buffer,
    #              global_step_counter, global_step_lock,
    #              agent_configs)

    # valid_thread(device, args,
    #              reinforce_configs, annunciator_configs,
    #              src_vocab, trg_vocab, data_set,
    #              global_agent, global_replay_buffer,
    #              global_step_counter, global_step_lock,
    #              agent_configs)

    # build multi thread for learning and validation
    process = []
    for rank in range(args.n):
        print("initialize training thread on cuda:%d" % (rank+1))
        p=mp.Process(
            target=train_thread,
            args=(rank, "cuda:%d"%(rank+1), args,
                  reinforce_configs, annunciator_configs,
                  src_vocab, trg_vocab, data_set,
                  global_agent, global_replay_buffer,
                  global_step_counter, global_step_lock,
                  agent_configs)
        )
        p.start()
        process.append(p)
    # run the dev thread for initiation
    print("initialize dev thread on cuda:0")
    p = mp.Process(
        target=valid_thread,
        args=("cuda:0", args,
              reinforce_configs, annunciator_configs,
              src_vocab, trg_vocab, data_set,
              global_agent, global_replay_buffer,
              global_step_counter, global_step_lock,
              agent_configs)
    )
    p.start()
    process.append(p)
    for p in process:
        p.join()


def train_thread(rank, device, args,
                 reinforce_configs, annunciator_configs,
                 src_vocab, trg_vocab, data_set,
                 global_SACAgent, global_replay_buffer,
                 global_step_counter, global_step_lock,
                 local_agent_configs):
    """
    build training thread for a local SACAgent on gpu device
    provide parameter soft-updates for the global_models.
    """
    GlobalNames.USE_GPU = args.use_gpu
    GlobalNames.SEED = reinforce_configs["seed"]
    torch.manual_seed(GlobalNames.SEED + rank)
    # build local SACAgent
    rephraser_model_configs = local_agent_configs["rephraser_model_configs"]
    rephraser_optimizer_configs = local_agent_configs["rephraser_optimizer_configs"]
    with open(reinforce_configs["victim_configs"], "r") as victim_f:
        victim_configs = yaml.load(victim_f, Loader=yaml.FullLoader)

    _, _, limit_dist = load_or_extract_near_vocab(
        config_path=reinforce_configs["victim_configs"],
        model_path=reinforce_configs["victim_model"],
        init_perturb_rate=reinforce_configs["init_perturb_rate"],
        save_to=os.path.join(args.save_to, "near_vocab"),
        save_to_full=os.path.join(args.save_to, "full_near_vocab"),
        top_reserve=12,
        emit_as_id=True)
    local_agent = rephraser.SACAgent(
        device=device,
        d_word_vec=victim_configs["model_configs"]["d_word_vec"],
        d_model=rephraser_model_configs["d_model"],
        limit_dist=limit_dist,
        dropout=rephraser_model_configs["dropout"],
        learnable_temperature=rephraser_model_configs["learnable_temperature"],
        init_temperature=rephraser_model_configs["init_temperature"],
        rephraser_optimizer_configs=rephraser_optimizer_configs,
        rank=rank,
        save_to=os.path.join(args.save_to),
        num_kept_checkpoints=reinforce_configs["num_kept_checkpoints"]
    )
    local_summary_writer = SummaryWriter(
        log_dir=os.path.join(args.save_to, "train_env%d" % rank))

    # # this is the secondary buffer
    # local_replay_buffer = ReplayBuffer(
    #     max_sen_len=global_replay_buffer.max_sen_len,
    #     state_dim=victim_configs["model_configs"]["d_word_vec"],
    #     action_dim=victim_configs["model_configs"]["d_word_vec"],
    #     capacity=max(reinforce_configs["replay_buffer_capacity"]/20, 10000)
    # )

    # build environment (include annunciator update settings)
    reinforce_iterator = DataIterator(
        dataset=data_set, batch_size=reinforce_configs["batch_size"],
        use_bucket=True, buffer_size=reinforce_configs["buffer_size"],
        numbering=True)
    local_data_iterator = reinforce_iterator.build_generator()
    local_env = Translate_Env(
        reinforce_configs=reinforce_configs,
        annunciator_configs=annunciator_configs,
        src_vocab=src_vocab, trg_vocab=trg_vocab,
        data_iterator=local_data_iterator,
        save_to=args.save_to, device=device)

    episode_count = 0  # a batch of sentences as an episode & learning episodes
    local_step = global_step_counter.value  # initiate local agent update steps
    patience_t = annunciator_configs["patience"]
    trust_acc = 0.5
    while True:  # infinite loop of data set iterator (epoch)
        # we will continue with a new iterator with refreshed environments
        # whenever the last iterator breaks with "StopIteration"
        rephraser_iterator = reinforce_iterator.build_generator()
        local_env.reset_data_iter(rephraser_iterator)

        try:
            while True: # loop training episodes
                # the environment will be initiated by an actor as self learning
                local_env.reset()
                states, _, _, _ = local_env.get_state()
                # x_emb = local_env.reset(local_agent.actor)
                # x_emb = torch.from_numpy(x_emb)
                # if device != "cpu":
                #     x_emb = x_emb.to(device)
                annunciator_base_step = local_step
                episode_length = 0  # the rollout steps
                episode_rewards = np.array([0.] * states.shape[0])
                done = True  # whether the episode is finished & should reset episode with new batch

                while episode_length <= args.max_episode_lengths:  # loop TD learning rollouts
                    # check for the environment updates using current **global** agent
                    if episode_count % local_agent_configs["rephraser_update_steps"] == 0:
                        """ stop criterion:
                        until mse reaches the bound within patience.
                        otherwise the training thread stops
                        """
                        try:  # update environment
                            INFO("Update environment")
                            # sync from the global agent to avoid unnecessary locking.
                            local_agent.sync_from(global_SACAgent)
                            annunciator_base_step, trust_acc = local_env.update_annunciator(
                                local_agent,
                                annunciator_base_step,
                                min_update_steps=annunciator_configs["valid_freq"],
                                max_update_steps=annunciator_configs["annunciator_update_steps"],
                                accuracy_bound=annunciator_configs["acc_bound"],
                                overall_update_weight=1-trust_acc,
                                summary_writer=local_summary_writer)
                            # global_SACAgent.to("cpu").train()
                            # switch back to training mode
                            local_agent = local_agent.to(device).train()
                        except StopIteration:
                            INFO("finish one training epoch, reset data_iterator")
                            break
                        annunciator_base_step += 1  # a flag to label the scorer updates
                        if trust_acc < annunciator_configs["d_converged_bound"]:
                            # GAN target reached, scorer has reached its limit.
                            patience_t -= 1
                            INFO("scorer reached GAN convergence bound: %d times" % patience_t)
                        else:  # reset patience if scorer is refreshed
                            patience_t = annunciator_configs["patience"]

                    if patience_t == 0 or episode_count > args.max_episodes:
                        WARN("maximum patience & training step reached. Thread stop")
                        break

                    local_agent = local_agent.to(device).train()  # switch the agent to training mode
                    if done:  # can't create SARSA with only one step, start a new episode
                        INFO("sync from global agent")
                        local_agent.sync_from(global_SACAgent)
                        local_agent = local_agent.to(device).train()
                        local_env.reset()
                        done = False

                    states, masks, rephrase_positions, _ = local_env.get_state()
                    roll_out_rewards = np.array([0] * states.shape[0])  # n-step TD learning accumulation
                    for i in range(args.action_roll_steps):
                        # sample actions for buffer collection (start with fully random)
                        if global_replay_buffer.size() < local_agent_configs["rephraser_learning_batch"]:
                            # fast random actions in the **action_space** as initial exploration
                            actions = local_env.limit_dist * np.random.sample([states.shape[0], states.shape[2]])
                            actions = torch.tensor(actions,dtype=torch.float32, device=device)
                        else:  # explore and exploit
                            local_agent = local_agent.train(False)  # set to validation mode  .detach().cpu().numpy()
                            label_emb = rephraser.slice_by_indices(states, rephrase_positions, device=device)
                            actions, _ = local_agent.actor.sample_normal(
                                states, x_pad_indicator=1.-masks, label_emb=label_emb, reparamization=False)
                        # able to update agent with current buffer
                        if global_replay_buffer.size() >= local_agent_configs["rephraser_learning_batch"]:
                            local_agent = local_agent.train()
                            local_agent.update_local_net(
                                local_agent_configs,
                                global_replay_buffer,
                                target_critic=global_SACAgent.critic.to(device),
                                update_step=local_step,
                                discount_factor=reinforce_configs["gamma"],
                                summary_writer=local_summary_writer,
                                update_trust_region=trust_acc
                            )
                            global_SACAgent.critic.to("cpu")

                            if local_step % local_agent_configs["soft_update_freq"] == 0:
                                with global_step_lock:  # update global model
                                    INFO("update global model by local model on: %s"%device)
                                    global_step_counter.value += 1
                                    local_agent.soft_update_target_net(global_SACAgent, reinforce_configs["tau"])

                                    if global_step_counter.value % reinforce_configs["save_freq"] == 0:
                                        INFO("save global model %d by device: %s"%(global_step_counter.value, device))
                                        with global_SACAgent.soft_update_lock:
                                            global_SACAgent.save_model(global_step_counter.value)
                                            # save final global model (as object)
                                            if trust_acc < annunciator_configs["d_converged_bound"]:
                                                global_SACAgent.save_model(save_to_final=args.save_to)

                            local_step += 1
                            INFO("finish local update by %2f"%trust_acc)

                        _, rewards, _ = local_env.step(actions)
                        new_states, _, new_rephrase_positions, survival_signals = local_env.get_state()

                        roll_out_rewards = [roll_out_rewards[t]+rewards[t] for t in range(states.shape[0])]
                        episode_rewards = [episode_rewards[t]+rewards[t] for t in range(states.shape[0])]
                        survival_signals_cpu = survival_signals.cpu().numpy().tolist()
                        survival_and_no_max = np.array([0 if episode_length == local_env.sent_len[t]
                                                        else survival_signals_cpu[t][0]
                                                        for t in range(states.shape[0])])
                        INFO("step %d of episode:%d, r:%f, device:%s"%(i, episode_count, np.array(episode_rewards).mean(), actions.device))
                        # add the tensor variables to the replay buffer,
                        # learn by rollout rewards or episode reward accumulated

                        # update the overall signals for rollout
                        if 1 in survival_signals:
                            done = False
                        else:
                            done = True
                        if local_env.index == states.shape[1]-1:
                            # note that the environment index will increase beyond the shape
                            done = True
                        if done:  # break from roll-out
                            break

                            # print(states.shape, masks.shape, actions.shape, rephrase_positions.shape, torch.tensor(roll_out_rewards).shape,
                        #   survival_signals.shape, torch.tensor(survival_and_no_max).shape)
                        INFO("write buffer begin")
                        #cleanse unecessary experience to save time
                        invalid_sample = []
                        if local_env.index>min(local_env.sent_len):
                            for i in range(states.shape[0]):  # for every record
                                if local_env.index > local_env.sent_len[i]:
                                    invalid_sample.append(i)

                        global_replay_buffer.add(
                            states=states,
                            masks=masks,
                            actions=actions,
                            rephrase_positions=rephrase_positions,
                            rewards=torch.tensor(rewards, dtype=torch.float32).unsqueeze(dim=-1),
                            survival_signals=survival_signals,
                            survive_and_no_max=torch.tensor(survival_and_no_max, dtype=torch.float32).unsqueeze(dim=-1),
                            invalid_sample=invalid_sample
                        )
                        INFO("write buffer end, size=%d"%global_replay_buffer.size())
                        episode_length += 1
                        states = new_states
                        rephrase_positions = new_rephrase_positions

                    if done:  # break from current episodes
                        episode_count += 1
                        break
                if episode_count > args.max_episodes:
                    return
        except StopIteration:
            INFO("finish one training epoch, reset data_iterator")
            continue


def valid_thread(device, args,
                 reinforce_configs, annunciator_configs,
                 src_vocab, trg_vocab, data_set,
                 global_SACAgent, global_replay_buffer,
                 global_step_counter, global_step_lock,
                 local_agent_configs):
    """
    build validation thread to check global SACAgent tests and update buffer
    """
    GlobalNames.USE_GPU = args.use_gpu
    GlobalNames.SEED = reinforce_configs["seed"]
    torch.manual_seed(GlobalNames.SEED)

    # build local agent and sync from the global model to avoid deadlock
    rephraser_model_configs = local_agent_configs["rephraser_model_configs"]
    rephraser_optimizer_configs = local_agent_configs["rephraser_optimizer_configs"]
    with open(reinforce_configs["victim_configs"], "r") as victim_f:
        victim_configs = yaml.load(victim_f, Loader=yaml.FullLoader)
    _, _, limit_dist = load_or_extract_near_vocab(
        config_path=reinforce_configs["victim_configs"],
        model_path=reinforce_configs["victim_model"],
        init_perturb_rate=reinforce_configs["init_perturb_rate"],
        save_to=os.path.join(args.save_to, "near_vocab"),
        save_to_full=os.path.join(args.save_to, "full_near_vocab"),
        top_reserve=12,
        emit_as_id=True)
    local_agent = rephraser.SACAgent(
        device=device,
        d_word_vec=victim_configs["model_configs"]["d_word_vec"],
        d_model=rephraser_model_configs["d_model"],
        limit_dist=limit_dist,
        dropout=rephraser_model_configs["dropout"],
        learnable_temperature=rephraser_model_configs["learnable_temperature"],
        init_temperature=rephraser_model_configs["init_temperature"],
        rephraser_optimizer_configs=rephraser_optimizer_configs,
        rank="dev",
        save_to=os.path.join(args.save_to),
        num_kept_checkpoints=reinforce_configs["num_kept_checkpoints"]
    )
    local_summary_writer = SummaryWriter(
        log_dir=os.path.join(args.save_to, "valid_env"))

    # build an independent dev env (annunciator)
    reinforce_iterator = DataIterator(
        dataset=data_set, batch_size=reinforce_configs["batch_size"],
        use_bucket=True, buffer_size=reinforce_configs["buffer_size"],
        numbering=True)
    local_data_iterator = reinforce_iterator.build_generator()
    local_env = Translate_Env(
        reinforce_configs=reinforce_configs,
        annunciator_configs=annunciator_configs,
        src_vocab=src_vocab, trg_vocab=trg_vocab,
        data_iterator=local_data_iterator,
        save_to=args.save_to, device=device)

    local_step = 0
    episode_count = 0  # determines the global sync
    update_signal = False # whether update the environment annunciator
    while True: # inifinite loop of dataset iterator
        rephraser_iterator = reinforce_iterator.build_generator()
        local_env.reset_data_iter(rephraser_iterator)
        try:
            annunciator_base_step = local_step
            while True:  # inifinite loop of the episodes
                # check for the environment annunciator updates;
                with global_step_lock:
                    if global_step_counter.value % (local_agent_configs["rephraser_update_steps"]) == 0:
                        update_signal=True
                if update_signal:
                    try:
                        INFO("update valid env annunciator")
                        update_signal=False
                        annunciator_base_step, trust_acc = local_env.update_annunciator(
                            local_agent,
                            annunciator_base_step,
                            min_update_steps=annunciator_configs["valid_freq"],
                            max_update_steps=annunciator_configs["annunciator_update_steps"],
                            accuracy_bound=annunciator_configs["acc_bound"],
                            summary_writer=local_summary_writer)
                        # global_SACAgent.to("cpu").train()
                        local_agent = local_agent.to(device).train(False)
                    except StopIteration:
                        INFO("finish one training epoch, reset data_iterator")
                        break
                # prepare a new episode
                local_agent.sync_from(global_SACAgent)
                local_agent = local_agent.to(device).train(False)
                local_env.reset()
                states, masks, rephrase_positions, _ = local_env.get_state()
                episode_rewards = np.array([0.] * states.shape[0])
                episode_length = 0
                done = False
                while not done:
                    states, masks, rephrase_positions, _ = local_env.get_state()
                    label_embs = rephraser.slice_by_indices(states, rephrase_positions, device=device)
                    # directly use the expectation as validation
                    mu, _ = local_agent.actor.forward(
                        states, x_pad_indicator=1.-masks, label_emb=label_embs
                    )

                    actions = mu * local_agent.actor.action_range

                    _, rewards, _ = local_env.step(actions)
                    episode_rewards = [episode_rewards[t]+rewards[t] for t in range(states.shape[0])]
                    new_states, _, _, survival_signals = local_env.get_state()
                    survival_signals_cpu = survival_signals.cpu().numpy().tolist()
                    survival_and_no_max = np.array([0 if episode_length == local_env.sent_len[t]
                                                    else survival_signals_cpu[t][0]
                                                    for t in range(states.shape[0])])
                    INFO("step %d survive and no max:%d, %s"%(episode_length, survival_and_no_max.sum(), actions.device))

                    if 1 in survival_signals:
                        done = False
                    else:
                        episode_count += 1
                        done = True
                    if local_env.index == states.shape[1]-1:
                        # note that the environment index will increase beyond the shape
                        done = True
                    if done:  # skip the final step
                        break
                    INFO("write buffer begin")
                    invalid_sample = []
                    if local_env.index>min(local_env.sent_len):
                        for i in range(states.shape[0]):  # for every record
                            if local_env.index > local_env.sent_len[i]:
                                invalid_sample.append(i)
                    # print(states.shape, masks.shape, actions.shape, rephrase_positions.shape, torch.tensor(roll_out_rewards).shape,
                    #   survival_signals.shape, torch.tensor(survival_and_no_max).shape)
                    global_replay_buffer.add(
                        states=states,
                        masks=masks,
                        actions=actions,
                        rephrase_positions=rephrase_positions,
                        rewards=torch.tensor(rewards, dtype=torch.float32).unsqueeze(dim=-1),
                        survival_signals=survival_signals,
                        survive_and_no_max=torch.tensor(survival_and_no_max, dtype=torch.float32).unsqueeze(dim=-1),
                        invalid_sample=invalid_sample
                    )
                    INFO("write buffer end, size=%d"%global_replay_buffer.size())
                    local_step += 1
                    episode_length += 1
                INFO("mean episode rewards on batch:%f"%np.array(episode_rewards).mean())
                local_summary_writer.add_scalar(
                    "valid_rewards",
                    scalar_value=np.array(episode_rewards).mean(), global_step=local_step)

        except StopIteration:
            INFO("finish one validation epoch, reset data_iterator")
            continue


if __name__ == "__main__":
    run()

# default actor threads as 1 # load reinforce configs # the Global variable of  USE_GPU is mainly used for environments # build vocabulary and data iterator for env # we build the parallel data sets and iterate inside a thread # collect range of action space: # build global SACAgent for the final policy (on cpu) # load global ckp (only for the AC parameters) if needed # save the initial model # test_for_throughput(global_replay_buffer) # make global objects shared memory # collect available devices and distribute env on the available gpu # initialize global parameters for the current training trial # "i is the type code for c_int" # train_thread(0, device, args, #              reinforce_configs, annunciator_configs, #              src_vocab, trg_vocab, data_set, #              global_agent, global_replay_buffer, #              global_step_counter, global_step_lock, #              agent_configs) # valid_thread(device, args, #              reinforce_configs, annunciator_configs, #              src_vocab, trg_vocab, data_set, #              global_agent, global_replay_buffer, #              global_step_counter, global_step_lock, #              agent_configs) # build multi thread for learning and validation # run the dev thread for initiation build training thread for a local SACAgent on gpu device     provide parameter soft-updates for the global_models. # build local SACAgent # # this is the secondary buffer # local_replay_buffer = ReplayBuffer( #     max_sen_len=global_replay_buffer.max_sen_len, #     state_dim=victim_configs["model_configs"]["d_word_vec"], #     action_dim=victim_configs["model_configs"]["d_word_vec"], #     capacity=max(reinforce_configs["replay_buffer_capacity"]/20, 10000) # ) # build environment (include annunciator update settings) # a batch of sentences as an episode & learning episodes # initiate local agent update steps # infinite loop of data set iterator (epoch) # we will continue with a new iterator with refreshed environments # whenever the last iterator breaks with "StopIteration" # loop training episodes # the environment will be initiated by an actor as self learning # x_emb = local_env.reset(local_agent.actor) # x_emb = torch.from_numpy(x_emb) # if device != "cpu": #     x_emb = x_emb.to(device) # the rollout steps # whether the episode is finished & should reset episode with new batch # loop TD learning rollouts # check for the environment updates using current **global** agent stop criterion:                         until mse reaches the bound within patience.                         otherwise the training thread stops # update environment # sync from the global agent to avoid unnecessary locking. # global_SACAgent.to("cpu").train() # switch back to training mode # a flag to label the scorer updates # GAN target reached, scorer has reached its limit. # reset patience if scorer is refreshed # switch the agent to training mode # can't create SARSA with only one step, start a new episode # n-step TD learning accumulation # sample actions for buffer collection (start with fully random) # fast random actions in the **action_space** as initial exploration # explore and exploit # set to validation mode  .detach().cpu().numpy() # able to update agent with current buffer # update global model # save final global model (as object) # add the tensor variables to the replay buffer, # learn by rollout rewards or episode reward accumulated # update the overall signals for rollout # note that the environment index will increase beyond the shape # break from roll-out # print(states.shape, masks.shape, actions.shape, rephrase_positions.shape, torch.tensor(roll_out_rewards).shape, #   survival_signals.shape, torch.tensor(survival_and_no_max).shape) #cleanse unecessary experience to save time # for every record # break from current episodes build validation thread to check global SACAgent tests and update buffer # build local agent and sync from the global model to avoid deadlock # build an independent dev env (annunciator) # determines the global sync # whether update the environment annunciator # inifinite loop of dataset iterator # inifinite loop of the episodes # check for the environment annunciator updates; # global_SACAgent.to("cpu").train() # prepare a new episode # directly use the expectation as validation # note that the environment index will increase beyond the shape # skip the final step # for every record # print(states.shape, masks.shape, actions.shape, rephrase_positions.shape, torch.tensor(roll_out_rewards).shape, #   survival_signals.shape, torch.tensor(survival_and_no_max).shape)