from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from scipy import misc
import sys
import os
import argparse
#import tensorflow as tf
import numpy as np
import mxnet as mx
import random
import cv2
import sklearn
from sklearn.decomposition import PCA
from time import sleep
from easydict import EasyDict as edict
from mtcnn_detector import MtcnnDetector
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'src', 'common'))
import face_image
import face_preprocess

def do_flip(data):
  for idx in xrange(data.shape[0]):
    data[idx,:,:] = np.fliplr(data[idx,:,:])

def get_model(ctx, image_size, model_str, layer):
  _vec = model_str.split(',')
  assert len(_vec)==2
  prefix = _vec[0]
  epoch = int(_vec[1])
  print('loading',prefix, epoch)
  sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
  all_layers = sym.get_internals()
  sym = all_layers[layer+'_output']
  model = mx.mod.Module(symbol=sym, context=ctx, label_names = None)
  #model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0], image_size[1]))], label_shapes=[('softmax_label', (args.batch_size,))])
  model.bind(data_shapes=[('data', (1, 3, image_size[0], image_size[1]))])
  model.set_params(arg_params, aux_params)
  return model

class FaceModel:
  def __init__(self, args):
    self.args = args
    ctx = mx.gpu(args.gpu)
    _vec = args.image_size.split(',')
    assert len(_vec)==2
    image_size = (int(_vec[0]), int(_vec[1]))
    self.model = None
    self.ga_model = None
    if len(args.model)>0:
      self.model = get_model(ctx, image_size, args.model, 'fc1')
    if len(args.ga_model)>0:
      self.ga_model = get_model(ctx, image_size, args.ga_model, 'fc1')

    self.threshold = args.threshold
    self.det_minsize = 50
    self.det_threshold = [0.6,0.7,0.8]
    #self.det_factor = 0.9
    self.image_size = image_size
    mtcnn_path = os.path.join(os.path.dirname(__file__), 'mtcnn-model')
    if args.det==0:
      detector = MtcnnDetector(model_folder=mtcnn_path, ctx=ctx, num_worker=1, accurate_landmark = True, threshold=self.det_threshold)
    else:
      detector = MtcnnDetector(model_folder=mtcnn_path, ctx=ctx, num_worker=1, accurate_landmark = True, threshold=[0.0,0.0,0.2])
    self.detector = detector

  def get_input(self, face_img):
    ret = self.detector.detect_face(face_img, det_type = self.args.det)
    if ret is None:
      return None
    bbox, points = ret
    if bbox.shape[0]==0:
      return None
    #bbox = bbox[0,0:4]
    #bbx selection
    bbw = bbox[:, 2] - bbox[:, 0] + 1
    bbh = bbox[:, 3] - bbox[:, 1] + 1
    area = bbw * bbh
    ind = np.where(area == np.max(area))
    bbox = bbox[ind[0][0],0:4]
    ###
    #points = points[0,:].reshape((2,5)).T
    points = points[ind[0][0],:].reshape((2,5)).T
    #print(bbox)
    #print(points)
    nimg = face_preprocess.preprocess(face_img, bbox, points, image_size='112,112')
    nimg = cv2.cvtColor(nimg, cv2.COLOR_BGR2RGB)
    aligned = np.transpose(nimg, (2,0,1))
    return aligned

  def get_input_center_box(self, face_img):
    ret = self.detector.detect_face(face_img, det_type = self.args.det)
    if ret is None:
      return None
    bbox, points = ret
    if bbox.shape[0]==0:
      return None

    ##bbox = bbox[0,0:4]
    ##bbx selection

    #bbw = bbox[:, 2] - bbox[:, 0] + 1
    #bbh = bbox[:, 3] - bbox[:, 1] + 1
    #area = bbw * bbh
    #ind = np.where(area == np.max(area))
    #bbox = bbox[ind[0][0],0:4]

    nrof_faces = bbox.shape[0]
    det = bbox[:,0:4]
    img_size = np.asarray(face_img.shape)[0:2]

    if nrof_faces>1:
        bounding_box_size = (det[:,2]-det[:,0])*(det[:,3]-det[:,1])
        img_center = img_size / 2
        offsets = np.vstack([ (det[:,0]+det[:,2])/2-img_center[1], (det[:,1]+det[:,3])/2-img_center[0]  ])
        offset_dist_squared = np.sum(np.power(offsets,2.0),0)
        index = np.argmax(bounding_box_size-offset_dist_squared*2.0) # some extra weight on the centering
        det = det[index,:]
        points = points[index,:]
    det = np.squeeze(det)
    bbox = det
    #bb = np.zeros(4, dtype=np.int32)
    #bb[0] = np.maximum(det[0]-args.margin/2, 0)
    #bb[1] = np.maximum(det[1]-args.margin/2, 0)
    #bb[2] = np.minimum(det[2]+args.margin/2, img_size[1])
    #bb[3] = np.minimum(det[3]+args.margin/2, img_size[0])

    #cropped = img[bb[1]:bb[3],bb[0]:bb[2],:]
    #scaled = misc.imresize(cropped, (args.image_size, args.image_size), interp='bilinear')
    #nrof_successfully_aligned += 1
    #misc.imsave(output_filename, scaled)
    #text_file.write('%s %d %d %d %d\n' % (output_filename, bb[0], bb[1], bb[2], bb[3]))
    ###
    #points = points[0,:].reshape((2,5)).T
    points = points.reshape((2,5)).T
    #print(bbox)
    #print(points)
    nimg = face_preprocess.preprocess(face_img, bbox, points, image_size='112,112')
    #nimg = cv2.cvtColor(nimg, cv2.COLOR_BGR2RGB)
    #aligned = np.transpose(nimg, (2,0,1))
    return nimg #aligned

  def get_input_center_box_return_bbx(self, face_img):
    ret = self.detector.detect_face(face_img, det_type = self.args.det)
    if ret is None:
      return None,None
    bbox, points = ret
    if bbox.shape[0]==0:
      return None,None

    ##bbox = bbox[0,0:4]
    ##bbx selection

    #bbw = bbox[:, 2] - bbox[:, 0] + 1
    #bbh = bbox[:, 3] - bbox[:, 1] + 1
    #area = bbw * bbh
    #ind = np.where(area == np.max(area))
    #bbox = bbox[ind[0][0],0:4]

    nrof_faces = bbox.shape[0]
    det = bbox[:,0:4]
    img_size = np.asarray(face_img.shape)[0:2]

    if nrof_faces>1:
        bounding_box_size = (det[:,2]-det[:,0])*(det[:,3]-det[:,1])
        img_center = img_size / 2
        offsets = np.vstack([ (det[:,0]+det[:,2])/2-img_center[1], (det[:,1]+det[:,3])/2-img_center[0]  ])
        offset_dist_squared = np.sum(np.power(offsets,2.0),0)
        index = np.argmax(bounding_box_size-offset_dist_squared*2.0) # some extra weight on the centering
        det = det[index,:]
        points = points[index,:]
    det = np.squeeze(det)
    bbox = det
    #bb = np.zeros(4, dtype=np.int32)
    #bb[0] = np.maximum(det[0]-args.margin/2, 0)
    #bb[1] = np.maximum(det[1]-args.margin/2, 0)
    #bb[2] = np.minimum(det[2]+args.margin/2, img_size[1])
    #bb[3] = np.minimum(det[3]+args.margin/2, img_size[0])

    #cropped = img[bb[1]:bb[3],bb[0]:bb[2],:]
    #scaled = misc.imresize(cropped, (args.image_size, args.image_size), interp='bilinear')
    #nrof_successfully_aligned += 1
    #misc.imsave(output_filename, scaled)
    #text_file.write('%s %d %d %d %d\n' % (output_filename, bb[0], bb[1], bb[2], bb[3]))
    ###
    #points = points[0,:].reshape((2,5)).T
    points = points.reshape((2,5)).T
    #print(bbox)
    #print(points)
    nimg = face_preprocess.preprocess(face_img, bbox, points, image_size='112,112')
    nimg = cv2.cvtColor(nimg, cv2.COLOR_BGR2RGB)
    aligned = np.transpose(nimg, (2,0,1))
    return aligned, bbox #aligned

 #get the largest bbox
  def get_input_out_128(self, face_img):
    ret = self.detector.detect_face(face_img, det_type = self.args.det)
    if ret is None:
      return None
    bbox, points = ret
    if bbox.shape[0]==0:
      return None
    #bbox = bbox[0,0:4]
    #bbx selection
    bbw = bbox[:, 2] - bbox[:, 0] + 1
    bbh = bbox[:, 3] - bbox[:, 1] + 1
    area = bbw * bbh
    ind = np.where(area == np.max(area))
    bbox = bbox[ind[0][0],0:4]
    ###
    #points = points[0,:].reshape((2,5)).T
    points = points[ind[0][0],:].reshape((2,5)).T
    #print(bbox)
    #print(points)
    nimg = face_preprocess.preprocess_out_128(face_img, bbox, points, image_size='112,112')
    #nimg = cv2.cvtColor(nimg, cv2.COLOR_BGR2RGB)
    #aligned = np.transpose(nimg, (2,0,1))
    return nimg #aligned

  def get_input_out_128_all_bbx(self, face_img):
    ret = self.detector.detect_face(face_img, det_type = self.args.det)
    if ret is None:
      return []
    bboxs, points = ret
    if bboxs.shape[0]==0:
      return []

    nimgs = []

    N = bboxs.shape[0]

    for i in range(N):
        bbox = bboxs[i]
        bbox = bbox[0:4]

        #bbx selection
        #bbw = bbox[:, 2] - bbox[:, 0] + 1
        #bbh = bbox[:, 3] - bbox[:, 1] + 1
        #area = bbw * bbh
        #ind = np.where(area == np.max(area))
        #bbox = bbox[ind[0][0],0:4]
        ###
        point = points[i]
        point = point.reshape((2,5)).T
        #print(bbox)
        #print(points)
        nimg = face_preprocess.preprocess_out_128(face_img, bbox, point, image_size='112,112')
        #nimg = cv2.cvtColor(nimg, cv2.COLOR_BGR2RGB)
        #aligned = np.transpose(nimg, (2,0,1))
        nimgs.append(nimg)
    return nimgs #aligned

  def get_input_out_112(self, face_img):
    ret = self.detector.detect_face(face_img, det_type = self.args.det)
    if ret is None:
      return None
    bbox, points = ret
    if bbox.shape[0]==0:
      return None
    #bbox = bbox[0,0:4]
    #bbx selection
    bbw = bbox[:, 2] - bbox[:, 0] + 1
    bbh = bbox[:, 3] - bbox[:, 1] + 1
    area = bbw * bbh
    ind = np.where(area == np.max(area))
    bbox = bbox[ind[0][0],0:4]
    ###
    #points = points[0,:].reshape((2,5)).T
    points = points[ind[0][0],:].reshape((2,5)).T
    #print(bbox)
    #print(points)
    nimg = face_preprocess.preprocess(face_img, bbox, points, image_size='112,112')
    #nimg = cv2.cvtColor(nimg, cv2.COLOR_BGR2RGB)
    #aligned = np.transpose(nimg, (2,0,1))
    return nimg #aligned

  def get_feature(self, aligned):
    input_blob = np.expand_dims(aligned, axis=0)
    data = mx.nd.array(input_blob)
    db = mx.io.DataBatch(data=(data,))
    self.model.forward(db, is_train=False)
    embedding = self.model.get_outputs()[0].asnumpy()
    embedding = sklearn.preprocessing.normalize(embedding).flatten()

    return embedding

  def get_feature_batch(self, aligned):
    if(len(aligned.shape) == 3):
        return get_feature(aligned)
    else:
        input_blob = aligned

    data = mx.nd.array(input_blob)
    db = mx.io.DataBatch(data=(data,))
    self.model.forward(db, is_train=False)
    embedding = self.model.get_outputs()[0].asnumpy()
    embedding = sklearn.preprocessing.normalize(embedding).flatten()

    return embedding

  def get_ga(self, aligned):
    input_blob = np.expand_dims(aligned, axis=0)
    data = mx.nd.array(input_blob)
    db = mx.io.DataBatch(data=(data,))
    self.ga_model.forward(db, is_train=False)
    ret = self.ga_model.get_outputs()[0].asnumpy()
    g = ret[:,0:2].flatten()
    gender = np.argmax(g)
    a = ret[:,2:202].reshape( (100,2) )
    a = np.argmax(a, axis=1)
    age = int(sum(a))

    return gender, age


# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import os
import numpy as np
import torch
from maskrcnn_benchmark.modeling import registry
from torch import nn
from torch.nn import functional as F

from maskrcnn_benchmark.layers import smooth_l1_loss, kl_div_loss, entropy_loss, soft_cross_entropy_loss, \
    Label_Smoothing_Regression
