import math
import collections
import numpy as np

from seaqube.nlp.context2vec.common.defs import Toks, SENT_COUNTS_FILENAME, WORD_COUNTS_FILENAME


def read_batch(f, batchsize, word2index):
    batch = []
    while len(batch) < batchsize:
        line = f.readline()
        if not line: break
        sent_words = line.strip().lower().split()
        assert(len(sent_words) > 1)
        sent_inds = []
        for word in sent_words:
            if word in word2index:
                ind = word2index[word]
            else:
                ind = word2index['<UNK>']
            sent_inds.append(ind)
        batch.append(sent_inds)
    return batch

class SentenceReaderDict:
    def __init__(self, prepared_corpus: dict, trimfreq, batchsize):
        '''
        Initialize reader.
        :param path: input directory
        :param trimfreq: treat all words with lower frequency than trimfreq as unknown words
        :param batchsize: the size of the minibatch that will be read in every iteration
        '''
        self.prepared_corpus = prepared_corpus
        self.batchsize = batchsize
        self.trimmed_word2count, self.word2index, self.index2word = self.read_and_trim_vocab(trimfreq)
        self.total_words = sum(self.trimmed_word2count.values())
        self.fds = []


    def open(self):
        self.fds = []
        for word_len, count in self.prepared_corpus['sent_counts'].items():
            batches = int(math.ceil(float(count) / self.batchsize))
            fd = self.prepared_corpus['batches'][word_len] # load the sentences for this current length
            self.fds = self.fds + [fd] * batches
        #np.random.seed(1034)
        np.random.shuffle(self.fds)

    def close(self):
        pass

    def read_and_trim_vocab(self, trimfreq):
        trimmed_word2count = collections.Counter()
        index2word = {Toks.UNK: '<UNK>', Toks.BOS: '<BOS>', Toks.EOS: '<EOS>'}
        word2index = {'<UNK>': Toks.UNK, '<BOS>': Toks.BOS, '<EOS>': Toks.EOS}
        unknown_counts = 0

        for word, count in self.prepared_corpus['word_counts'].items():
            if count >= trimfreq and word.lower() != '<unk>' and word.lower() != '<rw>':
                ind = len(word2index)
                word2index[word] = ind
                index2word[ind] = word
                trimmed_word2count[ind] = count
            else:
                unknown_counts += count
        trimmed_word2count[word2index['<BOS>']] = 0
        trimmed_word2count[word2index['<EOS>']] = 0
        trimmed_word2count[word2index['<UNK>']] = unknown_counts

        return trimmed_word2count, word2index, index2word

    def next_batch(self):
        for fd in self.fds:
            batch = self.read_batch(fd, self.batchsize, self.word2index)
            yield batch

    def read_batch(self, sentences: list, batchsize, word2index):
        batch = []
        sentences_loader = iter(sentences)

        while len(batch) < batchsize:
            try:
                sent_words = next(sentences_loader)
                assert (len(sent_words) > 1)
                sent_inds = []
                for word in sent_words:
                    if word in word2index:
                        ind = word2index[word]
                    else:
                        ind = word2index['<UNK>']
                    sent_inds.append(ind)
                batch.append(sent_inds)

            except StopIteration:
                break

        return batch

class SentenceReaderDir(object):
    '''
    Reads a batch of sentences at a time from a corpus directory in random order.
    Assumes that the sentences are split into different files in the directory according to their word lengths.
    '''
    
    sent_counts_filename = SENT_COUNTS_FILENAME
    word_counts_filename = WORD_COUNTS_FILENAME

    def __init__(self, path, trimfreq, batchsize):
        '''
        Initialize reader.
        :param path: input directory
        :param trimfreq: treat all words with lower frequency than trimfreq as unknown words
        :param batchsize: the size of the minibatch that will be read in every iteration
        '''
        self.path = path
        self.batchsize = batchsize
        self.trimmed_word2count, self.word2index, self.index2word = self.read_and_trim_vocab(trimfreq)
        self.total_words = sum(self.trimmed_word2count.values())
        self.fds = []
        
    def open(self):
        self.fds = []
        with open(self.path+'/'+self.sent_counts_filename) as f:
            for line in f:
                [filename, count] = line.strip().split()
                batches = int(math.ceil(float(count) / self.batchsize))
                fd = open(self.path+'/'+filename, 'r')
                self.fds = self.fds + [fd]*batches
        np.random.seed(1034)
        np.random.shuffle(self.fds)
    

    def close(self):
        fds_set = set(self.fds)
        for f in fds_set:
            f.close()
                        
            
    def read_and_trim_vocab(self, trimfreq):
        word2count = collections.Counter()
        with open(self.path+'/'+self.word_counts_filename) as f:
            for line in f:
                [word, count] = line.strip().lower().split()
                word2count[word] = int(count)
    
        trimmed_word2count = collections.Counter()
        index2word = {Toks.UNK:'<UNK>', Toks.BOS:'<BOS>', Toks.EOS:'<EOS>'}
        word2index = {'<UNK>':Toks.UNK, '<BOS>':Toks.BOS, '<EOS>':Toks.EOS}
        unknown_counts = 0
        for word, count in word2count.items():
            if count >= trimfreq and word.lower() != '<unk>' and word.lower() != '<rw>':    
                ind = len(word2index)
                word2index[word] = ind
                index2word[ind] = word
                trimmed_word2count[ind] = count
            else:
                unknown_counts += count
        trimmed_word2count[word2index['<BOS>']] = 0
        trimmed_word2count[word2index['<EOS>']] = 0
        trimmed_word2count[word2index['<UNK>']] = unknown_counts
        
        return trimmed_word2count, word2index, index2word

       
    def next_batch(self):                
        for fd in self.fds:
            batch = read_batch(fd, self.batchsize, self.word2index)
            yield batch
 
 
 
if __name__ == '__main__':
    import sys   
    reader = SentenceReaderDir(sys.argv[1],int(sys.argv[2]),int(sys.argv[3]))
    
    for i in range(2):
        print('epoc', i)
        reader.open()
        i = 0
        j = 0
        for batch in reader.next_batch():
            if i < 3:
                print(batch)
                print()
            i += 1
            j += len(batch)
        print('batches', i)
        print('sents', j)
        reader.close()

