
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

import os
import time
import json
import numpy as np
from collections import defaultdict
from speaker import Speaker
from mbert import mBERT
from visual import Visual

from utils import read_vocab,write_vocab,build_vocab,Tokenizer,padding_idx,timeSince, read_img_features, read_object_constraints
import utils
from env import R2RBatch
from agent import Seq2SeqAgent
from eval import Evaluation
from param import args
from tqdm import tqdm

import warnings
warnings.filterwarnings("ignore")


from tensorboardX import SummaryWriter
from transformers import BertTokenizer

from transformers import BertModel, BertConfig, AdamW, get_linear_schedule_with_warmup


log_dir = 'snap/%s' % args.name
if not os.path.exists(log_dir):
    os.makedirs(log_dir)

TRAIN_VOCAB = 'tasks/R2R/data/train_vocab.txt'
TRAINVAL_VOCAB = 'tasks/R2R/data/trainval_vocab.txt'

if args.feature_size == 2048:
    IMAGENET_FEATURES = 'img_features/ResNet-152-imagenet.tsv'
elif args.feature_size == 512:
    IMAGENET_FEATURES = 'img_features/CLIP-ViT-B-32-views.tsv'
PLACE365_FEATURES = 'img_features/ResNet-152-places365.tsv'

if args.features == 'imagenet':
    features = IMAGENET_FEATURES

if args.fast_train:
    name, ext = os.path.splitext(features)
    features = name + "-fast" + ext

feedback_method = args.feedback # teacher or sample

print(args)


def train_speaker(train_env, tok, n_iters, log_every=500, val_envs={}):
    writer = SummaryWriter(logdir=log_dir)
    listner = Seq2SeqAgent(train_env, "", tok, args.maxAction)
    speaker = Speaker(train_env, listner, tok)

    if args.fast_train:
        log_every = 40

    best_bleu = defaultdict(lambda: 0)
    best_loss = defaultdict(lambda: 1232)
    for idx in range(0, n_iters, log_every):
        interval = min(log_every, n_iters - idx)

        # Train for log_every interval
        speaker.env = train_env
        speaker.train(interval)   # Train interval iters

        print()
        print("Iter: %d" % idx)

        # Evaluation
        for env_name, (env, evaluator) in val_envs.items():
            if 'train' in env_name: # Ignore the large training set for the efficiency
                continue

            print("............ Evaluating %s ............." % env_name)
            speaker.env = env
            path2inst, loss, word_accu, sent_accu = speaker.valid()
            path_id = next(iter(path2inst.keys()))
            print("Inference: ", tok.decode_sentence(path2inst[path_id]))
            print("GT: ", evaluator.gt[str(path_id)]['instructions'])
            bleu_score, precisions = evaluator.bleu_score(path2inst)

            # Tensorboard log
            writer.add_scalar("bleu/%s" % (env_name), bleu_score, idx)
            writer.add_scalar("loss/%s" % (env_name), loss, idx)
            writer.add_scalar("word_accu/%s" % (env_name), word_accu, idx)
            writer.add_scalar("sent_accu/%s" % (env_name), sent_accu, idx)
            writer.add_scalar("bleu4/%s" % (env_name), precisions[3], idx)

            # Save the model according to the bleu score
            if bleu_score > best_bleu[env_name]:
                best_bleu[env_name] = bleu_score
                print('Save the model with %s BEST env bleu %0.4f' % (env_name, bleu_score))
                speaker.save(idx, os.path.join(log_dir, 'state_dict', 'best_%s_bleu' % env_name))

            if loss < best_loss[env_name]:
                best_loss[env_name] = loss
                print('Save the model with %s BEST env loss %0.4f' % (env_name, loss))
                speaker.save(idx, os.path.join(log_dir, 'state_dict', 'best_%s_loss' % env_name))

            # Screen print out
            print("Bleu 1: %0.4f Bleu 2: %0.4f, Bleu 3 :%0.4f,  Bleu 4: %0.4f" % tuple(precisions))