from maskrcnn_benchmark.modeling.utils import cat
from .model_msg_passing import IMPContext
from .model_vtranse import VTransEFeature
from .model_vctree import VCTreeLSTMContext
from .model_motifs import LSTMContext, FrequencyBias, SpuerFrequencyBias
from .model_motifs_with_attribute import AttributeLSTMContext
from .model_transformer import TransformerContext, ScaledDotProductAttention, MultiHeadAttention
from .model_transformer_super import TransformerContext as TransformerContextSuper
from .utils_relation import layer_init, get_box_info, get_box_pair_info
from maskrcnn_benchmark.data import get_dataset_statistics
from scipy.special import softmax as spysoftmax
from maskrcnn_benchmark.layers.gcn.gcn_layers import GCN
from maskrcnn_benchmark.layers.gcn._utils import adj_normalize
import json

@registry.ROI_RELATION_PREDICTOR.register("TransformerPredictor")
class TransformerPredictor(nn.Module):
    def __init__(self, config, in_channels):
        super(TransformerPredictor, self).__init__()
        self.attribute_on = config.MODEL.ATTRIBUTE_ON
        # load parameters
        self.num_obj_cls = config.MODEL.ROI_BOX_HEAD.NUM_CLASSES
        self.num_att_cls = config.MODEL.ROI_ATTRIBUTE_HEAD.NUM_ATTRIBUTES
        self.num_rel_cls = config.MODEL.ROI_RELATION_HEAD.NUM_CLASSES
        self.devices = config.MODEL.DEVICE
        assert in_channels is not None
        num_inputs = in_channels
        self.use_vision = config.MODEL.ROI_RELATION_HEAD.PREDICT_USE_VISION
        self.use_bias = config.MODEL.ROI_RELATION_HEAD.PREDICT_USE_BIAS

        # load class dict
        statistics = get_dataset_statistics(config)
        obj_classes, rel_classes, att_classes = statistics['obj_classes'], statistics['rel_classes'], statistics[
            'att_classes']
        assert self.num_obj_cls == len(obj_classes)
        assert self.num_att_cls == len(att_classes)
        assert self.num_rel_cls == len(rel_classes)
        # module construct
        self.context_layer = TransformerContext(config, obj_classes, rel_classes, in_channels)

        # post decoding
        self.hidden_dim = config.MODEL.ROI_RELATION_HEAD.CONTEXT_HIDDEN_DIM
        self.pooling_dim = config.MODEL.ROI_RELATION_HEAD.CONTEXT_POOLING_DIM
        self.post_emb = nn.Linear(self.hidden_dim, self.hidden_dim * 2)
        self.post_cat = nn.Linear(self.hidden_dim * 2, self.pooling_dim)
        self.rel_compress = nn.Linear(self.pooling_dim, self.num_rel_cls)
        self.ctx_compress = nn.Linear(self.hidden_dim * 2, self.num_rel_cls)

        # initialize layer parameters
        layer_init(self.post_emb, 10.0 * (1.0 / self.hidden_dim) ** 0.5, normal=True)
        layer_init(self.rel_compress, xavier=True)
        layer_init(self.ctx_compress, xavier=True)
        layer_init(self.post_cat, xavier=True)

        if self.pooling_dim != config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM:
            self.union_single_not_match = True
            self.up_dim = nn.Linear(config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM, self.pooling_dim)
            layer_init(self.up_dim, xavier=True)
        else:
            self.union_single_not_match = False

        if self.use_bias:
            # convey statistics into FrequencyBias to avoid loading again
            self.freq_bias = FrequencyBias(config, statistics)

    def forward(self, proposals, rel_pair_idxs, rel_labels, rel_binarys, roi_features, union_features, logger=None):
        """
        Returns:
            obj_dists (list[Tensor]): logits of object label distribution
            rel_dists (list[Tensor])
            rel_pair_idxs (list[Tensor]): (num_rel, 2) index of subject and object
            union_features (Tensor): (batch_num_rel, context_pooling_dim): visual union feature of each pair
        """
        if self.attribute_on:
            obj_dists, obj_preds, att_dists, edge_ctx = self.context_layer(roi_features, proposals, logger)
        else:
            obj_dists, obj_preds, edge_ctx = self.context_layer(roi_features, proposals, logger)

        # post decode
        edge_rep = self.post_emb(edge_ctx)
        edge_rep = edge_rep.view(edge_rep.size(0), 2, self.hidden_dim)
        head_rep = edge_rep[:, 0].contiguous().view(-1, self.hidden_dim)
        tail_rep = edge_rep[:, 1].contiguous().view(-1, self.hidden_dim)

        num_rels = [r.shape[0] for r in rel_pair_idxs]
        num_objs = [len(b) for b in proposals]
        assert len(num_rels) == len(num_objs)

        head_reps = head_rep.split(num_objs, dim=0)
        tail_reps = tail_rep.split(num_objs, dim=0)
        obj_preds = obj_preds.split(num_objs, dim=0)

        # from object level feature to pairwise relation level feature
        prod_reps = []
        pair_preds = []
        for pair_idx, head_rep, tail_rep, obj_pred in zip(rel_pair_idxs, head_reps, tail_reps, obj_preds):
            prod_reps.append(torch.cat((head_rep[pair_idx[:, 0]], tail_rep[pair_idx[:, 1]]), dim=-1))
            pair_preds.append(torch.stack((obj_pred[pair_idx[:, 0]], obj_pred[pair_idx[:, 1]]), dim=1))
        prod_rep = cat(prod_reps, dim=0)
        pair_pred = cat(pair_preds, dim=0)

        ctx_gate = self.post_cat(prod_rep)

        # use union box and mask convolution
        if self.use_vision:
            if self.union_single_not_match:
                visual_rep = ctx_gate * self.up_dim(union_features)
            else:
                visual_rep = ctx_gate * union_features

        rel_dists = self.rel_compress(visual_rep) + self.ctx_compress(prod_rep)

        # use frequence bias
        if self.use_bias:
            freq_dists_bias = self.freq_bias.index_with_labels(pair_pred)
            freq_dists_bias = F.dropout(freq_dists_bias, 0.3, training=self.training)
            rel_dists = rel_dists + freq_dists_bias

        obj_dists = obj_dists.split(num_objs, dim=0)
        rel_dists = rel_dists.split(num_rels, dim=0)

        add_losses = {}

        if self.attribute_on:
            att_dists = att_dists.split(num_objs, dim=0)
            return (obj_dists, att_dists), rel_dists, add_losses
        else:
            return obj_dists, rel_dists, add_losses

@registry.ROI_RELATION_PREDICTOR.register("TransformerSuperPredictor")
class TransformerSuperPredictor(nn.Module):
    def __init__(self, config, in_channels):
        super(TransformerSuperPredictor, self).__init__()
        self.attribute_on = config.MODEL.ATTRIBUTE_ON
        # load parameters
        self.num_obj_cls = config.MODEL.ROI_BOX_HEAD.NUM_CLASSES
        self.num_att_cls = config.MODEL.ROI_ATTRIBUTE_HEAD.NUM_ATTRIBUTES
        self.num_rel_cls = config.MODEL.ROI_RELATION_HEAD.NUM_CLASSES

        assert in_channels is not None
        num_inputs = in_channels
        self.use_vision = config.MODEL.ROI_RELATION_HEAD.PREDICT_USE_VISION
        self.use_bias = config.MODEL.ROI_RELATION_HEAD.PREDICT_USE_BIAS

        # load class dict
        statistics = get_dataset_statistics(config)
        obj_classes, rel_classes, att_classes = statistics['obj_classes'], statistics['rel_classes'], statistics[
            'att_classes']
        assert self.num_obj_cls == len(obj_classes)
        assert self.num_att_cls == len(att_classes)
        assert self.num_rel_cls == len(rel_classes)

        # module construct
        self.context_layer = TransformerContext(config, obj_classes, rel_classes, in_channels)

        # post decoding
        self.hidden_dim = config.MODEL.ROI_RELATION_HEAD.CONTEXT_HIDDEN_DIM
        self.pooling_dim = config.MODEL.ROI_RELATION_HEAD.CONTEXT_POOLING_DIM
        self.post_emb = nn.Linear(self.hidden_dim, self.hidden_dim * 2)
        self.post_cat = nn.Linear(self.hidden_dim * 2, self.pooling_dim)
        self.rel_compress = nn.Linear(self.pooling_dim, self.num_rel_cls)
        self.ctx_compress = nn.Linear(self.hidden_dim * 2, self.num_rel_cls)
        self.devices = config.MODEL.DEVICE
        # initialize layer parameters
        layer_init(self.post_emb, 10.0 * (1.0 / self.hidden_dim) ** 0.5, normal=True)
        layer_init(self.post_cat, xavier=True)
        layer_init(self.rel_compress, xavier=True)
        layer_init(self.ctx_compress, xavier=True)

        if self.pooling_dim != config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM:
            self.union_single_not_match = True
            self.up_dim = nn.Linear(config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM, self.pooling_dim)
            layer_init(self.up_dim, xavier=True)
        else:
            self.union_single_not_match = False

        if self.use_bias:
            # convey statistics into FrequencyBias to avoid loading again
            self.freq_bias = FrequencyBias(config, statistics)
        # Clean Head
        self.val_alpha = config.MODEL.ROI_RELATION_HEAD.VAL_ALPHA
        self.rel_compress_clean = nn.Linear(self.pooling_dim, self.num_rel_cls)
        self.ctx_compress_clean = nn.Linear(self.hidden_dim * 2, self.num_rel_cls)
        layer_init(self.rel_compress_clean, xavier=True)
        layer_init(self.ctx_compress_clean, xavier=True)

        if self.use_bias:
            # convey statistics into FrequencyBias to avoid loading again
            self.freq_bias_clean = FrequencyBias(config, statistics)

    def forward(self, proposals, rel_pair_idxs, rel_labels, rel_binarys, roi_features, union_features, logger=None):
        """
        Returns:
            obj_dists (list[Tensor]): logits of object label distribution
            rel_dists (list[Tensor])
            rel_pair_idxs (list[Tensor]): (num_rel, 2) index of subject and object
            union_features (Tensor): (batch_num_rel, context_pooling_dim): visual union feature of each pair
        """
        if self.attribute_on:
            obj_dists, obj_preds, att_dists, edge_ctx = self.context_layer(roi_features, proposals, logger)
        else:
            obj_dists, obj_preds, edge_ctx = self.context_layer(roi_features, proposals, logger)

        # post decode
        edge_rep = self.post_emb(edge_ctx)
        edge_rep = edge_rep.view(edge_rep.size(0), 2, self.hidden_dim)
        head_rep = edge_rep[:, 0].contiguous().view(-1, self.hidden_dim)
        tail_rep = edge_rep[:, 1].contiguous().view(-1, self.hidden_dim)

        num_rels = [r.shape[0] for r in rel_pair_idxs]
        num_objs = [len(b) for b in proposals]
        assert len(num_rels) == len(num_objs)

        head_reps = head_rep.split(num_objs, dim=0)
        tail_reps = tail_rep.split(num_objs, dim=0)
        obj_preds = obj_preds.split(num_objs, dim=0)

        # from object level feature to pairwise relation level feature
        prod_reps = []
        pair_preds = []
        for pair_idx, head_rep, tail_rep, obj_pred in zip(rel_pair_idxs, head_reps, tail_reps, obj_preds):
            prod_reps.append(torch.cat((head_rep[pair_idx[:, 0]], tail_rep[pair_idx[:, 1]]), dim=-1))
            pair_preds.append(torch.stack((obj_pred[pair_idx[:, 0]], obj_pred[pair_idx[:, 1]]), dim=1))
        prod_rep = cat(prod_reps, dim=0)
        pair_pred = cat(pair_preds, dim=0)

        ctx_gate = self.post_cat(prod_rep)

        # use union box and mask convolution
        if self.use_vision:
            if self.union_single_not_match:
                visual_rep = ctx_gate * self.up_dim(union_features)
            else:
                visual_rep = ctx_gate * union_features

        rel_dists_noisy = self.rel_compress(visual_rep) + self.ctx_compress(prod_rep)
        # use frequence bias
        if self.use_bias:
            freq_dists_bias = self.freq_bias.index_with_labels(pair_pred)
            freq_dists_bias = F.dropout(freq_dists_bias, 0.3, training=self.training)
            rel_dists_noisy = rel_dists_noisy + freq_dists_bias

        rel_dists_clean = self.rel_compress_clean(visual_rep) + self.ctx_compress_clean(prod_rep)
        if self.use_bias:
            freq_dists_bias_clean = self.freq_bias_clean.index_with_labels(pair_pred)
            freq_dists_bias_clean = F.dropout(freq_dists_bias_clean, 0.3, training=self.training)
            rel_dists_clean = rel_dists_clean + freq_dists_bias_clean

        rel_dists = rel_dists_clean
        if not self.training:
            rel_dists = (1.0 - self.val_alpha) * rel_dists + self.val_alpha * rel_dists_noisy

        add_losses = {}
        obj_dists = obj_dists.split(num_objs, dim=0)
        rel_dists = rel_dists.split(num_rels, dim=0)

        if self.attribute_on:
            att_dists = att_dists.split(num_objs, dim=0)
            return (obj_dists, att_dists), rel_dists, add_losses
        else:
            return obj_dists, rel_dists, add_losses

