# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: kgtk
#     language: python
#     name: kgtk
# ---

import os
import pandas as pd

# # Data Augmentation With Wikidata
# In this notebook we show how to use KGTK to augment the Rotten Tomatoes movie dataset from Kaggle file with movie and actor awards from Wikidata. 
#
# Download the `rotten_tomatoes_movies.csv` from https://www.kaggle.com/stefanoleone992/rotten-tomatoes-movies-and-critic-reviews-dataset?select=rotten_tomatoes_movies.csv. Once you download the file, define an enviroment variable to record where you put the file.

# **Edit the cell below to record the path where you put the downloaded file**

os.environ['RT'] = "/Users/pedroszekely/Downloads/rotten_tomatoes_movies.csv"

# !head -1 $RT

# **Some preliminaries to set up KGTK to access the Wikidata files. You need to get the files from our download site (instructions coming soon)**

# + jupyter={"source_hidden": true}
# Parameters

# Folder on local machine where to create the output and temporary folders
output_path = "/Users/pedroszekely/Downloads/kypher"

# The names of the output and temporary folders
output_folder = "wikidata_os_v5"
temp_folder = "temp.wikidata_os_v5"

# The location of input Wikidata files
wikidata_folder = "/Volumes/GoogleDrive/Shared drives/KGTK/datasets/wikidata-20200803-v4/"
wikidata_folder = "/Users/pedroszekely/Downloads/kypher/wikidata_os_v1/"
# The wikidata_os files can be downloaded from https://drive.google.com/drive/folders/1V6oAQKmwQ4LJnrBai-uv5gHWphFSCt50?usp=sharing

# Location of the cache database for kypher
cache_path = "/Users/pedroszekely/Downloads/kypher/temp.useful_wikidata_files_v4"
# cache_path = "/Users/pedroszekely/Downloads/kypher/temp.useful_wikidata_files_v4/wikidata.sqlite3.db"
# Whether to delete the cache database
delete_database = False

# shortcuts to commands
kgtk = "time kgtk --debug"
# kgtk = "kgtk --debug"

# + jupyter={"outputs_hidden": true, "source_hidden": true}
# The names of files in the KGTK Wikidata distirbution that we will use in this notebook.
file_names = {
    "claims": "claims.tsv.gz",
    "label": "labels.en.tsv.gz",
    "alias": "aliases.en.tsv.gz",
    "description": "descriptions.en.tsv.gz",
    "item": "claims.wikibase-item.tsv.gz",
    "qualifiers": "qualifiers.tsv.gz",
    "sitelinks": "sitelinks.tsv.gz",
    "qualifiers_time": "qualifiers.time.tsv.gz",
    "property_datatypes": "metadata.property.datatypes.tsv.gz",
    "isa": "derived.isa.tsv.gz",
    "p279star": "derived.P279star.tsv.gz",
    "p279": "derived.P279.tsv.gz",
    "p31": "derived.P31.tsv.gz"
}

# We will define environment variables to hold the full paths to the files as we will use them in the shell commands
kgtk_environment_variables = []

os.environ['WIKIDATA'] = wikidata_folder
kgtk_environment_variables.append('WIKIDATA')

for key, value in file_names.items():
    variable = key.upper()
    os.environ[variable] = wikidata_folder + value
    kgtk_environment_variables.append(variable)
    
# KGTK creates a SQLite database to index the knowledge graph.
if cache_path:
    os.environ['STORE'] = "{}/wikidata.sqlite3.db".format(cache_path)
else:
    os.environ['STORE'] = "{}/{}/wikidata.sqlite3.db".format(output_path, temp_folder)
kgtk_environment_variables.append('STORE')

# We will create many temporary files, so set up a folder for outputs and one for the temporary files.
os.environ['TEMP'] = "{}/{}".format(output_path, temp_folder) 
os.environ['OUT'] = "{}/{}".format(output_path, output_folder) 
kgtk_environment_variables.append('TEMP')
kgtk_environment_variables.append('OUT')

# Envronment variables with shortcuts to the commands we use often
os.environ['kgtk'] = kgtk
os.environ['kypher'] = "time kgtk --debug query --graph-cache " + os.environ['STORE']
os.environ['kypher'] = "kgtk query --graph-cache " + os.environ['STORE']
kgtk_environment_variables.append('kgtk')
kgtk_environment_variables.append('kypher')

# We'll save the current working directory so we can call into other example notebooks later
os.environ["EXAMPLES_DIR"] = os.getcwd()
kgtk_environment_variables.append('EXAMPLES_DIR')

