# -*- coding: utf-8 -*-
"""data.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1TZWO9-LOzuimvdA1GsIXUBvJN0Au9zYj
"""

import numpy as np
import pandas as pd
from nltk.classify import NaiveBayesClassifier
from nltk.corpus import subjectivity
from nltk.sentiment import SentimentAnalyzer
from nltk.sentiment.util import *
import matplotlib.pyplot as mlpt
import csv

##pip install fix_yahoo_finance
##pip install yfinance
import yfinance as yf
import datetime as dt
import snscrape.modules.twitter as sntwitter
##pip install vaderSentiment
##pip install pyspark
import re
import warnings
import pyspark as spark
import pandas as pd
warnings.filterwarnings('ignore')
from pyspark.sql.types import *
from pyspark.sql import SQLContext
from pyspark.sql.functions import col,udf,monotonically_increasing_id,unix_timestamp,round,avg
from vaderSentiment.vaderSentiment import SentimentIntensityAnalyzer
# TODO: import dependencies
import sys
class DataLoader:
    def __init__(self, stock_id, start, end):
        self.start = start
        self.end = end
        self.stock_id = stock_id
        self.twitter_data = []
        self.stock_price = []
        print("import data from {} to {}".format(self.start, self.end))

    def load_twitter_dataid(self, save = False, savepath = 'Tweets.csv'):
        #Import Tweets from Twitter
        arg1 = '$ %s since:%s until:%s'%(self.stock_id, self.start, self.end)
        for i,tweet in enumerate(sntwitter.TwitterSearchScraper(arg1,'corona lang:en').get_items()):
            self.twitter_data.append([tweet.date,  tweet.content]) #tweet.id
        
        self.twitter_data = pd.DataFrame(self.twitter_data, columns=['Date', 'Tweets'])
        data = self.twitter_data
        cdata=pd.DataFrame(columns=['Date','Tweets'])
        index=0
        for index,row in data.iterrows():
            stre=row["Tweets"]
            my_new_string = re.sub('[^ a-zA-Z0-9]', '', stre)
            cdata.sort_index()
            cdata.at[index,'Date']=row["Date"]
            cdata.at[index,'Tweets']=my_new_string
            index=index+1
        self.twitter_data = cdata
        if save:
            self.twitter_data.to_csv(savepath)
        return self.twitter_data

    def load_stock_price(self, save=False, savepath = 'Prices.csv'):
        #Import Stock Market Data from Yahoo! Finance
        data = yf.download(self.stock_id, start=self.start, end=self.end)
        stock_price = pd.DataFrame(data=data)
        stock_price.to_csv(savepath)
        stock_price=pd.read_csv(savepath)
        price=stock_price.drop(['Open', 'High', 'Low', 'Close', 'Volume'], axis=1)
        price.columns = ["Date",'Prices']   
        self.stock_price = price
        return self.stock_price

class Sentiment:
  def __init__(self, stock_data, tweets_data):
    self.stock_data = stock_data
    self.tweets_data = tweets_data
    self.tweets_senti=[]
    self.tweets_senti_price=[]
    
  def senti_analysis(self, save = False, savepath = 'Tweets_Senti_Price.csv'):
    tweets=self.tweets_data
    tweets['Date'] = tweets['Date'].astype('datetime64[ns]').dt.date
    sc = spark.SparkContext()
    sql = spark.SQLContext(sc)
    FullDataTw=sql.createDataFrame(tweets).dropna()
    FullDataTw.select(monotonically_increasing_id().alias("rowId"),"*")
    #setting column names of Twitter dataset
    CleanDF = FullDataTw.withColumnRenamed('Tweets', 'Tweet')
    CleanDF = FullDataTw.withColumnRenamed('Date', 'Date_Time')
    analyser = SentimentIntensityAnalyzer()
    def senti_score_udf(sentence):
      snt = analyser.polarity_scores(sentence)
      return ([snt['neg'], snt['neu'], snt['pos'], snt['compound']])
    func_udf2 = udf(senti_score_udf, ArrayType(FloatType()))
    CleanDF = CleanDF.withColumn('p_neg', func_udf2(CleanDF['Tweets'])[0])
    CleanDF = CleanDF.withColumn('p_neu', func_udf2(CleanDF['Tweets'])[1])
    CleanDF = CleanDF.withColumn('p_pos', func_udf2(CleanDF['Tweets'])[2])
    CleanDF = CleanDF.withColumn('p_comp', func_udf2(CleanDF['Tweets'])[3])
    FinalTw = CleanDF.drop('Unnamed: 0')
    FinalTw.registerTempTable("temp")
    FinalTw_avg = sql.sql("SELECT Date_Time As DateTime,AVG(p_neg) as P_Neg,AVG(p_neu) as P_Neu,AVG(p_pos) as P_Pos,AVG(p_comp) as P_Comp FROM temp GROUP BY Date_Time")
    tw=FinalTw_avg.toPandas()
    tw.sort_values('DateTime', inplace=True)
    tw=tw.reset_index(drop=True)
    self.tweets_senti = tw
    tw['Prices']=""
    # combine sentiment data and stock price data together in one dataframe
    price=self.stock_data
    for i in range (0,len(tw)):
      for j in range (0,len(price)):
        get_tweet_date=tw["DateTime"].iloc[i]
        get_stock_date=price.Date.iloc[j]
        if(str(get_stock_date)==str(get_tweet_date)):
            tw.at[i,'Prices']=price["Prices"].iloc[j]
            break
    data=tw[tw["Prices"]!=""].reset_index(drop=True)
    self.tweets_senti_price = data
    if save:
      self.tweets_senti_price.to_csv(savepath)
    return self.tweets_senti_price


class Obtain_Data: 
  def __init__(self, stock_id, start, end):
    self.stock_id=stock_id
    self.start=start
    self.end=end
    self.dataloader = DataLoader(stock_id=stock_id, start=start, end=end)
    self.stock_data=self.dataloader.load_stock_price(save=True)
    self.tweets_data=self.dataloader.load_twitter_dataid(save=True)
    self.sentiment = Sentiment(self.stock_data,self.tweets_data)
    self.final_dataset=self.sentiment.senti_analysis(save=True)

  def getSplitData(self):
    print("Raw Stock Data Size: ", self.stock_data.shape,"Raw Tweets Data Size: ", self.tweets_data.shape)
    print("Sentiment Data Size",self.final_dataset)

  def getdata(self):
    return self.final_dataset

