import sys
import json
import hashlib
import gc
from operator import *
import shlex

from pyspark import StorageLevel
from pyspark.sql import SQLContext
from pyspark.sql.functions import *
from pyspark.sql.types import *

import numpy as np

from subjectivity_clues import clues


def expect(name, var, expected, op=eq):
    if op(var, expected):
        log('[checkpoint] {} = {}'.format(name, expected))
    else:
        log('[error] {} = {}'.format(name, expected))
        raise Exception(name)


def log(message):
    log_file = 'sample_subjectivity_tweets.log'
    with open(log_file, 'a') as f:
        f.write(message)
        f.write('\n')
        f.flush()
        f.close()
    print message


def to_json(name, jsons):
    filename = '{}.json'.format(name)
    with open(filename, 'w') as f:
        for j in jsons:
            f.write(j)
            f.write('\n')


def to_csv(name, jsons):
    filename = '{}.csv'.format(name)
    with open(filename, 'w') as f:
        for tweet in jsons:
            t = json.loads(tweet)
            body = t['body'].replace('\n', ' ').replace('\r', '').replace('"', '""')
            f.write('"{}",{},{},"{}"\n'.format(t['id'], t['verb'], t['postedTime'], body))


def sample(rdd, size, seed):
    items = rdd.collect()
    rand = np.random.RandomState(seed)
    sampled = rand.choice(items, size=size, replace=False)
    expect('sampled', len(set(sampled)), size)
    return sampled.tolist()


def sha(name, ext='json'):
    BUF_SIZE = 65536
    filename = '{}.{}'.format(name, ext)

    sha1 = hashlib.sha1()
    with open(filename, 'rb') as f:
        while True:
            data = f.read(BUF_SIZE)
            if not data:
                break
            sha1.update(data)

    return sha1.hexdigest()


def read_and_parse_clues():
    DEFAULT_FILENAME = os.getcwd() + os.sep + 'subjectivity_clues' + os.sep + 'subjclueslen1-HLTEMNLP05.tff'

    lines = None
    with open(DEFAULT_FILENAME, 'r') as f:
        lines = f.readlines()

    clues = dict()
    for l in lines:
        clue = dict(token.split('=') for token in shlex.split(l))
        word = clue['word1']
        clues[word] = clue

    return clues


def calculate_relevant(lexicons, sentence):
    PRIORPOLARITY = {
        'positive': 1,
        'negative': -1,
        'both': 0,
        'neutral': 0
    }

    TYPE = {
        'strongsubj': 2,
        'weaksubj': 1
    }

    total_score = 0

    for w in sentence.split(' '):
        if w not in lexicons.keys():
            continue

        total_score += PRIORPOLARITY[lexicons[w]['priorpolarity']] * TYPE[lexicons[w]['type']]

    return total_score


# Make sure Python uses UTF-8 as tweets contains emoticon and unicode
reload(sys)
sys.setdefaultencoding('utf-8')

# Use SQLContext for better support
sqlContext = SQLContext(sc)

# Define storage level
DISK_ONLY_2 = StorageLevel(True, False, False, False, 2)
MEMORY_AND_DISK = StorageLevel(True, True, False, False, 1)

# Read GNIP's JSON file
directory = "tweets"
datasets = sqlContext.read.json(directory)
log('# Completed reading JSON files')

# Check checksum count
file_count = datasets.where(datasets['verb'].isNull()).count()
expect('file_count', file_count, 21888)

# Check post count
all_posts = datasets.where(datasets['verb'] == 'post')
all_posts_count = all_posts.count()
expect('all_posts_count', all_posts_count, 1570398)

# Check share count
all_shares = datasets.where(datasets['verb'] == 'share')
all_shares_count = all_shares.count()
expect('all_shares_count', all_shares_count, 1112590)

# Check dataset count
info_dataset = datasets.select('info')
info_dataset.registerTempTable('info')
all_tweets_count = info_dataset.select('info.activity_count').groupBy().sum('activity_count').collect()[0][0]
expect('all_tweets_count', all_tweets_count, 2682988)
expect('all_tweets_count', all_tweets_count, all_posts_count + all_shares_count)
log('# Completed validating tweets count')

