import os
import torch.nn as nn
import torch
from torch.autograd import Variable
import numpy as np
import logging
from six.moves import cPickle as pickle

from core.config import cfg
from modeling.get_dataset_counts_rel import get_rel_counts


logger = logging.getLogger(__name__)


class FrequencyBias(nn.Module):
    """
    The goal of this is to provide a simplified way of computing
    P(predicate | obj1, obj2, img).
    """

    def __init__(self, ds_name, eps=1e-3):
        super(FrequencyBias, self).__init__()

        if ds_name.find('vg') >= 0:
            ds_name = 'vg'
        elif ds_name.find('vrd') >= 0:
            ds_name = 'vrd'
        elif ds_name.find('gvqa') >= 0:
            ds_name = 'gvqa'
        else:
            raise NotImplementedError

        if cfg.MODEL.USE_OVLP_FILTER:
            must_overlap = True
        else:
            must_overlap = False
        fg_matrix, bg_matrix = get_rel_counts(ds_name, must_overlap=must_overlap)
        bg_matrix += 1
        fg_matrix[:, :, 0] = bg_matrix

        pred_dist = np.log(fg_matrix / (fg_matrix.sum(2)[:, :, None] + 1e-08) + eps) # (59k, 59k, 29k)

        self.num_objs = pred_dist.shape[0]
        pred_dist = torch.FloatTensor(pred_dist).view(-1, pred_dist.shape[2])

        self.rel_baseline = nn.Embedding(pred_dist.size(0), pred_dist.size(1))
        self.rel_baseline.weight.data = pred_dist
        
        logger.info('Frequency bias tables loaded.')

    def rel_index_with_labels(self, labels):
        """
        :param labels: [batch_size, 2] 
        :return: 
        """
        return self.rel_baseline(labels[:, 0] * self.num_objs + labels[:, 1])
