from actions import make_action
from utils import merge_tracks
from utils import split_tracks
from utils import tracklet_IoU
import sys
import matplotlib
import matplotlib.pyplot as plt
import cv2
from scipy.interpolate import spline
from utils import log_histogram
import gc
from data_providers import make_data_provider
from models import make_model
from label_storage import LabelStorage
import numpy as np
from utils import IoU
from utils import hypotheses_IoU
from utils import interpolate_tracklet
from utils import Hypothesis
import logging
from utils import Detection
import os
import tensorflow as tf
import pickle
import argparse
import time
from tqdm import tqdm
from copy import deepcopy
import json
import glob
import shutil
import matplotlib.pyplot as plt
from sklearn.metrics.pairwise import pairwise_distances

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
tf.set_random_seed(1)

# List of all flags

base_label_config = {
    "features"           : "bdif-bbox-conf-soc3-intp",
    "score_fun"          : "idf_1",
    "output_fun"         : "iou_list",
    "app_dir"            : ""
}

base_data_provider_config = {
    "name"         : "DukeDataProvider",
    "frames"       : "(1,70000,71000)",
    "frequency"     : 0.33,
    "scene_size"   : 6,
    "_camera_shift": 1000000
}

base_model_config = {
    "name" : "BidiRNNIoUPredictorModel",
    "embedding_size"   : 300,
    "rnn_cell_size"    : 300,
    "dropout"          : 0.0,
    "l2_norm"          : 0.0,
    "batch_size"       : 100,
    "lr"               : 1e-3,
    "samples_per_epoch": 100000,
    "_eval_batch_size" : 100,
    "_save_every_secs" : 6000000,
    "predict_ious"     : 1,
    "predict_labels"   : 1,
    "predict_bboxes"   : 1,
    "layers"           : 1
}

base_final_solution_config = {
    "name"             : "GreedyFinalSolution",
    "score_cutoff"     : 2.5,
    "iou_cutoff"       : 0.30,
    "no_overlap_policy": 1
}

base_nms_config = {
    "name"            : "SimpleScoreAndNMSHypotheses",
    "pairwise_min_iou": 1e-9,
    "pairwise_max_dt": 3,
    "nms_option": "start"
}

TestConfig = {
    "label_config"               : deepcopy(base_label_config),
    "train_data_provider_config" : deepcopy(base_data_provider_config),
    "input_data_provider_config" : deepcopy(base_data_provider_config),
    "model_config"               : deepcopy(base_model_config),
    "final_solution_config"      : deepcopy(base_final_solution_config),
    "nms_config"                 : deepcopy(base_nms_config),
    "mode"                       : "test",
    "logging"                    : "consol",
    "experiment_name"            : "mot17_dpm_but3",
    "hardmine"                   : 1,
    "eval_runner_cnt"            : 0,
    "datapath"                   : "",
    "target_metric"              : "IDF"
}

parser = argparse.ArgumentParser()
for k, v in TestConfig.items():
    if type(v) is dict:
        for k2, v2 in v.items():
            parser.add_argument('--%s.%s' % (k, k2), default=v2, type=type(v2))
    else:
        parser.add_argument('--%s' % k, default=v, type=type(v))
args = parser.parse_args()
for arg in vars(args):
    k = arg
    v = getattr(args, k)
    print(k, "->", v)
    if '.' in k:
        config_name, config_var = k.split('.')
        TestConfig[config_name][config_var] = v
    else:
        TestConfig[k] = v

logger = logging.getLogger('main')
logger.setLevel(logging.INFO)
ch = logging.StreamHandler()
ch.setLevel(logging.INFO)
formatter = logging.Formatter('%(levelname)s: %(name)s\n'
                              '%(asctime)s: %(message)s')
ch.setFormatter(formatter)
log_path = "runs/%s/logs" % TestConfig["experiment_name"]
if not os.path.exists(log_path):
    os.makedirs(log_path)
fh = logging.FileHandler(os.path.join(log_path, '%s.log' % os.getpid()), "w")
fh.setLevel(logging.INFO)
fh.setFormatter(logging.Formatter('%(levelname)10s    %(name)60s     '
                                  '%(asctime)30s    %(message)s'))

if "console" in TestConfig['logging'].split("-"):
    logger.addHandler(ch)
if "file" in TestConfig['logging'].split("-"):
    logger.addHandler(fh)

# Main procedure for running the model on the data.

def run_once(model,
             label_config,
             data_provider_config,
             nms_config,
             final_solution_config,
             debug=False):

    data_provider = make_data_provider(data_provider_config)

    detections = data_provider.detections()

    nms_action = make_action(nms_config)
    select_action = make_action(final_solution_config)

    LabelStorage.instance = None
    LabelStorage(label_config, data_provider)
    max_det_confidence = LabelStorage.instance.max_det_confidence

    # Operating under assumption of one continuous scene here!
    start_time = min([detection.time for detection in detections])
    finish_time = max([detection.time for detection in detections])
    final_tracks = []

    # Shifting by 1/3 of the batch length on every step.
    # In particular, in batch [1, 2, 3] we assume that
    # results in the first third (1) are already fixed
    # results in the second third (2) will be fixed after this batch
    # results in the third third (3) will be discarded and recomputed in the next batch.
    batch_len = data_provider.scene_size
    batch_shift = batch_len // 3

    all_observed_hypos = []
    all_selected_hypos = []

    for batch_start in tqdm(range(start_time, finish_time + 1, batch_shift)):

        logger.info("Starting batch @ %d, last @ %d",
                    batch_start, finish_time)

        batch_end = batch_start + batch_len

        # First let's gather all the detection in the current batch.
        # That includes all detection from the detector
        # + all detections obtained as the result of tracking from earlier batches
        # (remember, we shift the detections by the shift that is output by the model)

        batch_detections = [detection for detection in detections
                            if batch_start <= detection.time < batch_end]

        previous_batch = []
        for track in final_tracks:
            if track[-1].time >= batch_start:
                keep_len = min(len(track), track[-1].time - batch_start + 1)
                for det in track[-keep_len:]:
                    bbox = np.zeros((len(batch_detections), 4),
                                    dtype=np.float32)
                    bbox[:, :] = det.bbox
                    dbox = np.asarray([d.bbox
                                       for d
                                       in batch_detections]).reshape((-1, 4))
                    dtime = np.asarray([d.time
                                       for d
                                       in batch_detections])
                    ious = IoU(bbox, dbox)

                    tokeep = np.where(np.bitwise_or(
                        ious < select_action.iou_cutoff,
                        dtime != det.time))[0]
                    batch_detections = [
                        batch_detections[idx]
                        for idx in tokeep]

                if track[-1].time == batch_start + batch_shift - 1:
                    previous_batch.append(track[-keep_len:])

        for track in previous_batch:
            batch_detections.append(track[-1])

        # Generate a set of all hypotheses with MHT tracker

        hypotheses, all_hypos = nms_action.do(detections=batch_detections,
                                              model=model,
                                              previous_batch=previous_batch,
                                              scene_start=batch_start,
                                              scene_end=batch_end)
        all_observed_hypos += all_hypos

        # Greedily select the best set of hypotheses

        tracks, selected_hypos = \
            select_action.do(hypotheses=hypotheses,
                             previous_batch=previous_batch)

        all_selected_hypos += selected_hypos

        if batch_start == start_time and batch_end >= finish_time:
            final_tracks = tracks
            break

        # Merge current tracks with the results obtained from previous batches.
        taken = set()
        for track in final_tracks:
            if track[-1].time == batch_start + batch_shift - 1:
                mid = len(taken)
                taken.add(mid)
                tmp = tracks[mid]
                while len(track) > 0 and track[-1].time >= tmp[0].time:
                    track.pop()
                for det in tmp:
                    if len(track) == 0:
                        track.append(det)
                    else:
                        track += interpolate_tracklet(track[-1], det)[1:]

        # Add new trajectories that were not there in the previous batch.
        for mid in range(len(tracks)):
            if mid not in taken:
                if tracks[mid][0].time < batch_start + 2 * batch_shift or \
                                batch_end >= finish_time:
                    tmp = tracks[mid]
                    final_tracks.append(tmp)

        # Cut at 2/3 of the batch
        if batch_end <= finish_time:
            for track in final_tracks:
                while len(track) > 0 and\
                    track[-1].time >= batch_start + 2 * batch_shift:
                    track.pop()
        else:
            break

        final_tracks = [t for t in final_tracks if len(t) > 0]

        # Add the interpolated detections that were output by the model to the global set of detections.
        for track in final_tracks:
            ptr = len(track) - 1
            while ptr >= 0 and track[ptr].time >= batch_start:
                if track[ptr].time >= batch_start + batch_shift:
                    track[ptr].interpolated = False
                    track[ptr].confidence = max_det_confidence
                    detections.append(track[ptr])
                ptr -= 1

    return [t for t in final_tracks if len(t) > 0],\
           all_observed_hypos,\
           all_selected_hypos

# Simple utility for logging arguments, and the code for logging arguments
def tostring(dic):
    out = ""
    for k in sorted(dic.keys()):
        if not k.startswith("_"):
            if len(out) > 0:
                out += ','
            out += k + '=' + str(dic[k])
    return out

label_json = tostring(TestConfig["label_config"])
model_json = tostring(TestConfig["model_config"])
train_json = tostring(TestConfig["train_data_provider_config"])
valid_json = tostring(TestConfig["input_data_provider_config"])

# Paths to all necessary resources (created if necessary)

gendata_stopfile = os.path.join("runs",
                                TestConfig["experiment_name"],
                                "gendata_stopfile")
ckpt_dir = os.path.join("runs", TestConfig["experiment_name"], "checkpoints")
datagen_ckpt_dir = \
    os.path.join("runs", TestConfig["experiment_name"], "datagen_checkpoints")
best_model_dir = os.path.join("runs", TestConfig["experiment_name"],
                              "best_model")
summ_dir = os.path.join("runs", TestConfig["experiment_name"],
                        "summaries", TestConfig["mode"])
data_dir = os.path.join("runs", TestConfig["experiment_name"], "data")

if TestConfig["datapath"] == "":
    dataset_dir = os.path.join("runs", TestConfig["experiment_name"], "dataset")
else:
    dataset_dir = TestConfig["datapath"]

for dir in [ckpt_dir, summ_dir, best_model_dir,
            data_dir, datagen_ckpt_dir, dataset_dir]:
    if not os.path.exists(dir):
        os.makedirs(dir)

with open(os.path.join("runs", TestConfig["experiment_name"], "config"), "a")\
        as f:
    line = json.dumps(TestConfig, indent=4, sort_keys=True)
    f.write(line)
    f.write("\n")

