"""
Visualises weighted list of words wherein size of each word indicates frequency and implied importance
"""

#-------------------------------------------#
# Import dependencies                       #
#-------------------------------------------#

import time
import numpy as np 
import pandas as pd 
import matplotlib.pyplot as plt

from datetime import datetime, timedelta
from PIL import Image
from wordcloud import WordCloud
from IPython.display import Image as im

#-------------------------------------------#
# Global variables                          #
#-------------------------------------------#

IMAGE = '../figures/ladyliberty.png'
FILENAME = '../datasets/cleaned_en_only_small_iranian.csv'
OUTPUT = '../figures/wordcloud.png'
CHUNKSIZE = 10000

#-------------------------------------------#
# D-types                                   #
#-------------------------------------------#

thirtyone_d_types = {                                                                       # Version 1.0 datasets, i.e. 'ira' and 'iranian', have a tweet_language column and so have 31
    'tweetid':                  'str',                                                      # IDs are numeric in the unhashed datasets and alphanumeric in the hashed datasets
    'userid':                   'str',                                                      # IDs are numeric in the unhashed datasets and alphanumeric in the hashed datasets
    'user_display_name':        'str',
    'user_screen_name':         'str',
    'user_reported_location':   'str',
    'user_profile_description': 'str',
    'user_profile_url':         'str',
    'follower_count':           'str',
    'following_count':          'str',
    'account_creation_date':    'str',
    'account_language':         'str',
    'tweet_language':           'str',
    'tweet_text':               'str',
    'tweet_time':               'str',
    'tweet_client_name':        'str',
    'in_reply_to_tweetid':      'str',
    'in_reply_to_userid':       'str',
    'quoted_tweet_tweetid':     'str',
    'is_retweet':               'bool',
    'retweet_userid':           'str',
    'retweet_tweetid':          'str',
    'latitude':                 'str',                                                      # parse as string, for general purpose, and then check later if it is numeric
    'longitude':                'str',                                                      # parse as string, for general purpose, and then check later if it is numeric
    'quote_count':              'str',
    'reply_count':              'str',
    'like_count':               'str',
    'retweet_count':            'str',
    'hashtags':                 'str',
    'urls':                     'str',
    'user_mentions':            'str',
    'poll_choices':             'str'
}

thirty_d_types = {                                                                          # Version 1.1 datasets do not have a tweet_language column and so have only 30 columns
    'tweetid':                  'str',                                                      # IDs are numeric in the unhashed datasets and alphanumeric in the hashed datasets
    'userid':                   'str',                                                      # IDs are numeric in the unhashed datasets and alphanumeric in the hashed datasets
    'user_display_name':        'str',
    'user_screen_name':         'str',
    'user_reported_location':   'str',
    'user_profile_description': 'str',
    'user_profile_url':         'str',
    'follower_count':           'str',
    'following_count':          'str',
    'account_creation_date':    'str',
    'account_language':         'str',
    'tweet_text':               'str',
    'tweet_time':               'str',
    'tweet_client_name':        'str',
    'in_reply_to_tweetid':      'str',
    'in_reply_to_userid':       'str',
    'quoted_tweet_tweetid':     'str',
    'is_retweet':               'bool',
    'retweet_userid':           'str',
    'retweet_tweetid':          'str',
    'latitude':                 'str',                                                      # parse as string, for general purpose, and then check later if it is numeric
    'longitude':                'str',                                                      # parse as string, for general purpose, and then check later if it is numeric
    'quote_count':              'str',
    'reply_count':              'str',
    'like_count':               'str',
    'retweet_count':            'str',
    'hashtags':                 'str',
    'urls':                     'str',
    'user_mentions':            'str',
    'poll_choices':             'str'
}

dates_to_parse = [                                                                          # column(s) that needs to be converted to datetime dtype
    'account_creation_date',
    'tweet_time'
    ]

#-------------------------------------------#
# Functions                                 #
#-------------------------------------------#
def get_chunk_tweet_string(batch_no, chunk):
    """
    Gets the tweets from a single chunk and converts them into a string of words
    """

    tweet_df = chunk
    num_tweets = len(tweet_df)

    print("Beginning processing of ",num_tweets," tweets at: ",str(datetime.now()))

    tweet_words = []
    for i in range(num_tweets):
       tweet = tweet_df.iloc[i]['tweet_text']
       tweet = tweet.split()
       tweet = ",".join(tweet)
       tweet_words.append(tweet)
       tweet_string = ",".join(tweet_words)

       if (i % 10000 == 0):                                                                # tweet checkpoint
           print(str(i) + " tweets processed")

    print("Finished processing of",num_tweets,"of",str(batch_no*int(num_tweets))," tweets at: ",str(datetime.now()))

    return tweet_string

def read_in(dataset):
    print("Reading in dataset at: ", str(datetime.now()))
    dataset_reader = pd.read_csv(
                            dataset,
                            header=0,
                            dtype=thirtyone_d_types,
                            chunksize=CHUNKSIZE,                                            # read in CSV in chunks of this size
                            parse_dates=dates_to_parse,                                     # parse certain columns as datetime dtype
                            keep_default_na=False                                           # do not parse rows with default NaN values
                            )
    batch_no = 1                                                                            # keeps track of how many chunks have been processed
    chunk_tweet_strings = []
    print("Entering loop through reader...")
    for chunk in dataset_reader:                                                            # process each DataFrame
        single_chunk_tweet_string = get_chunk_tweet_string(batch_no, chunk)
        chunk_tweet_strings.append(single_chunk_tweet_string)
        tweet_string = ",".join(chunk_tweet_strings)
        batch_no +=1
    return tweet_string

def generate_word_cloud(image, string, OUTPUT):
    """
    Function to take in tweet string and generate a word cloud based on image
    """

    mask = np.array(Image.open(IMAGE))

    wc = WordCloud(
                   background_color="white", 
                   max_words=2000, 
                   mask=mask
                   )
    wc.generate(string)

    f = plt.figure(figsize=(50,50))
    f.add_subplot(1,2, 1)
    plt.imshow(
               mask,
               cmap=plt.cm.gray, 
               interpolation='bilinear'
               )
    plt.title(
              'Original Stencil', 
              size=40
              )
    plt.axis("off")

    f.add_subplot(1,2, 2)
    plt.imshow(wc, interpolation='bilinear')
    plt.axis("off")
    f.savefig(OUTPUT)

#-------------------------------------------#
# Main                                      #
#-------------------------------------------#

def main(): 
    start_time = time.time()                                                                # program start time

    tweet_string = read_in(FILENAME)
    generate_word_cloud(IMAGE, tweet_string, OUTPUT)
    
    runtime = "Execution: %s seconds" % timedelta(seconds=time.time() - start_time)         # calculate runtime
    
    print(runtime)

#-------------------------------------------#
# Call Main()                               #
#-------------------------------------------#

if __name__ == "__main__":
    main()                                                                                