from pyspark import SparkContext
from pyspark import SparkConf
from pyspark.sql import Row
from pyspark.sql import SparkSession
from pyspark.sql.functions import explode
from pyspark.sql.functions import col
from pyspark.sql.functions import dense_rank
from pyspark.sql.functions import collect_set
from pyspark.sql.window import Window
from pyspark.streaming import StreamingContext
import json
import time

'''
Test for Spark Stream to be used with a Tweet Tester, to see if we can using SQL to parse the tweets from the twitter stream, and stop gracefully
Apparently there is some dark magic to stopping spark streaming gracefully and executing one last batch calculation
Consumes the tweets sampled by the Tweet Tester app.

On Terminal 1 run,
$ python TweetTester10-GracefulStreaming.py

On Terminal 2 run,
$ python SparkTester10-GracefulStreaming.py

as the tweet tester streams tweets into the spark streaming app, we should see ...

+-------------------+----------+--------------+
|topic              |num_tweets|num_tweets_rnk|
+-------------------+----------+--------------+
|RHOC               |3         |1             |
|BTSARMY            |2         |2             |
|BTSxELLEN          |2         |3             |
|Education          |2         |4             |
|IoT                |2         |5             |
|MisuseOfLaw_Exposed|2         |6             |
|ออฟกัน             |2         |7             |
|เบบี๋              |2         |8             |
|세븐틴                |2         |9             |
|안전ㅅㅏ설토ㅌㅗㅅㅏㅇㅣ트      |2         |10            |
+-------------------+----------+--------------+


+--------------+---------------------+
|user          |top_topics_of_user   |
+--------------+---------------------+
|damn00131     |[안전ㅅㅏ설토ㅌㅗㅅㅏㅇㅣ트]      |
|cco5minbjc    |[IoT]                |
|FHFactorHouse |[RHOC]               |
|Yun_Miho      |[BTSxELLEN, BTSARMY] |
|BTtaehyeong   |[세븐틴]                |
|brunlii       |[Education]          |
|suttirat1     |[เบบี๋, ออฟกัน]      |
|_MINGYUKIM    |[세븐틴]                |
|beang_beang   |[เบบี๋, ออฟกัน]      |
|taesbaemixed  |[BTSxELLEN, BTSARMY] |
|Vintage_Beso  |[RHOC]               |
|DivyamVerma117|[MisuseOfLaw_Exposed]|
|HardTechTV    |[IoT]                |
|ccorona89     |[RHOC]               |
+--------------+---------------------+


+--------------+---------------------+
|mention       |top_topics_of_mention|
+--------------+---------------------+
|AshramGaushala|[MisuseOfLaw_Exposed]|
|mingukie9731  |[안전ㅅㅏ설토ㅌㅗㅅㅏㅇㅣ트]      |
|SBSNOW        |[세븐틴]                |
|AtthaphanP    |[เบบี๋, ออฟกัน]      |
|IsntDaveOne   |[RHOC]               |
|PeggySulahian |[RHOC]               |
|shannoncoulter|[RHOC]               |
|TheEllenShow  |[BTSxELLEN, BTSARMY] |
|THR           |[RHOC]               |
|AsaramBapuJi  |[MisuseOfLaw_Exposed]|
|pledis_17     |[세븐틴]                |
|theinnocence_J|[세븐틴]                |
|BTS_twt       |[BTSxELLEN, BTSARMY] |
+--------------+---------------------+

See 
* https://stackoverflow.com/questions/35093336/how-to-stop-spark-streaming-when-the-data-source-has-run-out
* https://stackoverflow.com/questions/24374620/python-loop-to-run-for-certain-amount-of-seconds
'''

my_port = 5555
my_host = "spark-2-1"
my_app_name = "hw9_py_test"
my_spark_master = "local[2]" #local StreamingContext with two working threads
my_spark_singleton_name = "sparkSessionSingletonInstance"
my_batch_interval = 5 #batch sampling interval, seconds
my_top_n_topics = 3 #top n most frequently-occurring hashtags among all tweets during the sampling period, or at each sampling period
my_total_run_time = 30 #script run duration, seconds


def getSparkSessionInstance(sparkConf):
    if (my_spark_singleton_name not in globals()):
        globals()[my_spark_singleton_name] = SparkSession\
            .builder\
            .config(conf=sparkConf)\
            .getOrCreate()
    return globals()[my_spark_singleton_name]

    