@registry.ROI_RELATION_PREDICTOR.register("TransformerTransferPredictor")
class TransformerTransferPredictor(nn.Module):
    def __init__(self, config, in_channels):
        super(TransformerTransferPredictor, self).__init__()
        self.attribute_on = config.MODEL.ATTRIBUTE_ON
        # load parameters
        self.num_obj_cls = config.MODEL.ROI_BOX_HEAD.NUM_CLASSES
        self.num_att_cls = config.MODEL.ROI_ATTRIBUTE_HEAD.NUM_ATTRIBUTES
        self.num_rel_cls = config.MODEL.ROI_RELATION_HEAD.NUM_CLASSES
        self.dropout_rate = config.MODEL.ROI_RELATION_HEAD.TRANSFORMER.DROPOUT_RATE

        assert in_channels is not None
        num_inputs = in_channels
        self.use_vision = config.MODEL.ROI_RELATION_HEAD.PREDICT_USE_VISION
        self.use_bias = config.MODEL.ROI_RELATION_HEAD.PREDICT_USE_BIAS
        self.with_knowdist = False
        self.devices = config.MODEL.DEVICE
        self.with_transfer = config.MODEL.ROI_RELATION_HEAD.WITH_TRANSFER_CLASSIFIER
        self.with_cleanclf = config.MODEL.ROI_RELATION_HEAD.WITH_CLEAN_CLASSIFIER
        # load class dict
        statistics = get_dataset_statistics(config)
        obj_classes, rel_classes, att_classes = statistics['obj_classes'], statistics['rel_classes'], statistics[
            'att_classes']
        assert self.num_obj_cls == len(obj_classes)
        assert self.num_att_cls == len(att_classes)
        assert self.num_rel_cls == len(rel_classes)
        self.val_alpha = config.MODEL.ROI_RELATION_HEAD.VAL_ALPHA

        # module construct
        self.context_layer = TransformerContext(config, obj_classes, rel_classes, in_channels)

        # post decoding
        self.hidden_dim = config.MODEL.ROI_RELATION_HEAD.CONTEXT_HIDDEN_DIM
        self.pooling_dim = config.MODEL.ROI_RELATION_HEAD.CONTEXT_POOLING_DIM
        self.post_emb = nn.Linear(self.hidden_dim, self.hidden_dim * 2)
        layer_init(self.post_emb, 10.0 * (1.0 / self.hidden_dim) ** 0.5, normal=True)
        self.post_cat = nn.Linear(self.hidden_dim * 2, self.pooling_dim)
        layer_init(self.post_cat, xavier=True)

        if self.pooling_dim != config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM:
            self.union_single_not_match = True
            self.up_dim = nn.Linear(config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM, self.pooling_dim)
            layer_init(self.up_dim, xavier=True)
        else:
            self.union_single_not_match = False

        # initialize layer parameters
        self.rel_compress = nn.Linear(self.pooling_dim, self.num_rel_cls)
        self.ctx_compress = nn.Linear(self.hidden_dim * 2, self.num_rel_cls)
        layer_init(self.rel_compress, xavier=True)
        layer_init(self.ctx_compress, xavier=True)
        if self.use_bias:
            # convey statistics into FrequencyBias to avoid loading again
            self.freq_bias = FrequencyBias(config, statistics)

        # the transfer classifier
        if self.with_cleanclf:
            self.rel_compress_clean = nn.Linear(self.pooling_dim, self.num_rel_cls)
            self.ctx_compress_clean = nn.Linear(self.hidden_dim * 2, self.num_rel_cls)
            layer_init(self.rel_compress_clean, xavier=True)
            layer_init(self.ctx_compress_clean, xavier=True)
            # self.gcns_rel_clean = GCN(self.pooling_dim, self.pooling_dim, self.dropout_rate)
            # self.gcns_ctx_clean = GCN(self.hidden_dim * 2, self.hidden_dim * 2, self.dropout_rate)
            self.freq_bias_clean = FrequencyBias(config, statistics)
        if self.with_transfer:
            #pred_adj_np = np.load('./misc/conf_mat_adj_mat.npy')
            print("Using Confusion Matrix Transfer!")
            pred_adj_np = np.load('./misc/conf_mat_freq_train.npy')
            # pred_adj_np = 1.0 - pred_adj_np
            pred_adj_np[0, :] = 0.0
            pred_adj_np[:, 0] = 0.0
            pred_adj_np[0, 0] = 1.0
            # adj_i_j means the baseline outputs category j, but the ground truth is i.
            pred_adj_np = pred_adj_np / (pred_adj_np.sum(-1)[:, None] + 1e-8)
            pred_adj_np = adj_normalize(pred_adj_np)
            self.pred_adj_nor = torch.from_numpy(pred_adj_np).float().to(self.devices)
            # self.pred_adj_layer_clean = nn.Linear(self.num_rel_cls, self.num_rel_cls, bias=False)
            # #layer_init(self.pred_adj_layer_clean, xavier=True)
            # with torch.no_grad():
            #     self.pred_adj_layer_clean.weight.copy_(torch.eye(self.num_rel_cls,dtype=torch.float), non_blocking=True)
                #self.pred_adj_layer_clean.weight.copy_(self.pred_adj_nor, non_blocking=True)

    def forward(self, proposals, rel_pair_idxs, rel_labels, rel_binarys, roi_features, union_features, logger=None):
        """
        Returns:
            obj_dists (list[Tensor]): logits of object label distribution
            rel_dists (list[Tensor])
            rel_pair_idxs (list[Tensor]): (num_rel, 2) index of subject and object
            union_features (Tensor): (batch_num_rel, context_pooling_dim): visual union feature of each pair
        """
        if self.attribute_on:
            obj_dists, obj_preds, att_dists, edge_ctx = self.context_layer(roi_features, proposals, logger)
        else:
            obj_dists, obj_preds, edge_ctx = self.context_layer(roi_features, proposals, logger)

        # post decode
        edge_rep = self.post_emb(edge_ctx)
        edge_rep = edge_rep.view(edge_rep.size(0), 2, self.hidden_dim)
        head_rep = edge_rep[:, 0].contiguous().view(-1, self.hidden_dim)
        tail_rep = edge_rep[:, 1].contiguous().view(-1, self.hidden_dim)

        num_rels = [r.shape[0] for r in rel_pair_idxs]
        num_objs = [len(b) for b in proposals]
        assert len(num_rels) == len(num_objs)

        head_reps = head_rep.split(num_objs, dim=0)
        tail_reps = tail_rep.split(num_objs, dim=0)
        obj_preds = obj_preds.split(num_objs, dim=0)

        # from object level feature to pairwise relation level feature
        prod_reps = []
        pair_preds = []
        for pair_idx, head_rep, tail_rep, obj_pred in zip(rel_pair_idxs, head_reps, tail_reps, obj_preds):
            prod_reps.append(torch.cat((head_rep[pair_idx[:, 0]], tail_rep[pair_idx[:, 1]]), dim=-1))
            pair_preds.append(torch.stack((obj_pred[pair_idx[:, 0]], obj_pred[pair_idx[:, 1]]), dim=1))
        prod_rep = cat(prod_reps, dim=0)
        pair_pred = cat(pair_preds, dim=0)

        ctx_gate = self.post_cat(prod_rep)

        # use union box and mask convolution
        if self.use_vision:
            if self.union_single_not_match:
                visual_rep = ctx_gate * self.up_dim(union_features)
            else:
                visual_rep = ctx_gate * union_features

        rel_dists_general = self.rel_compress(visual_rep) + self.ctx_compress(prod_rep)
        if self.use_bias:
            freq_dists_bias = self.freq_bias.index_with_labels(pair_pred)
            freq_dists_bias = F.dropout(freq_dists_bias, 0.3, training=self.training)
            rel_dists_general = rel_dists_general + freq_dists_bias
        rel_dists = rel_dists_general
        # the transfer classifier
        if self.with_cleanclf:
            rel_dists_clean = self.rel_compress_clean(visual_rep) + self.ctx_compress_clean(prod_rep)
            if self.use_bias:
                freq_dists_bias_clean = self.freq_bias_clean.index_with_labels(pair_pred)
                freq_dists_bias_clean = F.dropout(freq_dists_bias_clean, 0.3, training=self.training)
                rel_dists_clean = rel_dists_clean + freq_dists_bias_clean
            rel_dists = rel_dists_clean

        if self.with_transfer:
            rel_dists = (self.pred_adj_nor @ rel_dists.T).T

        add_losses = {}
        # if self.with_knowdist:
        #     rel_dists_specific_soft = F.log_softmax(rel_dists, -1)
        #     rel_dists_general_soft = F.softmax(rel_dists_general, -1)
        #     add_losses['know_dist_kl'] = self.kd_alpha * self.kl_loss(rel_dists_specific_soft, rel_dists_general_soft)

        obj_dists = obj_dists.split(num_objs, dim=0)
        rel_dists = rel_dists.split(num_rels, dim=0)

        if self.attribute_on:
            att_dists = att_dists.split(num_objs, dim=0)
            return (obj_dists, att_dists), rel_dists, add_losses
        else:
            return obj_dists, rel_dists, add_losses

@registry.ROI_RELATION_PREDICTOR.register("IMPPredictor")
class IMPPredictor(nn.Module):
    def __init__(self, config, in_channels):
        super(IMPPredictor, self).__init__()
        self.num_obj_cls = config.MODEL.ROI_BOX_HEAD.NUM_CLASSES
        self.num_rel_cls = config.MODEL.ROI_RELATION_HEAD.NUM_CLASSES
        self.use_bias = False

        assert in_channels is not None

        self.context_layer = IMPContext(config, self.num_obj_cls, self.num_rel_cls, in_channels)

        # post decoding
        self.hidden_dim = config.MODEL.ROI_RELATION_HEAD.CONTEXT_HIDDEN_DIM
        self.pooling_dim = config.MODEL.ROI_RELATION_HEAD.CONTEXT_POOLING_DIM

        if self.pooling_dim != config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM:
            self.union_single_not_match = True
            self.up_dim = nn.Linear(config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM, self.pooling_dim)
            layer_init(self.up_dim, xavier=True)
        else:
            self.union_single_not_match = False

        # freq
        if self.use_bias:
            statistics = get_dataset_statistics(config)
            self.freq_bias = FrequencyBias(config, statistics)

    def forward(self, proposals, rel_pair_idxs, rel_labels, rel_binarys, roi_features, union_features, logger=None):
        """
        Returns:
            obj_dists (list[Tensor]): logits of object label distribution
            rel_dists (list[Tensor])
            rel_pair_idxs (list[Tensor]): (num_rel, 2) index of subject and object
            union_features (Tensor): (batch_num_rel, context_pooling_dim): visual union feature of each pair
        """

        if self.union_single_not_match:
            union_features = self.up_dim(union_features)

        # encode context infomation
        obj_dists, rel_dists = self.context_layer(roi_features, proposals, union_features, rel_pair_idxs, logger)

        num_objs = [len(b) for b in proposals]
        num_rels = [r.shape[0] for r in rel_pair_idxs]
        assert len(num_rels) == len(num_objs)

        if self.use_bias:
            obj_preds = obj_dists.max(-1)[1]
            obj_preds = obj_preds.split(num_objs, dim=0)

            pair_preds = []
            for pair_idx, obj_pred in zip(rel_pair_idxs, obj_preds):
                pair_preds.append(torch.stack((obj_pred[pair_idx[:, 0]], obj_pred[pair_idx[:, 1]]), dim=1))
            pair_pred = cat(pair_preds, dim=0)

            rel_dists = rel_dists + self.freq_bias.index_with_labels(pair_pred.long())

        obj_dists = obj_dists.split(num_objs, dim=0)
        rel_dists = rel_dists.split(num_rels, dim=0)

        # we use obj_preds instead of pred from obj_dists
        # because in decoder_rnn, preds has been through a nms stage
        add_losses = {}

        return obj_dists, rel_dists, add_losses

