import time
import numpy as np
import utilities


# Training model
def train(train_file, settings, model_folder, path_dictionary):
    print('initialize the algorithm ...')
    ml_ope = MLOPE(settings['num_terms'], settings['num_topics'], settings['alpha'],
                   settings['tau0'], settings['kappa'], settings['iter_infer'])

    # Start
    print('-->START<--')
    i = 0
    time_start = time.time()
    while i < settings['iter_train']:
        i += 1
        print('\n***iter_train:%d***\n' % (i))
        datafp = open(train_file, 'r')
        j = 0
        while True:
            j += 1
            (wordids, wordcts) = utilities.read_minibatch_list_frequencies(datafp, settings['batch_size'])
            # Stop condition
            if len(wordids) == 0:
                break

            print('---num_minibatch:%d---' % (j))
            ml_ope.static_online(wordids, wordcts)

        datafp.close()

    # write settings
    print('write setting ...')
    file_name_setting = '%s/setting.txt' % model_folder
    utilities.write_setting(settings, file_name_setting)

    # write final model to file
    print('write model ...')
    file_name_beta = '%s/beta_final.dat' % model_folder
    utilities.write_topics(ml_ope.beta, file_name_beta)
    time_end = time.time()

    # write topic
    file_name_topic = '%s/topic.txt' % model_folder
    dictionary = utilities.read_dictionary(path_dictionary)
    utilities.write_topics2(file_name_beta, file_name_topic, dictionary, num_topic=settings['num_topics'])
    # Finish
    print('-------------------------DONE-------------------------', (time_end - time_start))


class MLOPE:
    def __init__(self, num_terms, num_topics, alpha, tau0, kappa, iter_infer):
        """
        Arguments:
            num_terms: Number of unique terms in the corpus (length of the vocabulary).
            num_topics: Number of topics shared by the whole corpus.
            alpha: Hyperparameter for prior on topic mixture theta.
            tau0: A (positive) learning parameter that downweights early iterations.
            kappa: Learning rate: exponential decay rate should be between
                   (0.5, 1.0] to guarantee asymptotic convergence.
            iter_infer: Number of iterations of FW algorithm

        Note that if you pass the same set of all documents in the corpus every time and
        set kappa=0 this class can also be used to do batch OPE.
        """
        self.num_terms = num_terms
        self.num_topics = num_topics
        self.alpha = alpha
        self.tau0 = tau0
        self.kappa = kappa
        self.INF_MAX_ITER = iter_infer
        self.updatect = 1

        # Initializable beta (topics)
        self.beta = np.random.rand(self.num_topics, self.num_terms) + 1e-10
        # import my_MLOPE.utilities as utilities
        # self.beta = utilities.init_beta('F:\\Data\\Important Data\\3.11.2016\\define_kws.txt', 'F:\\Data\\Important Data\\3.11.2016\\dictionary3.txt', self.num_topics, self.num_terms, 0.01)

        # print(self.beta[0], '\n')
        beta_norm = self.beta.sum(axis=1)
        # print(beta_norm)
        self.beta /= beta_norm[:, np.newaxis]
        # print(self.beta[0])

    def static_online(self, wordids, wordcts):
        """
        First does an E step on the mini-batch given in wordids and
        wordcts, then uses the result of that E step to update the
        topics in M step.

        Arguments:
        batch_size: Number of documents of the mini-batch.
        wordids: A list whose each element is an array (terms), corresponding to a document.
                 Each element of the array is index of a unique term, which appears in the document,
                 in the vocabulary.
        wordcts: A list whose each element is an array (frequency), corresponding to a document.
                 Each element of the array says how many time the corresponding term in wordids appears
                 in the document.
        Returns time the E and M steps have taken and the list of topic mixtures of all documents in the mini-batch.
        """
        batch_size = len(wordids)
        # E_step
        start1 = time.time()
        theta = self.e_step(batch_size, wordids, wordcts)
        end1 = time.time()

        # M_step
        start2 = time.time()
        self.m_step(batch_size, wordids, wordcts, theta)
        end2 = time.time()
        return end1 - start1, end2 - start2, theta

    def e_step(self, batch_size, wordids, wordcts):
        """
        Does e step

        Returns topic mixtures theta.
        """
        # Declare theta of minibatch
        theta = np.zeros((batch_size, self.num_topics))
        # Inference
        for d in range(batch_size):
            thetad = self.infer_doc(wordids[d], wordcts[d])
            theta[d, :] = thetad

        return (theta)

    def infer_doc(self, ids, cts):
        """
        Does inference for a document using Online MAP Estimation algorithm.

        Arguments:
        ids: an element of wordids, corresponding to a document.
        cts: an element of wordcts, corresponding to a document.

        Returns inferred theta.
        """
        # locate cache memory
        beta = self.beta[:, ids]
        # Initialize theta randomly
        theta = np.random.rand(self.num_topics) + 1
        theta /= sum(theta)

        x = np.dot(theta, beta)
        # Loop
        T = [1, 0]
        for l in range(1, self.INF_MAX_ITER):
            # Pick fi uniformly
            T[np.random.randint(2)] += 1
            # Select a vertex with the largest value of
            # derivative of the function F
            df = T[0] * np.dot(beta, cts / x) + T[1] * (self.alpha - 1) / theta
            index = np.argmax(df)
            alpha = 1.0 / (l + 1)
            # Update theta
            theta *= 1 - alpha
            theta[index] += alpha
            # Update x
            x = x + alpha * (beta[index, :] - x)

        return theta

    def m_step(self, batch_size, wordids, wordcts, theta):
        """
        Does m step: update global variables beta.
        """
        # Compute intermediate beta which is denoted as "unit beta"
        beta = np.zeros((self.num_topics, self.num_terms), dtype=float)
        for d in range(batch_size):
            beta[:, wordids[d]] += np.outer(theta[d], wordcts[d])
        # Check zeros index
        beta_sum = beta.sum(axis=0)
        ids = np.where(beta_sum != 0)[0]
        unit_beta = beta[:, ids]
        # Nomalize the intermediate beta
        unit_beta_norm = unit_beta.sum(axis=1)
        unit_beta /= unit_beta_norm[:, np.newaxis]
        # Update beta
        rhot = pow(self.tau0 + self.updatect, -self.kappa)
        self.rhot = rhot
        self.beta *= (1 - rhot)
        self.beta[:, ids] += unit_beta * rhot
        self.updatect += 1