import pandas as pd
import tensorflow as tf
import nltk 
from nltk.corpus import stopwords
from sklearn.preprocessing import LabelEncoder
from tensorflow.keras.layers.experimental.preprocessing import TextVectorization
import re
import string
from tensorflow.keras import layers
from tensorflow import keras
from tensorflow.keras import losses
from matplotlib import pyplot as plt
from sklearn.decomposition import PCA
import plotly.express as px 
from plotly.io import write_html

train_url = "https://github.com/PhilChodrow/PIC16b/blob/master/datasets/fake_news_train.csv?raw=true"
df = pd.read_csv(train_url)
le = LabelEncoder()
df["fake"] = le.fit_transform(df["fake"])
num_fake = len(df["fake"].unique())

nltk.download('stopwords')
stopword = stopwords.words('english')

def make_dataset(df):
  
  df['title'] = df['title'].apply(lambda x: ' '.join([word for word in x.split() if word not in (stopword)]))
  df['text'] = df['text'].apply(lambda x: ' '.join([word for word in x.split() if word not in (stopword)]))

  ds = tf.data.Dataset.from_tensor_slices((
      {"title": df[["title"]],
       "text": df[["text"]]},
      {"fake": df["fake"]}
    ))
  return ds

ds = make_dataset(df)

ds = ds.shuffle(buffer_size = len(ds))

train_size = int(0.8*len(ds))
val_size = int(0.2*len(ds))

train = ds.take(train_size).batch(20)
val = ds.skip(train_size).take(val_size).batch(20)

print(len(ds),sum(df["fake"] == 1),sum(df["fake"] == 0))
11740/22449

#preparing a text vectorization layer for tf model
size_vocabulary = 2000

def standardization(input_data):
    lowercase = tf.strings.lower(input_data)
    no_punctuation = tf.strings.regex_replace(lowercase,
                                  '[%s]' % re.escape(string.punctuation),'')
    return no_punctuation 

title_vectorize_layer = TextVectorization(
    standardize=standardization,
    max_tokens=size_vocabulary, # only consider this many words
    output_mode='int',
    output_sequence_length=500) 

title_vectorize_layer.adapt(train.map(lambda x, y: x["title"]))

title_input = keras.Input(
    shape=(1,),
    name = "title",
    dtype = "string"
)

title_features = title_vectorize_layer(title_input)
title_features = layers.Embedding(size_vocabulary, output_dim = 2, name="embedding")(title_features)
title_features = layers.Dropout(0.2)(title_features)
title_features = layers.GlobalAveragePooling1D()(title_features)
title_features = layers.Dropout(0.2)(title_features)
title_features = layers.Dense(32, activation='relu')(title_features)

output = layers.Dense(num_fake, name="fake")(title_features) 

model1 = keras.Model(
    inputs = title_input,
    outputs = output
)

keras.utils.plot_model(model1)

model1.compile(optimizer="adam",
        loss = losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=["accuracy"])

history = model1.fit(train, 
            validation_data=val,
            epochs = 30,
            verbose = False)

plt.plot(history.history["accuracy"], label = "training")
plt.plot(history.history["val_accuracy"], label = "validation")
plt.gca().set(xlabel = "epoch", ylabel = "accuracy")
plt.legend()

text_input = keras.Input(
    shape=(1,),
    name = "text", # same name as the dictionary key in the dataset
    dtype = "string"
)

text_features = title_vectorize_layer(text_input)
text_features = layers.Embedding(size_vocabulary, output_dim = 2, name="embedding")(text_features)
text_features = layers.Dropout(0.2)(text_features)
text_features = layers.GlobalAveragePooling1D()(text_features)
text_features = layers.Dropout(0.2)(text_features)
text_features = layers.Dense(32, activation='relu')(text_features)

output = layers.Dense(num_fake, name="fake")(text_features) 

model2 = keras.Model(
    inputs = text_input,
    outputs = output
)

keras.utils.plot_model(model2)

model2.compile(optimizer="adam",
        loss = losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=["accuracy"])

history = model2.fit(train, 
            validation_data=val,
            epochs = 30,
            verbose = False)

plt.plot(history.history["accuracy"], label = "training")
plt.plot(history.history["val_accuracy"], label = "validation")
plt.gca().set(xlabel = "epoch", ylabel = "accuracy")
plt.legend()

main = layers.concatenate([title_features, text_features], axis = 1) # Combine the layers of the first two models
main = layers.Dense(32, activation='relu')(main)

output = layers.Dense(num_fake, name="fake")(main) 
model3 = keras.Model(
    inputs = [title_input, text_input],
    outputs = output
)

keras.utils.plot_model(model3)

model3.compile(optimizer="adam",
              loss = losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=["accuracy"])
history = model3.fit(train, 
                    validation_data=val,
                    epochs = 30,
                    verbose = False)

plt.plot(history.history["accuracy"], label = "training")
plt.plot(history.history["val_accuracy"], label = "validation")
plt.gca().set(xlabel = "epoch", ylabel = "accuracy")
plt.legend()

test_url = "https://github.com/PhilChodrow/PIC16b/blob/master/datasets/fake_news_test.csv?raw=true"
df_test = pd.read_csv(test_url)
ds = make_dataset(df_test)

ds = ds.shuffle(buffer_size = len(ds))
test = ds.batch(20)
model3.evaluate(test)

weights = model3.get_layer('embedding2').get_weights()[0] # get the weights from the embedding layer
vocab = title_vectorize_layer.get_vocabulary()                # get the vocabulary from our data prep for later

pca = PCA(n_components=2)
weights = pca.fit_transform(weights)

embedding_df = pd.DataFrame({
    'word' : vocab, 
    'x0'   : weights[:,0],
    'x1'   : weights[:,1]
})

fig = px.scatter(embedding_df, 
                 x = "x0", 
                 y = "x1", 
                 size = [2]*len(embedding_df),
                # size_max = 2,
                 hover_name = "word")
write_html(fig, "embedding.html")