import fire
from pyspark.sql import SparkSession
from pyspark.sql import functions as f
from similarities_helper import build_tfidf_dataset, build_weekly_tfidf_dataset, select_topN_subreddits


def _tfidf_wrapper(func, inpath, outpath, topN, term_colname, exclude):
    spark = SparkSession.builder.getOrCreate()

    df = spark.read.parquet(inpath)

    df = df.filter(~ f.col(term_colname).isin(exclude))

    include_subs = select_topN_subreddits(topN)

    df = func(df, include_subs, term_colname)

    df.write.parquet(outpath,mode='overwrite',compression='snappy')

    spark.stop()

def tfidf(inpath, outpath, topN, term_colname, exclude):
    return _tfidf_wrapper(build_tfidf_dataset, inpath, outpath, topN, term_colname, exclude)

def tfidf_weekly(inpath, outpath, topN, term_colname, exclude):
    return _tfidf_wrapper(build_weekly_tfidf_dataset, inpath, outpath, topN, term_colname, exclude)

def tfidf_authors(outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_authors.parquet',
                  topN=25000):

    return tfidf("/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
                 outpath,
                 topN,
                 'author',
                 ['[deleted]','AutoModerator']
                 )

def tfidf_terms(outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_terms.parquet',
                topN=25000):

    return tfidf("/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
                 outpath,
                 topN,
                 'term',
                 []
                 )

def tfidf_authors_weekly(outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_authors.parquet',
                  topN=25000):

    return tfidf_weekly("/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
                 outpath,
                 topN,
                 'author',
                 ['[deleted]','AutoModerator']
                 )

def tfidf_terms_weekly(outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_terms.parquet',
                topN=25000):

    return tfidf_weekly("/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
                 outpath,
                 topN,
                 'term',
                 []
                 )


if __name__ == "__main__":
    fire.Fire({'authors':tfidf_authors,
               'terms':tfidf_terms,
               'authors_weekly':tfidf_authors_weekly,
               'terms_weekly':tfidf_terms_weekly})