@registry.ROI_RELATION_PREDICTOR.register("MotifPredictor")
class MotifPredictor(nn.Module):
    def __init__(self, config, in_channels):
        super(MotifPredictor, self).__init__()
        self.attribute_on = config.MODEL.ATTRIBUTE_ON
        self.num_obj_cls = config.MODEL.ROI_BOX_HEAD.NUM_CLASSES
        self.num_att_cls = config.MODEL.ROI_ATTRIBUTE_HEAD.NUM_ATTRIBUTES
        self.num_rel_cls = config.MODEL.ROI_RELATION_HEAD.NUM_CLASSES

        assert in_channels is not None
        num_inputs = in_channels
        self.use_vision = config.MODEL.ROI_RELATION_HEAD.PREDICT_USE_VISION
        self.use_bias = config.MODEL.ROI_RELATION_HEAD.PREDICT_USE_BIAS

        # load class dict
        statistics = get_dataset_statistics(config)
        obj_classes, rel_classes, att_classes = statistics['obj_classes'], statistics['rel_classes'], statistics[
            'att_classes']
        assert self.num_obj_cls == len(obj_classes)
        assert self.num_att_cls == len(att_classes)
        assert self.num_rel_cls == len(rel_classes)
        # init contextual lstm encoding
        if self.attribute_on:
            self.context_layer = AttributeLSTMContext(config, obj_classes, att_classes, rel_classes, in_channels)
        else:
            self.context_layer = LSTMContext(config, obj_classes, rel_classes, in_channels)

        # post decoding
        self.hidden_dim = config.MODEL.ROI_RELATION_HEAD.CONTEXT_HIDDEN_DIM
        self.pooling_dim = config.MODEL.ROI_RELATION_HEAD.CONTEXT_POOLING_DIM
        self.post_emb = nn.Linear(self.hidden_dim, self.hidden_dim * 2)
        self.post_cat = nn.Linear(self.hidden_dim * 2, self.pooling_dim)
        self.rel_compress = nn.Linear(self.pooling_dim, self.num_rel_cls, bias=True)

        # initialize layer parameters
        layer_init(self.post_emb, 10.0 * (1.0 / self.hidden_dim) ** 0.5, normal=True)
        layer_init(self.post_cat, xavier=True)
        layer_init(self.rel_compress, xavier=True)

        if self.pooling_dim != config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM:
            self.union_single_not_match = True
            self.up_dim = nn.Linear(config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM, self.pooling_dim)
            layer_init(self.up_dim, xavier=True)
        else:
            self.union_single_not_match = False

        if self.use_bias:
            # convey statistics into FrequencyBias to avoid loading again
            self.freq_bias = FrequencyBias(config, statistics)

        self.with_clean_classifier = config.MODEL.ROI_RELATION_HEAD.WITH_CLEAN_CLASSIFIER

        if self.with_clean_classifier:
            if self.pooling_dim != config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM:
                self.union_single_not_match = True
                self.up_dim_clean = nn.Linear(config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM, self.pooling_dim)
                layer_init(self.up_dim_clean, xavier=True)
            else:
                self.union_single_not_match = False
            self.post_cat_clean = nn.Linear(self.hidden_dim * 2, self.pooling_dim)
            self.rel_compress_clean = nn.Linear(self.pooling_dim, self.num_rel_cls, bias=True)
            layer_init(self.post_cat_clean, xavier=True)
            layer_init(self.rel_compress_clean, xavier=True)
            if self.use_bias:
                # convey statistics into FrequencyBias to avoid loading again
                self.freq_bias_clean = FrequencyBias(config, statistics)
            self.with_transfer = config.MODEL.ROI_RELATION_HEAD.WITH_TRANSFER_CLASSIFIER
            if self.with_transfer:
                self.devices = config.MODEL.DEVICE
                print("!!!!!!!!!With Confusion Matrix Channel!!!!!")
                pred_adj_np = np.load('./misc/conf_mat_freq_train.npy')
                # pred_adj_np = 1.0 - pred_adj_np
                pred_adj_np[0, :] = 0.0
                pred_adj_np[:, 0] = 0.0
                pred_adj_np[0, 0] = 1.0
                # adj_i_j means the baseline outputs category j, but the ground truth is i.
                pred_adj_np = pred_adj_np / (pred_adj_np.sum(-1)[:, None] + 1e-8)
                pred_adj_np = adj_normalize(pred_adj_np)
                self.pred_adj_nor = torch.from_numpy(pred_adj_np).float().to(self.devices)

    def forward(self, proposals, rel_pair_idxs, rel_labels, rel_binarys, roi_features, union_features, logger=None):
        """
        Returns:
            obj_dists (list[Tensor]): logits of object label distribution
            rel_dists (list[Tensor])
            rel_pair_idxs (list[Tensor]): (num_rel, 2) index of subject and object
            union_features (Tensor): (batch_num_rel, context_pooling_dim): visual union feature of each pair
        """

        # encode context infomation
        if self.attribute_on:
            obj_dists, obj_preds, att_dists, edge_ctx = self.context_layer(roi_features, proposals, logger)
        else:
            obj_dists, obj_preds, edge_ctx, _ = self.context_layer(roi_features, proposals, logger)

        # post decode
        edge_rep = self.post_emb(edge_ctx)
        edge_rep = edge_rep.view(edge_rep.size(0), 2, self.hidden_dim)
        head_rep = edge_rep[:, 0].contiguous().view(-1, self.hidden_dim)
        tail_rep = edge_rep[:, 1].contiguous().view(-1, self.hidden_dim)

        num_rels = [r.shape[0] for r in rel_pair_idxs]
        num_objs = [len(b) for b in proposals]
        assert len(num_rels) == len(num_objs)

        head_reps = head_rep.split(num_objs, dim=0)
        tail_reps = tail_rep.split(num_objs, dim=0)
        obj_preds = obj_preds.split(num_objs, dim=0)

        prod_reps = []
        pair_preds = []
        for pair_idx, head_rep, tail_rep, obj_pred in zip(rel_pair_idxs, head_reps, tail_reps, obj_preds):
            prod_reps.append(torch.cat((head_rep[pair_idx[:, 0]], tail_rep[pair_idx[:, 1]]), dim=-1))
            pair_preds.append(torch.stack((obj_pred[pair_idx[:, 0]], obj_pred[pair_idx[:, 1]]), dim=1))
        prod_rep = cat(prod_reps, dim=0)
        pair_pred = cat(pair_preds, dim=0)
        prod_rep = self.post_cat(prod_rep)

        if self.use_vision:
            if self.union_single_not_match:
                prod_rep = prod_rep * self.up_dim(union_features)
            else:
                prod_rep = prod_rep * union_features

        rel_dists = self.rel_compress(prod_rep)

        if self.use_bias:
            freq_dists_bias = self.freq_bias.index_with_labels(pair_pred.long())
            freq_dists_bias = F.dropout(freq_dists_bias, 0.3, training=self.training)
            rel_dists = rel_dists + freq_dists_bias

        if self.with_clean_classifier:
            prod_rep_clean = cat(prod_reps, dim=0)
            prod_rep_clean = self.post_cat_clean(prod_rep_clean)
            if self.use_vision:
                if self.union_single_not_match:
                    prod_rep_clean = prod_rep_clean * self.up_dim_clean(union_features)
                else:
                    prod_rep_clean = prod_rep_clean * union_features

            rel_dists_clean = self.rel_compress_clean(prod_rep_clean)
            if self.use_bias:
                freq_dists_bias_clean = self.freq_bias_clean.index_with_labels(pair_pred.long())
                freq_dists_bias_clean = F.dropout(freq_dists_bias_clean, 0.3, training=self.training)
                rel_dists_clean = rel_dists_clean + freq_dists_bias_clean
            if self.with_transfer:
                rel_dists_clean = (self.pred_adj_nor @ rel_dists_clean.T).T

            rel_dists = rel_dists_clean

        obj_dists = obj_dists.split(num_objs, dim=0)
        rel_dists = rel_dists.split(num_rels, dim=0)

        # we use obj_preds instead of pred from obj_dists
        # because in decoder_rnn, preds has been through a nms stage
        add_losses = {}

        if self.attribute_on:
            att_dists = att_dists.split(num_objs, dim=0)
            return (obj_dists, att_dists), rel_dists, add_losses
        else:
            return obj_dists, rel_dists, add_losses

@registry.ROI_RELATION_PREDICTOR.register("VCTreePredictor")
class VCTreePredictor(nn.Module):
    def __init__(self, config, in_channels):
        super(VCTreePredictor, self).__init__()
        self.attribute_on = config.MODEL.ATTRIBUTE_ON
        self.num_obj_cls = config.MODEL.ROI_BOX_HEAD.NUM_CLASSES
        self.num_att_cls = config.MODEL.ROI_ATTRIBUTE_HEAD.NUM_ATTRIBUTES
        self.num_rel_cls = config.MODEL.ROI_RELATION_HEAD.NUM_CLASSES

        assert in_channels is not None
        num_inputs = in_channels

        # load class dict
        statistics = get_dataset_statistics(config)
        obj_classes, rel_classes, att_classes = statistics['obj_classes'], statistics['rel_classes'], statistics[
            'att_classes']
        assert self.num_obj_cls == len(obj_classes)
        assert self.num_att_cls == len(att_classes)
        assert self.num_rel_cls == len(rel_classes)
        # init contextual lstm encoding
        self.context_layer = VCTreeLSTMContext(config, obj_classes, rel_classes, statistics, in_channels)

        # post decoding
        self.hidden_dim = config.MODEL.ROI_RELATION_HEAD.CONTEXT_HIDDEN_DIM
        self.pooling_dim = config.MODEL.ROI_RELATION_HEAD.CONTEXT_POOLING_DIM
        self.post_emb = nn.Linear(self.hidden_dim, self.hidden_dim * 2)
        self.post_cat = nn.Linear(self.hidden_dim * 2, self.pooling_dim)

        # learned-mixin
        # self.uni_gate = nn.Linear(self.pooling_dim, self.num_rel_cls)
        # self.frq_gate = nn.Linear(self.pooling_dim, self.num_rel_cls)
        self.ctx_compress = nn.Linear(self.pooling_dim, self.num_rel_cls)
        # self.uni_compress = nn.Linear(self.pooling_dim, self.num_rel_cls)
        # layer_init(self.uni_gate, xavier=True)
        # layer_init(self.frq_gate, xavier=True)
        layer_init(self.ctx_compress, xavier=True)
        # layer_init(self.uni_compress, xavier=True)

        # initialize layer parameters
        layer_init(self.post_emb, 10.0 * (1.0 / self.hidden_dim) ** 0.5, normal=True)
        layer_init(self.post_cat, xavier=True)

        if self.pooling_dim != config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM:
            self.union_single_not_match = True
            self.up_dim = nn.Linear(config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM, self.pooling_dim)
            layer_init(self.up_dim, xavier=True)
        else:
            self.union_single_not_match = False

        self.freq_bias = FrequencyBias(config, statistics)

        self.with_clean_classifier = config.MODEL.ROI_RELATION_HEAD.WITH_CLEAN_CLASSIFIER

        if self.with_clean_classifier:
            if self.pooling_dim != config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM:
                self.union_single_not_match = True
                self.up_dim_clean = nn.Linear(config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM, self.pooling_dim)
                layer_init(self.up_dim_clean, xavier=True)
            else:
                self.union_single_not_match = False
            self.post_cat_clean = nn.Linear(self.hidden_dim * 2, self.pooling_dim)
            self.ctx_compress_clean = nn.Linear(self.pooling_dim, self.num_rel_cls, bias=True)
            layer_init(self.post_cat_clean, xavier=True)
            layer_init(self.ctx_compress_clean, xavier=True)
            # convey statistics into FrequencyBias to avoid loading again
            self.freq_bias_clean = FrequencyBias(config, statistics)

            self.with_transfer = config.MODEL.ROI_RELATION_HEAD.WITH_TRANSFER_CLASSIFIER
            if self.with_transfer:
                self.devices = config.MODEL.DEVICE
                print("!!!!!!!!!With Confusion Matrix Channel!!!!!")
                pred_adj_np = np.load('./misc/conf_mat_freq_train.npy')
                # pred_adj_np = 1.0 - pred_adj_np
                pred_adj_np[0, :] = 0.0
                pred_adj_np[:, 0] = 0.0
                pred_adj_np[0, 0] = 1.0
                # adj_i_j means the baseline outputs category j, but the ground truth is i.
                pred_adj_np = pred_adj_np / (pred_adj_np.sum(-1)[:, None] + 1e-8)
                pred_adj_np = adj_normalize(pred_adj_np)
                self.pred_adj_nor = torch.from_numpy(pred_adj_np).float().to(self.devices)

    def forward(self, proposals, rel_pair_idxs, rel_labels, rel_binarys, roi_features, union_features, logger=None):
        """
        Returns:
            obj_dists (list[Tensor]): logits of object label distribution
            rel_dists (list[Tensor])
            rel_pair_idxs (list[Tensor]): (num_rel, 2) index of subject and object
            union_features (Tensor): (batch_num_rel, context_pooling_dim): visual union feature of each pair
        """

        # encode context infomation
        obj_dists, obj_preds, edge_ctx, binary_preds = self.context_layer(roi_features, proposals, rel_pair_idxs,
                                                                          logger)

        # post decode
        edge_rep = F.relu(self.post_emb(edge_ctx))
        edge_rep = edge_rep.view(edge_rep.size(0), 2, self.hidden_dim)
        head_rep = edge_rep[:, 0].contiguous().view(-1, self.hidden_dim)
        tail_rep = edge_rep[:, 1].contiguous().view(-1, self.hidden_dim)

        num_rels = [r.shape[0] for r in rel_pair_idxs]
        num_objs = [len(b) for b in proposals]
        assert len(num_rels) == len(num_objs)

        head_reps = head_rep.split(num_objs, dim=0)
        tail_reps = tail_rep.split(num_objs, dim=0)
        obj_preds = obj_preds.split(num_objs, dim=0)

        prod_reps = []
        pair_preds = []
        for pair_idx, head_rep, tail_rep, obj_pred in zip(rel_pair_idxs, head_reps, tail_reps, obj_preds):
            prod_reps.append(torch.cat((head_rep[pair_idx[:, 0]], tail_rep[pair_idx[:, 1]]), dim=-1))
            pair_preds.append(torch.stack((obj_pred[pair_idx[:, 0]], obj_pred[pair_idx[:, 1]]), dim=1))
        prod_rep = cat(prod_reps, dim=0)
        pair_pred = cat(pair_preds, dim=0)

        prod_rep = self.post_cat(prod_rep)

        # learned-mixin Gate
        # uni_gate = torch.tanh(self.uni_gate(self.drop(prod_rep)))
        # frq_gate = torch.tanh(self.frq_gate(self.drop(prod_rep)))

        if self.union_single_not_match:
            union_features = self.up_dim(union_features)

        ctx_dists = self.ctx_compress(prod_rep * union_features)
        # uni_dists = self.uni_compress(self.drop(union_features))
        frq_dists = self.freq_bias.index_with_labels(pair_pred.long())
        frq_dists = F.dropout(frq_dists, 0.3, training=self.training)
        rel_dists = ctx_dists + frq_dists
        # rel_dists = ctx_dists + uni_gate * uni_dists + frq_gate * frq_dists
        if self.with_clean_classifier:
            prod_rep_clean = cat(prod_reps, dim=0)
            prod_rep_clean = self.post_cat_clean(prod_rep_clean)
            if self.union_single_not_match:
                union_features = self.up_dim_clean(union_features)

            ctx_dists_clean = self.ctx_compress_clean(prod_rep_clean * union_features)
            # uni_dists = self.uni_compress(self.drop(union_features))
            frq_dists_clean = self.freq_bias_clean.index_with_labels(pair_pred.long())
            frq_dists_clean = F.dropout(frq_dists_clean, 0.3, training=self.training)
            rel_dists_clean = ctx_dists_clean + frq_dists_clean
            if self.with_transfer:
                rel_dists_clean = (self.pred_adj_nor @ rel_dists_clean.T).T
            rel_dists = rel_dists_clean
        obj_dists = obj_dists.split(num_objs, dim=0)
        rel_dists = rel_dists.split(num_rels, dim=0)

        # we use obj_preds instead of pred from obj_dists
        # because in decoder_rnn, preds has been through a nms stage
        add_losses = {}

        if self.training:
            binary_loss = []
            if binary_preds[0].requires_grad:
                for bi_gt, bi_pred in zip(rel_binarys, binary_preds):
                    bi_gt = (bi_gt > 0).float()
                    binary_loss.append(F.binary_cross_entropy_with_logits(bi_pred, bi_gt))
                add_losses["binary_loss"] = sum(binary_loss) / len(binary_loss)

        return obj_dists, rel_dists, add_losses

