import os

BASE_DIR = os.getcwd() + '/..'
RAWDATA_DIR = BASE_DIR + '/rawdata'

raw_script = RAWDATA_DIR + '/love_actually.txt'
cast_csv = RAWDATA_DIR + '/love_actually_cast.csv'

from pandas import DataFrame
import pandas as pd
import numpy as np

# read in the script
with open(raw_script, 'r', encoding='utf8') as s:
    raw_df = DataFrame({'raw': s.readlines()})
raw_df[:10]

len(raw_df)

# filter out new/empty lines and lines annotated as songs
lines = DataFrame(raw_df[(raw_df.raw.str.strip() != '') & ~(raw_df.raw.str.contains('\(song\)'))])

len(lines)

lines[:10]

# mutate: add new columns to annotate scene markers and then calculate scene numbers for each line, based on markers.

lines['is_scene'] = np.where(lines.raw.str.contains("\[ Scene #"), True, False)
# fortunately, numpy has a cumsum method like R
lines['scene'] = np.cumsum(np.where(lines.raw.str.contains("\[ Scene #"), True, False))

lines[:10]

# clean up now that we're done with the is_scene column
lines = lines[~lines.is_scene] # filter out scene markers
lines = lines.drop('is_scene', axis=1) # remove is_scene column

lines[:10]

# take raw and partition by ':'
# if partition exists, left side represents the speaker, right side the dialogue of the speaker
# if partition does not exist, entire line is dialogue
raw_partitioned = lines.raw.str.rpartition(':')

lines['speaker'] = raw_partitioned[0]
lines['dialogue'] = raw_partitioned[2]

# add new column with line numbers
lines['line'] = np.cumsum(lines.speaker != '')

lines[:15]

# collapse dialogue belonging to the same line into one row:
    # group by scene and line (do not index by group labels, else we'll get a weird multi-level index at the end)
    # aggregate the dialogue across these groups (and aggregate speaker, but each line should have one speaker anyway)
lines = lines.groupby(['scene', 'line'], as_index=False).aggregate({'dialogue': lambda x: x.str.cat().strip(), 'speaker': lambda x: x.str.cat()})

lines[:10]

# read in the character -> actor mapping
with open(cast_csv, 'r') as c:
    cast = pd.read_csv(c)

lines = lines.merge(cast)
lines['character'] = lines['speaker'] + ' (' + lines['actor'] + ')'

# clean up the output
lines = lines.sort_values('line').reset_index(drop=True)

lines[:15]

by_speaker_scene = lines.groupby(['scene', 'character'], as_index=False).size().reset_index(name='n')

by_speaker_scene[:10]

speaker_scene_matrix = pd.pivot_table(by_speaker_scene, index=['character'],
   columns=['scene'], fill_value=0, aggfunc=np.size)

speaker_scene_matrix.shape

speaker_scene_matrix

# Normalizing so that the # of scenes for each character adds up to 1 (http://varianceexplained.org/r/love-actually-network/#fn:hclust)

# in pandas, axis specifies the axis **along which** to compute
# - so we're summing along the columns (i.e. row sum) and then dividing that sum along rows ('index' in pandas)
norm = speaker_scene_matrix.div(speaker_scene_matrix.sum(axis='columns'), axis='index')

from scipy.spatial.distance import pdist
from scipy.cluster import hierarchy

# cityblock = Manhattan distance
# complete finds similar clusters
h = hierarchy.complete(pdist(norm, metric='cityblock'))

import matplotlib.pyplot as plt

plt.figure(figsize=(12, 4))
dn = hierarchy.dendrogram(h, labels=norm.index.values, leaf_rotation=90)
plt.ylabel('cityblock distance')
plt.xlabel('hierarchy.complete(pdist(norm, metric=\'cityblock\'))')
plt.show()

scenes = by_speaker_scene.groupby('scene').filter(lambda g: len(g) > 1)

# Temporarily suppressing SettingWithCopyWarning below
original_setting = pd.get_option('chained_assignment')
pd.set_option('chained_assignment', None)

# now that we've removed scene numbers, update them so that they are contiguous integers
old_to_new_scenes = {o:n+1 for n,o in enumerate(scenes.scene.unique())}
scenes.scene = scenes.scene.apply(lambda x: old_to_new_scenes[x])

