# coding=utf-8
from __future__ import division
import time

from thrift.transport import TSocket
from thrift.protocol import TBinaryProtocol
from thrift.transport import TTransport
from .hbase import Hbase
import pymongo

SUBSTITUTION_TOKEN = 'SUB'

class NgramService(object):
    
    @staticmethod
    def _is_subst(ngram):
        return SUBSTITUTION_TOKEN in set(ngram.split())

    @classmethod
    def configure(cls, substitutions, mongo_host='localhost', hbase_host=('localhost', '9090')):
        cls.substitutions = sorted(substitutions)
        
        cls.m_client = pymongo.MongoClient(host=mongo_host)
        cls.m_1grams = cls.m_client['1grams']['default']
        cls.m_2grams = cls.m_client['2grams']['default']
        cls.m_3grams = cls.m_client['3grams']['default']
    
        # HBASE
        cls.h_transport = TTransport.TBufferedTransport(TSocket.TSocket(*hbase_host))
        protocol = TBinaryProtocol.TBinaryProtocolAccelerated(cls.h_transport)
        cls.h_client = Hbase.Client(protocol)
        cls.h_transport.open()
        cls.h_rate = 0
        cls.h_start = time.time()
        
        cls.substitution_counts = dict([(subst, cls.get_freq(subst)[subst]) for subst in cls.substitutions])

    @classmethod
    def hbase_count(cls, table, ngram):
        """
        :rtype: int
        """
        cls.h_rate += 1
        time_diff = time.time() - cls.h_start
        if time_diff > 30:
            print "HBase req rate:", cls.h_rate/time_diff, "r/s"
            cls.start = time.time()
            cls.rate = 0
        try:
            res = cls.h_client.get(table, ngram, "ngram:cnt", None)
            return long(res[0].value)
        except (ValueError, IndexError):
            return 0

    @classmethod
    def get_freq(cls, ngram):
        """Get ngram frequency from Google Ngram corpus"""
        split_len = len(ngram.split())
        if NgramService._is_subst(ngram):
            if split_len == 1:
                return cls.substitution_counts
            if split_len == 2:
                res = cls.m_2grams.find_one({'ngram': ngram})
            elif split_len == 3:
                res = cls.m_3grams.find_one({'ngram': ngram})
            else:
                raise Exception('%d-grams are not supported yet' % split_len)
            try:
                counts = res['count']
                res = dict([ (ngram.replace(SUBSTITUTION_TOKEN, subst), long(counts.get(subst, 0))) for subst in cls.substitutions])
            except:
                res = dict([ (ngram.replace(SUBSTITUTION_TOKEN, subst), 0) for subst in cls.substitutions])
        else:
            if split_len == 1:
                try:
                    count = cls.m_1grams.find_one({'ngram': ngram})['count']
                except:
                    count = 0
            elif split_len == 2:
                count = cls.hbase_count('ngrams2', ngram)
            else:
                raise Exception('%d-grams are not supported' % split_len)
            res = {ngram: count}
        return res

@@ -51,16 +51,52 @@
 ort time
+%0Afrom collections import defaultdict
 %0A%0Afrom t
@@ -402,24 +402,16 @@
 et(ngram
-.split()
 )%0A%0A    @
@@ -1179,22 +1179,25 @@
  dict(%5B(
+(
 subst,
+),
  cls.get
@@ -1881,18 +1881,16 @@
 lit_
-len = len(
+ngram = 
 ngra
@@ -1898,16 +1898,52 @@
 .split()
+%0A        split_len = len(split_ngram
 )%0A      
@@ -1966,32 +1966,38 @@
 e._is_subst(
+split_
 ngram):%0A
             
@@ -1976,32 +1976,94 @@
 t(split_ngram):%0A
+            sub_index = split_ngram.index(SUBSTITUTION_TOKEN)%0A
             if s
@@ -2479,95 +2479,96 @@
     
-    res = dict(%5B (ngram.replace(SUBSTITUTION_TOKEN, subst), long(counts.get(subst, 0)))
+except:%0A                counts = defaultdict(lambda: 0)%0A            res = %7B%7D%0A           
  for
@@ -2586,34 +2586,33 @@
 ls.substitutions
-%5D)
+:
 %0A            exc
@@ -2612,122 +2612,149 @@
     
-except:%0A                res = dict(%5B (ngram.replace(SUBSTITUTION_TOKEN, subst), 0) for subst in cls.substitutions%5D
+    cur_ngram = split_ngram%5B:%5D%0A                cur_ngram%5Bsub_index%5D = subst%0A                res%5Btuple(cur_ngram)%5D = long(counts.get(subst, 0)
 )%0A  
@@ -3150,21 +3150,36 @@
  res = %7B
-ngram
+tuple(ngram.split())
 : count%7D

Update ngram_service.py