@registry.ROI_RELATION_PREDICTOR.register("CausalAnalysisPredictor")
class CausalAnalysisPredictor(nn.Module):
    def __init__(self, config, in_channels):
        super(CausalAnalysisPredictor, self).__init__()
        self.cfg = config
        self.attribute_on = config.MODEL.ATTRIBUTE_ON
        self.spatial_for_vision = config.MODEL.ROI_RELATION_HEAD.CAUSAL.SPATIAL_FOR_VISION
        self.num_obj_cls = config.MODEL.ROI_BOX_HEAD.NUM_CLASSES
        self.num_rel_cls = config.MODEL.ROI_RELATION_HEAD.NUM_CLASSES
        self.fusion_type = config.MODEL.ROI_RELATION_HEAD.CAUSAL.FUSION_TYPE
        self.separate_spatial = config.MODEL.ROI_RELATION_HEAD.CAUSAL.SEPARATE_SPATIAL
        self.use_vtranse = config.MODEL.ROI_RELATION_HEAD.CAUSAL.CONTEXT_LAYER == "vtranse"
        self.effect_type = config.MODEL.ROI_RELATION_HEAD.CAUSAL.EFFECT_TYPE

        assert in_channels is not None
        num_inputs = in_channels

        # load class dict
        statistics = get_dataset_statistics(config)
        obj_classes, rel_classes = statistics['obj_classes'], statistics['rel_classes']
        assert self.num_obj_cls == len(obj_classes)
        assert self.num_rel_cls == len(rel_classes)
        # init contextual lstm encoding
        if config.MODEL.ROI_RELATION_HEAD.CAUSAL.CONTEXT_LAYER == "motifs":
            self.context_layer = LSTMContext(config, obj_classes, rel_classes, in_channels)
        elif config.MODEL.ROI_RELATION_HEAD.CAUSAL.CONTEXT_LAYER == "vctree":
            self.context_layer = VCTreeLSTMContext(config, obj_classes, rel_classes, statistics, in_channels)
        elif config.MODEL.ROI_RELATION_HEAD.CAUSAL.CONTEXT_LAYER == "vtranse":
            self.context_layer = VTransEFeature(config, obj_classes, rel_classes, in_channels)
        else:
            print('ERROR: Invalid Context Layer')

        # post decoding
        self.hidden_dim = config.MODEL.ROI_RELATION_HEAD.CONTEXT_HIDDEN_DIM
        self.pooling_dim = config.MODEL.ROI_RELATION_HEAD.CONTEXT_POOLING_DIM

        if self.use_vtranse:
            self.edge_dim = self.pooling_dim
            self.post_emb = nn.Linear(self.hidden_dim, self.pooling_dim * 2)
            self.ctx_compress = nn.Linear(self.pooling_dim, self.num_rel_cls, bias=False)
        else:
            self.edge_dim = self.hidden_dim
            self.post_emb = nn.Linear(self.hidden_dim, self.hidden_dim * 2)
            self.post_cat = nn.Sequential(*[nn.Linear(self.hidden_dim * 2, self.pooling_dim),
                                            nn.ReLU(inplace=True), ])
            self.ctx_compress = nn.Linear(self.pooling_dim, self.num_rel_cls)
        self.vis_compress = nn.Linear(self.pooling_dim, self.num_rel_cls)

        if self.fusion_type == 'gate':
            self.ctx_gate_fc = nn.Linear(self.pooling_dim, self.num_rel_cls)
            layer_init(self.ctx_gate_fc, xavier=True)

        # initialize layer parameters
        layer_init(self.post_emb, 10.0 * (1.0 / self.hidden_dim) ** 0.5, normal=True)
        if not self.use_vtranse:
            layer_init(self.post_cat[0], xavier=True)
            layer_init(self.ctx_compress, xavier=True)
        layer_init(self.vis_compress, xavier=True)

        assert self.pooling_dim == config.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM

        # convey statistics into FrequencyBias to avoid loading again
        self.freq_bias = FrequencyBias(config, statistics)

        # add spatial emb for visual feature
        if self.spatial_for_vision:
            self.spt_emb = nn.Sequential(*[nn.Linear(32, self.hidden_dim),
                                           nn.ReLU(inplace=True),
                                           nn.Linear(self.hidden_dim, self.pooling_dim),
                                           nn.ReLU(inplace=True)
                                           ])
            layer_init(self.spt_emb[0], xavier=True)
            layer_init(self.spt_emb[2], xavier=True)

        self.label_smooth_loss = Label_Smoothing_Regression(e=1.0)

        # untreated average features
        self.effect_analysis = config.MODEL.ROI_RELATION_HEAD.CAUSAL.EFFECT_ANALYSIS
        self.average_ratio = 0.0005

        self.register_buffer("untreated_spt", torch.zeros(32))
        self.register_buffer("untreated_conv_spt", torch.zeros(self.pooling_dim))
        self.register_buffer("avg_post_ctx", torch.zeros(self.pooling_dim))
        self.register_buffer("untreated_feat", torch.zeros(self.pooling_dim))

    def pair_feature_generate(self, roi_features, proposals, rel_pair_idxs, num_objs, obj_boxs, logger,
                              ctx_average=False):
        # encode context infomation
        obj_dists, obj_preds, edge_ctx, binary_preds = self.context_layer(roi_features, proposals, rel_pair_idxs,
                                                                          logger, ctx_average=ctx_average)
        obj_dist_prob = F.softmax(obj_dists, dim=-1)

        # post decode
        edge_rep = self.post_emb(edge_ctx)
        edge_rep = edge_rep.view(edge_rep.size(0), 2, self.edge_dim)
        head_rep = edge_rep[:, 0].contiguous().view(-1, self.edge_dim)
        tail_rep = edge_rep[:, 1].contiguous().view(-1, self.edge_dim)
        # split
        head_reps = head_rep.split(num_objs, dim=0)
        tail_reps = tail_rep.split(num_objs, dim=0)
        obj_preds = obj_preds.split(num_objs, dim=0)
        obj_prob_list = obj_dist_prob.split(num_objs, dim=0)
        obj_dist_list = obj_dists.split(num_objs, dim=0)
        ctx_reps = []
        pair_preds = []
        pair_obj_probs = []
        pair_bboxs_info = []
        for pair_idx, head_rep, tail_rep, obj_pred, obj_box, obj_prob in zip(rel_pair_idxs, head_reps, tail_reps,
                                                                             obj_preds, obj_boxs, obj_prob_list):
            if self.use_vtranse:
                ctx_reps.append(head_rep[pair_idx[:, 0]] - tail_rep[pair_idx[:, 1]])
            else:
                ctx_reps.append(torch.cat((head_rep[pair_idx[:, 0]], tail_rep[pair_idx[:, 1]]), dim=-1))
            pair_preds.append(torch.stack((obj_pred[pair_idx[:, 0]], obj_pred[pair_idx[:, 1]]), dim=1))
            pair_obj_probs.append(torch.stack((obj_prob[pair_idx[:, 0]], obj_prob[pair_idx[:, 1]]), dim=2))
            pair_bboxs_info.append(get_box_pair_info(obj_box[pair_idx[:, 0]], obj_box[pair_idx[:, 1]]))
        pair_obj_probs = cat(pair_obj_probs, dim=0)
        pair_bbox = cat(pair_bboxs_info, dim=0)
        pair_pred = cat(pair_preds, dim=0)
        ctx_rep = cat(ctx_reps, dim=0)
        if self.use_vtranse:
            post_ctx_rep = ctx_rep
        else:
            post_ctx_rep = self.post_cat(ctx_rep)

        return post_ctx_rep, pair_pred, pair_bbox, pair_obj_probs, binary_preds, obj_dist_prob, edge_rep, obj_dist_list

    def forward(self, proposals, rel_pair_idxs, rel_labels, rel_binarys, roi_features, union_features, logger=None):
        """
        Returns:
            obj_dists (list[Tensor]): logits of object label distribution
            rel_dists (list[Tensor])
            rel_pair_idxs (list[Tensor]): (num_rel, 2) index of subject and object
            union_features (Tensor): (batch_num_rel, context_pooling_dim): visual union feature of each pair
        """
        num_rels = [r.shape[0] for r in rel_pair_idxs]
        num_objs = [len(b) for b in proposals]
        obj_boxs = [get_box_info(p.bbox, need_norm=True, proposal=p) for p in proposals]

        assert len(num_rels) == len(num_objs)

        post_ctx_rep, pair_pred, pair_bbox, pair_obj_probs, binary_preds, obj_dist_prob, edge_rep, obj_dist_list = self.pair_feature_generate(
            roi_features, proposals, rel_pair_idxs, num_objs, obj_boxs, logger)

        if (not self.training) and self.effect_analysis:
            with torch.no_grad():
                avg_post_ctx_rep, _, _, avg_pair_obj_prob, _, _, _, _ = self.pair_feature_generate(roi_features,
                                                                                proposals,
                                                                                rel_pair_idxs,
                                                                                num_objs, obj_boxs,
                                                                                logger,
                                                                                ctx_average=True)

        if self.separate_spatial:
            union_features, spatial_conv_feats = union_features
            post_ctx_rep = post_ctx_rep * spatial_conv_feats

        if self.spatial_for_vision:
            post_ctx_rep = post_ctx_rep * self.spt_emb(pair_bbox)

        rel_dists = self.calculate_logits(union_features, post_ctx_rep, pair_pred, use_label_dist=False)
        rel_dist_list = rel_dists.split(num_rels, dim=0)

        add_losses = {}
        # additional loss
        if self.training:
            rel_labels = cat(rel_labels, dim=0)

            # binary loss for VCTree
            if binary_preds is not None:
                binary_loss = []
                for bi_gt, bi_pred in zip(rel_binarys, binary_preds):
                    bi_gt = (bi_gt > 0).float()
                    binary_loss.append(F.binary_cross_entropy_with_logits(bi_pred, bi_gt))
                add_losses["binary_loss"] = sum(binary_loss) / len(binary_loss)

            # branch constraint: make sure each branch can predict independently
            add_losses['auxiliary_ctx'] = F.cross_entropy(self.ctx_compress(post_ctx_rep), rel_labels)
            if not (self.fusion_type == 'gate'):
                add_losses['auxiliary_vis'] = F.cross_entropy(self.vis_compress(union_features), rel_labels)
                add_losses['auxiliary_frq'] = F.cross_entropy(self.freq_bias.index_with_labels(pair_pred.long()),
                                                              rel_labels)

            # untreated average feature
            if self.spatial_for_vision:
                self.untreated_spt = self.moving_average(self.untreated_spt, pair_bbox)
            if self.separate_spatial:
                self.untreated_conv_spt = self.moving_average(self.untreated_conv_spt, spatial_conv_feats)
            self.avg_post_ctx = self.moving_average(self.avg_post_ctx, post_ctx_rep)
            self.untreated_feat = self.moving_average(self.untreated_feat, union_features)

        elif self.effect_analysis:
            with torch.no_grad():
                # untreated spatial
                if self.spatial_for_vision:
                    avg_spt_rep = self.spt_emb(self.untreated_spt.clone().detach().view(1, -1))
                # untreated context
                avg_ctx_rep = avg_post_ctx_rep * avg_spt_rep if self.spatial_for_vision else avg_post_ctx_rep
                avg_ctx_rep = avg_ctx_rep * self.untreated_conv_spt.clone().detach().view(1,
                                                                                -1) if self.separate_spatial else avg_ctx_rep
                # untreated visual
                avg_vis_rep = self.untreated_feat.clone().detach().view(1, -1)
                # untreated category dist
                avg_frq_rep = avg_pair_obj_prob

            if self.effect_type == 'TDE':  # TDE of CTX
                rel_dists = self.calculate_logits(union_features, post_ctx_rep, pair_obj_probs) - self.calculate_logits(
                    union_features, avg_ctx_rep, pair_obj_probs)
            elif self.effect_type == 'NIE':  # NIE of FRQ
                rel_dists = self.calculate_logits(union_features, avg_ctx_rep, pair_obj_probs) - self.calculate_logits(
                    union_features, avg_ctx_rep, avg_frq_rep)
            elif self.effect_type == 'TE':  # Total Effect
                rel_dists = self.calculate_logits(union_features, post_ctx_rep, pair_obj_probs) - self.calculate_logits(
                    union_features, avg_ctx_rep, avg_frq_rep)
            else:
                assert self.effect_type == 'none'
                pass
            rel_dist_list = rel_dists.split(num_rels, dim=0)

        return obj_dist_list, rel_dist_list, add_losses

    def moving_average(self, holder, input):
        assert len(input.shape) == 2
        with torch.no_grad():
            holder = holder * (1 - self.average_ratio) + self.average_ratio * input.mean(0).view(-1)
        return holder

    def calculate_logits(self, vis_rep, ctx_rep, frq_rep, use_label_dist=True, mean_ctx=False):
        if use_label_dist:
            frq_dists = self.freq_bias.index_with_probability(frq_rep)
        else:
            frq_dists = self.freq_bias.index_with_labels(frq_rep.long())

        if mean_ctx:
            ctx_rep = ctx_rep.mean(-1).unsqueeze(-1)
        vis_dists = self.vis_compress(vis_rep)
        ctx_dists = self.ctx_compress(ctx_rep)

        if self.fusion_type == 'gate':
            ctx_gate_dists = self.ctx_gate_fc(ctx_rep)
            union_dists = ctx_dists * torch.sigmoid(vis_dists + frq_dists + ctx_gate_dists)
            # union_dists = (ctx_dists.exp() * torch.sigmoid(vis_dists + frq_dists + ctx_constraint) + 1e-9).log()    # improve on zero-shot, but low mean recall and TDE recall
            # union_dists = ctx_dists * torch.sigmoid(vis_dists * frq_dists)                                          # best conventional Recall results
            # union_dists = (ctx_dists.exp() + vis_dists.exp() + frq_dists.exp() + 1e-9).log()                        # good zero-shot Recall
            # union_dists = ctx_dists * torch.max(torch.sigmoid(vis_dists), torch.sigmoid(frq_dists))                 # good zero-shot Recall
            # union_dists = ctx_dists * torch.sigmoid(vis_dists) * torch.sigmoid(frq_dists)                           # balanced recall and mean recall
            # union_dists = ctx_dists * (torch.sigmoid(vis_dists) + torch.sigmoid(frq_dists)) / 2.0                   # good zero-shot Recall
            # union_dists = ctx_dists * torch.sigmoid((vis_dists.exp() + frq_dists.exp() + 1e-9).log())               # good zero-shot Recall, bad for all of the rest

        elif self.fusion_type == 'sum':
            union_dists = vis_dists + ctx_dists + frq_dists
        else:
            print('invalid fusion type')

        return union_dists

    def binary_ce_loss(self, logits, gt):
        batch_size, num_cat = logits.shape
        answer = torch.zeros((batch_size, num_cat), device=gt.device).float()
        answer[torch.arange(batch_size, device=gt.device), gt.long()] = 1.0
        return F.binary_cross_entropy_with_logits(logits, answer) * num_cat

    def fusion(self, x, y):
        return F.relu(x + y) - (x - y) ** 2

