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

############ Wall Street Journal ############
# import news
wsj = pd.read_csv("D:\\Dropbox\\China Foreign Share Discount\\WSJ\\All\\factiva0712.csv")

wsj['text'] = wsj['lp']+ '/n/n' + wsj['td']

wsj = wsj.dropna(subset=['text'])

wsj['text'] = [remove_punctuation(text) for text in wsj['text']]        # remove punctuations
tokenizer = RegexpTokenizer(r'\w+')
wsj['text'] = wsj['text'].apply(lambda x: tokenizer.tokenize(x.lower()))        # tokenize
wsj['text'] = wsj['text'].apply(lambda x: word_lemmatizer(x))       # remove stop words

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

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

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

# political news flag
wsj['political'] = wsj['ns'].str.contains('gcat : Political/General News', flags=re.IGNORECASE, regex=True)
wsj['gpir'] = wsj['ns'].str.contains('gpir', flags=re.IGNORECASE, regex=True)

# other news flag
wsj['econ'] = wsj['ns'].str.contains('ecat : Economic News', flags=re.IGNORECASE, regex=True)
wsj['corp'] = wsj['ns'].str.contains('ccat : Corporate/Industrial News', flags=re.IGNORECASE, regex=True)
wsj['commodity'] = wsj['ns'].str.contains('mcat : Commodity/Financial Market News', flags=re.IGNORECASE, regex=True)

# mainly China flag
wsj['re_split'] = wsj['re'].str.split('|')
wsj.loc[:, 're_china'] = wsj.re_split.map(lambda x: 'china' in x[0])
wsj.loc[:, 're_first'] = wsj.re_split.map(lambda x: x[0])

'''
wsj1 = wsj[['pd', 'political', 'gpir', 'econ', 'corp', 'commodity', 're_china', 're_first', 'pos', 'neg', 'pos_lm', 'neg_lm', 'wordcount']]
wsj1.to_csv("sentiment_wsj0712.csv")
'''