# AUTOGENERATED! DO NOT EDIT! File to edit: 01_descriptive_analysis.ipynb (unless otherwise specified).

__all__ = ['describer', 'describer', 'describer', 'describer', 'describer', 'describer', 'describer', 'describer',
           'describer', 'describer']

# Cell
from newstrends import loader

from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import pandas as pd
from collections import Counter
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import matplotlib.pyplot as plt

import multiprocessing
from gensim.models import Word2Vec
from time import time
import os
import math

# Cell
class describer(loader.article_holder):
    "inherit everything from article_holder including init"

    subclass="describer"
    vectorizer=None


# Cell
class describer(describer):
    "Add in vectorize"

    def fitVectorizer(self, vectorizer:CountVectorizer=CountVectorizer, ngram_range=(1,2), max_features=10000):
        try:
            _ = self.df[:1].quickReplace
        except:
            raise Exception("No article data found")

        self.vectorizer=vectorizer(stop_words=self.stopwords, ngram_range=ngram_range, max_features=max_features).fit(self.df.quickReplace)

    def getTopNWords(self, topN = 10, lastDate=None, window=None, source=None):
        " get topN important words for each publication "

        "check if properly formatted"
        if type(self.df) != pd.core.frame.DataFrame:
            raise Exception("Dataframe not loaded")
        if self.vectorizer == None:
            raise Exception("No vectorizer found")


        "Get Dataframe for source and time period"
        sources=source
        if source==None:
            sources=[x for x in self.df.source.unique()]
        df = self.df[self.df.source.isin(sources)]

        "get counts of features from count vectorizer"
        X = self.vectorizer.transform(df.quickReplace)
        vocab = list(self.vectorizer.get_feature_names())
        counts = X.sum(axis=0).A1
        counts = Counter(dict(zip(vocab, counts)))

        return counts.most_common(10)


# Cell
class describer(describer):
    "Now having cooccurances could be nice"

    def generateCoOccurances(self, pubList = ["newyorktimes", "foxnews", "washingtonpost", "cnn", "breitbart", "abcnews", "dailycaller"], verbose=False, topK:"int<100" = 20):
        " get cooccurances of terms in my df, up to 100"

        if type(self.df) != pd.core.frame.DataFrame:
            raise Exception("Dataframe not loaded")
        if self.vectorizer == None:
            raise Exception("No vectorizer found")

        vectorizer = CountVectorizer(stop_words=self.stopwords, max_features=10000).fit(self.df.quickReplace)

        # get the transformed DF
        X = vectorizer.transform(self.df.quickReplace)
        X[X > 0] = 1

        coOccurance = (X.T * X)
        coOccurance.setdiag(0)
        d = coOccurance.todense()

        checkLength = topK*2
        if checkLength > 100:
            checkLength = 100

        top_prs = np.dstack(np.unravel_index(np.argpartition(d.ravel(),-checkLength)[:,-checkLength:],d.shape))[0]

        vals = []
        keys = vectorizer.get_feature_names()
        for pair in top_prs:
            newEntry = [keys[pair[0]], keys[pair[1]]]
            if newEntry not in vals:
                vals.append(newEntry)
            if len(vals) >= topK:
                break

        #So now for each day for each time period I want to math out the co-occurances!
        return vals


# Cell
class describer(describer):
    "update describer with SOT"

    def similarityOverTime(self, endogenous:str, exogenous:list=[], lastNScrapes:int=6, scope=10):
        "caluculate tfidf most similar to publisher"

        vectorizer = TfidfVectorizer(stop_words = self.stopwords, max_features=10000)
        vectorizer.fit(self.df.quickReplace)
        self.df.date = pd.to_datetime(self.df.date)

        recent_dates = sorted(self.df.date.unique())
        today = self.df.date.max()

        cosims = pd.DataFrame()

        for i in range(1,scope+1):
            df = self.df[self.df.date == recent_dates[-i]]
            endo = vectorizer.transform(df[df.source==endogenous].quickReplace).sum(axis=0)

            for exo in exogenous:
                exog = vectorizer.transform(df[df.source==exo].quickReplace).sum(axis=0)
                cosims.at[exo, recent_dates[-i]] = cosine_similarity(endo, exog)
        #for testing purposes
        self.cosims = cosims
        return cosims


# Cell
class describer(describer):
    "upgrade describer with mSOT"

    def meanSimilarityOverTime(self, endogenous:str, exogenous:list, maxLag=24, timeFrame=240):
        "calcultes TFIDF for endo, each exo using TF for doc as a page and IDF over entire corpus"
        vectorizer = TfidfVectorizer(stop_words = self.stopwords, max_features=10000)
        vectorizer.fit(self.df.quickReplace)
        self.df.date = pd.to_datetime(self.df.date)

        recent_dates = sorted(self.df.date.unique())
        today = self.df.date.max()

        lags = pd.DataFrame()

        print("beginning loops")

        for lag in range(maxLag):
            print(lag)
            cosims = pd.DataFrame()
            for i in range(1,timeFrame+1):
                try:
                    endo = vectorizer.transform(self.df[(self.df.source==endogenous) & (self.df.date == recent_dates[-i])].quickReplace).sum(axis=0)
                    for exoPub in exogenous:
                        exog = vectorizer.transform(self.df[(self.df.source==exoPub) & (self.df.date == recent_dates[-(i+lag)])].quickReplace).sum(axis=0)
                        cosims.at[recent_dates[-i], exoPub] = cosine_similarity(endo, exog)
                except:
                    pass
            for exoPub in exogenous:
                lags.at[lag, exoPub] = cosims[exoPub].mean()
            self.lags = lags        #for testing purposes

        return lags


