import sys
import argparse
import numpy as np
from pyspark import SparkContext

def toLowerCase(s):

    return s.lower()

def stripNonAlpha(s):
    """ Remove non alphabetic characters. E.g. 'B:a,n+a1n$a' becomes 'Banana' """
    return ''.join([c for c in s if c.isalpha()])


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description = 'Text Analysis through TFIDF computation',formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('mode', help='Mode of operation',choices=['TF','IDF','TFIDF','SIM','TOP'])
    parser.add_argument('input', help='Input file or list of files.')
    parser.add_argument('output', help='File in which output is stored')
    parser.add_argument('--master',default="local[20]",help="Spark Master")
    parser.add_argument('--idfvalues',type=str,default="idf", help='File/directory containing IDF values. Used in TFIDF mode to compute TFIDF')
    parser.add_argument('--other',type=str,help = 'Score to which input score is to be compared. Used in SIM mode')
    args = parser.parse_args()

    sc = SparkContext(args.master, 'Text Analysis')


    if args.mode=='TF':
        # Read text file at args.input, compute TF of each term,
        # and store result in file args.output. All terms are first converted to
        # lowercase, and have non alphabetic characters removed
        # (i.e., 'Ba,Na:Na.123' and 'banana' count as the same term). Empty strings, i.e., ""
        # are also removed
        lines=sc.textFile(sys.argv[2])
        lines.flatMap(lambda s: s.split()) \
             .map(lambda word:(stripNonAlpha(toLowerCase((word)) , 1))\
             .filter(lambda word: word[0]!='')\
             .reduceByKey(lambda x, y: x + y) \
             .saveAsTextFile(sys.argv[3])
        pass

    if args.mode=='TOP':
        # Read file at args.input, comprizing strings representing pairs of the form (TERM,VAL),
        # where TERM is a string and VAL is a numeric value. Find the pairs with the top 20 values,
        # and store result in args.output
        lines=sc.textFile(sys.argv[2])
        words=lines.flatMap(lambda s: s.split()) \
             .map(lambda word:(stripNonAlpha(toLowerCase((word)) , 1))\
             .filter(lambda word: word[0]!='')\
             .reduceByKey(lambda x, y: x + y) \
             .takeOrdered(20, lambda (key, value): -value)
        f=open(sys.argv[3], 'w')
        for x in words:
                f.write("(%s, %f)\n" % x)
        f.close()
        pass

    if args.mode=='IDF':
        # Read list of files from args.input, compute IDF of each term,
        # and store result in file args.output.  All terms are first converted to
        # lowercase, and have non alphabetic characters removed
        # (i.e., 'Ba,Na:Na.123' and 'banana' count as the same term). Empty strings ""
        # are removed
        alltext=sc.wholeTextFiles(sys.argv[2])
        doc=alltext.count()
        idf_word=alltext.flatMapValues(lambda s : s.split())\
                        .map(lambda word: (stripNonAlpha(toLowerCase((word[1])), word[0]))\
                        .filter(lambda word: word[0]!= '')\
                        .distinct()\
                        .map(lambda word: (word[0], 1))\
                        .reduceByKey(lambda x, y:x + y)\
                        .map(lambda word: (word[0], np.log(doc/word[1])))\
                        .saveAsTextFile(sys.argv[3])
        pass

    if args.mode=='TFIDF':
        # Read  TF scores from file args.input the IDF scores from file args.idfvalues,
        # compute TFIDF score, and store it in file args.output. Both input files contain
        # strings representing pairs of the form (TERM,VAL),
        # where TERM is a lowercase letter-only string and VAL is a numeric value.
       TF=sc.textFile(args.input).map(eval)
       IDF=sc.textFile(args.idfvalues).map(eval)
       TFIDF=TF.join(IDF)\
               .map(lambda word: (word[0], word[1][0] * word[1][1]))\
               .saveAsTextFile(args.output)
              #.takeOrdered(20, lambda (key, value): -value)
       #f=open(args.output, 'w')
       #for x in TFIDF:
            #f.write("(%s, %f)\n" % x)
       #f.close()

       pass

    if args.mode=='SIM':
        # Read  scores from file args.input the scores from file args.other,
        # compute the cosine similarity between them, and store it in file args.output. Both input files contain
        # strings representing pairs of the form (TERM,VAL),
        # where TERM is a lowercase, letter-only string and VAL is a numeric value.
        tfidfa=sc.textFile(args.input).map(eval)
        tfidfb=sc.textFile(args.other).map(eval)
        numer=tfidfa.join(tfidfb).map(lambda x: x[1][0] * x[1][1]).reduce(lambda x, y: x + y)
        deno=np.sqrt(tfidfa.map(lambda x: x[1] * x[1]).reduce(lambda x, y: x + y)*tfidfb.map(lambda x: x[1] * x[1]).reduce(lambda x, y: x + y))
        cos = numer/deno
        f = open(args.output, 'w')
        f.write("%f" % cos)
        f.close()
        pass