# enforce categorical ordering by character similarities found by our hierarchical clustering earlier (rather than alphabetical)
scenes.character = pd.Categorical(scenes.character).reorder_categories(dn['ivl'])

# Turn SettingWithCopyWarning back on
pd.set_option('chained_assignment', original_setting)

%matplotlib inline
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import matplotlib.collections as collections
import matplotlib.colors as colors
import seaborn as sns

def get_ylim_for_scene(ss, character_ordering, s):
    max_index = 0
    min_index = len(ss.scene.unique()) + 1
    for c in ss.loc[ss.scene == s].character:
        c_index = np.where(character_ordering == c)[0][0]
        max_index = c_index if c_index > max_index else max_index
        min_index = c_index if c_index < min_index else min_index
    return min_index, max_index

num_scenes = len(scenes.scene.unique())
num_characters = len(scenes.character.unique())

plt.figure(figsize=(6, 8))

sns.set_style('whitegrid')
timeline = sns.stripplot(x='scene', y='character', color='black', data=scenes, marker='D')
timeline.xaxis.set_major_locator(ticker.MultipleLocator(5))
timeline.set(xlim=(0,num_scenes + 1))

# inspired by: https://stackoverflow.com/a/40896356
vlines = []
for i in range(num_scenes+1):
    y_min, y_max = get_ylim_for_scene(scenes, np.array(dn['ivl']), i)
    pair=[(i, y_min), (i, y_max)]
    vlines.append(pair)

# horizontal grid lines
hlines = []
for i in range(num_characters+1):
    pair=[(0, i), (num_scenes+1, i)]
    hlines.append(pair)

hlinecoll = collections.LineCollection(hlines, colors=colors.to_rgb('lightgrey'), linewidths=1)
timeline.add_collection(hlinecoll)

vlinecoll = collections.LineCollection(vlines, colors=colors.to_rgb('black'), linewidths=1.1)
timeline.add_collection(vlinecoll)

plt.show()

non_airport_scenes = speaker_scene_matrix.loc[:, speaker_scene_matrix.sum(axis='index') < 10].reindex(dn['ivl'])

cooccur = np.dot(non_airport_scenes, non_airport_scenes.T)

plt.figure()
sns.heatmap(cooccur, annot=True, fmt='d', cmap='Reds', xticklabels=dn['ivl'], yticklabels=dn['ivl'])
plt.show()

import networkx as nx
import pylab

def show_graph_with_labels(adjacency_matrix, mylabels):
    plt.figure(figsize=(12, 10))
    rows, cols = np.where(adjacency_matrix > 0)
    edges = zip(rows.tolist(), cols.tolist())
    gr = nx.Graph()
    gr.add_edges_from(edges)
    nx.draw_networkx(gr, node_size=400, labels=mylabels, with_labels=True)
    limits=pylab.axis('off') # turn off axis
    plt.show()

show_graph_with_labels(cooccur, dict(zip(list(range(len(dn['ivl']))), dn['ivl'])))

speaker_scene_matrix.columns = speaker_scene_matrix.columns.droplevel() # drop 'n'

speaker_scene_matrix.shape

scenes.shape

#!conda install rpy2
# + Add C:\..\Anaconda3\Library\mingw-w64\bin to PATH (https://github.com/ContinuumIO/anaconda-issues/issues/1129#issuecomment-289413749)

# https://stackoverflow.com/a/11586672
# http://pandas.pydata.org/pandas-docs/version/0.19.2/r_interface.html#rpy-updating

from rpy2.robjects import r, pandas2ri, numpy2ri
pandas2ri.activate()
numpy2ri.activate()

r_speaker_scene_matrix = numpy2ri.py2ri(speaker_scene_matrix.values)
r_scenes = pandas2ri.py2ri(scenes)

# bug fix for py2ri_pandasseries, which calls do_slot_assign before FactorVector for category dtype and breaks things:
# https://github.com/randy3k/rpy2/blob/master/rpy/robjects/pandas2ri.py
from rpy2.rinterface import StrSexpVector
r_scenes[1].do_slot_assign('levels', StrSexpVector(scenes.character.cat.categories))

r.assign("speaker_scene_matrix", r_speaker_scene_matrix)
r.assign("scenes", r_scenes)
r("save(speaker_scene_matrix, scenes, file='love_actually_data.rda')")