# Training code
if TestConfig["mode"] == 'train':

    data_provider = make_data_provider(TestConfig["train_data_provider_config"])
    LabelStorage(label_config=TestConfig["label_config"],
                 data_provider=data_provider)

    # Create a model and possibly load it from the checkpoints
    model = make_model(TestConfig["model_config"], "infer",
                       tf.Session(), ckpt_dir)
    fw = tf.summary.FileWriter(summ_dir, graph=model.sess.graph)
    try:
        model.load_model()
        model.save_model()
    except:
        model.sess.run(tf.global_variables_initializer())
        model.save_model()

    # Read all the data that was prepared for training
    labeled_hypotheses = [[] for _ in range(10)]
    fls = os.listdir(dataset_dir)
    use_data_ptr = False
    if fls[0].endswith("data_ptr"):
        use_data_ptr = True
        print("Reading from data_ptr file")
        with open(os.path.join(dataset_dir, fls[0]), "r") as f:
            fls = [line.rstrip() for line in f.readlines()]
        print("First file %s" % fls[0])

    cur_hypotheses = []
    for path in fls:
        if not use_data_ptr:
            cur_hypotheses = pickle.load(
                open(os.path.join(dataset_dir, path), "rb"))
        else:
            cur_hypotheses = pickle.load(open(path, "rb"))

        for lh in cur_hypotheses:
            gtlen = np.sum(lh.labels)
            sqlen = 0
            mtlen = 0
            for idx in range(len(lh.tracklet)):
                if lh.tracklet[idx] is not None:
                    cam_id, _ = data_provider.cam_and_time(lh.tracklet[idx])
                    sqlen += 1.0
                    if lh.ious[idx] > 0.5:
                        mtlen += 1.0
            idf = 2 * mtlen * 1. / (gtlen + sqlen)
            bucket = max(0, min(9, int(idf * 10)))
            lh.idf = idf

            # clearing features
            delattr(lh, "features")
            for idx in range(len(lh.tracklet)):
                if lh.tracklet[idx] is not None:
                    lh.tracklet[idx].features = None

            labeled_hypotheses[bucket].append(lh)
        print("Loaded %d from %s" % (len(cur_hypotheses), path))

    epoch = 0

    while True:
        np.random.seed(epoch)
        min_bucket_size = min([len(x) for x in labeled_hypotheses])

        # Hardmining to select the balanced dataset
        for idx in range(10):

            # First sort at random
            order = np.arange(len(labeled_hypotheses[idx]))
            np.random.shuffle(order)

            # Then take min_bucket_size x HARDMINE
            # and select min_bucket_size hardest
            to_take = min(len(order), min_bucket_size * TestConfig["hardmine"])
            _, losses = model.train_epoch(labeled_hypotheses[idx][:to_take], 0,
                                          do_train=False)
            order = np.argsort(np.asarray(losses))
            order = order[::-1]

            labeled_hypotheses[idx] = [labeled_hypotheses[idx][idx2]
                                       for idx2 in order]

        cur_dataset = []
        for idx in range(10):
            cur_dataset += [labeled_hypotheses[idx][idx2]
                            for idx2
                            in range(min(min_bucket_size,
                                         len(labeled_hypotheses[idx])))]

        # Train one epoch and wait a bit to ensure that all evaluation runners
        # had the time to run on the current checkpoint (otherwise a large backlog
        # might be created)
        model.train_epoch(cur_dataset, epoch, do_train=True, do_save=True)
        print("Done epoch %d" % epoch)
        epoch += 1
        time.sleep(120)

# Generating the dataset - MASTER
if TestConfig["mode"] == 'gen_dataset':

    data_provider = make_data_provider(TestConfig["train_data_provider_config"])
    detections = data_provider.detections()
    ground_truth = data_provider.ground_truth()

    all_det_times = np.unique(np.asarray([
        det.time for det in detections]))

    LabelStorage(label_config=TestConfig["label_config"],
                 data_provider=data_provider)

    # Create and possibly load a model
    model = make_model(TestConfig["model_config"], "infer",
                       tf.Session(), datagen_ckpt_dir)
    fw = tf.summary.FileWriter(summ_dir, graph=model.sess.graph)
    try:
        model.load_model()
        model.save_model()
    except:
        model.sess.run(tf.global_variables_initializer())
        model.save_model()
    epoch = 0

    total_deleted_ever = 0

    size_history = []
    labeled_hypotheses = [[] for _ in range(10)]
    while(True):
        gc.collect()
        step = model.sess.run(model.global_step)

        # Read all the data generated by the slaves.
        # The slaves create two filex (X.pickle and X.pickle_flag)
        # As long as there is a pickle_flag, it is safe to read
        # the pickle file (the file is finished and can be read)

        fls = os.listdir(data_dir)
        added_hypotheses = []
        for path_flag in fls:
            if path_flag.endswith("_flag"):
                path = path_flag[:-5]
                cur_hypotheses = pickle.load(
                    open(os.path.join(data_dir, path), "rb"))
                added_hypotheses += cur_hypotheses
                os.remove(os.path.join(data_dir, path))
                os.remove(os.path.join(data_dir, path_flag))

        # computing IDF of each of the read hypotheses
        # to divide them into 10 buckets according to IDF
        for lh in added_hypotheses:
            gtlen = np.sum(lh.labels)
            sqlen = 0
            mtlen = 0
            for idx in range(len(lh.tracklet)):
                if lh.tracklet[idx] is not None:
                    sqlen += 1.0
                    if lh.ious[idx] > 0.5:
                        mtlen += 1.0
            idf = 2 * mtlen * 1. / (gtlen + sqlen)
            bucket = max(0, min(9, int(idf * 10)))
            lh.idf = idf
            labeled_hypotheses[bucket].append(lh)

        # If not enough training data was read, just leeep
        min_bucket_size = min([len(x) for x in labeled_hypotheses])
        if  min_bucket_size == 0:
            print("Not enough training data so far, waiting")
            time.sleep(30)
            continue

        np.random.seed(epoch)
        to_delete = []
        all_losses = []
        for idx in range(10):
            # Compute losses for every data sample and keep only the hardest ones
            _, losses = model.train_epoch(labeled_hypotheses[idx], 0,
                                          do_train=False)
            all_losses.append(losses)
            order = np.argsort(np.asarray(losses))
            order = order[::-1]

            #order = np.arange(len(labeled_hypotheses[idx]))
            #np.random.shuffle(order)

            labeled_hypotheses[idx] = [labeled_hypotheses[idx][idx2]
                                       for idx2 in order]
            # To avoid memory limit exceeded
            if len(labeled_hypotheses[idx]) > 20000:
                to_delete += [h for h in labeled_hypotheses[idx][20000:]]
                labeled_hypotheses[idx] = labeled_hypotheses[idx][:20000]

        min_bucket_size = min([len(x) for x in labeled_hypotheses])

        cur_dataset = []
        for idx in range(10):
            cur_dataset += [labeled_hypotheses[idx][idx2]
                            for idx2
                            in range(min(min_bucket_size,
                                         len(labeled_hypotheses[idx])))]

        # Report the summaries for the data generation process.

        sl = sum([len(x) for x in labeled_hypotheses])

        summary = tf.Summary()
        summary.value.add(tag='data idf',
                          histo=log_histogram(
                              np.asarray([lh.idf
                                          for lhypos in labeled_hypotheses
                                          for lh in lhypos]),
                              bins=10))
        summary.value.add(tag='mean loss',
                          simple_value=np.mean(np.concatenate(all_losses)))
        summary.value.add(tag='min bucket', simple_value=min_bucket_size)
        summary.value.add(tag="train data size", simple_value=len(cur_dataset))
        summary.value.add(tag="active data size", simple_value=sl)
        summary.value.add(tag="deleted data size", simple_value=len(to_delete))
        total_deleted_ever += len(to_delete)
        summary.value.add(tag="observed data size",
                          simple_value=sl + total_deleted_ever)
        fw.add_summary(summary, step)
        fw.flush()

        # Save the data samples
        pickle.dump(to_delete,
                    open(os.path.join(dataset_dir,
                                      "deleted_%d.pickle" % step), "wb"))

        pickle.dump([x for lhypos in labeled_hypotheses for x in lhypos],
                    open(os.path.join(dataset_dir, "remaining.pickle"),
                         "wb"))

        # Train the model on the hardest samples taken.
        model.train_epoch(cur_dataset, epoch, do_train=True, do_save=True)
        print("Done epoch %d" % epoch)
        epoch += 1
        size_history.append(sl + total_deleted_ever)
        # If there was a significant dataset growth, restarting the model training.
        if size_history[-1] >= 1.5 * size_history[0]:
            print("Dataset has grown at least 1.5x, clearing the model")
            step = model.sess.run(model.global_step)
            model.sess.run(tf.global_variables_initializer())
            model.sess.run(model.global_step.assign(step))
            size_history = size_history[-1:]

        # If there was less than 5% growth in 10 steps, stopping the training.
        if os.path.exists(gendata_stopfile) or\
                (len(size_history) > 10 and\
                                 size_history[-11] * 1.05 > size_history[-1]):
            print("Dataset growth less than 5% in 10 steps or stopfile, ending")
            break

# Generating data - SLAVE
if TestConfig["mode"] == 'gen_data':

    data_provider = make_data_provider(TestConfig["input_data_provider_config"])
    detections = data_provider.detections()
    ground_truth = data_provider.ground_truth()

    det_by_time = {}
    for det in detections:
        if det.time not in det_by_time.keys():
            det_by_time[det.time] = []
        det_by_time[det.time].append(det)

    start_time = min([detection.time for detection in detections])
    finish_time = max([detection.time for detection in detections])
    batch_len = data_provider.scene_size

    LabelStorage(label_config=TestConfig["label_config"],
                 data_provider=data_provider)

    model = make_model(TestConfig["model_config"], "infer",
                       tf.Session(), ckpt_dir)

    nms_action = make_action(TestConfig["nms_config"])

    all_hypotheses_ever = set()

    last_step = -1
    while True:
        # Try to load the model from the latest checkpoint
        try:
            model.load_model(datagen_ckpt_dir)
            new_step = model.sess.run(model.global_step)
        except:
            time.sleep(30)
            continue
        # If it was already observed, sleep
        if new_step == last_step:
            time.sleep(30)
            continue

        last_step = new_step
        last_ckpt = tf.train.latest_checkpoint(datagen_ckpt_dir)

        labeled_all_hypotheses = []

        # Precompute batches
        batches = []
        for batch_start in range(start_time, finish_time + 1, batch_len):
            batch_end = batch_start + batch_len
            batches.append((batch_start, batch_end))

        np.random.seed(new_step)
        order = np.arange(len(batches))
        batches = [batches[idx] for idx in order]

        for bid in tqdm(range(len(batches))):

            # Run the model on every batch
            # If the new model checkpoint is available, restart the process

            batch_start, batch_end = batches[bid]
            newest_ckpt = tf.train.latest_checkpoint(datagen_ckpt_dir)
            if newest_ckpt != last_ckpt:
                print("Model was updated, restarting")
                break

            #LabelStorage.instance.storage = {}
            batch_detections = []
            for when in range(batch_start, batch_end):
                if when in det_by_time.keys():
                    batch_detections += det_by_time[when]

            if len(batch_detections) == 0:
                continue

            # Good_hypotheses are the ones selected by MHT
            # All_hypotheses are all hypotheses observed by the MHT
            # We use all hypotheses for training.

            good_hypotheses, all_hypotheses =\
                nms_action.do(detections=batch_detections,
                              model=model,
                              previous_batch=[],
                              scene_start=batch_start,
                              scene_end=batch_end)

            out = []
            for h in all_hypotheses:
                if hash(h) not in all_hypotheses_ever:
                    all_hypotheses_ever.add(hash(h))
                    out.append(h)
            all_hypotheses = out

            # Compute the labels corresponging to all hypotheses

            lhypos = LabelStorage.instance.label_hypotheses(all_hypotheses,
                                                            mode="two")
            LabelStorage.instance.get_hypo_features(lhypos)
            labeled_all_hypotheses += lhypos

            if batch_end >= finish_time:
                break

        # Save all of the observed hypotheses with the corresponding labels
        # for the master to consume

        print("Step %d: + %d hypotheses = %d total" % (
            new_step, len(labeled_all_hypotheses), len(all_hypotheses_ever)))

        path = os.path.join(data_dir, "%d_%s.pickle" % (new_step, os.getpid()))
        pickle.dump(labeled_all_hypotheses, open(path, "wb"))
        with open(path + "_flag", "w") as f:
            f.write("Done creating labeled hypotheses")