def train(train_env, train_eval, tok, n_iters, log_every=500, val_envs={}, aug_env=None):
    writer = SummaryWriter(logdir=log_dir)
    listner = Seq2SeqAgent(train_env, "", tok, args.maxAction, eval=train_eval)

    speaker = None
    if args.self_train:
        speaker = Speaker(train_env, listner, tok)
        if args.speaker is not None:
            print("Load the speaker from %s." % args.speaker)
            speaker.load(args.speaker)

    start_iter = 0

    if args.load_encoder is not None:
        print("LOAD THE encoder from %s" % args.load_encoder)
        listner.load_encoder(os.path.join(args.load_encoder))

    if args.load_visual is not None and not args.visual_random:
        print("LOAD THE visual encoder from %s" % args.load_visual)
        listner.load_visual(os.path.join(args.load_visual))

    if args.load is not None:
        print("LOAD THE listener from %s" % args.load)
        start_iter = listner.load(os.path.join(args.load))
        listner.encoder_scheduler = get_linear_schedule_with_warmup(listner.encoder_optimizer, 0.2 * args.iters, args.iters,
                                                                 last_epoch=start_iter)
        listner.encoder_pos_scheduler = get_linear_schedule_with_warmup(listner.encoder_pos_optimizer, 0.2 * args.iters,
                                                                     args.iters, last_epoch=start_iter)

    start = time.time()

    best_val = {'val_seen': {"accu": 0., "state":"", 'update':False},
                'val_unseen': {"accu": 0., "state":"", 'update':False}}
    best_val_ndtw =  {'val_seen': {"accu": 0., "state":"", 'update':False},
                'val_unseen': {"accu": 0., "state":"", 'update':False}}
    if args.fast_train:
        log_every = 40

    for idx in range(start_iter, start_iter+n_iters, log_every):
        listner.logs = defaultdict(list)
        interval = min(log_every, n_iters-idx)
        iter = idx + interval

        # Train for log_every interval
        if aug_env is None:     # The default training process
            if args.accumulate_grad:
                listner.env = train_env
                for _ in range(interval):
                    listner.zero_grad()
                    listner.accumulate_gradient(feedback_method)
                    listner.accumulate_gradient(feedback_method)
                    listner.optim_step()
            else:
                listner.env = train_env
                listner.train(interval, feedback=feedback_method)   # Train interval iters
        else:
            if args.accumulate_grad:
                for _ in range(interval // 2):
                    listner.zero_grad()
                    listner.env = train_env

                    # Train with GT data
                    args.ml_weight = 0.2
                    listner.accumulate_gradient(feedback_method)
                    listner.env = aug_env

                    # Train with Back Translation
                    args.ml_weight = 0.6        # Sem-Configuration
                    listner.accumulate_gradient(feedback_method, speaker=speaker)
                    listner.optim_step()
            else:
                for _ in range(interval // 2):
                    # Train with GT data
                    listner.env = train_env
                    args.ml_weight = 0.2
                    listner.train(1, feedback=feedback_method)

                    # Train with Back Translation
                    listner.env = aug_env
                    args.ml_weight = 0.6
                    listner.train(1, feedback=feedback_method, speaker=speaker)

        # Log the training stats to tensorboard
        total = max(sum(listner.logs['total']), 1)
        length = max(len(listner.logs['critic_loss']), 1)
        critic_loss = sum(listner.logs['critic_loss']) / total #/ length / args.batchSize
        entropy = sum(listner.logs['entropy']) / total #/ length / args.batchSize
        predict_loss = sum(listner.logs['us_loss']) / max(len(listner.logs['us_loss']), 1)
        dis_loss = sum(listner.logs['dis_loss']) / total
        con_loss = sum(listner.logs['con_loss']) / total
        writer.add_scalar("loss/critic", critic_loss, idx)
        writer.add_scalar("policy_entropy", entropy, idx)
        writer.add_scalar("loss/unsupervised", predict_loss, idx)
        writer.add_scalar("total_actions", total, idx)
        writer.add_scalar("max_length", length, idx)
        writer.add_scalar("discriminator", dis_loss, idx)
        writer.add_scalar("loss/contrastive", con_loss, idx )
        print("total_actions", total)
        print("max_length", length)

        # Run validation
        loss_str = ""
        for env_name, (env, evaluator) in val_envs.items():
            listner.env = env

            # Get validation loss under the same conditions as training
            iters = None if args.fast_train or env_name != 'train' else 20     # 20 * 64 = 1280

            # Get validation distance from goal under test evaluation conditions
            listner.test(use_dropout=False, feedback='argmax', iters=iters)
            result = listner.get_results()
            score_summary, _ = evaluator.score(result)
            loss_str += ", %s " % env_name
            for metric,val in score_summary.items():
                if metric in ['success_rate']:
                    writer.add_scalar("accuracy/%s" % env_name, val, idx)
                    if env_name in best_val:
                        if val > best_val[env_name]['accu']:
                            best_val[env_name]['accu'] = val
                            best_val[env_name]['update'] = True
                if metric in ['ndtw']:
                    writer.add_scalar("ndtw/%s" % env_name, val, idx)
                    if env_name in best_val_ndtw:
                        if val > best_val_ndtw[env_name]['accu']:
                            best_val_ndtw[env_name]['accu'] = val
                            best_val_ndtw[env_name]['update'] = True
                loss_str += ', %s: %.3f' % (metric, val)

        for env_name in best_val:
            if best_val[env_name]['update']:
                best_val[env_name]['state'] = 'Iter %d %s' % (iter, loss_str)
                best_val[env_name]['update'] = False
                listner.save(idx, os.path.join("snap", args.name, "state_dict", "best_%s" % (env_name)))

        for env_name in best_val_ndtw:
            if best_val_ndtw[env_name]['update']:
                best_val_ndtw[env_name]['state'] = 'Iter %d %s' % (iter, loss_str)
                best_val_ndtw[env_name]['update'] = False
                listner.save(idx, os.path.join("snap", args.name, "state_dict", "best_%s_ndtw" % (env_name)))

        print(('%s (%d %d%%) %s' % (timeSince(start, float(iter)/n_iters),
                                             iter, float(iter)/n_iters*100, loss_str)))

        if iter % 5000 == 0:
            print("BEST RESULT TILL NOW")
            for env_name in best_val:
                print(env_name, best_val[env_name]['state'])
            for env_name in best_val_ndtw:
                print(env_name, best_val_ndtw[env_name]['state'])

        if iter % 50000 == 0:
            listner.save(idx, os.path.join("snap", args.name, "state_dict", "Iter_%06d" % (iter)))

    listner.save(idx, os.path.join("snap", args.name, "state_dict", "LAST_iter%d" % (idx)))


def valid(train_env, tok, val_envs={}):
    agent = Seq2SeqAgent(train_env, "", tok, args.maxAction)

    print("Loaded the listener model at iter %d from %s" % (agent.load(args.load), args.load))

    for env_name, (env, evaluator) in val_envs.items():
        if env_name == 'train':
            continue
        agent.logs = defaultdict(list)
        agent.env = env

        iters = None
        agent.test(use_dropout=False, feedback='argmax', iters=iters)
        result = agent.get_results()

        # if env_name == 'val_unseen':
        #     evaluator.compare_path(result)

        sr = []
        ndtw = []
        sdtw = []
        spl = []
        weights = []
        if env_name != 'test':
            if args.eval_environment:
                score_summary, scores = evaluator.score_environment(result)
                loss_str = "Env name: %s" % env_name
                for scan, value in score_summary.items():
                    loss_str = loss_str + '/n' + scan
                    for metric, val in value.items():
                        if metric == 'success_rate':
                            sr.append(val)
                        elif metric == 'ndtw':
                            ndtw.append(val)
                        elif metric == 'sdtw':
                            sdtw.append(val)
                        elif metric == 'spl':
                            spl.append(val)
                        loss_str += ', %s: %.4f' % (metric, val)
                    weights.append(len(scores[scan]['nav_errors']))
                print(loss_str)
                mean_sr = np.average(np.array(sr), weights=weights)
                std_sr = np.average((np.array(sr) - mean_sr) ** 2, weights=weights)
                mean_ndtw = np.average(np.array(ndtw), weights=weights)
                std_ndtw = np.average((np.array(ndtw) - mean_ndtw) ** 2, weights=weights)
                mean_sdtw = np.average(np.array(sdtw), weights=weights)
                std_sdtw = np.average((np.array(sdtw) - mean_sdtw) ** 2, weights=weights)
                mean_spl = np.average(np.array(spl), weights=weights)
                std_spl = np.average((np.array(spl) - mean_spl) ** 2, weights=weights)
                print("SR:", mean_sr, std_sr ** (1 / 2))
                print("SPL:", mean_spl, std_spl ** (1 / 2))
                print("NDTW:", mean_ndtw, std_ndtw ** (1 / 2))
                print("SDTW:", mean_sdtw, std_sdtw ** (1 / 2))
                print(weights)
            else:
                score_summary, _ = evaluator.score(result)
                loss_str = "Env name: %s" % env_name
                for metric,val in score_summary.items():
                    loss_str += ', %s: %.4f' % (metric, val)
                print(loss_str)

        # if args.submit:
        if True:
            json.dump(
                result,
                open(os.path.join(log_dir, "submit_%s.json" % env_name), 'w'),
                sort_keys=True, indent=4, separators=(',', ': ')
            )


def beam_valid(train_env, tok, val_envs={}):
    listener = Seq2SeqAgent(train_env, "", tok, args.maxAction)

    speaker = Speaker(train_env, listener, tok)
    if args.speaker is not None:
        print("Load the speaker from %s." % args.speaker)
        speaker.load(args.speaker)

    print("Loaded the listener model at iter % d" % listener.load(args.load))

    final_log = ""
    for env_name, (env, evaluator) in val_envs.items():
        listener.logs = defaultdict(list)
        listener.env = env

        listener.beam_search_test(speaker)
        results = listener.results

        def cal_score(x, alpha, avg_speaker, avg_listener):
            speaker_score = sum(x["speaker_scores"]) * alpha
            if avg_speaker:
                speaker_score /= len(x["speaker_scores"])
            # normalizer = sum(math.log(k) for k in x['listener_actions'])
            normalizer = 0.
            listener_score = (sum(x["listener_scores"]) + normalizer) * (1-alpha)
            if avg_listener:
                listener_score /= len(x["listener_scores"])
            return speaker_score + listener_score

        if args.param_search:
            # Search for the best speaker / listener ratio
            interval = 0.01
            logs = []
            for avg_speaker in [False, True]:
                for avg_listener in [False, True]:
                    for alpha in np.arange(0, 1 + interval, interval):
                        result_for_eval = []
                        for key in results:
                            result_for_eval.append({
                                "instr_id": key,
                                "trajectory": max(results[key]['paths'],
                                                  key=lambda x: cal_score(x, alpha, avg_speaker, avg_listener)
                                                  )['trajectory']
                            })
                        score_summary, _ = evaluator.score(result_for_eval)
                        for metric,val in score_summary.items():
                            if metric in ['success_rate']:
                                print("Avg speaker %s, Avg listener %s, For the speaker weight %0.4f, the result is %0.4f" %
                                      (avg_speaker, avg_listener, alpha, val))
                                logs.append((avg_speaker, avg_listener, alpha, val))
            tmp_result = "Env Name %s\n" % (env_name) + \
                    "Avg speaker %s, Avg listener %s, For the speaker weight %0.4f, the result is %0.4f\n" % max(logs, key=lambda x: x[3])
            print(tmp_result)
            # print("Env Name %s" % (env_name))
            # print("Avg speaker %s, Avg listener %s, For the speaker weight %0.4f, the result is %0.4f" %
            #       max(logs, key=lambda x: x[3]))
            final_log += tmp_result
            print()
        else:
            avg_speaker = True
            avg_listener = True
            alpha = args.alpha

            result_for_eval = []
            for key in results:
                result_for_eval.append({
                    "instr_id": key,
                    "trajectory": [(vp, 0, 0) for vp in results[key]['dijk_path']] + \
                                  max(results[key]['paths'],
                                   key=lambda x: cal_score(x, alpha, avg_speaker, avg_listener)
                                  )['trajectory']
                })
            # result_for_eval = utils.add_exploration(result_for_eval)
            score_summary, _ = evaluator.score(result_for_eval)

            if env_name != 'test':
                loss_str = "Env Name: %s" % env_name
                for metric, val in score_summary.items():
                    if metric in ['success_rate']:
                        print("Avg speaker %s, Avg listener %s, For the speaker weight %0.4f, the result is %0.4f" %
                              (avg_speaker, avg_listener, alpha, val))
                    loss_str += ",%s: %0.4f " % (metric, val)
                print(loss_str)
            print()

            if args.submit:
                json.dump(
                    result_for_eval,
                    open(os.path.join(log_dir, "submit_%s.json" % env_name), 'w'),
                    sort_keys=True, indent=4, separators=(',', ': ')
                )
    print(final_log)


def setup():
    torch.manual_seed(1)
    torch.cuda.manual_seed(1)
    # dist.init_process_group(backend='nccl', world_size=2, init_method='...')
    # Check for vocabs
    if not os.path.exists(TRAIN_VOCAB):
        write_vocab(build_vocab(splits=['train']), TRAIN_VOCAB)
    if not os.path.exists(TRAINVAL_VOCAB):
        write_vocab(build_vocab(splits=['train','val_seen','val_unseen']), TRAINVAL_VOCAB)


def cleanup():
    dist.destroy_process_group()


def train_val():
    ''' Train on the training set, and validate on seen and unseen splits. '''
    # args.fast_train = True
    setup()
    # Create a batch training environment that will also preprocess text
    vocab = read_vocab(TRAIN_VOCAB)
    # tok = Tokenizer(vocab=vocab, encoding_length=args.maxInput)
    tok = BertTokenizer.from_pretrained('bert-base-multilingual-cased', cache_dir="./cache/")
    # encoding = tok("Right now you're facing towards a fireplace. Now turn right, there are two open door in front of you, move towards the open door which is to your right and exit the room through the door. Now slightly turn left, there is an open door in front of you, move towards the door and enter in to the room. Now you are facing towards a sofa, move forward and stand in front of the sofa and it is your end point.", padding='max_length', truncation=False, max_length=args.maxInput)
    # print(tok.convert_ids_to_tokens(encoding['input_ids']))
    # exit()

    feat_dict = read_img_features(features)

    featurized_scans = set([key.split("_")[0] for key in list(feat_dict.keys())])

    train_env = R2RBatch(feat_dict, batch_size=args.batchSize, splits=['train'], tokenizer=tok)
    train_eval = Evaluation(['train'], featurized_scans, tok)
    from collections import OrderedDict

    val_env_names = ['val_unseen', 'val_seen']

    if not args.beam:
        val_env_names.append("train")

    val_envs = OrderedDict(
        ((split,
          (R2RBatch(feat_dict, batch_size=args.batchSize, splits=[split], tokenizer=tok),
           Evaluation([split], featurized_scans, tok))
          )
         for split in val_env_names
         )
    )

    if args.submit:
        val_envs['test'] = (R2RBatch(feat_dict, batch_size=args.batchSize, splits=['test'], tokenizer=tok), None)

    if args.train == 'listener':
        train(train_env, train_eval, tok, args.iters, val_envs=val_envs)
    elif args.train == 'validlistener':
        if args.beam:
            beam_valid(train_env, tok, val_envs=val_envs)
        else:
            valid(train_env, tok, val_envs=val_envs)
    elif args.train == 'speaker':
        train_speaker(train_env, tok, args.iters, val_envs=val_envs)
    elif args.train == 'validspeaker':
        valid_speaker(tok, val_envs)
    else:
        assert False

    # cleanup()


def valid_speaker(tok, val_envs):
    import tqdm
    listner = Seq2SeqAgent(None, "", tok, args.maxAction)
    speaker = Speaker(None, listner, tok)
    speaker.load(args.load)

    for env_name, (env, evaluator) in val_envs.items():
        if env_name == 'train':
            continue
        print("............ Evaluating %s ............." % env_name)
        speaker.env = env
        path2inst, loss, word_accu, sent_accu = speaker.valid(wrapper=tqdm.tqdm)
        path_id = next(iter(path2inst.keys()))
        print("Inference: ", tok.decode_sentence(path2inst[path_id]))
        print("GT: ", evaluator.gt[path_id]['instructions'])
        pathXinst = list(path2inst.items())
        name2score = evaluator.lang_eval(pathXinst, no_metrics={'METEOR'})
        score_string = " "
        for score_name, score in name2score.items():
            score_string += "%s_%s: %0.4f " % (env_name, score_name, score)
        print("For env %s" % env_name)
        print(score_string)
        print("Average Length %0.4f" % utils.average_length(path2inst))


def train_val_augment():
    """
    Train the listener with the augmented data
    """
    setup()

    # Create a batch training environment that will also preprocess text
    vocab = read_vocab(TRAIN_VOCAB)
    # tok = Tokenizer(vocab=vocab, encoding_length=args.maxInput)
    tok = BertTokenizer.from_pretrained('bert-base-multilingual-cased')

    # Load the env img features
    feat_dict = read_img_features(features)
    featurized_scans = set([key.split("_")[0] for key in list(feat_dict.keys())])

    # Load the augmentation data
    aug_path = args.aug

    # Create the training environment
    train_env = R2RBatch(feat_dict, batch_size=args.batchSize,
                         splits=['train'], tokenizer=tok)
    aug_env   = R2RBatch(feat_dict, batch_size=args.batchSize,
                         splits=[aug_path], tokenizer=tok, name='aug')

    # Printing out the statistics of the dataset
    stats = train_env.get_statistics()
    print("The training data_size is : %d" % train_env.size())
    print("The average instruction length of the dataset is %0.4f." % (stats['length']))
    print("The average action length of the dataset is %0.4f." % (stats['path']))
    stats = aug_env.get_statistics()
    print("The augmentation data size is %d" % aug_env.size())
    print("The average instruction length of the dataset is %0.4f." % (stats['length']))
    print("The average action length of the dataset is %0.4f." % (stats['path']))

    # Setup the validation data
    val_envs = {split: (R2RBatch(feat_dict, batch_size=args.batchSize, splits=[split],
                                 tokenizer=tok), Evaluation([split], featurized_scans, tok))
                for split in ['train', 'val_seen', 'val_unseen']}

    # Start training
    train(train_env, tok, args.iters, val_envs=val_envs, aug_env=aug_env)


def train_mbert():
    setup()

    tok = BertTokenizer.from_pretrained('bert-base-multilingual-cased')

    feat_dict = read_img_features(features)
    featurized_scans = set([key.split("_")[0] for key in list(feat_dict.keys())])

    train_env = R2RBatch(feat_dict, batch_size=args.batchSize, splits=['train'], tokenizer=tok)

    val_envs = {split: (R2RBatch(feat_dict, batch_size=args.batchSize, splits=[split],
                                 tokenizer=tok), Evaluation([split], featurized_scans, tok))
                for split in ['train', 'val_seen', 'val_unseen']}

    mbert(train_env, tok, args.iters, val_envs=val_envs)


def mbert(train_env, tok, n_iters, log_every=100, val_envs={}):
    writer = SummaryWriter(logdir=log_dir)
    encoder = mBERT(train_env, "", tok, args.maxAction)

    start_iter = 0

    start = time.time()

    #best_loss = defaultdict(lambda: 10000000)

    best_loss = {'val_seen': {"loss": 10000000000., "state": ""},
                'val_unseen': {"loss": 10000000000., "state": ""},
                 'train': {"loss": 1000000000., "state": ""}}

    for idx in range(start_iter, start_iter + n_iters, log_every):
        encoder.logs = defaultdict(list)
        interval = min(log_every, n_iters - idx)
        iter = idx + interval

        encoder.env = train_env
        encoder.train(interval, feedback=feedback_method)

        # Log the training stats to tensorboard
        con_loss = sum(encoder.logs['con_loss']) / (args.batchSize * log_every)

        writer.add_scalar("loss/contrastive", con_loss, idx)

        # Run validation
        loss_str = ""

        for env_name, (env, evaluator) in val_envs.items():
            encoder.env = env

            # Get validation loss under the same conditions as training
            iters = None if env_name != 'train' else 20  # 20 * 64 = 1280

            # Get validation distance from goal under test evaluation conditions
            result = encoder.test(use_dropout=False, iters=iters)
            # result = encoder.loss
            loss_str += ", %s " % env_name
            loss_str += ', %s: %.3f' % ("contrastive loss", result)

            writer.add_scalar("loss/%s" % (env_name), result, idx)

            if result < best_loss[env_name]["loss"]:
                best_loss[env_name]["loss"] = result
                best_loss[env_name]['state'] = 'Iter %d %s' % (iter, loss_str)
                print('Save the model with %s BEST env loss %0.4f' % (env_name, result))
                encoder.save(idx, os.path.join(log_dir, 'state_dict', 'best_%s_loss' % env_name))

        print(('%s (%d %d%%) %s' % (timeSince(start, float(iter) / n_iters),
                                    iter, float(iter) / n_iters * 100, loss_str)))

        if iter % 1000 == 0:
            print("BEST RESULT TILL NOW")
            for env_name in best_loss:
                print(env_name, best_loss[env_name]['state'])

        if iter % 50000 == 0:
            encoder.save(idx, os.path.join("snap", args.name, "state_dict", "Iter_%06d" % (iter)))

    encoder.save(idx, os.path.join("snap", args.name, "state_dict", "LAST_iter%d" % (idx)))

def train_visual():
    setup()

    tok = BertTokenizer.from_pretrained('bert-base-multilingual-cased', cache_dir="./cache/")

    feat_dict = read_img_features(features)
    featurized_scans = set([key.split("_")[0] for key in list(feat_dict.keys())])

    if args.objects_constraints:
        objects = read_object_constraints(args.objects_constraints)
        train_env = R2RBatch(feat_dict, batch_size=args.batchSize, splits=['train'], tokenizer=tok, objects=objects)

        val_envs = {split: (R2RBatch(feat_dict, batch_size=args.batchSize, splits=[split],
                                     tokenizer=tok, objects=objects), Evaluation([split], featurized_scans, tok))
                    for split in ['train', 'val_seen', 'val_unseen']}
    else:
        train_env = R2RBatch(feat_dict, batch_size=args.batchSize, splits=['train'], tokenizer=tok)

        val_envs = {split: (R2RBatch(feat_dict, batch_size=args.batchSize, splits=[split],
                                     tokenizer=tok), Evaluation([split], featurized_scans, tok))
                    for split in ['train', 'val_seen', 'val_unseen']}

    visual(train_env, tok, args.iters, val_envs=val_envs)


def visual(train_env, tok, n_iters, log_every=100, val_envs={}):
    writer = SummaryWriter(logdir=log_dir)
    encoder = Visual(train_env, "", tok, args.maxAction)

    start_iter = 0

    start = time.time()

    if args.load_encoder is not None:
        print("LOAD THE encoder from %s" % args.load_encoder)
        encoder.load_encoder(os.path.join(args.load_encoder))

    if args.load is not None and not args.visual_random:
        print("LOAD THE visual encoder from %s" % args.load)
        start_iter = encoder.load(os.path.join(args.load))

    best_loss = {'val_seen': {"loss": 10000000000., "state": ""},
                'val_unseen': {"loss": 10000000000., "state": ""},
                 'train': {"loss": 1000000000., "state": ""}}

    for idx in range(start_iter, start_iter + n_iters, log_every):
        encoder.logs = defaultdict(list)
        interval = min(log_every, n_iters - idx)
        iter = idx + interval

        encoder.env = train_env
        encoder.train(interval, feedback=feedback_method)

        # Log the training stats to tensorboard
        con_loss = sum(encoder.logs['con_loss']) / (args.batchSize * log_every)

        writer.add_scalar("loss/contrastive", con_loss, idx)

        # Run validation
        loss_str = ""

        for env_name, (env, evaluator) in val_envs.items():
            encoder.env = env

            # Get validation loss under the same conditions as training
            iters = None if env_name != 'train' else 20  # 20 * 64 = 1280

            # Get validation distance from goal under test evaluation conditions
            result = encoder.test(use_dropout=False, iters=iters)
            # result = encoder.loss
            loss_str += ", %s " % env_name
            loss_str += ', %s: %.3f' % ("contrastive loss", result)

            writer.add_scalar("loss/%s" % (env_name), result, idx)

            if result < best_loss[env_name]["loss"]:
                best_loss[env_name]["loss"] = result
                best_loss[env_name]['state'] = 'Iter %d %s' % (iter, loss_str)
                print('Save the model with %s BEST env loss %0.4f' % (env_name, result))
                encoder.save(idx, os.path.join(log_dir, 'state_dict', 'best_%s_loss' % env_name))

        print(('%s (%d %d%%) %s' % (timeSince(start, float(iter) / n_iters),
                                    iter, float(iter) / n_iters * 100, loss_str)))

        if iter % 1000 == 0:
            print("BEST RESULT TILL NOW")
            for env_name in best_loss:
                print(env_name, best_loss[env_name]['state'])

        if iter % 50000 == 0:
            encoder.save(idx, os.path.join("snap", args.name, "state_dict", "Iter_%06d" % (iter)))

    encoder.save(idx, os.path.join("snap", args.name, "state_dict", "LAST_iter%d" % (idx)))


if __name__ == "__main__":
    if args.train in ['speaker', 'rlspeaker', 'validspeaker',
                      'listener', 'validlistener']:
        # mp.spawn(train_val,
        #          args=(2,),
        #          nprocs=2,
        #          join=True)
        train_val()
    elif args.train == 'auglistener':
        train_val_augment()
    elif args.train == 'mbert':
        train_mbert()
    elif args.train == 'visual':
        train_visual()
    else:
        assert False

    # features = 'img_features/ResNet-152-imagenet.tsv'
    # feat_dict = read_img_features(features)
    # featurized_scans = set([key.split("_")[0] for key in list(feat_dict.keys())])
    # tok = BertTokenizer.from_pretrained('bert-base-multilingual-cased')
    # evaluator = Evaluation(['val_unseen'], featurized_scans, tok)
    #
    # with open("snap/agent42/submit_val_unseen.json") as f:
    #     results_baseline = json.load(f)
    # f.close()
    #
    # with open("snap/agent61/submit_val_unseen.json") as f:
    #     results_clear = json.load(f)
    # f.close()
    #
    # assert len(results_baseline) == len(results_clear)
    #
    # results_baseline_dict = dict()
    # for result in results_baseline:
    #     results_baseline_dict[result['instr_id']] = [result['trajectory'], result['path_id']]
    #
    # results_clear_dict = dict()
    # for result in results_clear:
    #     results_clear_dict[result['instr_id']] = [result['trajectory'], result['path_id']]
    #
    # results_baseline_ordered = []
    # results_clear_ordered = []
    #
    # for k, v in results_baseline_dict.items():
    #     results_baseline_ordered.append({'instr_id': k, 'trajectory': v[0], 'path_id': v[1]})
    #     clear_result = results_clear_dict[k]
    #     results_clear_ordered.append({'instr_id': k, 'trajectory': clear_result[0], 'path_id': clear_result[1]})
    #
    # indexes = [i for i in range(len(results_baseline_ordered))]
    # sr_baseline = []
    # ndtw_baseline = []
    # sdtw_baseline = []
    # spl_baseline = []
    # sr_clear = []
    # ndtw_clear = []
    # sdtw_clear = []
    # spl_clear = []
    # ndtw_count = 0
    # sr_count = 0
    # # score_baseline = []
    # # score_clear = []
    # # score_baseline, _ = evaluator.score(results_baseline_ordered)
    # # score_clear, _ = evaluator.score(results_clear_ordered)
    # for i in tqdm(range(1000)):
    #     samples = np.random.choice(indexes, len(results_baseline_ordered))
    #     baseline = [results_baseline_ordered[i] for i in samples]
    #     clear = [results_clear_ordered[i] for i in samples]
    #     score_baseline, _ = evaluator.score(baseline)
    #     score_clear, _ = evaluator.score(clear)
    #
    #     if score_clear['ndtw'] > score_baseline['ndtw']:
    #         ndtw_count += 1
    #     if score_clear['success_rate'] > score_baseline['success_rate']:
    #         sr_count += 1
    #
    #     # for metric, val in score_baseline.items():
    #     #     if metric == 'success_rate':
    #     #         sr_baseline.append(val)
    #     #     elif metric == 'ndtw':
    #     #         ndtw_baseline.append(val)
    #     #     elif metric == 'sdtw':
    #     #         sdtw_baseline.append(val)
    #     #     elif metric == 'spl':
    #     #         spl_baseline.append(val)
    #     #
    #     # for metric, val in score_clear.items():
    #     #     if metric == 'success_rate':
    #     #         sr_clear.append(val)
    #     #     elif metric == 'ndtw':
    #     #         ndtw_clear.append(val)
    #     #     elif metric == 'sdtw':
    #     #         sdtw_clear.append(val)
    #     #     elif metric == 'spl':
    #     #         spl_clear.append(val)
    #
    #     print("ndtw",ndtw_count/(i+1))
    #     print("sr",sr_count/(i+1))
    #
    # # sr_baseline = np.array(sr_baseline)
    # # ndtw_baseline = np.array(ndtw_baseline)
    # # sdtw_baseline = np.array(sdtw_baseline)
    # # spl_baseline = np.array(spl_baseline)
    # # sr_clear = np.array(sr_clear)
    # # ndtw_clear = np.array(ndtw_clear)
    # # sdtw_clear = np.array(sdtw_clear)
    # # spl_clear = np.array(spl_clear)
    # #
    # # sr_diff = sr_clear - sr_baseline
    # # ndtw_diff = ndtw_clear - ndtw_baseline
    # # sdtw_diff = sdtw_clear - sdtw_baseline
    # # spl_diff = spl_clear - spl_baseline
    # #
    # # sr_mean = np.mean(sr_diff)
    # # sr_std = np.std(sr_diff)
    # # ndtw_mean = np.mean(ndtw_diff)
    # # ndtw_std = np.std(ndtw_diff)
    # # sdtw_mean = np.mean(sdtw_diff)
    # # sdtw_std = np.std(sdtw_diff)
    # # spl_mean = np.mean(spl_diff)
    # # spl_std = np.std(spl_diff)
    # #
    # # print("SR:", sr_mean, sr_std)
    # # print("NDTW", ndtw_mean, ndtw_std)
    # # print("SDTW:", sdtw_mean, sdtw_std)
    # # print("SPL", spl_mean, spl_std)
    #
    # print("ndtw", ndtw_count / 5000)
    # print("sr", sr_count / 5000)