def make_roi_relation_predictor(cfg, in_channels):
    func = registry.ROI_RELATION_PREDICTOR[cfg.MODEL.ROI_RELATION_HEAD.PREDICTOR]
    return func(cfg, in_channels)


"""Get predictions directly from XGBoost model, within ms2pip framework."""

import hashlib
import logging
import os
import urllib.request
from itertools import islice

import numpy as np
import xgboost as xgb

from ms2pip.exceptions import (InvalidModificationFormattingError,
                               InvalidXGBoostModelError,
                               UnknownModificationError)
from ms2pip.ms2pipC import AMINO_ACID_IDS, ms2pip_pyx

logger = logging.getLogger("ms2pip")

def process_peptides_xgb(peprec, model_params, ptm_ids, model_dir, num_cpu=1):
    """Get predictions for peptides directly from XGBoost model."""
    feature_vectors, mzs = _get_ms2pip_data_for_xgb(peprec, model_params, ptm_ids)
    feature_vectors = xgb.DMatrix(feature_vectors)

    num_ions = (peprec["peptide"].str.len() - 1).to_list()
    peptide_lengths = peprec["peptide"].str.len().to_list()
    charges = peprec["charge"].to_list()
    spec_ids = peprec["spec_id"].to_list()

    preds_list = []
    for ion_type, model_file in model_params["xgboost_model_files"].items():
        # Get predictions from XGBoost model
        bst = xgb.Booster({"nthread": num_cpu})
        bst.load_model(os.path.join(model_dir, model_file))
        preds = bst.predict(feature_vectors)

        # Reshape into arrays for each peptide
        preds = _split_list_by_lengths(preds, num_ions)
        if ion_type in ["x", "y", "y2", "z"]:
            preds = [np.array(x[::-1], dtype=np.float32) for x in preds]
        elif ion_type in ["a", "b", "b2", "c"]:
            preds = [np.array(x, dtype=np.float32) for x in preds]
        else:
            raise ValueError(f"Unsupported ion_type: {ion_type}")
        preds_list.append(preds)

    predictions = [list(t) for t in zip(*preds_list)]

    # List of objects with `get` method is expected, use spoofer class
    return [
        _MultiprocessingResultSpoofer(
            (mzs, predictions, None, peptide_lengths, charges, spec_ids)
        )
    ]

def _get_ms2pip_data_for_xgb(peprec, model_params, ptm_ids):
    """Get feature vectors and mz values for all peptides in self.data."""
    peaks_version = model_params["peaks_version"]

    vector_list = []
    mz_list = []
    for row in peprec.to_dict(orient="records"):

        peptide = np.array(
            [0] + [AMINO_ACID_IDS[x] for x in row["peptide"].replace("L", "I")] + [0],
            dtype=np.uint16,
        )
        modpeptide = apply_mods(peptide, row["modifications"], ptm_ids)
        charge = row["charge"]

        vector_list.append(
            np.array(
                ms2pip_pyx.get_vector(peptide, modpeptide, charge), dtype=np.uint16
            )
        )
        mzs = ms2pip_pyx.get_mzs(modpeptide, peaks_version)
        mz_list.append([np.array(m, dtype=np.float32) for m in mzs])

    feature_vectors = np.vstack(vector_list)

    return feature_vectors, mz_list

def _split_list_by_lengths(list_in, lengths):
    list_in = iter(list_in)
    return [list(islice(list_in, elem)) for elem in lengths]

class _MultiprocessingResultSpoofer:
    """Spoof result structure of multiprocessing, for direct XGB predictions."""

    def __init__(self, contents):
        self.contents = contents

    def get(self):
        return self.contents

def apply_mods(peptide, mods, PTMmap):
    """
    Takes a peptide sequence and a set of modifications. Returns the modified
    version of the peptide sequence, c- and n-term modifications. This modified
    version are hard coded in ms2pipfeatures_c.c for now.
    """
    modpeptide = np.array(peptide[:], dtype=np.uint16)

    if mods != "-":
        l = mods.split("|")
        if len(l) % 2 != 0:
            raise InvalidModificationFormattingError(mods)
        for i in range(0, len(l), 2):
            tl = l[i + 1]
            if tl in PTMmap:
                modpeptide[int(l[i])] = PTMmap[tl]
            else:
                raise UnknownModificationError(tl)

    return modpeptide

def check_model_presence(model, model_hash, model_dir):
    """Check whether xgboost model is downloaded."""
    filename = os.path.join(model_dir, model)
    if not os.path.isfile(filename):
        return False
    return check_model_integrity(filename, model_hash)

def download_model(model, model_hash, model_dir):
    """Download the xgboost model from the Genesis server."""
    if not os.path.isdir(model_dir):
        os.mkdir(model_dir)
    filename = os.path.join(model_dir, model)

    logger.info(f"Downloading {model} to {filename}...")
    urllib.request.urlretrieve(
        os.path.join("http://genesis.ugent.be/uvpublicdata/ms2pip/", model),
        filename
    )
    if not check_model_integrity(filename, model_hash):
        raise InvalidXGBoostModelError()

def check_model_integrity(filename, model_hash):
    """Check that models are correctly downloaded."""
    sha1_hash = hashlib.sha1()
    with open(filename, "rb") as modelfile:
        while True:
            chunk = modelfile.read(16 * 1024)
            if not chunk:
                break
            sha1_hash.update(chunk)
    if sha1_hash.hexdigest() == model_hash:
        return True
    else:
        logger.warn("Model hash not recognised")
        return False

def validate_requested_xgb_model(xgboost_model_files, xgboost_model_hashes, model_dir):
    """Validate requestes xgboost models, and download if neccessary"""

    for _, model_file in xgboost_model_files.items():
        if not check_model_presence(
            model_file, xgboost_model_hashes[model_file], model_dir
        ):
            download_model(
                model_file,
                xgboost_model_hashes[model_file],
                model_dir
            )

def initialize_xgb_models(xgboost_model_files, model_dir, nthread) -> dict:
    """Initialize xgboost models and return them in a dict wit ion types as keys"""
    xgb.set_config(verbosity=0)

    xgboost_models = {}
    for ion_type in xgboost_model_files.keys():
        xgb_model = xgb.Booster({"nthread": nthread})
        xgb_model.load_model(
            os.path.join(model_dir, xgboost_model_files[ion_type])
        )
        xgboost_models[ion_type] = xgb_model

    return xgboost_models


"""
model_interface.py
"""

######################
# Import modules
######################

#import librosa
#import sounddevice as sd
import argparse
import numpy as np
#import tensorflow as tf
import tflite_runtime.interpreter as tflite
import io
import time
import csv

from utils import *

######################
# YAMNetInterface
######################