# Evaluation process - MASTER
if args.mode.startswith('continuous_eval_joiner'):

    writer = tf.summary.FileWriter(summ_dir)
    best_metric = -100
    while True:
        # Find all the statistics generated by the slaves
        dir = str.replace(summ_dir, "joiner", "runner")

        fls = glob.glob("%s/*.txt" % dir)
        fls.sort(key=os.path.getmtime)

        if len(fls) == 0:
            time.sleep(10)
            continue

        # If there are as many summaries for the oldest checkpoint
        # as there are runners, gather all the statistics, average it,
        # and save to tensorboard
        min_ckpt = fls[0].split("/")[-1].split("_")[0]
        ctr = sum([file.split("/")[-1].split("_")[0] == min_ckpt
                   for file in fls])
        print("Found min_ckpt %s cnt %d" % (min_ckpt, ctr))
        if ctr == args.eval_runner_cnt:
            data = {}
            sum_wght = 0
            for file in fls:
                if not file.split("/")[-1].split("_")[0] == min_ckpt:
                    continue
                with open(file, "r") as f:
                    lines = [line.rstrip() for line in f.readlines()]
                os.remove(file)

                keys = [line.split(" ")[0] for line in lines]
                vals = [float(line.split(" ")[1]) for line in lines]
                for k, v in zip(keys, vals):
                    if k == "GTLEN":
                        sum_wght += v
                        cur_wght = v
                        continue
                    if k not in data.keys():
                        data[k] = 0
                    data[k] += v * cur_wght

            summary = tf.Summary()
            for k, v in data.items():
                summary.value.add(tag=k, simple_value=v / sum_wght)
            writer.add_summary(summary, min_ckpt)
            writer.flush()

            # If the value for a given checkpoint is the best,
            # save the model in the best_model folder
            if data[TestConfig["target_metric"]] > best_metric:
                best_metric = data[TestConfig["target_metric"]]
                for file in glob.glob(r'%s/model-%s.*' % (ckpt_dir, min_ckpt)):
                    shutil.copy(file, best_model_dir)
                with open(os.path.join(best_model_dir, "checkpoint"), "w") as f:
                    f.write("model_checkpoint_path: \"model-%s\"\n" % min_ckpt)
                    f.write("all_model_checkpoint_paths: \"model-%s\"\n" % min_ckpt)
        else:
            # If something happened and there are no results for the oldest checkpoint
            # but it is already very stale, just ignore it and proceed to the newer ones.
            if len(fls) > args.eval_runner_cnt * 3:
                print("But too long list of files, removing tail")
                for file in fls:
                    if not file.split("/")[-1].split("_")[0] == min_ckpt:
                        continue
                    os.remove(file)

        time.sleep(10)
        continue
# Evaluation process - SLAVE
if args.mode.startswith('continuous_eval_runner'):
    final_json = tostring(TestConfig["final_solution_config"])

    old_step = -1
    data_provider = make_data_provider(TestConfig["input_data_provider_config"])

    while True:
        # Read the latest model
        LabelStorage.instance = None
        LabelStorage(label_config=TestConfig["label_config"],
                     data_provider=data_provider)
        tf.reset_default_graph()
        model = make_model(TestConfig["model_config"], "infer",
                           tf.Session(), ckpt_dir)
        try:
            if args.mode.endswith("best"):
                model.load_model(best_model_dir)
                print("Loading from best")
            else:
                model.load_model(ckpt_dir)
            new_step = model.sess.run(model.global_step)
        except:
            logger.info("Waiting for model checkpoint")
            time.sleep(10)
            continue

        if new_step == old_step:
            time.sleep(10)
            continue
        else:
            old_step = new_step
            logger.info("Evaluating @ step %d", new_step)

        # Run inference and save the tracks

        final_tracks, all_hypos, selected_hypos = \
            run_once(model,
                     TestConfig["label_config"],
                     TestConfig["input_data_provider_config"],
                     TestConfig["nms_config"],
                     TestConfig["final_solution_config"])

        # We are only interested in tracks longer than the length of the batch.
        final_tracks = [t for t in final_tracks
                        if len(t) >
                        TestConfig["input_data_provider_config"]["scene_size"]]

        # Evaluate the quality of tracking

        mota, idf = data_provider.evaluate(tracks=final_tracks)

        # To compute the summaries, also take a part of the observed hypotheses
        # and compute summaries on it

        if len(all_hypos) > 100000:
            all_hypos = all_hypos[:100000]
        labeled_hypos = LabelStorage.instance.label_hypotheses(all_hypos)

        summaries, _ = model.train_epoch(labeled_hypos, 0, do_train=False)

        # Save the summaries and the tracking results for evaluation master to gather

        path = os.path.join(summ_dir, "%d_%d.txt" % (new_step, os.getpid()))

        with open(path, "w") as f:
            f.write("GTLEN %d\n" % sum([len(x)
                                        for x
                                        in data_provider.ground_truth()]))
            f.write("MOTA %0.3f\n" % mota)
            f.write("IDF %0.3f\n" % idf)
            for k, v in summaries.items():
                f.write("%s %0.3f\n" % (k, v))

        if args.mode.endswith("best"):
            exit(0)

# Inference
if args.mode.startswith('infer'):

    data_provider = make_data_provider(TestConfig["input_data_provider_config"])
    min_det = min([det.time for det in data_provider.detections()])
    LabelStorage.instance = None
    LabelStorage(label_config=TestConfig["label_config"],
                 data_provider=data_provider)
    tf.reset_default_graph()
    model = make_model(TestConfig["model_config"], "infer",
                       tf.Session(), ckpt_dir)
    model.load_model(best_model_dir)

    final_tracks, all_hypos, selected_hypos = \
        run_once(model,
                 TestConfig["label_config"],
                 TestConfig["input_data_provider_config"],
                 TestConfig["nms_config"],
                 TestConfig["final_solution_config"])

    # We only consider in the final solution tracks that are longer than the length of the batch.
    cur_final_tracks = [t for t in final_tracks
                        if len(t) >
                        TestConfig["input_data_provider_config"]["scene_size"]]
    # Final results are saved in runs/$experiment_name/summaries/infer/tracks_$idx.txt
    data_provider.save_tracks(cur_final_tracks, os.path.join(summ_dir,
                                                             "tracks_%s.txt" % TestConfig["input_data_provider_config"]["frames"]))


#!/usr/bin/env python3

import logging
import numpy as np
import time
import torch

import cv2

from . import transform as transform

logger = logging.getLogger(__name__)

def compute_frames_distances(frames, frame_pairs, num_frames, dist_func):
    """
    This function computes the distance between all pairs of frame
    and store the result in matrix

    Args:
        frames:
        dist_func:
        frame_pairs:
        num_frames:

    Returns:
        frames_mse (np.array): matrix of distances between each pair
    """
    # if torch.cuda.is_available():
    #     frames = frames.cuda()
    frames_mse = np.zeros((num_frames, num_frames))
    for pairs in frame_pairs:
        t1, t2 = pairs
        f_t1, f_t2 = torch.chunk(frames[:,[t1,t2], :, :], chunks=2, dim=1)  # CxTxWxH
        mse_val = dist_func(f_t1, f_t2)
        frames_mse[t1, t2] = mse_val.item()
        frames_mse[t2, t1] = mse_val.item()

    return frames_mse

def retry_load_images(image_paths, retry=10, backend="pytorch"):
    """
    This function is to load images with support of retrying for failed load.

    Args:
        image_paths (list): paths of images needed to be loaded.
        retry (int, optional): maximum time of loading retrying. Defaults to 10.
        backend (str): `pytorch` or `cv2`.

    Returns:
        imgs (list): list of loaded images.
    """
    for i in range(retry):
        imgs = [cv2.imread(image_path) for image_path in image_paths]

        if all(img is not None for img in imgs):
            if backend == "pytorch":
                imgs = torch.as_tensor(np.stack(imgs))
            return imgs
        else:
            logger.warn("Reading failed. Will retry.")
            time.sleep(1.0)
        if i == retry - 1:
            raise Exception("Failed to load images {}".format(image_paths))

def get_sequence(center_idx, half_len, sample_rate, num_frames):
    """
    Sample frames among the corresponding clip.

    Args:
        center_idx (int): center frame idx for current clip
        half_len (int): half of the clip length
        sample_rate (int): sampling rate for sampling frames inside of the clip
        num_frames (int): number of expected sampled frames

    Returns:
        seq (list): list of indexes of sampled frames in this clip.
    """
    seq = list(range(center_idx - half_len, center_idx + half_len, sample_rate))

    for seq_idx in range(len(seq)):
        if seq[seq_idx] < 0:
            seq[seq_idx] = 0
        elif seq[seq_idx] >= num_frames:
            seq[seq_idx] = num_frames - 1
    return seq

def pack_pathway_output(cfg, frames):
    """
    Prepare output as a list of tensors. Each tensor corresponding to a
    unique pathway.
    Args:
        frames (tensor): frames of images sampled from the video. The
            dimension is `channel` x `num frames` x `height` x `width`.
    Returns:
        frame_list (list): list of tensors with the dimension of
            `channel` x `num frames` x `height` x `width`.
    """
    if cfg.MODEL.ARCH in cfg.MODEL.SINGLE_PATHWAY_ARCH:
        frame_list = [frames]
    elif cfg.MODEL.ARCH in cfg.MODEL.MULTI_PATHWAY_ARCH:
        fast_pathway = frames
        # Perform temporal sampling from the fast pathway.
        slow_pathway = torch.index_select(
            frames,
            1,
            torch.linspace(
                0, frames.shape[1] - 1, frames.shape[1] // cfg.SLOWFAST.ALPHA
            ).long(),
        )
        frame_list = [slow_pathway, fast_pathway]
    else:
        raise NotImplementedError(
            "Model arch {} is not in {}".format(
                cfg.MODEL.ARCH,
                cfg.MODEL.SINGLE_PATHWAY_ARCH + cfg.MODEL.MULTI_PATHWAY_ARCH,
            )
        )
    return frame_list

def spatial_sampling(
    frames,
    spatial_idx=-1,
    min_scale=256,
    max_scale=320,
    crop_size=224,
    random_horizontal_flip=True,
    inverse_uniform_sampling=False,
):
    """
    Perform spatial sampling on the given video frames. If spatial_idx is
    -1, perform random scale, random crop, and random flip on the given
    frames. If spatial_idx is 0, 1, or 2, perform spatial uniform sampling
    with the given spatial_idx.
    Args:
        frames (tensor): frames of images sampled from the video. The
            dimension is `num frames` x `height` x `width` x `channel`.
        spatial_idx (int): if -1, perform random spatial sampling. If 0, 1,
            or 2, perform left, center, right crop if width is larger than
            height, and perform top, center, buttom crop if height is larger
            than width.
        min_scale (int): the minimal size of scaling.
        max_scale (int): the maximal size of scaling.
        crop_size (int): the size of height and width used to crop the
            frames.
        inverse_uniform_sampling (bool): if True, sample uniformly in
            [1 / max_scale, 1 / min_scale] and take a reciprocal to get the
            scale. If False, take a uniform sample from [min_scale,
            max_scale].
    Returns:
        frames (tensor): spatially sampled frames.
    """
    assert spatial_idx in [-1, 0, 1, 2]
    if spatial_idx == -1:
        frames, _ = transform.random_short_side_scale_jitter(
            images=frames,
            min_size=min_scale,
            max_size=max_scale,
            inverse_uniform_sampling=inverse_uniform_sampling,
        )
        frames, _ = transform.random_crop(frames, crop_size)
        if random_horizontal_flip:
            frames, _ = transform.horizontal_flip(0.5, frames)
    else:
        # The testing is deterministic and no jitter should be performed.
        # min_scale, max_scale, and crop_size are expect to be the same.
        assert len({min_scale, max_scale, crop_size}) == 1
        frames, _ = transform.random_short_side_scale_jitter(
            frames, min_scale, max_scale
        )
        frames, _ = transform.uniform_crop(frames, crop_size, spatial_idx)
    return frames


