import re
import os
import shutil
import textwrap
import argparse
import ujson as json

from pyspark import SparkContext, SparkConf

from models import links, text

import logging
log = logging.getLogger() 

MODELS = [
    links.EntityCounts,
    links.EntityNameCounts,
    text.TermIndicies,
    text.TermFrequencies,
    text.TermDocumentFrequencies
]

class BuildModel(object):
    """ Wrapper for modules which extract models of entities or text using a corpus of linked documents """
    def __init__(self, **kwargs):
        self.corpus_path = kwargs.pop('corpus_path')
        self.output_path = kwargs.pop('output_path')
        self.sample = kwargs.pop('sample')
        self.sort = False

        modelcls = kwargs.pop('modelcls')
        model_args = {p:kwargs[p] for p in modelcls.__init__.__code__.co_varnames if p in kwargs}
        
        self.model_name = re.sub('([A-Z])', r' \1', modelcls.__name__).strip()

        log.info("Building %s...", self.model_name)
        self.model = modelcls(**model_args)

    def __call__(self):
        log.info('Processing corpus: %s ...', self.corpus_path)
        c = SparkConf().setAppName('Build %s' % self.model_name)

        log.info('Using spark master: %s', c.get('spark.master'))
        sc = SparkContext(conf=c)

        corpus = sc.textFile(self.corpus_path).map(json.loads)
        m = self.model.build(corpus)

        if self.sample > 0:
            if self.sort:
                m = m.map(lambda (k,v): (v,k)).sortByKey(False)
            print '\n'.join(str(i) for i in m.take(self.sample))
        elif self.output_path:
            log.info("Saving to: %s", self.output_path)
            if os.path.isdir(self.output_path):
                log.warn('Writing over output path: %s', self.output_path)
                shutil.rmtree(self.output_path)

            m = self.model.format(m)
            m.saveAsTextFile(self.output_path, 'org.apache.hadoop.io.compress.GzipCodec')

        log.info('Done.')

    @classmethod
    def add_arguments(cls, p):
        p.add_argument('corpus_path', metavar='CORPUS_PATH')
        p.add_argument('--save', dest='output_path', required=False, default=None, metavar='OUTPUT_PATH')
        p.add_argument('--sample', dest='sample', required=False, default=0, type=int, metavar='N')
        p.set_defaults(cls=cls)

        sp = p.add_subparsers()
        for modelcls in MODELS:
            name = modelcls.__name__
            help_str = modelcls.__doc__.split('\n')[0]
            desc = textwrap.dedent(modelcls.__doc__.rstrip())
            csp = sp.add_parser(name,
                                help=help_str,
                                description=desc,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
            modelcls.add_arguments(csp)
        return p

@@ -666,16 +666,56 @@
 ample')%0A
+        self.fmt = kwargs.pop('format')%0A
         
@@ -732,16 +732,16 @@
 = False%0A
-
 %0A       
@@ -1413,16 +1413,167 @@
 (corpus)
+%0A        m = self.model.format(m)%0A%0A        if self.fmt == 'redis':%0A            m = self.format_redis(m)%0A        else:%0A            m = m.map(json.dumps)
 %0A%0A      
@@ -2010,46 +2010,8 @@
 ath)
-%0A%0A            m = self.model.format(m)
 %0A   
@@ -2121,24 +2121,446 @@
 o('Done.')%0A%0A
+    def format_redis(self, model, namespace=None):%0A        cmd = '%5Cr%5Cn'.join(%5B%22*3%22, %22$3%22, %22SET%22, %22$%7B%7D%22, %22%7B%7D%22, %22$%7B%7D%22, %22%7B%7D%22%5D)+'%5Cr'%0A%0A        if namespace == None:%0A            namespace = ''%0A        else:%0A            namespace += ':'%0A%0A        return model%5C%0A            .map(lambda i: ((namespace+i%5B'_id'%5D.replace('%22','%5C%5C%22')).encode('utf-8'), json.dumps(i)))%5C%0A            .map(lambda (t, c): cmd.format(len(t), t, len(c), c))%0A%0A
     @classme
@@ -2754,32 +2754,32 @@
 ='OUTPUT_PATH')%0A
-
         p.add_ar
@@ -2862,16 +2862,127 @@
 ar='N')%0A
+        p.add_argument('--format', choices=%5B'json','redis'%5D, required=False, default='json', metavar='FORMAT')%0A
         

Support redis protocol output format for bulk insert of results