'''
This module does sequence parsing,
stores and writes the n-gram set 
to file.

@author: Gergely Acs <acs@crysys.hu>
'''

import ngram
from Utils import *
import os
#import pickle
from collections import Counter
from ProgressBar import *
from functools import cmp_to_key

def compare_grams(x, y):
    if len(x) == len(y):
        return cmp(x, y)

    return len(x) - len(y)


class NGramSet(Counter):

    def __init__(self, max_len, N_max = 5):
        Counter.__init__(self)
        self.N_max = N_max
        self.alphabet_size = 0
        self.max_len = max_len
        self.all_record_num = 0
        self.TERM = 0

    def load_dataset(self, in_file, dump_file):
        if not os.path.isfile(dump_file):
            print "File " + dump_file + " does not exist!\n"
            print "Creating " + dump_file
            self.parse_sequences(in_file)
            self.dump(dump_file)
        else:
            self.load_dump(dump_file)
            

    def load_dump(self, filename):
        file = open(filename)

        ngram_num = int(file.readline().strip())
        #self.alphabet_size = int(file.readline().strip())

        print "NGrams: ", ngram_num  
        print "Loading NGram file (%s, N=%d)..." % (filename, self.N_max)

        pbar = MyProgressBar('Loading dump', ngram_num)

        for (line_num, line) in enumerate(file):
            parts = line.strip().partition(':')
            tokens = parts[0].strip().split()
            ### NOTE: it is needed only if the file contains longer grams than N_max!!!!
            #if len(seq) > self.N_max:
            #    continue
            #####self[strToSeq(seq)] = float(parts[2].strip()) 
            self[strToSeq(tokens, dec=1)] = float(parts[2].strip()) 

            max_item = max(map(int, tokens)) - 1
            if self.alphabet_size < max_item:
                self.alphabet_size = max_item

            pbar.update(line_num + 1)

        pbar.finish()

        self.TERM = self.alphabet_size 

        print "Alphabet size:", self.alphabet_size

    # we assume that locations are numbered from 1 .. max (so we decrease each
    # location id with one)
    def parse_sequences(self, filename):
        print "Parsing sequence file (%s, N=%d)..." % (filename, self.N_max)
        file = open(filename)

        self.all_record_num = 0
        lines = []
        # First we check the alphabet
        for line in file:
            if line.startswith('#') or line.startswith('//') or line.startswith('%'):
                continue
            lines.append(line.strip().split()[:self.max_len])
            max_item = max(map(int, lines[-1]))
            if self.alphabet_size < max_item:
                self.alphabet_size = max_item

            self.all_record_num += 1

        self.TERM = self.alphabet_size 
        print "Alphabet size:", self.alphabet_size
        print "Termination code:", self.TERM
        print "Number of sequences:", self.all_record_num

        pbar = MyProgressBar('Parsing', self.all_record_num + 1)

        for (record, line) in enumerate(lines):
            seq = strToSeq(line, dec=1) + unichr(self.TERM)
            for i in range(1, self.N_max+1):
                G = ngram.NGram(N=i)
                self.update(G.ngrams(seq))

            pbar.update(record + 1)
        pbar.finish()

        file.close()

    def dump(self, filename):
        print "Creating ngram file (%s, N=%d)..." % (filename, self.N_max)

        file = open(filename, 'w')
        file.write(str(len(self)) + "\n")
        #file.write(str(self.alphabet_size) + "\n")

        pbar = MyProgressBar('Dumping', len(self))

        i = 0
        for gram in sorted(self.keys(), key=cmp_to_key(compare_grams)):
            ###file.write("%s : %f\n" % (" ".join(map(lambda x: str(ord(x)), gram)),self[gram]))
            ## NOTE: ord(x) + 1 should be in order to remain compatible with the input format
            file.write("%s : %f\n" % (" ".join(map(lambda x: str(ord(x)+1), gram)),self[gram]))
            i += 1
            pbar.update(i)

        pbar.finish()
        file.close()
        # pickle.dump(self.ngrams, file) 

        