# coding: utf-8
# Copyright (c) 2020 The MITRE Corporation.
"""
Inferencers to make predictions and analyze data using trained topic models.
"""

import json
import mxnet as mx
import numpy as np
import gluonnlp as nlp
import io
import os
import scipy
import umap
import logging
import pickle
from tmnt.modeling import BowVAEModel, CovariateBowVAEModel, SeqBowVED, MetricSeqBowVED
from tmnt.estimator import BowEstimator
from tmnt.data_loading import DataIterLoader, file_to_data, SparseMatrixDataIter
from tmnt.preprocess.vectorizer import TMNTVectorizer
from tmnt.distribution import HyperSphericalDistribution
from tmnt.utils.recalibrate import recalibrate_scores
from multiprocessing import Pool
from gluonnlp.data import BERTTokenizer, BERTSentenceTransform
from sklearn.datasets import load_svmlight_file
from functools import partial

MAX_DESIGN_MATRIX = 250000000

class BaseInferencer(object):
    """Base inference object for text encoding with a trained topic model.

    """
    def __init__(self, ctx):
        self.ctx = ctx

    def save(self, model_dir):
        raise NotImplementedError

    def encode_texts(self, intexts):
        raise NotImplementedError

    def get_top_k_words_per_topic(self, k):
        raise NotImplementedError

    def get_top_k_words_per_topic_per_covariate(self, k):
        raise NotImplementedError

class BowVAEInferencer(BaseInferencer):
    """
    """
    def __init__(self, estimator, pre_vectorizer=None):
        super().__init__(estimator.model.model_ctx)
        self.max_batch_size = 16
        self.vocab = estimator.model.vocabulary
        self.vectorizer = pre_vectorizer or TMNTVectorizer(initial_vocabulary=estimator.model.vocabulary)
        self.n_latent = estimator.model.n_latent
        self.model = estimator.model
        self.estimator = estimator
        if isinstance(estimator.model, CovariateBowVAEModel):
            self.covar_model = True
            self.n_covars = estimator.model.n_covars
            self.covar_net_layers = estimator.model.covar_net_layers
        else:
            self.covar_model = False

    @classmethod
    def from_saved(cls, param_file=None, config_file=None, vocab_file=None, model_dir=None, ctx=mx.cpu()):
        serialized_vectorizer_file = None
        if model_dir is not None:
            estimator = BowEstimator.from_saved(model_dir)
            serialized_vectorizer_file = os.path.join(model_dir,'vectorizer.pkl')
        else:
            param_file = os.path.join(model_dir,'model.params')
            vocab_file = os.path.join(model_dir,'vocab.json')
            config_file = os.path.join(model_dir,'model.config')
            estimator = BowEstimator.from_config(config=config_file, vocabulary=vocab_file, pretrained_param_file=param_file)
        estimator.initialize_with_pretrained()
        if os.path.exists(serialized_vectorizer_file):
            with open(serialized_vectorizer_file, 'rb') as fp:
                vectorizer = pickle.load(fp)
        else:
            vectorizer = None
        #model.load_parameters(str(param_file), allow_missing=False)
        return cls(estimator, pre_vectorizer=vectorizer)

    def save(self, model_dir: str) -> None:
        """
        Save model and vectorizer to disk

        Parameters:
            model_dir: Model directory to save parmaeters, config, vocabulary and vectorizer
        """
        self.estimator.write_model(model_dir)
        serialized_vector_file = os.path.join(model_dir, 'vectorizer.pkl')
        if self.vectorizer:
            with io.open(serialized_vector_file, 'wb') as fp:
                pickle.dump(self.vectorizer, fp)

    def get_model_details(self, sp_vec_file_or_X, y=None):
        if isinstance(sp_vec_file_or_X, str):
            data_csr, labels = load_svmlight_file(sp_vec_file, n_features=len(self.vocab))
        else:
            data_csr = sp_vec_file_or_X
            labels = y
        data_csr = mx.nd.sparse.csr_matrix(data_csr, dtype='float32')
        ## 1) K x W matrix of P(term|topic) probabilities
        w = self.model.decoder.collect_params().get('weight').data().transpose() ## (K x W)
        w_pr = mx.nd.softmax(w, axis=1)
        ## 2) D x K matrix over the test data of topic probabilities
        covars = labels if self.covar_model else None
        dt_matrix = self.encode_data(data_csr, covars, use_probs=True, target_entropy=2.0)
        ## 3) D-length vector of document sizes
        doc_lengths = data_csr.sum(axis=1)
        ## 4) vocab (in same order as W columns)
        ## 5) frequency of each word w_i \in W over the test corpus
        term_cnts = data_csr.sum(axis=0)
        return w_pr, dt_matrix, doc_lengths, term_cnts

    def get_pyldavis_details(self, sp_vec_file_or_X, y=None):
        w_pr, dt_matrix, doc_lengths, term_cnts = self.get_model_details(sp_vec_file_or_X, y=y)
        d1 = w_pr.asnumpy().tolist()
        d2 = list(map(lambda x: x.tolist(), dt_matrix))
        d3 = doc_lengths.asnumpy().tolist()
        d5 = term_cnts.asnumpy().tolist()
        d4 = list(map(lambda i: self.vocab.idx_to_token[i], range(len(self.vocab.idx_to_token))))
        d = {'topic_term_dists': d1, 'doc_topic_dists': d2, 'doc_lengths': d3, 'vocab': d4, 'term_frequency': d5 }
        return d

    def get_umap_embeddings(self, data, umap_metric='euclidean', use_probs=False, target_entropy=1.0):
        encs = self.encode_data(data, None, use_probs=use_probs, target_entropy=target_entropy)
        encs2 = np.array(encs)
        um = umap.UMAP(n_neighbors=4, min_dist=0.2, metric='euclidean')
        return um.fit_transform(encs2)

    def plot_to(self, embeddings, labels, f=None):
        import matplotlib.pyplot as plt
        plt.scatter(*embeddings.T, c=labels, s=0.8, alpha=0.9, cmap='coolwarm')
        if f is None:
            plt.show()
        else:
            plt.savefig(f)

    def export_full_model_inference_details(self, sp_vec_file, ofile):
        d = self.get_pyldavis_details(sp_vec_file)
        with io.open(ofile, 'w') as fp:
            json.dump(d, fp, sort_keys=True, indent=4)

    def encode_vec_file(self, sp_vec_file, use_probs=False):
        data_mat, labels = load_svmlight_file(sp_vec_file, n_features=len(self.vocab), zero_based=True)
        return self.encode_data(data_mat, labels, use_probs=use_probs), labels

    def encode_texts(self, texts, use_probs=True, include_bn=False):
        X, _ = self.vectorizer.transform(texts)
        encodings = self.encode_data(X, None, use_probs=use_probs, include_bn=include_bn)
        return encodings

    def _get_data_iterator(self, data_mat, labels):
        x_size = data_mat.shape[0] * data_mat.shape[1]
        if x_size <= MAX_DESIGN_MATRIX and isinstance(data_mat, scipy.sparse.csr.csr_matrix):
            data_mat = mx.nd.sparse.csr_matrix(data_mat, dtype='float32')
        elif isinstance(data_mat, mx.nd.NDArray):
            data_mat = mx.nd.array(data_mat, dtype='float32')
        batch_size = min(data_mat.shape[0], self.max_batch_size)
        last_batch_size = data_mat.shape[0] % batch_size
        covars = mx.nd.one_hot(mx.nd.array(labels, dtype='int'), self.n_covars) \
            if self.covar_model and labels[:-last_batch_size] is not None else None
        if last_batch_size < 1:
            data_to_iter = data_mat
        else:
            data_to_iter = data_mat[:-last_batch_size]
        if x_size > MAX_DESIGN_MATRIX:
            logging.info("Sparse matrix has total size = {}. Using Sparse Matrix data batcher.".format(x_size))
            if covars is None:
                covars = mx.nd.zeros(data_to_iter.shape[0])
            infer_iter = DataIterLoader(SparseMatrixDataIter(data_to_iter, covars, batch_size, last_batch_handle='discard',
                                                             shuffle=False))
        else:
            infer_iter = DataIterLoader(mx.io.NDArrayIter(data_to_iter, covars,
                                                      batch_size, last_batch_handle='discard', shuffle=False))
        return infer_iter, last_batch_size

    def encode_data(self, data_mat, labels=None, use_probs=True, include_bn=False, target_entropy=1.0):
        infer_iter, last_batch_size = self._get_data_iterator(data_mat, labels)
        encodings = []
        for _, (data,labels) in enumerate(infer_iter):
            data = data.as_in_context(self.ctx)
            if self.covar_model and labels is not None:
                labels = labels.as_in_context(self.ctx)
                encs = self.model.encode_data_with_covariates(data, labels, include_bn=include_bn)
            else:
                encs = self.model.encode_data(data, include_bn=include_bn)
            if use_probs:
                e1 = (encs - mx.nd.min(encs, axis=1).expand_dims(1)).astype('float64')
                encs = list(mx.nd.softmax(e1).asnumpy())
                encs = list(map(partial(recalibrate_scores, target_entropy=target_entropy), encs))
            else:
                encs = list(encs.astype('float64').asnumpy())
            encodings.extend(encs)
        ## handle the last batch explicitly as NDArrayIter doesn't do that for us
        if last_batch_size > 0:
            last_data = mx.nd.sparse.csr_matrix(data_mat[-last_batch_size:], dtype='float32')
            data = last_data.as_in_context(self.ctx)
            if self.covar_model and labels is not None:
                labels = mx.nd.one_hot(mx.nd.array(labels[-last_batch_size:], dtype='int'),
                                       self.n_covars).as_in_context(self.ctx)
                encs = self.model.encode_data_with_covariates(data, labels)
            else:
                encs = self.model.encode_data(data)
            if use_probs:
                e1 = (encs - mx.nd.min(encs, axis=1).expand_dims(1)).astype('float64')
                encs = list(mx.nd.softmax(e1).asnumpy())
                encs = list(map(partial(recalibrate_scores, target_entropy=target_entropy), encs))
            else:
                encs = list(encs.astype('float64').asnumpy())
            encodings.extend(encs)
        return encodings

    def get_likelihood_stats(self, data_mat, n_samples=50):
        ## Notes:
        ## Following ideas in the paper:
        ## Bayesian Autoencoders: Analysing and Fixing the Bernoulli likelihood for Out-of-Distribution Detection
        data_iter, last_batch_size = self._get_data_iterator(data_mat, None)
        all_stats = []
        for _, (data, labels) in enumerate(data_iter):
            elbos = []
            for s in range(0,n_samples):
                elbo, _,_,_,_,_,_ = self.model(data.as_in_context(self.ctx), labels)
                elbos.append(list(elbo.asnumpy()))
            wd_cnts = data.sum(axis=1).asnumpy()
            elbos_np = np.array(elbos) / (wd_cnts + 1)
            elbos_means = list(elbos_np.mean(axis=0))
            elbos_var   = list(elbos_np.var(axis=0))
            all_stats.extend(list(zip(elbos_means, elbos_var)))
        return all_stats

    def get_top_k_words_per_topic(self, k):
        sorted_ids = self.model.get_ordered_terms()
        topic_terms = []
        for t in range(self.n_latent):
            top_k = [ self.vocab.idx_to_token[int(i)] for i in list(sorted_ids[:k, t]) ]
            topic_terms.append(top_k)
        return topic_terms

    def get_top_k_words_per_topic_encoder(self, k, dataloader, sample_size=-1):
        sorted_ids = self.model.get_ordered_terms_encoder(dataloader, sample_size=sample_size)
        topic_terms = []
        for t in range(self.n_latent):
            top_k = [ self.vocab.idx_to_token[int(i)] for i in list(sorted_ids[:k, t]) ]
            topic_terms.append(top_k)
        return topic_terms

    def get_top_k_words_per_topic_per_covariate(self, k):
        n_topics = self.n_latent
        w = self.model.cov_decoder.cov_inter_decoder.collect_params().get('weight').data()
        n_covars = int(w.shape[1] / n_topics)
        topic_terms = []
        for i in range(n_covars):
            cv_i_slice = w[:, (i * n_topics):((i+1) * n_topics)]
            sorted_ids = cv_i_slice.argsort(axis=0, is_ascend=False)
            cv_i_terms = []
            for t in range(n_topics):
                top_k = [ self.vocab.idx_to_token[int(i)] for i in list(sorted_ids[:k, t].asnumpy()) ]
                cv_i_terms.append(top_k)
            topic_terms.append(cv_i_terms)
        return topic_terms

    def get_top_k_words_per_topic_over_scalar_covariate(self, k, min_v=0.0, max_v=1.0, step=0.1):
        raise NotImplemented

    def predict_text(self, txt, pred_threshold=0.5):
        X_csr, _      = self.vectorizer.transform(txt)
        X = mx.nd.sparse.csr_matrix(X_csr, dtype='float32')
        encodings = self.encode_data(X, None, use_probs=True, include_bn=False)
        preds     = self.model.predict(X).asnumpy()
        inv_map = [0] * len(self.vectorizer.label_map)
        for k in self.vectorizer.label_map:
            inv_map[self.vectorizer.label_map[k]] = k
        if not self.model.multilabel:
            bests = np.argmax(preds, axis=1)
            best_strs = [ inv_map[best] for best in bests ]
        else:
            best_strs = [ inv_map[i] for i in list(np.where(preds > pred_threshold)[0]) ]
        return best_strs, encodings