# Cell
import multiprocessing
from gensim.models import Word2Vec
from time import time
import os
import math

# Cell
class describer(describer):
    "upgrade describer with word2vec"

    def word2vec(self, path=".", fileName=""):

        train = self.df.drop_duplicates(subset="quickReplace")

        model = Word2Vec(size=100, workers=multiprocessing.cpu_count()-1)
        model.build_vocab(train.tokens)
        model.train(train.tokens, total_examples=model.corpus_count, epochs=model.epochs)

        if fileName != "":
            os.makedirs(path, exist_ok=True)
            model.save("{}/{}.model".format(path, fileName))

        self.w2v = model
        return model


# Cell
class describer(describer):
    "add getting w2v value"

    def getw2vvector(self, word):
        try:
            return self.w2v.wv[word]
        except:
            return 0


# Cell
class describer(describer):
    "compare cosine between vectors using one of the modelds"

    def w2vCosineSimilarity(self, list1, list2):
        "assumes a and b are pd.datafames of tokens (lists of lists)"

        a = [x for y in list1 for x in y]
        b = [x for y in list2 for x in y]

        firstVector = sum([self.getw2vvector(x) for x in a])
        secondVector = sum([self.getw2vvector(x) for x in b])

        return (self.w2v.wv.cosine_similarities(firstVector, [secondVector]))[0]

# Cell
class describer(describer):
    "update describer with similarity of w2v"

    def w2vSimilarityOverTime(self, endogenous:str, exogenous:list=[], lastNScrapes:int=6, scope=10, keywords=[]):
        "caluculate w2v similar to publisher"

        if type(self.getw2vvector("covid")) == int:
            raise Exception ("no w2v model")


        self.df.date = pd.to_datetime(self.df.date)

        recent_dates = sorted(self.df.date.unique())
        today = self.df.date.max()

        cosims = pd.DataFrame()

        #I'll pre-compute all my vectors first for eac

        for i in range(1,scope+1):
            df = self.df[self.df.date == recent_dates[-i]]
            if len(keywords) > 0:
                df = df[df.tokens.apply(lambda x: bool(set(x) & set(keywords)))]
            for exo in exogenous:
                try:
                    cosims.at[exo, recent_dates[-i]] = self.w2vCosineSimilarity(df[df.source==endogenous].tokens, df[df.source==exo].tokens)
                except:
                    pass
        #for testing purposes
        self.cosims = cosims
        return cosims


# Cell
class describer(describer):
    "upgrade describer with mSOT"

    def meanw2vSimilarityOverTime(self, endogenous:str, exogenous:list, maxLag=24, timeFrame=48, verbose=False, keywords=[]):
        "calcultes w2v for endo, compared to lagged exo; check for 48 timesteps (1 day)"
        self.df.date = pd.to_datetime(self.df.date)

        recent_dates = sorted(self.df.date.unique())
        today = self.df.date.max()

        lags = pd.DataFrame()

        dataframes = {}
        for i in range(1,timeFrame+maxLag+1):
            if verbose:
                if i%10==0:
                    print(i)
            df = self.df[self.df.date == recent_dates[-i]]

            if len(keywords) > 0:
                df = df[df.tokens.apply(lambda x: bool(set(x) & set(keywords)))]

            dataframes[recent_dates[-i]] = {}
            for publisher in [endogenous]+exogenous:
                a = [x for y in df[df.source==publisher].tokens for x in y]
                dataframes[recent_dates[-i]][publisher] = sum([self.getw2vvector(x) for x in a])

        for lag in range(maxLag):
            if verbose:
                print(lag)
            cosims = pd.DataFrame()
            for i in range(1,timeFrame+1):
                for exoPub in exogenous:
                    error_location ="vectors"
                    try:

                        firstVector = dataframes[recent_dates[-i]][endogenous]

                        error_location ="second vector"
                        secondVector = dataframes[recent_dates[-(i+lag)]][exoPub]

                        error_location="cosims"
                        cosims.at[recent_dates[-i], exoPub] = self.w2v.wv.cosine_similarities(firstVector, [secondVector])[0]
                        #cosims.at[recent_dates[-i], exoPub] = self.w2vCosineSimilarity(self.df[(self.df.source==endogenous) & (self.df.date == recent_dates[-i])].tokens,\
                                                                                      #self.df[(self.df.source==exoPub) & (self.df.date == recent_dates[-(i+lag)])].tokens)
                    except:
                        print("error {}".format(error_location))
                        continue


            for exoPub in exogenous:
                try:
                    lags.at[lag, exoPub] = cosims[exoPub].mean()
                except:
                    pass
            self.lags = lags        #for testing purposes

        return lags