class YAMNetInterface():
    """
    A class to interface with the YAMNet tflite model

    ...

    Attributes
    ----------
    TOP_K : int
        Number of top k predictions to track
    DURATION_S : int
        Length in seconds of audio clips
    MODE : str
        Two operating modes, a or b
    WATCH_SOUNDS : [str]
        List of sounds to watch for if in mode b
    MODEL_PATH : str
        Path for YAMNet tflite model
    SAMPLERATE_HZ: int
        Sample rate of audio, must be 16000 to work with YAMNet

    Methods
    -------
    predict(y)
        Takes np.array of sound wave and predicts a vector of softmax scores over classes
    scale_data_yamnet(y)
        Scales data between [-1,1], requirement of YAMNet
    predictMode(y)
        Wrapper function for prediction
    getTopK(yh)
        Takes predictions and return top-k classes
    load_class_names()
        Loads list of classes model was trained on and can predict
    getAudioModel()
        Utility to declare tflite model of YAMNet
    """

    def __init__(self,args):
        self.TOP_K = args.TOP_K
        self.DURATION_S = args.DURATION_S
        self.MODE = args.MODE
        self.WATCH_SOUNDS = args.WATCH_SOUNDS
        self.MODEL_PATH = "lite-model_yamnet_tflite_1.tflite"
        self.SAMPLERATE_HZ = 16000 # requirement of YAMNet

        # Load model
        self.interpreter, self.waveform_input_index, self.scores_output_index, self.embeddings_output_index, self.spectrogram_output_index = self.getAudioModel()

        # Load class names
        self.class_names = self.load_class_names()

    def predict(self,y):
        """ Perform prediction with YAMNet tflite model
        Args:
            y (np.array): sound wave in form of np.array with shape (length,)

        Returns:
            yh_k (list): top k classes from predictions
            yh_conf (list): corresponding softmax values for top predictions
        """
        # Scale between -1 and 1
        y = self.scale_data_yamnet(y)

        # Make a prediction
        scores, embeddings, spectrogram = self.predictModel(y)
        yh_k, yh_conf = self.getTopK(scores)
        return yh_k, yh_conf

    @staticmethod
    def scale_data_yamnet(y):
        """ Perform prediction with YAMNet tflite model
        Args:
            y (np.array): sound wave in form of np.array with shape (length,)

        Returns:
            y (np.array): sound wave scaled between [-1,1]
        """
        return 2.0 * (y - np.min(y)) / np.ptp(y) - 1

    def predictModel(self, y):
        """ Perform prediction with YAMNet tflite model
        Args:
            y (np.array): sound wave in form of np.array with shape (length,)

        Returns:
            scores (np.array): softmax values for all classes of YAMNet for given data
            embeddings (np.array): feature embedding of input data
            spectrogram (np.array): a spectrogram of the input
        """
        self.interpreter.resize_tensor_input(self.waveform_input_index, [len(y)], strict=True)
        self.interpreter.allocate_tensors()
        self.interpreter.set_tensor(self.waveform_input_index, y)
        self.interpreter.invoke()
        scores, embeddings, spectrogram = (
            self.interpreter.get_tensor(self.scores_output_index),
            self.interpreter.get_tensor(self.embeddings_output_index),
            self.interpreter.get_tensor(self.spectrogram_output_index),
        )
        return scores, embeddings, spectrogram

    def getTopK(self, yh):
        """ Perform prediction with YAMNet tflite model
        Args:
            yh (np.array): scores of prediction

        Returns:
            (list): top k class predictions from the input
        """
        yh = yh.mean(axis=0)
        yh_max_id = yh.argsort()[-self.TOP_K:][::-1]
        return [self.class_names[k] for k in yh_max_id], [yh[k] for k in yh_max_id]

    def load_class_names(self):
        """Returns list of class names corresponding to score vector."""
        with open("yamnet_class_map.csv", "r") as file:
            class_map_csv_text = file.read()
        class_map_csv = io.StringIO(class_map_csv_text)
        class_names = [
            display_name for (class_index, mid, display_name) in csv.reader(class_map_csv)
        ]
        class_names = class_names[1:]  # Skip CSV header
        return class_names

    def getAudioModel(self):
        """ Returns tflite model of YAMNet and dependencies """
        interpreter = tflite.Interpreter(model_path=self.MODEL_PATH)
        input_details = interpreter.get_input_details()
        waveform_input_index = input_details[0]["index"]
        output_details = interpreter.get_output_details()
        scores_output_index = output_details[0]["index"]
        embeddings_output_index = output_details[1]["index"]
        spectrogram_output_index = output_details[2]["index"]
        return (
            interpreter,
            waveform_input_index,
            scores_output_index,
            embeddings_output_index,
            spectrogram_output_index,
        )


import numpy as np
import pandas as pd
from .constants import *
from . import utils

def extract_base_intensity(maxed, peaks, cells, threshold_peaks):
    """Extract values from `maxed` where `peaks` exceeds `threshold_peaks`. Also extract
    cell label at read position from `cells` and i,j positions of reads.
    """
    # reads outside of cells get label 0
    read_mask = (peaks > threshold_peaks)
    values = maxed[:, :, read_mask].transpose([2, 0, 1])
    labels = cells[read_mask]
    positions = np.array(np.where(read_mask)).T

    return values, labels, positions

def format_bases(values, labels, positions, cycles, bases):
    """Arrange ?x?x? arrays of base intensity information into a dataframe in "long" format (one
    row per observation).
    """
    index = (CYCLE, cycles), (CHANNEL, bases)
    try:
        df = utils.ndarray_to_dataframe(values, index)
    except ValueError:
        print('failed to reshape extracted pixels to sequencing bases, writing empty table')
        return pd.DataFrame()

    df_positions = pd.DataFrame(positions, columns=[POSITION_I, POSITION_J])
    df = (df.stack([CYCLE, CHANNEL])
          .reset_index()
          .rename(columns={0: INTENSITY, 'level_0': READ})
          .join(pd.Series(labels, name=CELL), on=READ)
          .join(df_positions, on=READ)
          .sort_values([CELL, READ, CYCLE])
          )

    return df

def do_median_call(df_bases, cycles=12, channels=4, correction_only_in_cells=False):
    """Call reads from raw base signal using median correction. Use the
    `correction_within_cells` flag to specify if correction is based on reads within
    cells, or all reads.
    """
    if correction_only_in_cells:
        # first obtain transformation matrix W
        X_ = dataframe_to_values(df_bases.query('cell > 0'))
        _, W = transform_medians(X_.reshape(-1, channels))

        # then apply to all data
        X = dataframe_to_values(df_bases)
        Y = W.dot(X.reshape(-1, channels).T).T.astype(int)
    else:
        X = dataframe_to_values(df_bases)
        Y, W = transform_medians(X.reshape(-1, channels))

    df_reads = call_barcodes(df_bases, Y, cycles=cycles, channels=channels)

    return df_reads

def clean_up_bases(df_bases):
    """Sort. Pre-processing for `dataframe_to_values`.
    """
    return df_bases.sort_values([WELL, TILE, CELL, READ, CYCLE, CHANNEL])

def call_cells(df_reads):
    """Determine count of top barcodes
    """
    cols = [WELL, TILE, CELL]
    s = (df_reads
         .drop_duplicates([WELL, TILE, READ])
         .groupby(cols)[BARCODE]
         .value_counts()
         .rename('count')
         .sort_values(ascending=False)
         .reset_index()
         .groupby(cols)
         )

    return (df_reads
            .join(s.nth(0)[BARCODE].rename(BARCODE_0), on=cols)
            .join(s.nth(0)['count'].rename(BARCODE_COUNT_0), on=cols)
            .join(s.nth(1)[BARCODE].rename(BARCODE_1), on=cols)
            .join(s.nth(1)['count'].rename(BARCODE_COUNT_1), on=cols)
            .join(s['count'].sum() .rename(BARCODE_COUNT), on=cols)
            .assign(**{BARCODE_COUNT_0: lambda x: x[BARCODE_COUNT_0].fillna(0),
                       BARCODE_COUNT_1: lambda x: x[BARCODE_COUNT_1].fillna(0)})
            .drop_duplicates(cols)
            .drop([READ, BARCODE], axis=1)  # drop the read
            # drop the read coordinates
            .drop([POSITION_I, POSITION_J], axis=1)
            .filter(regex='^(?!Q_)')  # remove read quality scores
            .query('cell > 0')  # remove reads not in a cell
            )

def call_cells_mapping(df_reads, df_pool=None):
    """Determine count of top two barcodes for each cell. If `df_pool` is
    provided, barcodes that map to designed sequences will be prioritized.
    """
    if df_pool is not None:
        guide_info_cols = [SGRNA, GENE_SYMBOL, GENE_ID]
        # map reads
        df_mapped = (pd.merge(df_reads, df_pool[[PREFIX]],
                              how='left', left_on=BARCODE, right_on=PREFIX)
            .assign(mapped=lambda x: pd.notnull(x[PREFIX]))
            .drop(PREFIX, axis=1)
            )
        extra_col = ['mapped']
    else:
        df_mapped = df_reads
        extra_col = []

    # choose top 2 barcodes, priority given by (mapped,count)
    cols = [WELL, TILE, CELL]
    s = (df_mapped
        .drop_duplicates([WELL, TILE, READ])
        .groupby(cols + extra_col)[BARCODE]
        .value_counts()
        .rename('count')
        .reset_index()
        .sort_values(extra_col + ['count'], ascending=False)
        .groupby(cols)
        )

    df_cells = (df_reads
        .join(s.nth(0)[BARCODE].rename(BARCODE_0), on=cols)
        .join(s.nth(0)['count'].rename(BARCODE_COUNT_0), on=cols)
        .join(s.nth(1)[BARCODE].rename(BARCODE_1), on=cols)
        .join(s.nth(1)['count'].rename(BARCODE_COUNT_1), on=cols)
        .join(s['count'].sum() .rename(BARCODE_COUNT), on=cols)
        .assign(**{BARCODE_COUNT_0: lambda x: x[BARCODE_COUNT_0].fillna(0),
                    BARCODE_COUNT_1: lambda x: x[BARCODE_COUNT_1].fillna(0)})
        .drop_duplicates(cols)
        .drop([READ, BARCODE], axis=1)  # drop the read
        # drop the read coordinates
        .drop([POSITION_I, POSITION_J], axis=1)
        .filter(regex='^(?!Q_)')  # remove read quality scores
        .query('cell > 0')  # remove reads not in a cell
        )

    if df_pool is not None:
        # merge guide info
        df_cells = (pd.merge(df_cells, df_pool[[PREFIX] + guide_info_cols],
                             how='left', left_on=BARCODE_0, right_on=PREFIX)
            .rename({col: col + '_0' for col in guide_info_cols}, axis=1)
            .drop(PREFIX, axis=1)
            )
        df_cells = (pd.merge(df_cells, df_pool[[PREFIX] + guide_info_cols],
                             how='left', left_on=BARCODE_1, right_on=PREFIX)
            .rename({col: col + '_1' for col in guide_info_cols}, axis=1)
            .drop(PREFIX, axis=1)
            )

    return df_cells

def call_cells_mapping(df_reads,df_pool):
    """Determine count of top barcodes, barcodes prioritized if barcode maps to given pool design
    """
    guide_info_cols = [SGRNA,GENE_SYMBOL,GENE_ID]

    # map reads
    df_mapped = (pd.merge(df_reads,df_pool[[PREFIX]],how='left',left_on=BARCODE,right_on=PREFIX)
                 .assign(mapped = lambda x: pd.notnull(x[PREFIX]))
                 .drop(PREFIX,axis=1)
                )

    # choose top 2 barcodes, priority given by (mapped,count)
    cols = [WELL, TILE, CELL]
    s = (df_mapped
       .drop_duplicates([WELL, TILE, READ])
       .groupby(cols+['mapped'])[BARCODE]
       .value_counts()
       .rename('count')
       .reset_index()
       .sort_values(['mapped','count'],ascending=False)
       .groupby(cols)
        )

    df_cells = (df_reads
      .join(s.nth(0)[BARCODE].rename(BARCODE_0),       on=cols)
      .join(s.nth(0)['count'].rename(BARCODE_COUNT_0), on=cols)
      .join(s.nth(1)[BARCODE].rename(BARCODE_1),       on=cols)
      .join(s.nth(1)['count'].rename(BARCODE_COUNT_1), on=cols)
      .join(s['count'].sum() .rename(BARCODE_COUNT),   on=cols)
      .assign(**{BARCODE_COUNT_0: lambda x: x[BARCODE_COUNT_0].fillna(0),
                 BARCODE_COUNT_1: lambda x: x[BARCODE_COUNT_1].fillna(0)})
      .drop_duplicates(cols)
      .drop([READ, BARCODE], axis=1) # drop the read
      .drop([POSITION_I, POSITION_J], axis=1) # drop the read coordinates
      .filter(regex='^(?!Q_)') # remove read quality scores
      .query('cell > 0') # remove reads not in a cell
    )

    # merge guide info
    df_cells = (pd.merge(df_cells,df_pool[[PREFIX] + guide_info_cols],how='left',left_on=BARCODE_0,right_on=PREFIX)
                .rename({col:col+'_0' for col in guide_info_cols},axis=1)
                .drop(PREFIX,axis=1)
               )
    df_cells = (pd.merge(df_cells,df_pool[[PREFIX]+ guide_info_cols],how='left',left_on=BARCODE_1,right_on=PREFIX)
                .rename({col:col+'_1' for col in guide_info_cols},axis=1)
                .drop(PREFIX,axis=1)
               )

    return df_cells

def dataframe_to_values(df, value='intensity'):
    """Dataframe must be sorted on [cycle, channel].
    Returns N x cycles x channels.
    """
    cycles = df[CYCLE].value_counts()
    assert len(set(cycles)) == 1
    n_cycles = len(cycles)
    n_channels = len(df[CHANNEL].value_counts())
    x = np.array(df[value]).reshape(-1, n_cycles, n_channels)
    return x