# Remove post authored by @ChipotleTweet and news agencies
chipotle_tweet = 'id:twitter.com:141341662'
users_to_remove = [chipotle_tweet, 'id:twitter.com:759251', 'id:twitter.com:91478624', 'id:twitter.com:28785486',
                   'id:twitter.com:1652541', 'id:twitter.com:51241574', 'id:twitter.com:807095',
                   'id:twitter.com:34713362', 'id:twitter.com:3090733766', 'id:twitter.com:1367531',
                   'id:twitter.com:14293310', 'id:twitter.com:3108351', 'id:twitter.com:14173315',
                   'id:twitter.com:292777349', 'id:twitter.com:428333', 'id:twitter.com:624413',
                   'id:twitter.com:20562637', 'id:twitter.com:13918492', 'id:twitter.com:16184358',
                   'id:twitter.com:625697849', 'id:twitter.com:2467791', 'id:twitter.com:9763482',
                   'id:twitter.com:14511951', 'id:twitter.com:6017542', 'id:twitter.com:26574283',
                   'id:twitter.com:115754870']

all_posts_wo_specific_users = all_posts.where(~ col('actor.id').isin(users_to_remove))
all_posts_w_specific_users = all_posts.where(col('actor.id').isin(users_to_remove)).count()
expect('all_posts_wo_specific_users', all_posts_wo_specific_users.count(), all_posts_count - all_posts_w_specific_users)

# Remove share retweet of tweet by @ChipotleTweet and news agencies
all_shares_wo_specific_users = all_shares.where(~ col('object.actor.id').isin(users_to_remove))
all_shares_w_specific_users = all_shares.where(col('object.actor.id').isin(users_to_remove)).count()
expect('all_shares_wo_specific_users', all_shares_wo_specific_users.count(), all_shares_count - all_shares_w_specific_users)

# Generate tweets pool with only English tweet
tweets_pool = all_posts_wo_specific_users.unionAll(all_shares_wo_specific_users).filter("twitter_lang = 'en'")
tweets_pool.persist(MEMORY_AND_DISK)
tweets_pool_count = tweets_pool.count()
# Adding all post to all share will be greater than tweet pool because of non-English tweet
expected_tweets_pool_count = all_posts_count - all_posts_w_specific_users + \
                             all_shares_count - all_shares_w_specific_users
expect('tweets_pool_count', tweets_pool_count, expected_tweets_pool_count, op=lt)
log('# Completed constructing tweets pool')

# Check language of tweets
languages = tweets_pool.select('twitter_lang').distinct()
languages_count = languages.count()
language_check = languages.first()['twitter_lang']
expect('languages_count', languages_count, 1)
expect('language_check', language_check, 'en')
log('# Completed validating language variety')

# Take top 80% of tweets by length
tweets_pool_str_lengths = tweets_pool.select(length('body').alias('length')).rdd.map(lambda x: x.length).collect()
lengths_np = np.array(tweets_pool_str_lengths)
p = np.percentile(lengths_np, 20)

final_tweets_pool = tweets_pool.filter(length('body') >= p)
final_tweets_pool.persist(MEMORY_AND_DISK)
tweets_pool.unpersist(blocking=True)

final_tweets_pool_count = final_tweets_pool.count()
percentage_kept = float(final_tweets_pool_count) / tweets_pool_count
expect('percentage_kept', percentage_kept, 0.8, op=gt)
log('# Completed sampling top 80% of tweets by body length')

# Sampling
final_tweets_ids = final_tweets_pool.select(final_tweets_pool['id']).rdd.sortBy(lambda x: x.id).map(lambda x: x.id)

# Development tweets
dev_seed = 10102016
number_of_dev_samples = 3000
dev_posts = sample(final_tweets_ids, number_of_dev_samples, dev_seed)
dev_posts_count = len(dev_posts)
expect('dev_posts_count', dev_posts_count, number_of_dev_samples)
log('# Completed sampling dev tweets')

dev_posts_file = "dev_posts"
dev_posts_jsons = final_tweets_pool[final_tweets_pool['id'].isin(dev_posts)].toJSON().collect()
to_json(dev_posts_file, dev_posts_jsons)
to_csv(dev_posts_file, dev_posts_jsons)
expect('dev_posts_file', sha(dev_posts_file), '74447296831c8e3061fc0ee739f549c5b08b85a3')
expect('dev_posts_file', sha(dev_posts_file, ext='csv'), '6acfd1f8d238bc5d25d97d2c9e6f6b177699389a')
log('Exporting dev post to {}'.format(dev_posts_file))
log('# Completed exporting dev tweets')

del dev_posts_jsons
gc.collect()

# Find distinct set of tweets (unique body text)
post_pool = final_tweets_pool.where(final_tweets_pool['verb'] == 'post')
post_pool.persist(MEMORY_AND_DISK)
post_pool_ids = post_pool.select(post_pool['id']).rdd.sortBy(lambda x: x.id).map(lambda x: x.id).collect()
expect('post_pool', post_pool.count(), 1124935)