class SeqVEDInferencer(BaseInferencer):
    """Inferencer for sequence variational encoder-decoder models using BERT
    """
    def __init__(self, model, bert_vocab, max_length, bow_vocab=None, pre_vectorizer=None, ctx=mx.cpu()):
        super().__init__(ctx)
        self.model     = model
        self.bert_base = model.bert
        self.tokenizer = BERTTokenizer(bert_vocab)
        self.transform = BERTSentenceTransform(self.tokenizer, max_length, pair=False)
        self.bow_vocab = bow_vocab
        self.vectorizer = pre_vectorizer or TMNTVectorizer(initial_vocabulary=bow_vocab)

    @classmethod
    def from_saved(cls, param_file=None, config_file=None, vocab_file=None, model_dir=None, max_length=128, ctx=mx.cpu()):
        if model_dir is not None:
            param_file = os.path.join(model_dir, 'model.params')
            vocab_file = os.path.join(model_dir, 'vocab.json')
            config_file = os.path.join(model_dir, 'model.config')
            serialized_vectorizer_file = os.path.join(model_dir, 'vectorizer.pkl')
        with open(config_file) as f:
            config = json.loads(f.read())
        with open(vocab_file) as f:
            voc_js = f.read()
        if os.path.exists(serialized_vectorizer_file):
            with open(serialized_vectorizer_file, 'rb') as fp:
                vectorizer = pickle.load(fp)
        else:
            vectorizer = None
        bow_vocab = nlp.Vocab.from_json(voc_js)
        bert_base, vocab = nlp.model.get_model(config['bert_model_name'],
                                               dataset_name=config['bert_data_name'],
                                               pretrained=True, ctx=ctx, use_pooler=True,
                                               use_decoder=False, use_classifier=False)
        latent_dist_t = config['latent_distribution']['dist_type']
        n_latent    = config['n_latent']
        kappa       = config['latent_distribution']['kappa']
        num_classes = config['n_labels']
        classifier_dropout = config['classifier_dropout']
        pad_id      = vocab[vocab.padding_token]
        latent_dist = HyperSphericalDistribution(n_latent, kappa=kappa, ctx=ctx)
        model = SeqBowVED(bert_base, latent_dist=latent_dist, bow_vocab_size = len(bow_vocab), num_classes=num_classes,
                          dropout=classifier_dropout)
        model.load_parameters(str(param_file), allow_missing=False, ignore_extra=True)
        model.latent_dist.post_init(ctx) # need to call this after loading parameters now
        return cls(model, vocab, max_length, bow_vocab, pre_vectorizer=vectorizer, ctx=ctx)

    def _embed_sequence(self, ids, segs):
        embeddings = self.bert_base.word_embed(ids) + self.bert_base.token_type_embed(segs)
        return embeddings.transpose((1,0,2))

    def _encode_from_embedding(self, embeddings, lens):
        outputs, _ = self.bert_base.encoder(embeddings, valid_length=lens)
        outputs = outputs.transpose((1,0,2))
        # outputs should be (batch, seq_len, C) shaped now
        pooled_out = self.bert_base._apply_pooling(outputs)
        topic_encoding = self.model.latent_dist.get_mu_encoding(pooled_out)
        return topic_encoding

    def prep_text(self, txt):    # used for integrated gradients
        tokens = self.tokenizer(txt)
        ids, lens, segs = self.transform((txt,))
        return ( tokens,
                 mx.nd.array([ids], dtype='int32'),
                 mx.nd.array([lens], dtype='float32'),
                 mx.nd.array([segs], dtype='int32') )

    def encode_text(self, txt):
        tokens, ids, lens, segs = self.prep_text(txt)
        _, enc = self.model.bert(ids.as_in_context(self.ctx),
                                              segs.as_in_context(self.ctx), lens.as_in_context(self.ctx))
        topic_encoding = self.model.latent_dist.get_mu_encoding(enc)
        return topic_encoding, tokens

    def predict_text(self, txt):
        encoding, _ = self.encode_text(txt)
        return self.model.classifier(encoding)

    def get_likelihood_stats(self, txt, n_samples=50):
        tokens, ids, lens, segs = self.prep_text(txt)
        bow_vector = mx.nd.array(self.vectorizer.vectorizer.transform([txt]), dtype='float32').expand_dims(0)
        elbos = []
        for s in range(n_samples):
            elbo, _, _, _, _ = self.model(ids, segs, lens, bow_vector)
            elbos.append(list(elbo.asnumpy()))
        wd_cnts = bow_vector.sum().asnumpy()
        elbos_np = np.array(elbos) / (wd_cnts + 1)
        elbos_means = list(elbos_np.mean(axis=0))
        elbos_var   = list(elbos_np.var(axis=0))
        return elbos_means, elbos_var

    def encode_data(self, dataloader, use_probs=False):
        encodings = []
        for _, seqs in enumerate(dataloader):
            ids, lens, segs, _, _ = seqs
            _, encs = self.model.bert(ids.as_in_context(self.ctx),
                                      segs.as_in_context(self.ctx), lens.astype('float32').as_in_context(self.ctx))
            encs = self.model.latent_dist.get_mu_encoding(encs)
            if use_probs:
                e1 = (encs - mx.nd.min(encs, axis=1).expand_dims(1)).astype('float64')
                encs = list(mx.nd.softmax(e1).asnumpy())
                topic_encodings = map(recalibrate, encs)
            else:
                topic_encodings = list(encs.astype('float64').asnumpy())
            encodings.extend(topic_encodings)
        return encodings

    def get_top_k_words_per_topic(self, k):
        if self.bow_vocab:
            sorted_ids = self.model.get_top_k_terms(k)
            topic_terms = []
            for t in range(self.model.n_latent):
                top_k = [ self.bow_vocab.idx_to_token[int(i)] for i in list(sorted_ids[:k, t]) ]
                topic_terms.append(top_k)
            return topic_terms
        else:
            raise Exception("Bow vocabulary required for Inferencer in order to provide topic terms")

class MetricSeqVEDInferencer(SeqVEDInferencer):
    """Inferencer for sequence variational encoder-decoder models using BERT trained via Metric Learning
    """
    def __init__(self, model, bert_vocab, max_length, bow_vocab=None, ctx=mx.cpu()):
        super().__init__(model, bert_vocab, max_length, bow_vocab=bow_vocab, ctx=ctx)

    @classmethod
    def from_saved(cls, param_file=None, config_file=None, vocab_file=None, model_dir=None, max_length=128, ctx=mx.cpu()):
        if model_dir is not None:
            param_file = os.path.join(model_dir, 'model.params')
            vocab_file = os.path.join(model_dir, 'vocab.json')
            config_file = os.path.join(model_dir, 'model.config')
        with open(config_file) as f:
            config = json.loads(f.read())
        with open(vocab_file) as f:
            voc_js = f.read()
        bow_vocab = nlp.Vocab.from_json(voc_js)
        bert_base, vocab = nlp.model.get_model(config['bert_model_name'],
                                               dataset_name=config['bert_data_name'],
                                               pretrained=True, ctx=ctx, use_pooler=True,
                                               use_decoder=False, use_classifier=False) #, output_attention=True)
        latent_dist_t = config['latent_distribution']['dist_type']
        n_latent    = config['n_latent']
        kappa       = config['latent_distribution']['kappa']
        num_classes = config['n_labels']
        classifier_dropout = config['classifier_dropout']
        pad_id      = vocab[vocab.padding_token]
        latent_dist = HyperSphericalDistribution(n_latent, kappa=kappa, ctx=ctx)
        model = MetricSeqBowVED(bert_base, latent_dist=latent_dist, bow_vocab_size = len(bow_vocab), n_latent=n_latent,
                                dropout=classifier_dropout)
        model.load_parameters(str(param_file), allow_missing=False, ignore_extra=True)
        return cls(model, vocab, max_length, bow_vocab, ctx)


#  Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License").
#  You may not use this file except in compliance with the License.
#  A copy of the License is located at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  or in the "license" file accompanying this file. This file is distributed
#  on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
#  express or implied. See the License for the specific language governing
#  permissions and limitations under the License.

import logging
import module_to_import  # noqa
import os
import mxnet
import numpy

logger = logging.getLogger(__name__)

