#!/bin/python

from __future__ import print_function

from lm import LangModel
import random
import numpy as np

class Sampler:

    def __init__(self, lm, temp = 1.0):
        """Sampler for a given language model.

        Supports the use of temperature, i.e. how peaky we want to treat the
        distribution as. Temperature of 1 means no change, temperature <1 means
        less randomness (samples high probability words even more), and temp>1
        means more randomness (samples low prob words more than otherwise). See
        simulated annealing for what this means.
        """
        self.lm = lm
        self.temp = temp

    def sample_sentence(self, prefix = [], max_length = 100):
        """Sample a random sentence (list of words) from the language model.

        Samples words till either EOS symbol is sampled or max_length is reached.
        Does not make any assumptions about the length of the context.
        """
        assert type(prefix) == list
        sent = prefix
        word = self.sample_next(sent, False)
        while len(sent) <= max_length and word != "END_OF_SENTENCE":
            sent.append(word)
            word = self.sample_next(sent)
        return sent

    def sample_next(self, prev, incl_eos = True):
        """Samples a single word from context.

        Can be useful to debug the model, for example if you have a bigram model,
        and know the probability of X-Y should be really high, you can run
        sample_next([Y]) to see how often X get generated.

        incl_eos determines whether the space of words should include EOS or not.
        """
        wps = []
        tot = -np.inf # this is the log (total mass)
        for w in self.lm.vocab():
            if not incl_eos and w == "END_OF_SENTENCE":
                continue
            lp = self.lm.cond_logprob(w, prev, 0)
            wps.append([w, lp/self.temp])
            tot = np.logaddexp2(lp/self.temp, tot)
        p = random.random()
        word = random.choice(wps)[0]
        s = -np.inf # running mass
        for w,lp in wps:
            s = np.logaddexp2(s, lp)
            if p < pow(2, s-tot):
                word = w
                break
        return word

# if __name__ == "__main__":
#     from lm import Unigram
#     unigram = Unigram()
#     corpus = [
#         [ "sam", "i", "am" ]
#     ]
#     unigram.fit_corpus(corpus)
#     print(unigram.model)
#     sampler = Sampler(unigram)
#     for i in range(10):
#         print(i, ":", " ".join(str(x) for x in sampler.sample_sentence([])))