kgtk_environment_variables.sort()
for variable in kgtk_environment_variables:
    print("{}: \"{}\"".format(variable, os.environ[variable]))

# + jupyter={"source_hidden": true}
# %cd {output_path}
# -

# ### Import the Rotten Tomatoes movie file in KGTK
#
# 1. use the `cat` command to convert it to TSV format 
# 1. use the `rename-columns` command to rename the `rotten_tomatoes_link` column to `node1` so now we have a legal KGTK file
# 1. store the result in `rt-movies.tsv`

# !$kgtk cat --mode NONE -i $RT --input-format csv \
# / rename-columns --mode NONE --old-columns rotten_tomatoes_link --new-columns node1 \
# -o $TEMP/rt-movies.tsv

# !head -2 $TEMP/rt-movies.tsv

# ### Query Wikidata to get the count of awards for each movie in the Kaggle dataset 
#
# - `rt-movies.tsv`: the Rotten Tomatoes KGTK file
# - `claims`: the Wikidata file containing all claims (statements)
# - `rt: (movie_id)-[]->()`: assign to the `movie_id` variable all the movie ids in the the Rotten Tomatoes file
# - `claims: (movie)-[:P1258]->(rt_id)`: get the Rotten Tomatoes ids for all movies in Wikidata; note that the ids in Wikidata are strings
# - `claims: (movie)-[aid:P166]->(award)`: get the awards of the movies that have Rotten Tomatoes ids
# - `--where 'kgtk_unstringify(rt_id) = movie_id'`: do the join on the `movie_id` in the input file and the Rotten Tomatoes ids `rt_id` without the quotes
# - `--return 'distinct movie_id as node1, count(distinct aid) as award_count' `: return the movie ids and the count of awards
# - `-order-by 'award_count desc'`: order the results descending by count of awards

# %%time
# !$kypher -i $TEMP/rt-movies.tsv --as movies -i $CLAIMS --as claims \
# --match 'movies: (movie_id)-[]->(), claims: (movie)-[:P1258]->(rt_id), claims: (movie)-[aid:P166]->(award)' \
# --where 'kgtk_unstringify(rt_id) = movie_id' \
# --return 'distinct movie_id as node1, count(distinct aid) as award_count' \
# --order-by 'award_count desc' \
# -o $TEMP/rt-awards.tsv

# Look at the file we got

# !wc $TEMP/rt-awards.tsv

# !head $TEMP/rt-awards.tsv

# Add the award counts to the Kaggle file by 
#
# - doing a left `join` on the `node1` column 
# - `compact`ing the file to remove duplicate rows

# !$kgtk join  --mode NONE --left-file $TEMP/rt-movies.tsv --right-file $TEMP/rt-awards.tsv --left-join \
#   --left-file-join-columns node1 --right-file-join-columns node1 \
# / compact --mode NONE --columns node1 \
# -o $TEMP/rt-movies.awards.tsv

# Load in Pandas, the `award_count` column we added is at the end

pd.read_table(os.environ['TEMP']+"/rt-movies.awards.tsv", sep="\t").sort_values(by=['award_count'], ascending=False)

# ### Query Wikidata again to add the count of awards of the cast of the movies
# The query is similar to the previous one, except that we chain from the movies to the cast and from the cast to the awards.

# %%time
# !$kypher -i movies -i claims \
# --match 'movies: (movie_id)-[]->(), claims: (movie)-[:P1258]->(rt_id), claims: (movie)-[:P161]->(cast)-[cid:P166]->(award)' \
# --where 'kgtk_unstringify(rt_id) = movie_id' \
# --return 'distinct movie_id as node1, count(distinct cid) as cast_award_count' \
# --order-by 'cast_award_count desc' \
# -o $TEMP/rt-cast-awards.tsv

# Repeat the join/compact command to add the `cast_award_count` column to the file

# !$kgtk join --mode NONE --left-file $TEMP/rt-movies.awards.tsv --right-file $TEMP/rt-cast-awards.tsv --left-join \
#   --left-file-join-columns node1 --right-file-join-columns node1 \
# / compact --mode NONE --columns node1 \
# -o $TEMP/rt-movies.awards.movies.cast.tsv

df = pd.read_table(os.environ['TEMP']+"/rt-movies.awards.movies.cast.tsv", sep="\t")
df.sort_values(by=['award_count'], ascending=False)



StarcoderdataJupyterScriptsDedupFiltered