def train(hyperparameters=None,
          input_data_config=None,
          channel_input_dirs=None,
          output_data_dir=None,
          model_dir=None,
          num_gpus=None,
          num_cpus=None,
          hosts=None,
          current_host=None):

    logger.info("hyperparameters: {}".format(hyperparameters))
    logger.info("input_data_config: {}".format(input_data_config))
    logger.info("channel_input_dirs: {}".format(channel_input_dirs))
    logger.info("num_gpus: {}".format(num_gpus))
    logger.info("num_cpus: {}".format(num_cpus))
    logger.info("hosts: {}".format(hosts))
    logger.info("current_host: {}".format(current_host))
    logger.info("expected_channel_names: {}".format(hyperparameters["expected_channel_names"]))

    expected_channel_names = hyperparameters["expected_channel_names"]

    assert set(input_data_config.keys()) == set(expected_channel_names)
    assert set(channel_input_dirs.keys()) == set(expected_channel_names)

    expected_gpus = hyperparameters["expected_gpus"]
    assert num_gpus == int(expected_gpus)

    expected_cpus = hyperparameters["expected_cpus"]
    assert num_cpus == int(expected_cpus)

    assert hosts
    assert current_host
    assert current_host in hosts

    # Do something with mxnet
    if num_gpus:
        ctx = mxnet.gpu(0)
    else:
        ctx = mxnet.cpu()
    arr_one = mxnet.ndarray.full((100, 100), 1, ctx=ctx)
    arr_two = mxnet.ndarray.full((100, 100), 2, ctx=ctx)
    arr_sum = (arr_one + arr_two).asnumpy()
    expected = mxnet.ndarray.full((100, 100), 3, ctx=ctx).asnumpy()
    logger.info("Expected Sum: {}".format(expected))
    logger.info("Received Sum: {}".format(arr_sum))
    logger.info("Equal?: {}".format(numpy.array_equal(expected, arr_sum)))
    assert numpy.array_equal(expected, arr_sum), "{} != {}".format(expected, arr_sum)

    # Assert that we can open the input files and they contain something
    for channel in input_data_config:
        assert open(os.path.join(channel_input_dirs[channel], "data.txt")).read(), "Cannot read {}".format(channel)

    # Write out to the model dir and the output_data_dir
    with open(os.path.join(model_dir, 'model.txt'), 'w') as f:
        f.write("model data")
    with open(os.path.join(output_data_dir, 'output.txt'), 'w') as f:
        f.write("output data")

    logger.info("Done test!")


import logging
import multiprocessing
import os
import sys
import warnings
import xml.etree.ElementTree as ET
from itertools import repeat
from pathlib import Path

import astropy.units as u
import click
import matplotlib.pyplot as plt
import numpy as np
from astropy.coordinates import SkyCoord
from astropy.modeling import models
from astropy.table import Table

import cscripts
import ctools
import gammalib
import Plot_SED

###
# export CALDB="$GAMMAPY_DATA/cta-1dc/caldb"
#

log = logging.getLogger(__name__)

# path config
BASE_PATH = Path("../make.py").parent

AVAILABLE_MODELS = [
    "point-pwl",
    "point-ecpl",
    "point-log-parabola",
    "point-pwl2",
    "point-ecpl-3fgl",
    "point-ecpl-4fgl",
    "point-template",
    "diffuse-cube",
    "disk-pwl",
    "gauss-pwl",
    "gauss-pwlsimple",
    "point-pwlsimple",
    "disk-pwlsimple",
    "point-pwltest",
    "point-pwl-time",
]

DPI = 600
#########
# Set the configuration parameters

E_MIN = 0.3  # TeV
E_MAX = 100  # TeV
N_BIN = 50
N_PIX = 50
BINSZ = 0.02  # deg
CALDB = "1dc"
IRF = "South_z20_50h"
COORDSYS = "CEL"
PROJ = "CAR"

#######

# log.info("Starting...")

@click.group()
@click.option(
    "--log-level", default="INFO", type=click.Choice(["DEBUG", "INFO", "WARNING"])
)
@click.option("--show-warnings", is_flag=True, help="Show warnings?")
def cli(log_level, show_warnings):
    logging.basicConfig(level=log_level)

    if not show_warnings:
        warnings.simplefilter("ignore")

@cli.command("all", help="Run all steps")
@click.argument("model", type=click.Choice(list(AVAILABLE_MODELS)))
@click.option(
    "--obs_ids", default=1, nargs=1, help="Select a single observation", type=int
)
def all_cmd(model, obs_ids):
    xmlfile = BASE_PATH / f"ctools/models/{model}.xml"
    inmodels = gammalib.GModels(str(xmlfile))
    log.info(f"Reading {xmlfile}")

    mod = inmodels[0]
    name = mod.name()

    ra = mod.spatial().region().ra()
    dec = mod.spatial().region().dec()

    evt = BASE_PATH / f"data/models/{model}/"

    for obsid in np.arange(obs_ids):

        event_file = evt / f"events_1h_{obsid:04d}.fits.gz"
        log.info(f"Reading {event_file}")
        outmodel = BASE_PATH / f"ctools/results/models/{model}/results_{obsid:04d}.xml"
        log.info(f"Reading {outmodel}")

        log.info(f"Analysing data...")
        stackedPipeline(
            name=name,
            obsfile=str(event_file),
            l=ra,
            b=dec,
            emin=E_MIN,
            emax=E_MAX,
            enumbins=N_BIN,
            nxpix=N_PIX,
            nypix=N_PIX,
            binsz=BINSZ,
            coordsys=COORDSYS,
            proj=PROJ,
            caldb=CALDB,
            irf=IRF,
            debug=False,
            inmodel=xmlfile,
            outmodel=outmodel,
        )

        log.info(f"Plot results...")
        os.system(
            f"python Plot_SED.py -m bkg_cube.xml -n {model} -s {str(outmodel.parent)}/{str(outmodel.stem)}.fits -b {str(outmodel.parent)}/{str(outmodel.stem)}.txt -o {str(outmodel.parent)}/{str(outmodel.stem)}.png"
        )

@cli.command("Plot_distrib", help="Make pull-distributions")
@click.argument("model", type=click.Choice(list(AVAILABLE_MODELS)))
@click.option(
    "--obs_ids", default=1, nargs=1, help="Select a single observation", type=int
)
def plot_pull_distrib(model, obs_ids):

    sim_mod = BASE_PATH / f"ctools/models/{model}.xml"

    source = Plot_SED.get_target(sim_mod, model)
    spec_model = source[0].attrib.get("type")
    spat_model = source[1].attrib.get("type")
    if spec_model == "PowerLaw":
        spec_param = Plot_SED.get_PowerLaw_model(source[0])
        tab_spec = Table()
        tab_spec["amplitude"] = [spec_param.amplitude.value]
        tab_spec["index"] = [spec_param.alpha.value]
        ampl = []
        e_ampl = []
        index = []
        e_index = []
    if spat_model == "PointSource":
        spat_param = Plot_SED.get_PointSource_model(source[1])
        tab_spat = Table()
        tab_spat["ra"] = [float(spat_param[0])]
        tab_spat["dec"] = [float(spat_param[1])]
        ra = []
        e_ra = []
        dec = []
        e_dec = []

    for obsid in np.arange(obs_ids):

        fit_mod = BASE_PATH / f"ctools/results/models/{model}/results_{obsid:04d}.xml"
        source_fit = Plot_SED.get_target(fit_mod, model)
        spec_fit_model = source_fit[0].attrib.get("type")
        spat_fit_model = source_fit[1].attrib.get("type")

        if spec_fit_model == "PowerLaw":
            spec_fit_param = Plot_SED.get_PowerLaw_model_error(source_fit[0])
            ampl.append(float(spec_fit_param[0]))
            index.append(float(spec_fit_param[1]))
            e_ampl.append(float(spec_fit_param[2]))
            e_index.append(float(spec_fit_param[3]))

        if spat_fit_model == "PointSource":
            spat_fit_param = Plot_SED.get_PointSource_model_error(source_fit[1])
            ra.append(float(spat_fit_param[0]))
            dec.append(float(spat_fit_param[1]))
            e_ra.append(float(spat_fit_param[2]))
            e_dec.append(float(spat_fit_param[3]))

    if spec_fit_model == "PowerLaw":
        tab = Table()
        tab["amplitude"] = ampl
        tab["e_amplitude"] = e_ampl
        tab["index"] = index
        tab["e_index"] = e_index
        names = [name for name in tab.colnames if "e_" not in name]
        for name in names:
            pull = (tab[name] - tab_spec[name]) / tab["e_" + name]
            plt.hist(pull, bins=21, normed=True, range=(-5, 5))
            plt.xlim(-5, 5)
            plt.xlabel("(value - value_true) / error")
            plt.ylabel("PDF")
            plt.title(f"Pull distribution for {model}: {name} ")
            filename = f"ctools/results/models/{model}/pull-distribution-{name}.png"
            save_figure(filename)

    if spat_fit_model == "PointSource":
        tab = Table()
        tab["ra"] = ra
        tab["e_ra"] = e_ra
        tab["dec"] = dec
        tab["e_dec"] = e_dec
        names = [name for name in tab.colnames if "e_" not in name]
        for name in names:
            pull = (tab[name] - tab_spat[name]) / tab["e_" + name]
            plt.hist(pull, bins=21, normed=True, range=(-5, 5))
            plt.xlim(-5, 5)
            plt.xlabel("(value - value_true) / error")
            plt.ylabel("PDF")
            plt.title(f"Pull distribution for {model}: {name} ")
            filename = f"ctools/results/models/{model}/pull-distribution-{name}.png"
            save_figure(filename)