def transform_medians(X):
    """Estimate and correct differences in channel intensity and spectral overlap among sequencing
    channels. For each channel, find points where the largest signal is from that channel. Use the
    median of these points to define new basis vectors. Describe with linear transformation W
    so that W * X = Y, where Y is the corrected data.
    """

    def get_medians(X):
        arr = []
        for i in range(X.shape[1]):
            arr += [np.median(X[X.argmax(axis=1) == i], axis=0)]
        M = np.array(arr)
        return M

    M = get_medians(X).T
    M = M / M.sum(axis=0)
    W = np.linalg.inv(M)
    Y = W.dot(X.T).T.astype(int)
    return Y, W

def call_barcodes(df_bases, Y, cycles=12, channels=4):
    """Use the argmax over channels in corrected data `Y` to identify bases in each cycle. Result
    dataframe has one row per read and preserves any extra information from `df_bases`.
    """
    bases = sorted(set(df_bases[CHANNEL]))
    if any(len(x) != 1 for x in bases):
        raise ValueError('supplied weird bases: {0}'.format(bases))
    df_reads = df_bases.drop_duplicates([WELL, TILE, READ]).copy()
    df_reads[BARCODE] = call_bases_fast(Y.reshape(-1, cycles, channels), bases)
    Q = quality(Y.reshape(-1, cycles, channels))
    # needed for performance later
    for i in range(len(Q[0])):
        df_reads['Q_%d' % i] = Q[:, i]

    return (df_reads
            .assign(Q_min=lambda x: x.filter(regex='Q_\d+').min(axis=1))
            .drop([CYCLE, CHANNEL, INTENSITY], axis=1)
            )

def call_bases_fast(values, bases):
    """Apply argmax to `values` and form strings by indexing `bases` (e.g., "ACGT" if channel
    dimension in `values` is pre-sorted).
    """
    assert values.ndim == 3
    assert values.shape[2] == len(bases)
    calls = values.argmax(axis=2)
    calls = np.array(list(bases))[calls]
    return [''.join(x) for x in calls]

def quality(X):
    """Define an ad-hoc quality score per sequencing call based on the highest and
    second-highest channels. Adjusted empirically to give reasonable results with 4-color
    sequencing.
    """
    X = np.abs(np.sort(X, axis=-1).astype(float))
    Q = 1 - np.log(2 + X[..., -2]) / np.log(2 + X[..., -1])
    Q = (Q * 2).clip(0, 1)
    return Q

def reads_to_fastq(df_reads, microscope='MN', dataset='DS', flowcell='FC'):
    """Convert reads to fastq format, including quality, tile position, and x,y position.
    """
    def wrap(x): return '{' + x + '}'
    def join_fields(xs): return ':'.join(map(wrap, xs))

    a = '@{m}:{d}:{f}'.format(m=microscope, d=dataset, f=flowcell)
    b = join_fields([WELL, CELL, 'well_tile', READ, POSITION_I, POSITION_J])
    c = '\n{b}\n+\n{{phred}}'.format(b=wrap(BARCODE))
    fmt = a + b + c

    well_tiles = sorted(set(df_reads[WELL] + '_' + df_reads[TILE]))
    fields = [WELL, TILE, CELL, READ, POSITION_I, POSITION_J, BARCODE]

    Q = df_reads.filter(like='Q_').values

    reads = []
    for i, row in enumerate(df_reads[fields].values):
        d = dict(zip(fields, row))
        d['phred'] = ''.join(phred(q) for q in Q[i])
        d['well_tile'] = well_tiles.index(d[WELL] + '_' + d[TILE])
        reads.append(fmt.format(**d))

    return reads

def dataframe_to_fastq(df, file, dataset):
    """Dump reads to fastq format. May be useful for applying NGS-specific tools like FastQC.
    """
    s = '\n'.join(reads_to_fastq(df, dataset))
    with open(file, 'w') as fh:
        fh.write(s)
        fh.write('\n')

def phred(q):
    """Convert 0...1 to 0...30
    No ":".
    No "@".
    No "+".
    """
    n = int(q * 30 + 33)
    if n == 43:
        n += 1
    if n == 58:
        n += 1
    return chr(n)

def add_clusters(df_cells, barcode_col=BARCODE_0, radius=50,
                 verbose=True, ij=(POSITION_I, POSITION_J)):
    """Define spatial clusters based on barcode identity. Rows with the same value in `barcode_col`
    are assigned to a cluster if `ij` positions are within `radius`. Assigns cluster index -1 to
    clusters with only one cell.
    """
    from scipy.spatial import cKDTree
    import networkx as nx

    barcodes = df_cells[barcode_col].values

    I, J = ij
    kdt = cKDTree(np.array([df_cells[J], df_cells[I]]).T)
    num_cells = len(df_cells)

    if verbose:
        print(f'searching for clusters among {num_cells} {barcode_col} objects')
    pairs = kdt.query_pairs(radius)

    # detect clusters
    if len(pairs) == 0:
        clusters = []
    else:
        pairs = np.array(list(pairs))

        x = barcodes[pairs]
        y = x[:, 0] == x[:, 1]

        G = nx.Graph()
        G.add_edges_from(pairs[y])

        clusters = list(nx.connected_components(G))

    # fill in -1 for objects not in a cluster
    cluster_index = np.zeros(num_cells, dtype=int) - 1
    for i, c in enumerate(clusters):
        cluster_index[list(c)] = i

    df_cells = df_cells.copy()
    df_cells[CLUSTER] = cluster_index
    df_cells[CLUSTER_SIZE] = (df_cells
                              .groupby(CLUSTER)[barcode_col].transform('size'))

    # objects not in a cluster have cluster size 1
    df_cells.loc[df_cells[CLUSTER] == -1, CLUSTER_SIZE] = 1
    return df_cells

def index_singleton_clusters(clusters):
    """Replace cluster labels of -1 with ascending integers larger than the maximum cluster index.
    """
    clusters = clusters.copy()
    filt = clusters == -1
    n = clusters.max()
    clusters[filt] = range(n, n + len(filt))
    return clusters

def join_by_cell_location(df_cells, df_ph, max_distance=4):
    """Join together dataframes representing in situ sequencing data and phenotypic data on a
    per-cell basis. Datasets must already be spatially aligned; cells are matched within
    max_distance` (typically less than radius of cell nucleus) based on global position
    (columns `global_x`, `global_y`).

    Can speed up over independent fields of view with `utils.groupby_apply2`.
    """
    from scipy.spatial import cKDTree
    i_tree = df_ph['global_y']
    j_tree = df_ph['global_x']
    i_query = df_cells['global_y']
    j_query = df_cells['global_x']

    kdt = cKDTree(list(zip(i_tree, j_tree)))
    distance, index = kdt.query(list(zip(i_query, j_query)))
    cell_ph = df_ph.iloc[index]['cell'].pipe(list)
    cols_left = ['well', 'tile', 'cell_ph']
    cols_right = ['well', 'tile', 'cell']
    cols_ph = [c for c in df_ph.columns if c not in df_cells.columns]
    return (df_cells
            .assign(cell_ph=cell_ph, distance=distance)
            .query('distance < @max_distance')
            .join(df_ph.set_index(cols_right)[cols_ph], on=cols_left)
            )


from tqdm import tqdm
import gym
import numpy as np
from gym import spaces
from stable_baselines3 import PPO, A2C, DQN
from stable_baselines3.common.evaluation import evaluate_policy
from pettingzoo.mpe import simple_adversary_v2
from policies import static_policy, random_policy, follow, follow_non_goal_landmark_policy, follow_agent_closest_to_landmark_policy, compute_reward
import torch
import torch.nn as nn
import sys
from os.path import join, exists
from os import mkdir, unlink, listdir, getpid, getcwd
import torch
import numpy as np
from misc import load_parameters, flatten_parameters, save_checkpoint
from train_wm import VAE, MDRNNCell, Encoder, Decoder, MDRNN
from train_swm import ContrastiveSWM, EncoderMLP, TransitionGNN
import os

import tools

class SimpleAdversaryEnv(gym.Env):
    metadata = {'render.modes': ['human']}
    def __init__(self, world_model):
        super(SimpleAdversaryEnv, self).__init__()
        self.adversary_policy = follow_agent_closest_to_landmark_policy
        self.agent_policy = follow_non_goal_landmark_policy
        self.world_model = world_model
        self.adversary_eps = 0.2
        self.agent_eps = 0.2
        self.action_space = spaces.Discrete(5)
        self.observation_space = spaces.Box(low=-float('inf'), high=float('inf'), shape=(self.world_model.latent_dim,), dtype=np.float16)
        self._env = simple_adversary_v2.parallel_env(N=2, max_cycles=100, continuous_actions=False)
        self.observations = self._env.reset()
        self.actions_taken = None
    def step(self, action):
        self.world_model.update(self.actions_taken)
        actions = {}
        for agent in self._env.agents:
            if agent == 'adversary_0':
                actions[agent] = self.adversary_policy(self.observations[agent], eps=self.adversary_eps)
            elif self.agent_policy is not None and agent == 'agent_0':
                actions[agent] = self.agent_policy(self.observations[agent], eps=self.agent_eps)
            else:
                actions[agent] = action
        self.observations, rewards, dones, infos = self._env.step(actions)
        self.actions_taken = [actions[agent] for agent in ['adversary_0', 'agent_0', 'agent_1']]
        features = self.world_model.forward(self.observations['agent_1'])
        return features, rewards['agent_1'], dones['agent_1'], infos['agent_1']
    def reset(self):
        self.actions_taken = None
        self.observations = self._env.reset()
        return self.world_model.forward(self.observations['agent_1'])
    def render(self, mode='human'):
        return self._env.render(mode)
    def close(self):
        return self._env.close()

class WorldModel:
    def __init__(self, wm, setting, model_dir="models"):
        self.wm = wm
        self.device = torch.device('cpu')
        if wm == "wm":
            self.vae = torch.load(join(model_dir, setting, "vae.tar")).to(self.device)
            self.vae.eval()
            self.mdrnn = torch.load(join(model_dir, setting, "mdrnn.tar")).to(self.device)
            self.mdrnn.eval()
            self.latent_dim = self.vae.latent_dim + self.mdrnn.latent_dim
            self.temporal_latent = [torch.zeros(1, self.mdrnn.latent_dim).to(self.device) for _ in range(2)]
        elif wm == "swm":
            self.swm = torch.load(join(model_dir, setting, "swm.tar")).to(self.device)
            self.swm.eval()
            self.latent_dim = self.swm.latent_dim
    def forward(self, obs):
        obs = torch.from_numpy(obs).unsqueeze(0).to(self.device)
        if self.wm == "wm":
            _, self.spatial_latent, _ = self.vae(obs)
            return torch.cat((self.spatial_latent, self.temporal_latent[0]), dim=1).detach().numpy()
        elif self.wm == "swm":
            obs = self.swm.forward(obs)
            return torch.flatten(obs, start_dim=1).detach().numpy()
    def update(self, actions):
        if self.wm == "wm":
            if actions is None:
                self.temporal_latent = [torch.zeros(1, self.mdrnn.latent_dim).to(self.device)
                                        for _ in range(2)]
            else:
                actions = [torch.eye(5, device=self.device)[a] for a in actions]
                actions = torch.cat(actions, 0).unsqueeze(0)
                _, _, _, _, self.temporal_latent = self.mdrnn(actions, self.spatial_latent, self.temporal_latent)

def train_rl(
        world_model,
        setting,
        agent,
        train_timesteps=50_000,
        eval_episodes=100,
        model_dir="models",
        verbose=False):
    env = SimpleAdversaryEnv(WorldModel(world_model, setting))
    if agent == "ppo":
        model = PPO("MlpPolicy", env, verbose=0)
    elif agent == "a2c":
        model = A2C("MlpPolicy", env, verbose=0)
    else:
        raise Exception(f"Model {agent} not available.")
    model_dir = join(model_dir, setting)
    if verbose:
        print(f"Training {world_model}-{agent} at: {model_dir}")
    model.learn(total_timesteps=train_timesteps)
    mean, std = evaluate_policy(model, model.get_env(), n_eval_episodes=eval_episodes)
    if verbose:
        print(f"Saved {world_model}-{agent} at {model_dir}: {mean:.4f}")
    model.save(join(model_dir, f"{world_model}-{agent}"))

    return {
        f"{world_model}_agent_reward": mean
    }

if __name__ == "__main__":
    train_rl(world_model="wm", setting="random", agent="ppo")
    train_rl(world_model="swm", setting="random", agent="ppo")
    train_rl(world_model="wm", setting="spurious", agent="ppo")
    train_rl(world_model="swm", setting="spurious", agent="ppo")
    train_rl(world_model="wm", setting="expert", agent="ppo")
    train_rl(world_model="swm", setting="expert", agent="ppo")