import pandas as pd
import nltk
from nltk.tokenize import RegexpTokenizer
from sentiment import remove_punctuation, word_lemmatizer, sentiment

trump = pd.read_csv('D:\\Dropbox\\China Foreign Share Discount\\Trump twitter\\tweets_01-08-2021.csv', encoding="cp1252", parse_dates=['date'])

# do not count retweet
#tweets=trump.loc[trump['is_retweet']==0]
tweets = trump.loc[pd.notna(trump['text'])]

#create date variable
tweets['date'] = pd.to_datetime(tweets['date']).dt.date


# clean up the twitter data
# merge the same tweets together
separate=[]
for tweet in tweets['text']:
    if tweet.startswith('...') or tweet.endswith('...'):
        separate.append(1)
    else:
        separate.append(0)
tweets['separate'] = separate

index = tweets.index
tweets['drop'] = 0
for i, j in enumerate(index):
    if tweets['separate'][j] == 1:
        if i > 0 and tweets['separate'][index[i-1]] == 1 and tweets['date'][j] == tweets['date'][index[i-1]]:
            tweets['drop'][j] = 1
        else:
            k = 1
            while tweets['separate'][index[i+k]] == 1 and tweets['date'][j] == tweets['date'][index[i+k]]:
                tweets['text'][j] += " " + tweets['text'][index[i+k]]
                k = k+1

tweets = tweets.loc[tweets['drop']==0]

tweets_china = tweets.loc[tweets['text'].str.contains('China|Chinese|CHINA|CHINESE', regex=True)]        # China-related

tweets_china['text'] = [remove_punctuation(text) for text in tweets_china['text']]

# tokenizer
tokenizer = RegexpTokenizer(r'\w+')

tweets_china['text'] = tweets_china['text'].apply(lambda x: tokenizer.tokenize(x.lower()))

# lemmatizer
tweets_china['text'] = tweets_china['text'].apply(lambda x: word_lemmatizer(x))


# sentiment
positive = []
negative = []
positive_lm = []
negative_lm = []
wordcount = []

for tweet in tweets_china['text']:
    x1,y1,x2,y2,z = sentiment(tweet)
    positive.append(x1)
    negative.append(y1)
    positive_lm.append(x2)
    negative_lm.append(y2)
    wordcount.append(z)

tweets_china['pos'] = positive
tweets_china['neg'] = negative
tweets_china['pos_lm'] = positive_lm
tweets_china['neg_lm'] = negative_lm
tweets_china['wordcount'] = wordcount

'''
tweets_china1 = tweets_china[['date', 'pos', 'neg', 'pos_lm', 'neg_lm', 'wordcount']]
tweets_china1.to_csv("sentiment_twitter0404.csv")
'''