def stackedPipeline(
    name="Crab",
    obsfile="index.xml",
    l=0.01,
    b=0.01,
    emin=0.1,
    emax=100.0,
    enumbins=20,
    nxpix=200,
    nypix=200,
    binsz=0.02,
    coordsys="CEL",
    proj="CAR",
    caldb="prod2",
    irf="acdc1a",
    debug=False,
    inmodel="Crab",
    outmodel="results",
):
    """
    Simulation and stacked analysis pipeline

    Parameters
    ----------
    obs : `~gammalib.GObservations`
    Observation container
    ra : float, optional
    Right Ascension of counts cube centre (deg)
    dec : float, optional
    Declination of Region of counts cube centre (deg)
    emin : float, optional
    Minimum energy (TeV)
    emax : float, optional
    Maximum energy (TeV)
    enumbins : int, optional
    Number of energy bins
    nxpix : int, optional
    Number of pixels in X axis
    nypix : int, optional
    Number of pixels in Y axis
    binsz : float, optional
    Pixel size (deg)
    coordsys : str, optional
    Coordinate system
    proj : str, optional
    Coordinate projection
    debug : bool, optional
    Debug function
    """

    # Bin events into counts map
    bin = ctools.ctbin()
    bin["inobs"] = obsfile
    bin["ebinalg"] = "LOG"
    bin["emin"] = emin
    bin["emax"] = emax
    bin["enumbins"] = enumbins
    bin["nxpix"] = nxpix
    bin["nypix"] = nypix
    bin["binsz"] = binsz
    bin["coordsys"] = coordsys
    bin["proj"] = proj
    bin["xref"] = l
    bin["yref"] = b
    bin["debug"] = debug
    bin["outobs"] = "cntcube.fits"
    bin.execute()
    print("Datacube : done!")

    # Create exposure cube
    expcube = ctools.ctexpcube()
    # expcube['incube']=bin.obs()
    expcube["inobs"] = obsfile
    expcube["incube"] = "NONE"
    expcube["ebinalg"] = "LOG"
    expcube["caldb"] = caldb
    expcube["irf"] = irf
    expcube["emin"] = emin
    expcube["emax"] = emax
    expcube["enumbins"] = enumbins
    expcube["nxpix"] = nxpix
    expcube["nypix"] = nypix
    expcube["binsz"] = binsz
    expcube["coordsys"] = coordsys
    expcube["proj"] = proj
    expcube["xref"] = l
    expcube["yref"] = b
    expcube["debug"] = debug
    expcube["outcube"] = "cube_exp.fits"
    expcube.execute()
    print("Expcube : done!")

    # Create PSF cube
    psfcube = ctools.ctpsfcube()
    psfcube["inobs"] = obsfile
    psfcube["incube"] = "NONE"
    psfcube["ebinalg"] = "LOG"
    psfcube["caldb"] = caldb
    psfcube["irf"] = irf
    psfcube["emin"] = emin
    psfcube["emax"] = emax
    psfcube["enumbins"] = enumbins
    psfcube["nxpix"] = 10
    psfcube["nypix"] = 10
    psfcube["binsz"] = 1.0
    psfcube["coordsys"] = coordsys
    psfcube["proj"] = proj
    psfcube["xref"] = l
    psfcube["yref"] = b
    psfcube["debug"] = debug
    psfcube["outcube"] = "psf_cube.fits"
    psfcube.execute()
    print("Psfcube : done!")

    edispcube = ctools.ctedispcube()
    edispcube["inobs"] = obsfile
    edispcube["ebinalg"] = "LOG"
    edispcube["incube"] = "NONE"
    edispcube["caldb"] = caldb
    edispcube["irf"] = irf
    edispcube["xref"] = l
    edispcube["yref"] = b
    edispcube["proj"] = proj
    edispcube["coordsys"] = coordsys
    edispcube["binsz"] = 1.0
    edispcube["nxpix"] = 10
    edispcube["nypix"] = 10
    edispcube["emin"] = emin
    edispcube["emax"] = emax
    edispcube["enumbins"] = enumbins
    edispcube["outcube"] = "edisp_cube.fits"
    edispcube["debug"] = debug
    edispcube.execute()
    print("Edispcube : done!")

    # Create background cube
    bkgcube = ctools.ctbkgcube()
    bkgcube["inobs"] = obsfile
    bkgcube["incube"] = "cntcube.fits"
    bkgcube["caldb"] = caldb
    bkgcube["irf"] = irf
    bkgcube["debug"] = debug
    bkgcube["inmodel"] = str(inmodel)
    bkgcube["outcube"] = "bkg_cube.fits"
    bkgcube["outmodel"] = "bkg_cube.xml"
    bkgcube.execute()
    print("Bkgcube : done!")

    # Fix the instrumental background parameters
    bkgcube.models()["BackgroundModel"]["Prefactor"].fix()
    bkgcube.models()["BackgroundModel"]["Index"].fix()
    #
    #    # Attach background model to observation container
    bin.obs().models(bkgcube.models())
    #
    #
    #    # Set Exposure and Psf cube for first CTA observation
    #    # (ctbin will create an observation with a single container)
    bin.obs()[0].response(
        expcube.expcube(), psfcube.psfcube(), edispcube.edispcube(), bkgcube.bkgcube()
    )

    # Perform maximum likelihood fitting
    like = ctools.ctlike(bin.obs())
    #    like['inmodel']='bkg_cube.xml'
    like["edisp"] = True
    #    like['edispcube'] = 'edisp_cube.fits'
    #    like['expcube'] = 'cube_exp.fits'
    #    like['psfcube'] = 'psf_cube.fits'
    #    like['bkgcube'] = 'bkg_cube.fits'
    like["outmodel"] = str(outmodel)
    #    like['outcovmat']=inmodel+'_covmat.txt'
    like["debug"] = debug  # Switch this always on for results in console
    # like['statistic']='CSTAT'
    like.execute()
    print("Likelihood : done!")

    # Set the best-fit models (from ctlike) for the counts cube
    bin.obs().models(like.obs().models())

    # Obtain the best-fit butterfly
    try:
        butterfly = ctools.ctbutterfly(bin.obs())
        butterfly["srcname"] = name
        butterfly["inmodel"] = str(outmodel)
        butterfly["edisp"] = True
        butterfly["emin"] = emin
        butterfly["emax"] = emax
        butterfly["outfile"] = str(outmodel.parent) + "/" + str(outmodel.stem) + ".txt"
        butterfly["debug"] = debug  # Switch this always on for results in console
        # like['statistic']='CSTAT'
        butterfly.execute()
        print("Butterfly : done!")
    except:
        print("I COULDN'T CALCULATE THE BUTTERFLY....")

    # Extract the spectrum
    try:
        csspec = cscripts.csspec(bin.obs())
        csspec["srcname"] = name
        csspec["inmodel"] = str(outmodel)
        csspec["method"] = "AUTO"
        csspec["ebinalg"] = "LOG"
        csspec["emin"] = emin
        csspec["emax"] = emax
        csspec["enumbins"] = 10
        csspec["edisp"] = True
        csspec["outfile"] = str(outmodel.parent) + "/" + str(outmodel.stem) + ".fits"
        csspec["debug"] = debug  # Switch this always on for results in console
        csspec.execute()
        print("Csspec : done!")
    except:
        print("I COULDN'T CALCULATE THE SPECTRUM....")

    # Return
    return

def save_figure(filename):
    path = BASE_PATH / filename
    path.parent.mkdir(parents=True, exist_ok=True)
    log.info(f"Writing {path}")
    plt.savefig(path, dpi=DPI)
    plt.clf()
    plt.close()

if __name__ == "__main__":
    cli()

import numpy,import astropy,from astropy


import torch
from mmdet.core import bbox2result
from ..registry import DETECTORS
from .single_stage import SingleStageDetector
import numpy as np
import cv2
from mmcv.image import imread, imwrite
from mmdet.core import tensor2imgs, get_classes
import mmcv
from pycocotools import mask as maskUtils
from mmdet.core.post_processing import interplate

@DETECTORS.register_module
class DenseRepPointsMaskDetector(SingleStageDetector):

    def __init__(self,
                 backbone,
                 neck,
                 bbox_head,
                 train_cfg=None,
                 test_cfg=None,
                 pretrained=None):
        super(DenseRepPointsMaskDetector, self).__init__(backbone, neck, bbox_head, train_cfg,
                                                         test_cfg, pretrained)

        @property
        def with_mask(self):
            return True

    def forward_train(self,
                      img,
                      img_metas,
                      gt_bboxes,
                      gt_labels,
                      gt_bboxes_ignore=None,
                      gt_masks=None):
        x = self.extract_feat(img)
        outs = self.bbox_head(x, test=False)
        loss_inputs = outs + (gt_bboxes, gt_masks, gt_labels, img_metas, self.train_cfg)
        losses = self.bbox_head.loss(
            *loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore)
        return losses

    def simple_test(self, img, img_meta, rescale=False):
        x = self.extract_feat(img)
        outs = self.bbox_head(x, test=True)
        bbox_inputs = outs + (img_meta, self.test_cfg, rescale)
        bbox_list = self.bbox_head.get_bboxes(*bbox_inputs)

        _, det_points, det_masks, det_cls_labels = bbox_list[0]
        if det_points.shape[0] > 0:
            _det_points = det_points[:, :-1].reshape(det_points.shape[0], -1, 2)
            _det_masks = det_masks[:, :-1].reshape(det_masks.shape[0], -1, 1)
            _det_points = torch.cat([_det_points, _det_masks], dim=-1).reshape(det_points.shape[0], -1)
            det_points = torch.cat([_det_points, det_points[:, [-1]]], dim=-1)

        ori_shape = img_meta[0]['ori_shape']
        scale_factor = img_meta[0]['scale_factor']
        bbox_results = [
            bbox2result(det_bbox, det_labels, self.bbox_head.num_classes)
            for det_bbox, det_pts, det_masks, det_labels in bbox_list
        ][0]
        pts_results = [
            pts2result(det_points, det_cls_labels, self.bbox_head.num_classes)
        ][0]
        segm_results = [
            self.get_seg_masks(det_masks[:, :-1], det_pts[:, :-1], det_bbox, det_labels, self.test_cfg,
                               ori_shape, scale_factor, rescale)
            for i, (det_bbox, det_pts, det_masks, det_labels) in enumerate(bbox_list)
        ][0]

        if rescale:
            return bbox_results, segm_results
        else:
            # store pts_result if we want visualization
            return (bbox_results, segm_results), pts_results

    def get_seg_masks(self, pts_score, det_pts, det_bboxes, det_labels, test_cfg,
                      ori_shape, scale_factor, rescale=False):
        """Get segmentation masks from mask_pred and bboxes.

        Args:
            pts_score (Tensor or ndarray): shape (n, #class+1, num_pts).
                For single-scale testing, pts_pred is the direct output of
                model, whose type is Tensor, while for multi-scale testing,
                it will be converted to numpy array outside of this method.
            det_pts (Tensor): shape (n, num_pts)
            det_bboxes (Tensor): shape (n, 4/5)
            det_labels (Tensor): shape (n, )
            test_cfg (dict): rcnn testing config
            ori_shape: original image size

        Returns:
            list[list]: encoded masks
        """

        cls_segms = [[] for _ in range(self.bbox_head.num_classes - 1)]
        bboxes = det_bboxes.cpu().numpy()[:, :4]
        labels = det_labels.cpu().numpy() + 1

        if rescale:
            img_h, img_w = ori_shape[:2]
        else:
            img_h = np.round(ori_shape[0] * scale_factor).astype(np.int32)
            img_w = np.round(ori_shape[1] * scale_factor).astype(np.int32)

        scale_factor = 1.0

        for i in range(bboxes.shape[0]):
            bbox = (bboxes[i, :] / scale_factor).astype(np.int32)
            label = labels[i]
            w = max(bbox[2] - bbox[0] + 1, 1)
            h = max(bbox[3] - bbox[1] + 1, 1)
            _pts_score = pts_score[i]
            im_mask = np.zeros((img_h, img_w), dtype=np.uint8)
            im_pts = det_pts[i].clone()
            im_pts = im_pts.reshape(-1, 2)
            im_pts[:, 0] = (im_pts[:, 0] - bbox[0]) * (28 - 1) / (w - 1)
            im_pts[:, 1] = (im_pts[:, 1] - bbox[1]) * (28 - 1) / (h - 1)
            im_pts_score = _pts_score
            _corner_mask_pred_ = torch.zeros(2, 2)
            bbox_mask = interplate(_corner_mask_pred_, im_pts.cpu(), im_pts_score.cpu(), (h, w)).numpy()
            bbox_mask = (bbox_mask > test_cfg.get('mask_thr_binary', 0.5)).astype(np.uint8)
            im_mask[bbox[1]:bbox[1] + h, bbox[0]:bbox[0] + w] = bbox_mask
            rle = maskUtils.encode(
                np.array(im_mask[:, :, np.newaxis], order='F'))[0]
            cls_segms[label - 1].append(rle)
        return cls_segms

    def show_result(self,
                    data,
                    result,
                    img_norm_cfg,
                    dataset='coco',
                    score_thr=0.3,
                    out_file=None
                    ):
        bbox_seg_result, pts_result = result
        if isinstance(bbox_seg_result, tuple):
            bbox_result, segm_result = bbox_seg_result
        else:
            bbox_result, segm_result = bbox_seg_result, None
        img_tensor = data['img'][0]
        img_metas = data['img_meta'][0].data[0]
        imgs = tensor2imgs(img_tensor, **img_norm_cfg)
        assert len(imgs) == len(img_metas)

        if dataset is None:
            class_names = self.CLASSES
        elif isinstance(dataset, str):
            class_names = get_classes(dataset)
        elif isinstance(dataset, (list, tuple)):
            class_names = dataset
        else:
            raise TypeError(
                'dataset must be a valid dataset name or a sequence'
                ' of class names, not {}'.format(type(dataset)))

        for img, img_meta in zip(imgs, img_metas):
            h, w, _ = img_meta['img_shape']
            img_show = img[:h, :w, :]

            pts = np.vstack(pts_result)
            bboxes = np.vstack(bbox_result)
            # draw segmentation masks
            if segm_result is not None:
                segms = mmcv.concat_list(segm_result)
                inds = np.where(bboxes[:, -1] > score_thr)[0]
                for i in inds:
                    color_mask = np.random.randint(
                        0, 256, (1, 3), dtype=np.uint8)
                    mask = maskUtils.decode(segms[i]).astype(np.bool)
                    img_show[mask] = img_show[mask] * 0.5 + color_mask * 0.5
            # draw bounding boxes
            labels = [
                np.full(pts.shape[0], i, dtype=np.int32)
                for i, pts in enumerate(pts_result)
            ]
            labels = np.concatenate(labels)
            imshow_det_pts(
                img_show,
                bboxes,
                pts,
                labels,
                out_file=out_file,
                class_names=class_names,
                score_thr=score_thr)

