from .ranking import Ranking
from .interleaving_method import InterleavingMethod
import numpy as np
import copy

class Probabilistic(InterleavingMethod):
    '''
    Probabilistic Interleaving
    '''
    class Softmax(object):
        def __init__(self, tau, ranking):
            self.tau = tau
            self.ranking = ranking
            self.numerators = np.ndarray(len(ranking))
            self.doc_index = {}
            for r in range(len(ranking)):
                self.numerators[r] = 1.0 / (r+1) ** tau
                docid = ranking[r]
                if not docid in self.doc_index:
                    self.doc_index[docid] = set()
                self.doc_index[docid].add(r)
            self.denominator = np.sum(self.numerators)
            self._original_denominator = self.denominator
            self._non_zero_index = set(range(len(self.numerators)))

        def delete(self, docid):
            for idx in self.doc_index[docid]:
                self.denominator -= self.numerators[idx]
                self._non_zero_index.remove(idx)
            if not self.denominator > 0:
                self.denominator = 0

        def reset(self):
            self.denominator = self._original_denominator
            self._non_zero_index = set(range(len(self.numerators)))

        def sample(self):
            if self.denominator == 0:
                return None
            p = np.random.rand() * self.denominator
            cum = 0.0
            for i in self._non_zero_index:
                cum += self.numerators[i]
                if cum > p:
                    return self.ranking[i]
            return self.ranking[i]

    def __init__(self, tau, max_length, sample_num, *lists):
        '''
        tau: a parameter that determines the probability of documents
        max_length: the maximum length of resultant interleaving
        sample_num: If this is None, an interleaved ranking is generated
                    every time when `interleave` is called.
                    Otherwise, `sample_num` rankings are sampled in the
                    initialization, one of which is returned when `interleave`
                    is called.
        *lists: two lists of document IDs (no multileaving)
        '''
        self._softmaxs = {}
        for i, l in enumerate(lists):
            self._softmaxs[i] = self.Softmax(tau, l)
        super(Probabilistic, self).__init__(max_length, sample_num, *lists)

    def _sample(self, max_length, lists):
        '''
        Sample a ranking

        max_length: the maximum length of resultant interleaving
        *lists: lists of document IDs

        Return an instance of Ranking
        '''
        result = Ranking()
        ranker_indices = list(range(len(lists)))
        result.teams = {i: set() for i in ranker_indices}
        available_rankers = []

        while len(result) < max_length and len(ranker_indices) > 0:
            if len(available_rankers) == 0:
                available_rankers = list(ranker_indices)
                np.random.shuffle(available_rankers)
            ranker_idx = available_rankers.pop()
            docid = self._softmaxs[ranker_idx].sample()
            if docid is None:
                ranker_indices.remove(ranker_idx)
            else:
                result.append(docid)
                result.teams[ranker_idx].add(docid)
                for ranker_idx in ranker_indices:
                    if docid in self._softmaxs[ranker_idx].doc_index:
                        self._softmaxs[ranker_idx].delete(docid)

        # reset the state of softmax
        for i in self._softmaxs:
            self._softmaxs[i].reset()

        return result

    @classmethod
    def _compute_scores(cls, ranking, clicks):
        '''
        ranking: an instance of Ranking
        clicks: a list of indices clicked by a user

        Return a list of scores of each ranker.
        '''
        return {i: len([c for c in clicks if ranking[c] in ranking.teams[i]])
            for i in ranking.teams}


@@ -1667,19 +1667,21 @@
 _(self, 
-tau
+lists
 , max_le
@@ -1676,32 +1676,37 @@
 ists, max_length
+=None
 , sample_num, *l
@@ -1701,24 +1701,30 @@
 mple_num
-, *lists
+=None, tau=3.0
 ):%0A     
@@ -1742,56 +1742,24 @@
     
-tau: a parameter that determines the probability
+lists: two lists
  of 
@@ -1766,16 +1766,19 @@
 document
+ ID
 s%0A      
@@ -1827,32 +1827,152 @@
 ant interleaving
+.%0A                    If this is None (default), it is set to the minimum length%0A                    of the given lists.
 %0A        sample_
@@ -1991,16 +1991,26 @@
  is None
+ (default)
 , an int
@@ -2032,18 +2032,8 @@
 g is
- generated
 %0A   
@@ -2049,16 +2049,26 @@
         
+generated 
 every ti
@@ -2293,58 +2293,96 @@
     
-*lists: two lists of document IDs (no multileaving
+tau: a parameter that determines the probability of documents%0A             (default: 3.0
 )%0A  
@@ -2554,16 +2554,46 @@
 _init__(
+lists,%0A            max_length=
 max_leng
@@ -2610,16 +2610,19 @@
 _num
-, *lists
+=sample_num
 )%0A%0A 

Update the arguments of Probablistic