share_pool = final_tweets_pool.where(final_tweets_pool['verb'] == 'share')
share_pool.persist(MEMORY_AND_DISK)
expect('share_pool', share_pool.count(), 846141)

broadcast_post_ids = sc.broadcast(set(post_pool_ids))
unique_share_ids = share_pool.select(share_pool['id'], share_pool['object.id'].alias('object_id')).rdd.filter(lambda row: row['object_id'] not in broadcast_post_ids.value).map(lambda row: row.id).collect()
expect('unique_share_pool', len(unique_share_ids), 193006)
log('# Completed finding unique share tweet')

# Constructing distinct tweet pool
broadcast_unique_share_ids = sc.broadcast(unique_share_ids)
distinct_tweets_pool = final_tweets_pool.\
    select(final_tweets_pool['id'], final_tweets_pool['body']).\
    rdd.\
    filter(lambda row: row['id'] in broadcast_post_ids.value or row['id'] in broadcast_unique_share_ids.value)
distinct_tweets_pool.persist(MEMORY_AND_DISK)
distinct_tweets_count = distinct_tweets_pool.count()
expect('distinct_tweets_pool', distinct_tweets_count, 1124935 + 193006)

# Exclude development tweets
tweets_unsampled = distinct_tweets_pool.toDF().where(~ col('id').isin(dev_posts))
tweets_unsampled.persist(MEMORY_AND_DISK)
tweets_unsampled_count = tweets_unsampled.count()
# no. of dev intersect post pool: 1718, no. of share dev intersect unique share pool: 293
expect('tweets_unsampled', tweets_unsampled_count, 1124935 + 193006 - 1718 - 293)
log('# Completed constructing unsampled tweets')

# Calculate subjectivity
lexicons = read_and_parse_clues()
udfBodyToRelevant = udf(lambda body: calculate_relevant(lexicons, body), IntegerType())

tweets_lexicon = tweets_unsampled.withColumn('score', udfBodyToRelevant('body'))
tweets_lexicon.persist(MEMORY_AND_DISK)
log('# Completed constructing tweet lexicon')

# Take top and bottom
number_of_tweets_each = 1500
positive_tweets = tweets_lexicon.orderBy(desc('score')).take(number_of_tweets_each)
negative_tweets = tweets_lexicon.orderBy(asc('score')).take(number_of_tweets_each)

# Cut top and bottom via score for more deterministic sampling
min_positive_score = positive_tweets[-1]['score']
min_negative_score = negative_tweets[-1]['score']
expect('min_positive_score', min_positive_score, 7)
expect('min_negative_score', min_negative_score, -5)

positive_tweets = tweets_lexicon.filter('score > {}'.format(min_positive_score - 1)).orderBy(desc('score')).collect()
expect('positive_tweets', len(positive_tweets), 2012)
negative_tweets = tweets_lexicon.filter('score < {}'.format(min_negative_score + 1)).orderBy(asc('score')).collect()
expect('positive_tweets', len(negative_tweets), 1715)

positive_tweet_file = "positive_tweets"
positive_tweets_ids = map(lambda t: t['id'], positive_tweets)
positive_tweet_jsons = final_tweets_pool[final_tweets_pool['id'].isin(positive_tweets_ids)].toJSON().collect()
to_json(positive_tweet_file, positive_tweet_jsons)
to_csv(positive_tweet_file, positive_tweet_jsons)
log('Exporting positive tweets to {}'.format(positive_tweet_file))
log('# Completed exporting positive tweets')

expect('positive_tweet_file', sha(positive_tweet_file), 'cb2f8b691ccf3eae9846c67735f413a49befea28')
expect('positive_tweet_file', sha(positive_tweet_file, ext='csv'), 'd3d43ab4e03fdf106b9191f4e0161cfcde3f040e')

negative_tweet_file = "negative_tweets"
negative_tweet_ids = map(lambda t: t['id'], negative_tweets)
negative_tweet_jsons = final_tweets_pool[final_tweets_pool['id'].isin(negative_tweet_ids)].toJSON().collect()
to_json(negative_tweet_file, negative_tweet_jsons)
to_csv(negative_tweet_file, negative_tweet_jsons)
log('Exporting negative tweets to {}'.format(negative_tweet_file))
log('# Completed exporting negative tweets')

expect('negative_tweet_file', sha(negative_tweet_file), '086c43427078092e538a779b8b06a71341b8da48')
expect('negative_tweet_file', sha(negative_tweet_file, ext='csv'), 'd10a1a95156c28d844e9c4e668d766963c0636a4')