def pts2result(pts, labels, num_classes):
    """Convert detection results to a list of numpy arrays.

    Args:
        bboxes (Tensor): shape (n, 5)
        labels (Tensor): shape (n, )
        num_classes (int): class number, including background class

    Returns:
        list(ndarray): bbox results of each class
    """
    if pts.shape[0] == 0:
        return [
            np.zeros((0, pts.shape[1]), dtype=np.float32) for i in range(num_classes - 1)
        ]
    else:
        pts = pts.cpu().numpy()
        labels = labels.cpu().numpy()
        return [pts[labels == i, :] for i in range(num_classes - 1)]

def masks2result(masks, labels, num_classes):
    if masks.shape[0] == 0:
        return [
            np.zeros((0, masks.shape[1]), dtype=np.float32) for i in range(num_classes - 1)
        ]
    else:
        masks = masks.cpu().numpy()
        labels = labels.cpu().numpy()
        return [masks[labels == i, :] for i in range(num_classes - 1)]

def imshow_det_pts(img,
                   bboxes,
                   pts,
                   labels,
                   class_names=None,
                   score_thr=0,
                   bbox_color='green',
                   text_color='green',
                   thickness=1,
                   font_scale=0.5,
                   show=True,
                   win_name='',
                   wait_time=0,
                   out_file=None):
    """Draw bboxes and class labels (with scores) on an image.

    Args:
        img (str or ndarray): The image to be displayed.
        bboxes (ndarray): Bounding boxes (with scores), shaped (n, 4) or
            (n, 5).
        labels (ndarray): Labels of bboxes.
        class_names (list[str]): Names of each classes.
        score_thr (float): Minimum score of bboxes to be shown.
        bbox_color (str or tuple or :obj:`Color`): Color of bbox lines.
        text_color (str or tuple or :obj:`Color`): Color of texts.
        thickness (int): Thickness of lines.
        font_scale (float): Font scales of texts.
        show (bool): Whether to show the image.
        win_name (str): The window name.
        wait_time (int): Value of waitKey param.
        out_file (str or None): The filename to write the image.
    """
    assert pts.ndim == 2
    assert labels.ndim == 1
    assert pts.shape[0] == labels.shape[0]
    img = imread(img)

    if score_thr > 0:
        scores = pts[:, -1]
        inds = scores > score_thr
        pts = pts[inds, :]
        bboxes = bboxes[inds, :]
        labels = labels[inds]

    for pt, bbox, label in zip(pts, bboxes, labels):
        r_color = (np.random.randint(0, 256), np.random.randint(0, 256), np.random.randint(0, 256))
        pt_x = pt[0:-1:3]
        pt_y = pt[1:-1:3]
        pt_score = pt[2:-1:3]
        left_top = (int(bbox[0]), int(bbox[1]))
        right_bottom = (int(bbox[2]), int(bbox[3]))
        cv2.rectangle(
            img, left_top, right_bottom, r_color, thickness=thickness)
        for i in range(pt_x.shape[0]):
            r_fore_color = r_color
            r_back_color = [int(i / 2) for i in r_color]
            r_show_color = r_fore_color if pt_score[i] > 0.5 else r_back_color
            cv2.circle(img, (pt_x[i], pt_y[i]), 1, r_show_color, thickness=2)
        label_text = class_names[
            label] if class_names is not None else 'cls {}'.format(label)
        if len(pt) % 2 == 1:
            label_text += '|{:.02f}'.format(pt[-1])
        cv2.putText(img, label_text, (int(bbox[0]), int(bbox[1]) - 2),
                    cv2.FONT_HERSHEY_COMPLEX, font_scale, r_color)

    if show:
        imshow(img, win_name, wait_time)
    if out_file is not None:
        imwrite(img, out_file)

def imshow(img, win_name='', wait_time=0):
    """Show an image.

    Args:
        img (str or ndarray): The image to be displayed.
        win_name (str): The window name.
        wait_time (int): Value of waitKey param.
    """
    cv2.imshow(win_name, imread(img))
    cv2.waitKey(wait_time)

def make_heatmap(img_shape, pts, sigma=10):
    """ Make the ground-truth for landmark.
    img: the original color image
    labels: label with the Gaussian center(s) [[x0, y0],[x1, y1],...]
    sigma: sigma of the Gaussian.
    one_mask_per_point: masks for each point in different channels?
    """
    h, w = img_shape
    pts = torch.tensor(pts).cuda()
    heatmap = make_gaussian((h, w), pts, sigma=sigma)  # (N, H, W)
    heatmap = heatmap.max(0)[0]  # (H, W)
    return heatmap

def make_gaussian(size, pts, sigma=1):
    """ Make a square gaussian kernel.
    size: is the dimensions of the output gaussian
    sigma: is full-width-half-maximum, which
    can be thought of as an effective radius.
    """
    x = torch.arange(0, size[1], 1, dtype=torch.float32).cuda()
    y = torch.arange(0, size[0], 1, dtype=torch.float32).cuda()
    x = x.unsqueeze(0).unsqueeze(0)
    y = y.unsqueeze(0).unsqueeze(-1)
    x0 = pts[:, 0].unsqueeze(-1).unsqueeze(-1)
    y0 = pts[:, 1].unsqueeze(-1).unsqueeze(-1)
    return torch.exp(-4 * np.log(2) * ((x - x0) ** 2 + (y - y0) ** 2) / sigma ** 2)


import numpy as np
import torch
from dataloader.sampler import CategoriesSampler

def set_up_datasets(args):
    if args.dataset == 'cifar100':
        import dataloader.cifar100.cifar as Dataset
        args.base_class = 60
        args.num_classes=100
        args.way = 5
        args.shot = 5
        args.sessions = 9
    if args.dataset == 'cub200':
        import dataloader.cub200.cub200 as Dataset
        args.base_class = 100
        args.num_classes = 200
        args.way = 10
        args.shot = 5
        args.sessions = 11
    if args.dataset == 'mini_imagenet':
        import dataloader.miniimagenet.miniimagenet as Dataset
        args.base_class = 60
        args.num_classes=100
        args.way = 5
        args.shot = 5
        args.sessions = 9
    args.Dataset=Dataset
    return args

def get_dataloader(args,session):
    if session == 0:
        trainset, trainloader, testloader = get_base_dataloader(args)
    else:
        trainset, trainloader, testloader = get_new_dataloader(args)
    return trainset, trainloader, testloader

def get_base_dataloader(args):
    txt_path = "data/index_list/" + args.dataset + "/session_" + str(0 + 1) + '.txt'
    class_index = np.arange(args.base_class)
    if args.dataset == 'cifar100':

        trainset = args.Dataset.CIFAR100(root=args.dataroot, train=True, download=True,
                                         index=class_index, base_sess=True)
        testset = args.Dataset.CIFAR100(root=args.dataroot, train=False, download=False,
                                        index=class_index, base_sess=True)

    if args.dataset == 'cub200':
        trainset = args.Dataset.CUB200(root=args.dataroot, train=True,
                                       index=class_index, base_sess=True)
        testset = args.Dataset.CUB200(root=args.dataroot, train=False, index=class_index)

    if args.dataset == 'mini_imagenet':
        trainset = args.Dataset.MiniImageNet(root=args.dataroot, train=True,
                                             index=class_index, base_sess=True)
        testset = args.Dataset.MiniImageNet(root=args.dataroot, train=False, index=class_index)

    trainloader = torch.utils.data.DataLoader(dataset=trainset, batch_size=args.batch_size_base, shuffle=True,
                                              num_workers=8, pin_memory=True)
    testloader = torch.utils.data.DataLoader(
        dataset=testset, batch_size=args.test_batch_size, shuffle=False, num_workers=8, pin_memory=True)

    return trainset, trainloader, testloader

def get_base_dataloader_meta(args):
    txt_path = "data/index_list/" + args.dataset + "/session_" + str(0 + 1) + '.txt'
    class_index = np.arange(args.base_class)
    if args.dataset == 'cifar100':
        trainset = args.Dataset.CIFAR100(root=args.dataroot, train=True, download=True,
                                         index=class_index, base_sess=True)
        testset = args.Dataset.CIFAR100(root=args.dataroot, train=False, download=False,
                                        index=class_index, base_sess=True)

    if args.dataset == 'cub200':
        trainset = args.Dataset.CUB200(root=args.dataroot, train=True,
                                       index_path=txt_path)
        testset = args.Dataset.CUB200(root=args.dataroot, train=False,
                                      index=class_index)
    if args.dataset == 'mini_imagenet':
        trainset = args.Dataset.MiniImageNet(root=args.dataroot, train=True,
                                             index_path=txt_path)
        testset = args.Dataset.MiniImageNet(root=args.dataroot, train=False,
                                            index=class_index)

    # DataLoader(test_set, batch_sampler=sampler, num_workers=8, pin_memory=True)
    sampler = CategoriesSampler(trainset.targets, args.train_episode, args.episode_way,
                                args.episode_shot + args.episode_query)

    trainloader = torch.utils.data.DataLoader(dataset=trainset, batch_sampler=sampler, num_workers=args.num_workers,
                                              pin_memory=True)

    testloader = torch.utils.data.DataLoader(
        dataset=testset, batch_size=args.test_batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True)

    return trainset, trainloader, testloader

def get_new_dataloader(args,session):
    txt_path = "data/index_list/" + args.dataset + "/session_" + str(session + 1) + '.txt'
    if args.dataset == 'cifar100':
        class_index = open(txt_path).read().splitlines()
        trainset = args.Dataset.CIFAR100(root=args.dataroot, train=True, download=False,
                                         index=class_index, base_sess=False)
    if args.dataset == 'cub200':
        trainset = args.Dataset.CUB200(root=args.dataroot, train=True,
                                       index_path=txt_path)
    if args.dataset == 'mini_imagenet':
        trainset = args.Dataset.MiniImageNet(root=args.dataroot, train=True,
                                       index_path=txt_path)
    if args.batch_size_new == 0:
        batch_size_new = trainset.__len__()
        trainloader = torch.utils.data.DataLoader(dataset=trainset, batch_size=batch_size_new, shuffle=False,
                                                  num_workers=args.num_workers, pin_memory=True)
    else:
        trainloader = torch.utils.data.DataLoader(dataset=trainset, batch_size=args.batch_size_new, shuffle=True,
                                                  num_workers=args.num_workers, pin_memory=True)

    # test on all encountered classes
    class_new = get_session_classes(args, session)

    if args.dataset == 'cifar100':
        testset = args.Dataset.CIFAR100(root=args.dataroot, train=False, download=False,
                                        index=class_new, base_sess=False)
    if args.dataset == 'cub200':
        testset = args.Dataset.CUB200(root=args.dataroot, train=False,
                                      index=class_new)
    if args.dataset == 'mini_imagenet':
        testset = args.Dataset.MiniImageNet(root=args.dataroot, train=False,
                                      index=class_new)

    testloader = torch.utils.data.DataLoader(dataset=testset, batch_size=args.test_batch_size, shuffle=False,
                                             num_workers=args.num_workers, pin_memory=True)

    return trainset, trainloader, testloader

def get_session_classes(args,session):
    class_list=np.arange(args.base_class + session * args.way)
    return class_list