def main():
    host, port = my_host, my_port
    conf = SparkConf()
    conf.setMaster(my_spark_master)
    conf.setAppName(my_app_name)
    conf.set("spark.streaming.stopGracefullyOnShutdown","true")
    sc = SparkContext(conf=conf)
    sc.setLogLevel("OFF")
    ssc = StreamingContext(sc, my_batch_interval)

    # Create a socket stream on target ip:port and count the
    # words in input stream of \n delimited text (eg. generated by 'nc')
    tweets = ssc.socketTextStream(host, int(port))
    num_tweets_total = 0
    
    def process(time, rdd):
        '''
        helper function that calculates the stats for each sample interval (RDD)
        '''
        print("========= %s =========" % str(time))
        print("processing %s records in this mini-batch" % str(rdd.count()))
        spark = getSparkSessionInstance(rdd.context.getConf())
        df = spark.read.json(rdd)
        
        #todo: add check for the important keys entities.hashtags.text, user.screen_name, entities.user_mentions.screen_name
        
        # finds the top topics, using SQL
        df.createOrReplaceTempView("tweets")
        qry = ( "WITH step1 AS (SELECT explode(entities.hashtags.text) as topic FROM tweets)"
                ", topics AS (SELECT topic, count(*) as num_tweets FROM step1 GROUP BY topic)"
                ", topics_ranked AS (SELECT topic, num_tweets, dense_rank() over (order by num_tweets desc, topic asc) as num_tweets_rnk FROM topics)"
                " SELECT * FROM topics_ranked")
        topics_ranked = spark.sql(qry)
        top_3_topics_ranked = topics_ranked.filter(col("num_tweets_rnk") <= my_top_n_topics)
        top_3_topics_ranked.createOrReplaceTempView("top_3_topics_ranked")
        top_3_topics_ranked.persist()
        
        # finds the authors of popular topics, using SQL
        qry = ( "WITH topics_users AS (SELECT user.screen_name as user, explode(entities.hashtags.text) as topic FROM tweets)"
                ", top_topics_users AS (SELECT ta.user, ta.topic FROM topics_users ta JOIN top_3_topics_ranked tb ON tb.topic = ta.topic)"
                ", top_topics_users_list AS (SELECT user, collect_set(topic) as top_topics_of_user FROM top_topics_users GROUP BY user)"
                " SELECT * FROM top_topics_users_list")
        top_topics_users_list = spark.sql(qry)
        
        # finds the mentions from tweets with popular topics, using SQL
        qry = ( "WITH topics_tweets AS (SELECT id, explode(entities.hashtags.text) as topic FROM tweets)"
                ", top_topics_tweets AS (SELECT ta.id, ta.topic FROM topics_tweets ta JOIN top_3_topics_ranked tb ON tb.topic = ta.topic)"
                ", tweet_mentions AS (SELECT id, explode(entities.user_mentions.screen_name) as mention FROM tweets)"
                ", top_topics_mentions AS (SELECT ta.mention, tb.topic FROM tweet_mentions ta JOIN top_topics_tweets tb ON tb.id = ta.id)"
                ", top_topics_mentions_list AS (SELECT mention, collect_set(topic) as top_topics_of_mention FROM top_topics_mentions GROUP BY mention)"
                " SELECT * FROM top_topics_mentions_list")
        top_topics_mentions_list = spark.sql(qry)
        
        # show all the result dataframes
        top_3_topics_ranked.show(truncate=False)
        top_topics_users_list.show(truncate=False)
        top_topics_mentions_list.show(truncate=False)
    
    # calculates the stats for each sampling interval
    tweets.foreachRDD(process)
    
    #num_tweets_total = tweets.foreachRDD(lambda rdd : rdd.count())
    #tweets.foreachRDD(lambda rdd : num_tweets_total = rdd.count())
    #tweets.count()
    
    # def get_topics(tweet):
        # tweet_json = json.loads(tweet)
        # topics = []
        # try:
            # topics = tweet_json.entities.hash_tags.text
        # except:
            # topics = []
        # return topics
    
    # def get_top_topics(rdd):
        # finds the top topics, using SQL
        # spark = getSparkSessionInstance(rdd.context.getConf())
        # df = spark.read.json(rdd)
        # df.createOrReplaceTempView("tweets")
        # qry = ( "WITH step1 AS (SELECT explode(entities.hashtags.text) as topic FROM tweets)"
                # ", topics AS (SELECT topic, count(*) as num_tweets FROM step1 GROUP BY topic)"
                # ", topics_ranked AS (SELECT topic, num_tweets, dense_rank() over (order by num_tweets desc, topic asc) as num_tweets_rnk FROM topics)"
                # " SELECT * FROM topics_ranked")
        # topics_ranked = spark.sql(qry)
        # top_3_topics_ranked = topics_ranked.filter(col("num_tweets_rnk") <= my_top_n_topics)
        # top_3_topics_ranked.createOrReplaceTempView("top_3_topics_ranked")
        # top_3_topics_ranked.persist()
    
    # gets the stats for the overall counts
    # spark = getSparkSessionInstance(sc.getConf())
    # print(type(tweets))
    # top_3_topics_ranked = tweets.flatMap(get_topics)
    # top_3_topics_ranked.pprint()
    
    ssc.start()
    
    # stop causes an exception???
    ssc.awaitTermination(my_total_run_time)
    
    # t_end = time.time() + my_total_run_time
    # while time.time() < t_end:
        # pass
    try:
        #ssc.stop(stopSparkContext=True, stopGraceFully=True)
        ssc.stop(True, True)
    except:
        pass
    
    print("do final stats somehow")
    print("finished!")
    
if __name__ == "__main__":
    main()