from __future__ import division, print_function
from sys import platform as _platform
import matplotlib
import matplotlib.cm as colormap

if _platform == "linux" or _platform == "linux2":
    matplotlib.use('Agg')
import matplotlib.pylab as plt
from graph_tool.all import *
import os
import matplotlib.cm as colormap
from matplotlib.colors import ColorConverter as color_converter
import pandas as pd
from . import printing
# import printing
import random
import datetime
import copy
import shutil
import numpy as np
import operator
import math
from scipy.sparse.linalg.eigen.arpack import eigsh as largest_eigsh
import scipy.stats as stats
from scipy.sparse import csr_matrix, issparse, lil_matrix, dok_matrix
import sys
from scipy.stats import powerlaw, poisson
from collections import defaultdict
import traceback
from .basics import create_folder_structure
import powerlaw as fit_powerlaw

def print_f(*args, **kwargs):
    if 'class_name' not in kwargs:
        kwargs.update({'class_name': 'gt_tools'})
    printing.print_f(*args, **kwargs)

def add_vertex_property(g, fn, p_name='weight', vertex_id_to_vertex=None, vertex_id_dtype='int', property_dtype='int',
                        sep=None, comment='#', col=1):
    pmap = g.new_vertex_property(property_dtype)
    if property_dtype is 'int':
        property_type_mapper = int
    elif property_dtype is 'float':
        property_type_mapper = float
    else:
        property_type_mapper = str

    if vertex_id_dtype is 'int' or vertex_id_to_vertex is None:
        vertex_id_mapper = int
    elif vertex_id_dtype is 'float':
        vertex_id_mapper = float
    else:
        vertex_id_mapper = str
    get_vertex = g.vertex if vertex_id_to_vertex is None else (lambda x: vertex_id_to_vertex[x])
    with open(fn, 'r') as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith(comment):
                line = line.split(sep)
                v = get_vertex(vertex_id_mapper(line[0]))
                p = property_type_mapper(line[col])
                pmap[v] = p
    g.vp[p_name] = pmap

def load_edge_list(filename, directed=False, vertex_weights=None, edge_weights=None, vertex_id_name='NodeId',
                   vertex_id_dtype='int', sep=None, comment='#', multiple_edges_per_line=False, store=True,
                   try_load_gt=True):
    """
    Loads an edge-list into a Graph (graph-tool) object and returns the Graph object.

    Attributes
    ----------
    filename : str
        the filename of the edge-list
    directed : bool
        treat the edge-list directed (True) or not (False)
    vertex_weights : dict
        dict where keys are the name of the resulting property-map and values are dicts containing properties of the
        weight like data-type etc. (cf. add_vertex_property)
        important: dict containing the properties must contain 'filename' : "filename_of_mapping"
        Example:
            vertex_weights = dict()
            one_v_weight = dict()
            one_v_weight['filename'] = 'foo.file'
            one_v_weight['property_dtype'] = 'int'  # you can use anything available in add_vertex_property()
            vertex_weights['MyVertexWeight'] = one_v_weight

    edge_weights : dict
        dict where keys are the name of the resulting property_map and values are tuples containing column id where
        the weight is in the edge-list and the data-type.
        Example:
            weight_name : (0, 'int')
            creates an integer property-map containing the values found in the first column AFTER src, tar.
    vertex_id_name : str
        Name of property map where the original vertex-name/-id... should be stored
    vertex_id_dtype : str
        Data-type of the original vertex-name/-id
    sep : str or None
        separator which should be used to split the lines of the edge-list file
    comment : str
        lines starting with this string are treated as comments (ignored)
    multiple_edges_per_line : bool
        Do the lines contain multiple-edges. (e.g., src tar1 tar2 tar3)
        If this is true, edge_weights are not supported.
    store : bool
        Should the graph be stored in .gt format?
    try_load_gt : bool
        Should the function try to load a previously stored .gt file? ("filename.gt")
        Note if there where modifications in the original edge-list, the function will reload the edge-list.

    Notes
    -----
    -

    Examples
    --------
    -
     """
    store_fname = filename + '.gt'
    if os.path.isfile(store_fname) and try_load_gt:
        try:
            g = load_graph(store_fname)
        except:
            print('failed loading. recreate graph')
            os.remove(store_fname)
            return load_edge_list(filename, directed=directed, vertex_id_dtype=vertex_id_dtype, sep=sep,
                                  comment=comment)
        if 'mtime' in g.gp.keys() and os.path.isfile(filename) and g.gp['mtime'] != os.path.getmtime(filename):
            print('modified edge-list. recreate graph')
            os.remove(filename + '.gt')
            return load_edge_list(filename, directed=directed, vertex_id_dtype=vertex_id_dtype, sep=sep,
                                  comment=comment)
        else:
            return g
    else:
        g = Graph(directed=directed)
        nodeid_to_v = defaultdict(g.add_vertex)
        edge_list = []
        v_type = int

        edge_weights_dtypes = dict()
        edge_weights_pmaps = dict()
        edge_weights_names = dict()
        if edge_weights is not None:
            for name, (col, dtype) in edge_weights.iteritems():
                edge_weights_names[col] = name
                edge_weights_pmaps[col] = g.new_edge_property(dtype)
                if dtype == 'int':
                    edge_weights_dtypes[col] = int
                elif dtype == 'float':
                    edge_weights_dtypes[col] = float
                else:
                    edge_weights_dtypes[col] = float
        src_target_extractor = (
            lambda x: x[:2]) if edge_weights is not None or not multiple_edges_per_line else (
            lambda x: x)
        weights_extractor = (lambda x: None) if edge_weights is None else (
            lambda x: tuple([edge_weights_dtypes[idx](i) for
                             idx, i in enumerate(x[2:])]))
        edge_weights_list = list()
        with open(filename, 'r') as f:
            for line in filter(lambda x: not x.startswith(comment), map(lambda x: x.strip(), f)):
                line = line.split(sep)
                nodes = src_target_extractor(line)
                try:
                    nodes = map(v_type, nodes)
                except ValueError:
                    v_type = float
                    try:
                        nodes = map(v_type, nodes)
                    except ValueError:
                        v_type = str
                        nodes = map(v_type, nodes)
                try:
                    src = int(nodeid_to_v[nodes[0]])
                except IndexError:
                    continue
                dest = map(lambda x: int(nodeid_to_v[x]), nodes[1:])
                edges = [(src, d) for d in dest]
                edge_weights_list.append(weights_extractor(line))
                edge_list.extend(edges)
        g.add_edge_list(edge_list)
        if edge_weights is not None:
            edges = list(map(lambda x: g.edge(*x), edge_list))
            for col_idx, col_weights in enumerate(zip(*edge_weights_list)):
                current_pmap = edge_weights_pmaps[col_idx]
                for e, w in zip(edges, col_weights):
                    current_pmap[e] = w
                g.ep[edge_weights_names[col_idx]] = current_pmap
        if vertex_id_dtype is not None:
            node_id_pmap = g.new_vertex_property(vertex_id_dtype)
            for v_id, v in nodeid_to_v.iteritems():
                node_id_pmap[v] = v_id
            g.vp[vertex_id_name] = node_id_pmap
        g.gp['filename'] = g.new_graph_property('string', filename)
        g.gp['mtime'] = g.new_graph_property('object', os.path.getmtime(filename))

        if vertex_weights is not None:
            for name, props in vertex_weights.iteritems():
                kwargs = dict()
                kwargs['vertex_id_to_vertex'] = nodeid_to_v
                kwargs['vertex_id_dtype'] = vertex_id_dtype
                kwargs['sep'] = sep
                kwargs['comment'] = comment
                fn = props['filename']
                props.pop('filename', None)
                kwargs.update(props)
                add_vertex_property(g, fn, p_name=name, **kwargs)
        if store:
            g.save(filename + '.gt', fmt='gt')
    return g

def net_from_adj(mat, directed=True, parallel_edges=True):
    g = Graph(directed=directed)
    assert mat.shape[0] == mat.shape[1]
    if not issparse(mat):
        mat = csr_matrix(mat)
    elif not isinstance(mat, csr_matrix):
        mat = mat.tocsr()
    g.add_vertex(mat.shape[0])
    w = None
    if not parallel_edges:
        if np.issubdtype(mat.dtype, int):
            w = g.new_edge_property('int')
        else:
            w = g.new_edge_property('float')
    row_idx, col_idx = mat.nonzero()
    data = np.array(mat.data)
    if not directed:
        # diag-upper part only (including diag)
        row_idx, col_idx, data = zip(*[(r, c, d) for r, c, d in zip(row_idx, col_idx, data) if c <= r])
    if parallel_edges:
        row_idx = np.array([r for r, d in zip(row_idx, data) for i in range(int(d))])
        col_idx = np.array([c for c, d in zip(col_idx, data) for i in range(int(d))])
    g.add_edge_list(zip(col_idx, row_idx))
    if w:
        w.a = data
        g.ep['weights'] = w
    return g

def load_property(network, filename, type='int', resolve='NodeId', sep=None, line_groups=False):
    assert isinstance(network, Graph)
    type = type.lower()
    if type == 'str' or type == 'string':
        pmap = network.new_vertex_property('string')
        mapper = lambda x: str(x)
    else:
        pmap = network.new_vertex_property(type)
        if type == 'int':
            mapper = lambda x: int(x)
        else:
            mapper = lambda x: float(x)
    if resolve is not None:
        res_map = network.vp[resolve]
        resolve = {str(res_map[v]): v for v in network.vertices()}
    mapped_vertices = set()
    com_id = 0
    with open(filename, 'r') as f:
        for line in filter(lambda l_line: not l_line.startswith('#'), f):
            line = line.strip().split(sep)
            if line_groups:
                line_vertices = []
                appender = line_vertices.append
                for i in line:
                    try:
                        if resolve is None:
                            i = network.vertex(int(i))
                        else:
                            i = resolve[i]
                        appender(i)
                    except KeyError:
                        pass
                for v in line_vertices:
                    pmap[v] = mapper(com_id)
                    mapped_vertices.add(int(v))
                com_id += 1
            else:
                try:
                    v = network.vertex(int(line[0])) if resolve is None else resolve[line[0]]
                    pmap[v] = mapper(line[1])
                    mapped_vertices.add(int(v))
                except KeyError:
                    pass
    unmapped_v = set(map(int,network.vertices())) - mapped_vertices
    if unmapped_v:
        print(filename, 'contained no mapping for', len(unmapped_v) / network.num_vertices() * 100, '% of all vertices')
        print('unmapped vertices:', list(unmapped_v)[:100])
    return pmap

def get_graph_com_connectivity(g, com_map='com'):
    com_map = g.vp[com_map]
    intern_edges = 0
    between_edges = 0
    for e in g.edges():
        if com_map[e.source()] == com_map[e.target()]:
            intern_edges += 1
        else:
            between_edges += 1
    num_edges = (intern_edges + between_edges)
    intern_edges /= num_edges
    between_edges /= num_edges
    return intern_edges, between_edges

def check_aperiodic(g):
    if isinstance(g, str):
        a = adjacency(load_graph(g))
        name = g.rsplit('/')[-1].replace('.gt', '')
        print('aperiodic:', name)
    else:
        a = adjacency(g)
    b = a * a
    diag_two_sum = b.diagonal().sum()
    print('\tA*A diag sum:', int(diag_two_sum))
    b *= a
    diag_three_sum = b.diagonal().sum()
    print('\tA*A*A diag sum:', int(diag_three_sum))
    aper = bool(diag_two_sum) and bool(diag_three_sum)
    print('\taperiodic:', aper)
    return aper

class SBMGenerator():
    @staticmethod
    def gen_stoch_blockmodel(num_nodes=1000, blocks=5, self_con=.97, other_con=0.03, directed=False,
                             degree_seq='powerlaw', powerlaw_exp=2.4, num_links=None, loops=False, min_degree=1,
                             con_prob_matrix=None, increase_lcc_prob=True, parallel_edges=False,
                             node_pick_strat=('dist', 'dist'), scale_con_prop_matrix=True):
        g = Graph(directed=directed)
        com_pmap = g.new_vertex_property('int')
        if isinstance(blocks, list):
            num_nodes = sum(blocks)
            num_blocks = len(blocks)
            block_sizes = blocks
        else:
            num_blocks = blocks
            block_sizes = [int(num_nodes / num_blocks) for i in range(num_blocks)]
            num_unmapped_nodes = num_nodes % num_blocks
            # print('unmapped nodes', num_unmapped_nodes)
            if num_unmapped_nodes > 0:
                for i in range(num_unmapped_nodes):
                    block_sizes[i] += 1
        g.add_vertex(num_nodes)

        if con_prob_matrix is not None:
            assert con_prob_matrix.shape[0] == con_prob_matrix.shape[1]
            con_prob_matrix /= con_prob_matrix.sum()
        com_pmap.a = np.hstack([np.array([idx] * i) for idx, i in enumerate(block_sizes)]).flatten()
        g.vp['com'] = com_pmap
        other_con /= ((num_blocks - 1) if num_blocks > 1 else 1)

        prob_pmap = g.new_vertex_property('float')
        block_to_vertices = dict()
        block_to_cumsum = dict()

        if isinstance(degree_seq, (np.ndarray, list, tuple)) and len(degree_seq) == num_nodes:
            degree_seq = degree_seq.astype('float')
            fixed_deg_seq = True
        else:
            fixed_deg_seq = False
            if degree_seq == 'powerlaw':
                degree_seq = stats.zipf.rvs(powerlaw_exp, loc=min_degree, size=num_nodes).astype('float')
            elif degree_seq == 'random':
                degree_seq = np.random.random(size=num_nodes)
            elif degree_seq == 'exp':
                degree_seq = np.random.exponential(size=num_nodes)
            else:
                degree_seq = np.array([1.0] * num_nodes)
            degree_seq.sort()
        degree_seq /= degree_seq.sum()
        multiplier = min_degree/degree_seq.min()
        degree_seq *= multiplier
        # print(degree_seq)
        if num_links is None:
            # print('min degree:', degree_seq.min())
            num_links = int(np.round((degree_seq.round().sum() / 2)))
            print('#links not set. using:', num_links)
        # print(degree_seq)
        block_deg_seq_sum = dict()
        vertices_array = np.array(list(map(int, g.vertices())))
        degree_indices = set(range(g.num_vertices()))
        for i in range(num_blocks):
            mask = com_pmap.a == i
            block_num_vertices = mask.sum()
            block_to_vertices[i] = vertices_array[mask]
            if fixed_deg_seq:
                block_deg_seq_idx = block_to_vertices[i]
            else:
                block_deg_seq_idx = random.sample(degree_indices, block_num_vertices)
                degree_indices -= set(block_deg_seq_idx)
            block_deg_seq = degree_seq[block_deg_seq_idx]
            deg_seq_sum = block_deg_seq.sum()
            block_deg_seq /= deg_seq_sum
            cum_sum = np.cumsum(block_deg_seq)
            assert np.allclose(cum_sum[-1], 1)
            block_to_cumsum[i] = cum_sum
            block_deg_seq_sum[i] = deg_seq_sum
            prob_pmap.a[mask] = block_deg_seq
        if scale_con_prop_matrix:
            blocks_prob = list()
            for i in range(num_blocks):
                row = list()
                for j in range(num_blocks):
                    if con_prob_matrix is None:
                        if i == j:
                            val = self_con
                        else:
                            val = other_con
                    else:
                        val = con_prob_matrix.item((i, j))
                    row.append(val * block_deg_seq_sum[i] * block_deg_seq_sum[j])
                blocks_prob.append(np.array(row))
            blocks_prob = np.array(blocks_prob)
            blocks_prob /= blocks_prob.sum()
        else:
            blocks_prob = con_prob_matrix
        # print(blocks_prob)
        cum_sum = np.cumsum(blocks_prob)
        assert np.isclose(cum_sum[-1], 1)
        # print(cum_sum)
        if parallel_edges:
            edges = list()
            edges_adder = edges.append
        else:
            edges = set()
            edges_adder = edges.add

        if isinstance(node_pick_strat, str):
            node_pick_strat = (node_pick_strat, node_pick_strat)
        pick_funcs = dict()
        get_rnd_node = SBMGenerator.get_random_node
        inv_prob = SBMGenerator.inverse_prob
        pick_funcs['rnd'] = lambda b, s=None: random.choice(block_to_vertices[b])
        pick_funcs['dist'] = lambda b, s=None: block_to_vertices[b][get_rnd_node(block_to_cumsum[b])]
        pick_funcs['invdist'] = lambda b, s=None: block_to_vertices[b][
            get_rnd_node(inv_prob(block_to_cumsum[b]))]
        pick_funcs['dist_other_com_inv_dist'] = lambda b, s=None: \
            block_to_vertices[b][get_rnd_node(block_to_cumsum[b])] \
            if (s is None or s == b) else \
                block_to_vertices[b][get_rnd_node(inv_prob(block_to_cumsum[b]))]

        # more efficient way to pick both inverse to the dist
        if node_pick_strat[0] == node_pick_strat[1] == 'invdist':
            node_pick_strat[0] = node_pick_strat[1] = 'dist'
            block_to_cumsum = {key: inv_prob(val) for key, val in block_to_cumsum.items()}

        src_pick_func, dest_pick_func = map(lambda x: pick_funcs[x], node_pick_strat)

        get_one_rnd_block = SBMGenerator.get_one_random_block
        if increase_lcc_prob:
            for v in g.vertices():
                if directed or v.out_degree() == 0:
                    src_block = com_pmap[v]
                    init_len = len(edges)
                    while init_len == len(edges):
                        dest_b = get_one_rnd_block(cum_sum, num_blocks, src_block) \
                            if 'dist_other_com_inv_dist' not in node_pick_strat else src_block
                        dest_v = dest_pick_func(dest_b, src_block)
                        link = (int(v), dest_v)
                        is_loop = v == dest_v
                        if not is_loop:
                            if not directed:
                                link = tuple(sorted(link))
                            edges_adder(link)
                        elif loops:
                            edges_adder(link)

        get_rnd_blocks = SBMGenerator.get_random_blocks
        for link_idx in range(num_links - len(edges)):
            while True:
                #maybe switch to: get random node. identify block. get random dest-block.
                src_b, dest_b = get_rnd_blocks(cum_sum, num_blocks)
                src_v = src_pick_func(src_b, dest_b)
                dest_v = dest_pick_func(dest_b, src_b)
                link = (src_v, dest_v)
                is_loop = src_v == dest_v
                if not is_loop:
                    if not directed:
                        link = tuple(sorted(link))
                    edges_adder(link)
                    break
                elif loops:
                    edges_adder(link)
                    break
        if not isinstance(edges, list):
            edges = list(edges)
        g.add_edge_list(edges)
        return g

    @staticmethod
    def inverse_prob(p_cum_sum):
        p = np.hstack([p_cum_sum[0], np.ediff1d(p_cum_sum)])
        p = 1. / p
        p /= p.sum()
        return p.cumsum()

    @staticmethod
    def get_random_node(cum_sum):
        rand_num = np.random.random()
        idx = np.searchsorted(cum_sum, rand_num)
        if idx == len(cum_sum):
            idx -= 1
        return idx

    @staticmethod
    def get_random_blocks(cum_sum, num_blocks):
        rand_num = np.random.random()
        idx = np.searchsorted(cum_sum, rand_num)
        row = int(idx / num_blocks)
        col = idx % num_blocks
        #print(rand_num)
        #print(row, col)
        #print(cum_sum)
        return row, col

    @staticmethod
    def get_one_random_block(cum_sum, num_blocks, row):
        src_b = None
        dest_b = None
        while src_b is None or row != src_b:
            src_b, dest_b = SBMGenerator.get_random_blocks(cum_sum, num_blocks)
        return dest_b

    @staticmethod
    def gen_bow_tie_model(scc_size, out_size, in_size, con_prob_matrix=None, **kwargs):
        blocks = [scc_size, out_size, in_size]
        if con_prob_matrix is None:
            con_prob_matrix = np.zeros((3, 3))
            # in-comp
            con_prob_matrix[0, 0] = 0.001  # self
            con_prob_matrix[0, 1] = 0.1  # scc
            con_prob_matrix[0, 2] = 0.05  # out

            # scc
            con_prob_matrix[1, 1] = 5  # self
            con_prob_matrix[1, 2] = 0.1  # out

            con_prob_matrix[2, 2] = 0.01  # self
        np.set_printoptions(formatter=dict(float=lambda x: '%.5f' % x))
        print(con_prob_matrix / con_prob_matrix.sum())
        return SBMGenerator.gen_stoch_blockmodel(blocks=blocks, con_prob_matrix=con_prob_matrix,
                                                 directed=True, **kwargs)

    @staticmethod
    def analyse_graph(g, filename='output/net', draw_net=False):
        print(str(g))
        deg_map = g.degree_property_map('total')
        plt.close('all')
        ser = pd.Series(deg_map.a)
        ser.plot(kind='hist', bins=int(deg_map.a.max()), lw=0)
        plt.xlabel('degree')
        plt.ylabel('num nodes')
        res = fit_powerlaw.Fit(deg_map.a, discrete=True)
        print('powerlaw alpha:', res.power_law.alpha)
        print('powerlaw xmin:', res.power_law.xmin)
        plt.title('powerlaw alpha:' + str(res.power_law.alpha) + ' || powerlaw xmin:' + str(res.power_law.xmin))
        plt.savefig(filename + '_degdist.png', bbox_tight=True)
        plt.close('all')
        if draw_net:
            graph_draw(g, vertex_fill_color=g.vp['com'], output_size=(200, 200),
                       vertex_size=prop_to_size(deg_map, mi=2, ma=15, power=1.), output=filename + '_network.png',
                       bg_color=[1, 1, 1, 1])

def bow_tie(graph):
    assert graph.is_directed()
    largest_component = label_largest_component(graph)
    weakly_components = label_components(graph, directed=False)[0]
    largest_component_corresponding_weakly = list(
        set(weakly_components.a[np.nonzero(largest_component.a)[0]]))
    assert len(largest_component_corresponding_weakly) == 1
    largest_component_corresponding_weakly = largest_component_corresponding_weakly[0]
    wcc = (weakly_components.a == largest_component_corresponding_weakly)

    # Core, In and Out
    all_nodes = set(range(graph.num_vertices()))
    scc = set(np.nonzero(largest_component.a)[0])
    scc_node = random.sample(scc, 1)[0]
    graph_reversed = GraphView(graph, reversed=True)

    outc = np.nonzero(label_out_component(graph, scc_node).a)[0]
    inc = np.nonzero(label_out_component(graph_reversed, scc_node).a)[0]
    outc = set(outc) - scc
    inc = set(inc) - scc

    # Tubes, Tendrils and Other
    wcc = set(np.nonzero(wcc)[0])
    tube = set()
    out_tendril = set()
    in_tendril = set()
    other = all_nodes - wcc
    remainder = wcc - inc - outc - scc

    for idx, r in enumerate(remainder):
        print(idx + 1, '/', len(remainder), end='\r')
        predecessors = set(np.nonzero(label_out_component(graph_reversed, r).a)[0])
        successors = set(np.nonzero(label_out_component(graph, r).a)[0])
        if any(p in inc for p in predecessors):
            if any(s in outc for s in successors):
                tube.add(r)
            else:
                in_tendril.add(r)
        elif any(s in outc for s in successors):
            out_tendril.add(r)
        else:
            other.add(r)

    vp_bowtie = graph.new_vertex_property('string')
    for component, label in [
        (inc, 'IN'),
        (scc, 'SCC'),
        (outc, 'OUT'),
        (in_tendril, 'TL_IN'),
        (out_tendril, 'TL_OUT'),
        (tube, 'TUBE'),
        (other, 'OTHER')
    ]:
        for node in component:
            vp_bowtie[graph.vertex(node)] = label
    graph.vp['bowtie'] = vp_bowtie

    bow_tie = map(len, [inc, scc, outc, in_tendril, out_tendril, tube, other])
    assert sum(bow_tie) == graph.num_vertices()
    bow_tie = [100 * x / graph.num_vertices() for x in bow_tie]
    bow_tie = dict(IN=bow_tie[0], SCC=bow_tie[1], OUT=bow_tie[2], TL_IN=bow_tie[3], TL_OUT=bow_tie[4], TUBE=bow_tie[5],
                   OTHER=bow_tie[6])
    return bow_tie

# Generator Class works with GraphTool generators, as they provide more functionality than NetworkX Generators
class GraphGenerator():
    # init generator
    def __init__(self, num_nodes=5000, directed=False):
        self.directed = directed
        self.num_nodes = num_nodes
        self.graph = None
        self.node_id_map = None
        self.return_and_reset()

    def return_and_reset(self):
        result = self.graph
        self.graph = Graph(directed=self.directed)
        self.node_id_map = defaultdict(lambda x: self.graph.add_vertex())
        return result

    # start creating blockmodel graph
    def create_blockmodel_graph(self, blocks=7, connectivity=10, model="blockmodel-traditional"):
        def corr(a, b):
            if a == b:
                return 0.999
            else:
                return 0.001

        self.print_f("Starting to create Blockmodel Graph with {} nodes and {} blocks".format(self.num_nodes, blocks))

        self.graph, vertex_colors = random_graph(self.num_nodes, lambda: poisson(connectivity), directed=False, model=model, block_membership=lambda: random.randint(1, blocks),
                                                 vertex_corr=corr)
        self.graph.vertex_properties["colorsComm"] = vertex_colors
        return self.return_and_reset()

    def create_fully_connected_graph(self, size=1000, directed=False, self_edges=False):
        return self.create_stochastic_blockmodel_graph(blocks=1, size=size, directed=directed, self_edges=self_edges, self_block_connectivity=1.0, other_block_connectivity=1.0)

    @staticmethod
    def create_sbm_lined_up_matrix(blocks=10, self_block_connectivity=None, other_block_connectivity=None):
        if self_block_connectivity is None:
            self_block_connectivity = [0.9]
        elif isinstance(self_block_connectivity, (int, float)):
            self_block_connectivity = [self_block_connectivity]
        if other_block_connectivity is None:
            other_block_connectivity = [0.1]
        elif isinstance(other_block_connectivity, (int, float)):
            other_block_connectivity = [other_block_connectivity]
        connectivity_matrix = []
        blocks_range = range(blocks)
        for idx in blocks_range:
            row = []
            outer_prob = other_block_connectivity[idx % len(other_block_connectivity)]
            inner_prob = self_block_connectivity[idx % len(self_block_connectivity)]
            for jdx in blocks_range:
                if idx != jdx:
                    row.append(outer_prob / pow(abs(idx - jdx), 2))
                else:
                    row.append(inner_prob)
            connectivity_matrix.append(row)
        return connectivity_matrix

    # scale = None
    # scale = relative
    # scale = absolute
    def create_stochastic_blockmodel_graph(self, blocks=10, size=100, self_block_connectivity=0.9, other_block_connectivity=0.1, connectivity_matrix=None, directed=False,
                                           self_edges=False, power_exp=None, scale=None, plot_stat=False):
        size = size if isinstance(size, list) else [size]
        self_block_connectivity = self_block_connectivity if isinstance(self_block_connectivity, list) else [self_block_connectivity]
        other_block_connectivity = other_block_connectivity if isinstance(other_block_connectivity, list) else [other_block_connectivity]

        num_nodes = sum([size[i % len(size)] for i in range(blocks)])
        if power_exp is None:
            self.print_f("Starting to create Stochastic Blockmodel Graph with {} nodes and {} blocks".format(num_nodes, blocks))
        else:
            self.print_f("Starting to create degree-corrected (alpha=" + str(power_exp) + ") Stochastic Blockmodel Graph with {} nodes and {} blocks".format(num_nodes, blocks))
        self.print_f('convert/transform probabilities')
        blocks_range = range(blocks)
        block_sizes = np.array([size[i % len(size)] for i in blocks_range])

        # create connectivity matrix of self- and other-block-connectivity
        if connectivity_matrix is None:
            connectivity_matrix = []
            self.print_f('inner conn: ' + str(self_block_connectivity) + '\tother conn: ' + str(other_block_connectivity))
            for idx in blocks_range:
                row = []
                for jdx in blocks_range:
                    if idx == jdx:
                        row.append(self_block_connectivity[idx % len(self_block_connectivity)])
                    else:
                        if scale is not None:
                            prob = other_block_connectivity[idx % len(other_block_connectivity)] / (num_nodes - block_sizes[idx]) * block_sizes[jdx]
                            if directed:
                                row.append(prob)
                            else:
                                row.append(prob / 2)
                        else:
                            row.append(other_block_connectivity[idx % len(other_block_connectivity)])
                connectivity_matrix.append(row)

        # convert con-matrix to np.array
        if connectivity_matrix is not None and isinstance(connectivity_matrix, np.matrix):
            connectivity_matrix = np.asarray(connectivity_matrix)

        # convert con-matrix to np.array
        if connectivity_matrix is not None and not isinstance(connectivity_matrix, np.ndarray):
            connectivity_matrix = np.array(connectivity_matrix)

        self.print_f('conn mat')
        printing.print_matrix(connectivity_matrix)

        if scale == 'relative' or scale == 'absolute':
            new_connectivity_matrix = []
            for i in blocks_range:
                connectivity_row = connectivity_matrix[i, :] if connectivity_matrix is not None else None
                nodes_in_src_block = block_sizes[i]
                multp = 1 if scale == 'absolute' else (nodes_in_src_block * (nodes_in_src_block - 1))
                row_prob = [(connectivity_row[idx] * multp) / (nodes_in_src_block * (nodes_in_block - 1)) for idx, nodes_in_block in enumerate(block_sizes)]
                new_connectivity_matrix.append(np.array(row_prob))
            connectivity_matrix = np.array(new_connectivity_matrix)
            self.print_f(scale + ' scaled conn mat:')
            printing.print_matrix(connectivity_matrix)

        # create nodes and store corresponding block-id
        self.print_f('insert nodes')
        vertex_to_block = []
        appender = vertex_to_block.append
        colors = self.graph.new_vertex_property("float")
        for i in range(blocks):
            block_size = size[i % len(size)]
            for j in range(block_size):
                appender((self.graph.add_vertex(), i))
                node = vertex_to_block[-1][0]
                colors[node] = i

        # create edges
        get_rand = np.random.random
        add_edge = self.graph.add_edge

        self.print_f('create edge probs')
        degree_probs = defaultdict(lambda: dict())
        for vertex, block_id in vertex_to_block:
            if power_exp is None:
                degree_probs[block_id][vertex] = 1
            else:
                degree_probs[block_id][vertex] = math.exp(power_exp * np.random.random())

        tmp = dict()
        self.print_f('normalize edge probs')
        all_prop = []
        for block_id, node_to_prop in degree_probs.iteritems():
            sum_of_block_norm = 1 / sum(node_to_prop.values())
            tmp[block_id] = {key: val * sum_of_block_norm for key, val in node_to_prop.iteritems()}
            all_prop.append(tmp[block_id].values())
        degree_probs = tmp
        if plot_stat:
            plt.clf()
            plt.hist(all_prop, bins=15)
            plt.savefig("prop_dist.png", bbox_tight=True)
            plt.close('all')

        self.print_f('count edges between blocks')
        edges_between_blocks = defaultdict(lambda: defaultdict(int))
        for idx, (src_node, src_block) in enumerate(vertex_to_block):
            conn_mat_row = connectivity_matrix[src_block, :]
            for dest_node, dest_block in vertex_to_block:
                if get_rand() < conn_mat_row[dest_block]:
                    edges_between_blocks[src_block][dest_block] += 1

        self.print_f('create edges')
        for src_block, dest_dict in edges_between_blocks.iteritems():
            self.print_f(' -- Processing Block {}. Creating links to: {}'.format(src_block, dest_dict))
            for dest_block, num_edges in dest_dict.iteritems():
                self.print_f('   ++ adding {} edges to {}'.format(num_edges, dest_block))
                for i in range(num_edges):
                    # find src node
                    prob = np.random.random()
                    prob_sum = 0
                    src_node = None
                    for vertex, v_prob in degree_probs[src_block].iteritems():
                        prob_sum += v_prob
                        if prob_sum >= prob:
                            src_node = vertex
                            break
                    # find dest node
                    prob = np.random.random()
                    prob_sum = 0
                    dest_node = None
                    for vertex, v_prob in degree_probs[dest_block].iteritems():
                        prob_sum += v_prob
                        if prob_sum >= prob:
                            dest_node = vertex
                            break
                    if src_node is None or dest_node is None:
                        print('Error selecting node:', src_node, dest_node)
                    if self.graph.edge(src_node, dest_node) is None:
                        if self_edges or not src_node == dest_node:
                            add_edge(src_node, dest_node)
        self.graph.vertex_properties["colorsComm"] = colors
        return self.return_and_reset()

    def create_preferential_attachment(self, communities=10):
        self.graph = price_network(self.num_nodes, directed=False, c=0, gamma=1, m=1)
        self.graph.vertex_properties['colorsComm'] = community_structure(self.graph, 1000, communities)
        return self.return_and_reset()

    # add node to graph and check if node is in node_dict
    def add_node(self, node_id, further_mappings=None):
        v = self.node_id_map[node_id]
        self.graph.vp['NodeId'][v] = node_id
        if further_mappings is not None:
            assert isinstance(further_mappings, dict)
            for key, val in further_mappings.iteritems():
                self.graph.vp[key][v] = val
        return v

    def load_smw_collab_network(self, filename, communities=10):
        self.print_f("Creating Graph")
        id_prop = self.graph.new_vertex_property("int")
        self.graph.vp["label"] = self.graph.new_vertex_property("string")

        f = open(filename, "rb")
        for idx, line in enumerate(f):
            if idx % 1000 == 0:
                self.print_f("--> parsing line %d" % idx)
            split_line = line.strip("\n").split("\t")
            source_v = self.add_node(split_line[0], id_prop)
            if split_line[1] != "":
                target_v = self.add_node(split_line[1], id_prop)
                self.graph.add_edge(source_v, target_v)

        self.print_f("Detecting Communities")
        self.graph.vp['colorsComm'] = community_structure(self.graph, 1000, communities)
        remove_self_loops(self.graph)
        remove_parallel_edges(self.graph)
        return self.return_and_reset()

    @staticmethod
    def increment_neighbours(vertices, b):
        for n in vertices:
            b[int(n)] += 1

    # start creating random graph
    # NOTE:
    # If min_degree is too small, graph will be disconnected and consist of many smaller graphs!
    # This could make diffusion problematic!
    def create_random_graph(self, min_degree=2, max_degree=40, model="probabilistic", communities=10):
        # Function to sample edges between nodes!
        self.print_f('create random graph')

        def sample_k(min_val, max_val, k=None):
            accept = False
            while not accept:
                k = random.randint(min_val, max_val + 1)
                accept = random.random() < 1.0 / k
            return k

        self.graph = random_graph(self.num_nodes, lambda: sample_k(min_degree, max_degree), model=model, vertex_corr=lambda i, k: 1.0 / (1 + abs(i - k)), directed=self.directed,
                                  n_iter=100)
        self.graph.vp['colorsComm'] = community_structure(self.graph, 10000, max_degree / communities)
        return self.return_and_reset()

    # start loading  graph
    def create_karate_graph(self):
        self.graph = collection.data["karate"]
        # Removing descriptions and readme, as they screw with the GML parser of networkx!
        self.graph.gp['description'] = ''
        self.graph.gp['readme'] = ''
        # Calculating Colors and updating members
        self.graph.vp['colorsComm'] = community_structure(self.graph, 10000, 2)
        self.directed = self.graph.is_directed()
        self.num_nodes = self.graph.num_vertices()
        return self.return_and_reset()

    def loaded_post_action(self):
        self.directed = self.graph.is_directed()
        self.num_nodes = self.graph.num_vertices()
        self.print_f("Graph loaded with {} nodes and {} edges".format(self.graph.num_vertices(), self.graph.num_edges()))

    # load graph from gml
    def load_gml(self, fn):
        self.print_f("Loading GML")
        self.graph = load_graph(fn)
        self.loaded_post_action()
        return self.return_and_reset()

    # load graph from file
    def load_gt(self, fn):
        self.print_f("Loading GT")
        self.graph = load_graph(fn)
        self.loaded_post_action()
        return self.return_and_reset()

    @staticmethod
    def print_f(*args, **kwargs):
        kwargs.update({'class_name': 'GraphGenerator'})
        print_f(*args, **kwargs)

def calc_eigenvalues(graph, num_ev=100):
    num_ev = min(100, num_ev)
    print_f("Extracting adjacency matrix!")
    adj_mat = adjacency(graph, weight=None)
    print_f("Starting calculation of {} Eigenvalues".format(num_ev))
    evals_large_sparse, evecs_large_sparse = largest_eigsh(adj_mat, num_ev * 2, which='LM')
    print_f("Finished calculating Eigenvalues")
    weights = sorted([float(x) for x in evals_large_sparse], reverse=True)[:num_ev]
    graph.gp["top_eigenvalues"] = graph.new_graph_property("vector<float>", weights)
    return graph

def cleanup_graph(graph, largest_comp=True, parallel_edges=False, self_loops=False):
    if largest_comp:
        reduce_to_largest_component(graph)
    if not parallel_edges:
        remove_parallel_edges(graph)
    if not self_loops:
        remove_self_loops(graph)
    return graph

def reduce_to_largest_component(graph):
    print_f("Reducing graph to largest connected component!")
    l = label_largest_component(graph)
    graph = GraphView(graph, vfilt=l)
    graph.purge_vertices(in_place=True)
    return graph

def calc_vertex_properties(graph, max_iter_ev=1000, max_iter_hits=1000):
    print_f("Calculating PageRank")
    graph.vp["pagerank"] = pagerank(graph)

    print_f("Calculating Clustering Coefficient")
    graph.vp["clustercoeff"] = local_clustering(graph)

    print_f("Calculating Eigenvector Centrality")
    ev, ev_centrality = eigenvector(graph, weight=None, max_iter=max_iter_ev)
    graph.vp["evcentrality"] = ev_centrality

    print_f("Calculating HITS")
    eig, authorities, hubs = hits(graph, weight=None, max_iter=max_iter_hits)
    graph.vp["authorities"] = authorities
    graph.vp["hubs"] = hubs

    print_f("Calculating Degree Property Map")
    graph.vertex_properties["degree"] = graph.degree_property_map("total")
    return graph

'''
def fast_sd(g, src_ids=None, dest_ids=None, pairs=None, max_dist=10, loops=False):
    # print(g)
    all_vertices_ids = list(map(int, g.vertices()))
    if src_ids is None:
        src_ids = all_vertices_ids
    elif not hasattr(src_ids, '__iter__'):
        src_ids = [int(src_ids)]
    if dest_ids is None:
        dest_ids = all_vertices_ids
    elif not hasattr(dest_ids, '__iter__'):
        dest_ids = [int(dest_ids)]
    if pairs is not None:
        if isinstance(pairs, list):
            pairs = set(pairs)
        elif isinstance(pairs, tuple):
            pairs = {pairs}
    elif g.is_directed():
        pairs = {(src, dest) for src in src_ids for dest in dest_ids if src != dest}
    else:
        pairs = {(src, dest) for src in src_ids for dest in dest_ids if src < dest}

    assert isinstance(pairs, set)

    A = adjacency(g).astype(bool)
    M = A.copy()
    M_old = None
    distances = defaultdict(lambda: defaultdict(int))
    m_max = g.num_vertices() * (g.num_vertices() - 1)
    reachable = 0
    for current_distance in range(1, max_dist + 1):
        m_non_zero = M.nnz

        if m_non_zero == 0:
            break
        reachable += m_non_zero
        print(current_distance, reachable/m_max)
        #pairs_with_current_distance = pairs & set(zip(*m_non_zero))
        #for src, dest in pairs_with_current_distance:
        #    distances[dest][src] = current_distance
        #pairs -= pairs_with_current_distance
        if not pairs:
            # print('break at:', current_distance)
            break
        M_old = M.copy()
        M *= A
        M -= M_old
        M = M.astype(bool)
    if len(pairs):
        # print('#pairs without distance:', len(pairs))
        pass
    return distances
'''


import argparse
import datetime
import os
import shutil
import sys
import time

def migrate_to_goose(directory):

  new_migration = True
  first_iteration = True
  out_dir = os.path.join(directory, "output")
  datafix_dir = os.path.join(out_dir, "datafix")
  template_file_dir = os.path.join(directory, "template")
  template_file_path = os.path.join(template_file_dir, "create_goose_db_version_table_template.sql")
  current_filename = ""

  if not os.path.isdir(out_dir):
    os.system("mkdir %s" % out_dir)
  else:
    shutil.rmtree(out_dir)
    os.system("mkdir %s" % out_dir)

  # create dir for datafix file and create empty file
  if not os.path.isdir(datafix_dir):
    os.system("mkdir %s" % datafix_dir)
  else:
    shutil.rmtree(datafix_dir)
    os.system("mkdir %s" % datafix_dir)

  datafix_file_name = "migrate_to_goose_datafix.sql"
  datafix_file_path = os.path.join(datafix_dir, datafix_file_name)

  shutil.copyfile(template_file_path, datafix_file_path)

  up_migration_contents = ""
  down_migration_contents = ""

  for filename in sorted(os.listdir(directory)):
    f = os.path.join(directory, filename)

    if os.path.isfile(f):
      file_full_path = (os.path.join(os.path.dirname(os.path.abspath(__file__)), f))
      sliced_filename = filename[5:-4].replace(".down", "").replace(".up", "")

    # check if it is the same file we are processing
    if (sliced_filename != current_filename) and (not first_iteration):
      file_timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
      new_file_name_with_ext = file_timestamp + "_" + current_filename + ".sql"
      new_file_dir = os.path.join(out_dir, new_file_name_with_ext)

      with open(new_file_dir, 'w') as new_file:
        print("creating new migration file %s" % new_file_dir)
        # up migration
        new_file.write("-- +goose Up\n")
        new_file.write("-- +goose StatementBegin\n")
        new_file.write(up_migration_contents)
        new_file.write("\n-- +goose StatementEnd\n")
        new_file.write("\n")

        # down migration
        new_file.write("-- +goose Down\n")
        new_file.write("-- +goose StatementBegin\n")
        new_file.write(down_migration_contents)
        new_file.write("\n-- +goose StatementEnd\n")

        time.sleep(5)

      with open(datafix_file_path, 'a') as datafix_file:
        datafix_file.write("INSERT INTO goose_db_version (version_id, is_applied, tstamp) VALUES (%s, 1, NOW());\n" % file_timestamp)

      new_migration = True

    if ('.sql' in filename) and ('.down.' in filename):
      with open (file_full_path, "r") as down_mig_file:
        file_contents = down_mig_file.readlines()
        raw_file_contents = " ".join(map(str, file_contents))
        down_migration_contents = raw_file_contents

    elif ('.sql' in filename) and ('.up.' in filename):
      with open (file_full_path, "r") as up_mig_file:
        file_contents = up_mig_file.readlines()
        raw_file_contents = " ".join(map(str, file_contents))
        up_migration_contents = raw_file_contents

    else:
      # check if there is any migration file pending to be written

      print ("not a sql file - skipping migration of %s" % f)

    if first_iteration:
      first_iteration = False

    current_filename = sliced_filename

if __name__== "__main__":

  # read input args
  parser = argparse.ArgumentParser()
  parser.add_argument("-p", "--path", help = "folder path where the migrations are stored")
  args = parser.parse_args()

  if args.path:
    print("INFO: converting files under %s to goose format" % args.path)
    migrate_to_goose(args.path)
  else:
    print("ERROR: no input path supplied. aborting run")
    sys.exit()


"""Plotting code for social dimension figures."""
import sys
import os
sys.path.append("../..")
from commembed.jupyter import *
import commembed.linalg as linalg
import commembed.dimens as dimens
import commembed.plots as plots
import commembed.data as data
import numpy as np
import pandas as pd
import matplotlib as matplotlib
import matplotlib.pyplot as plt
import datetime
import seaborn as sns
import matplotlib.gridspec as gridspec
from matplotlib.colors import Normalize
#import commembed.data as data
#spark = data.spark_context()
import re

matplotlib.rcParams['mathtext.fontset'] = 'custom'

embedding = load_embedding('reddit', 'master')
dimen_list = dimens.load_dimen_list('final')
scores = dimens.score_embedding(embedding, dimen_list)
scores_nonpctl = scores

from tsmoothie.smoother import LowessSmoother

from matplotlib.patches import Polygon
from matplotlib import colors

cluster_rename_map = {
    "US cities and sports": "USA",
    "Non-US cities and places": "Non-USA",
    "Cars and motorcycles": "Cars",
    "Fitness and self-improvement": "Self-improvement",
    "Porn (heterosexual)": "Pornography",
    "Crytocurrencies": "Cryptocurrencies",

    "General interest": "General interest$^†$",
    "Gaming": "Gaming$^†$"
}

def many_densities_plot(dimens_to_plot, clustering_name="reddit_unmerged", percentilize=True, export_raw_data=False, single_column=False):

    clustering = load_clustering(clustering_name)

    cluster_rename = cluster_rename_map
    if single_column:
        cluster_rename["Recreational drugs"] = "Drugs"

    clustering["cluster_name"] = clustering["cluster_name"].map(lambda x: cluster_rename.get(x, x))

    clustering_copy = clustering.copy()
    clustering_copy["cluster_id"] = 999
    all_name = "All\\ communities" if not single_column else "All"
    clustering_copy["cluster_name"] = "${\\bf %s}$" % all_name
    clustering_with_all_communities = pd.concat([clustering, clustering_copy])

    # percentileize scores
    if percentilize:
        #scores = scores_nonpctl.apply(lambda x: np.digitize(x, np.percentile(x, np.arange(0, 100))))
        scores = scores_nonpctl.apply(lambda x: 100 * np.argsort(np.argsort(x)) / len(x))
    else:
        scores = scores_nonpctl

    # include all communities row?
    include_all_communities_row = True
    if include_all_communities_row:
        clustering_to_plot = clustering_with_all_communities
    else:
        clustering_to_plot = clustering

    cluster_id_name_size = clustering_to_plot.groupby('cluster_id').apply(lambda x: (x.name, x["cluster_name"].iloc[0], len(x)))
    cluster_id_name_size = sorted(cluster_id_name_size, key=lambda x: x[2])

    cluster_names = [x[1] + " (%d)" % x[2] for x in cluster_id_name_size]

    scatter_outliers_only = True
    show_density_estimate = True
    label_outliers =  not percentilize
    figure_size =  (9, (8/19.0)*len(cluster_names))

    if single_column:
        figure_size = (4.5, (8/19.0)*len(cluster_names))

    # end configuration

    plt.rcParams["image.composite_image"] = False

    to_plot = scores[dimens_to_plot]
    to_plot = to_plot.join(clustering_to_plot, on='community', how='inner')

    to_plot_nonpctl = scores_nonpctl[dimens_to_plot]
    to_plot_nonpctl = to_plot_nonpctl.join(clustering, on='community', how='inner')

    band_width = 0.1
    annotation_offset = 0.2
    density_n_bins = 40
    density_scaling_factor = 40 if percentilize else 0.25
    smooth_densities = percentilize

    smoother = LowessSmoother(smooth_fraction=0.2, iterations=1)

    cluster_ids = [x[0] for x in cluster_id_name_size]
    cluster_id_to_y = {cid: y for y, cid in enumerate(cluster_ids)}

    hide_outliers_for_clusters = ["Pornography"]
    hide_outliers_for_clusters = [clustering_to_plot[clustering_to_plot["cluster_name"] == cname]["cluster_id"].iloc[0] \
                                  for cname in hide_outliers_for_clusters]

    to_plot["y"] = to_plot["cluster_id"].map(cluster_id_to_y)# + (np.random.normal(size=len(to_plot)) * band_width) - (band_width / 2)

    fig, axs = plt.subplots(1, len(dimens_to_plot), figsize=figure_size, sharey=True)
    if not type(axs) is np.ndarray:
        axs = [axs]

    plt.subplots_adjust(wspace=0, hspace=0)

    glossary = []
    def annotate_outlier(ax, name, x, y, cluster_id, outliertype='minima',color=None):
        placement = {"minima":"top","maxima":"bottom"}[outliertype]

        glossary.append(name)

        ha = ('right' if outliertype == 'minima' else 'left')
        #ha = 'center'

        ax.annotate(name, xy=(x, y), size=8, ha=ha,
                    va=placement,
                    xytext=(x, y + annotation_offset * (-1 if placement == 'top' else 1)),
                   arrowprops=dict(color=color, arrowstyle="->"), c=color, zorder=999)

    metadata_rows = []

    cmap_name = {
        #"partisan": "coolwarm",
        #"gender": "PuOr",
        #"age": "PiYG"
    }

    for i, dimen in enumerate(dimens_to_plot):

        ax = axs[i]
        ax.set_frame_on(False)
        ax.tick_params(axis='both', length=0)
        ax.set_title(axis_name(dimen))

        density_bottoms = np.array(range(0, len(cluster_names)))
        tick_labels = []

        if i == 0:
            y_tick_offset = 0.2
            ax.set_yticks(density_bottoms + y_tick_offset)
            ax.set_yticklabels(cluster_names)
            ax.tick_params(axis='y', which='major', pad=25 if not single_column else 10)

        for b in density_bottoms:
            ax.axhline(b, color='#aaaaaa', linewidth=1, zorder=(-2 * b))

        ax.set_ylim(-0.1 if percentilize else -0.5, len(cluster_names) - 0.2)

        render_x = to_plot[dimen]
        x = to_plot_nonpctl[dimen]
        mean = np.mean(x)
        max_deviation = np.std(x)*3# np.percentile(np.abs(x - mean), 99)
        std_dev = np.std(x)
        render_thresh = std_dev * 2 if not percentilize else 0

        # cluster colormap
        cmap = axis_colormap(dimen)
        #if dimen == "partisan":
        #    cmap = matplotlib.colors.LinearSegmentedColormap.from_list("", ["#3A4CC0","#F7F7F6","#B30326"])
        cmap = matplotlib.cm.ScalarMappable(cmap=cmap)
        cmap.set_clim(mean - max_deviation, mean + max_deviation)

        if percentilize:
            # Need to generate colormap mapping percentiles to raw scores

            pctl_scores = scores[dimen].sort_values()
            raw_scores = (scores_nonpctl[dimen][pctl_scores.index] - (mean - max_deviation)) / (max_deviation * 2)

            class PercentileNormalize(colors.Normalize):
                def __init__(self, vmin=None, vmax=None, vcenter=None, clip=False):
                    self.vcenter = vcenter
                    colors.Normalize.__init__(self, vmin, vmax, clip)

                def __call__(self, value, clip=None):
                    #x, y = [self.vmin, self.vcenter, self.vmax], [0, 0.5, 1]
                    x, y = pctl_scores, raw_scores
                    return np.ma.masked_array(np.interp(value, x, y))

            cmap = matplotlib.cm.ScalarMappable(cmap=axis_colormap(dimen), norm=PercentileNormalize())

        means = to_plot.groupby("cluster_name").agg({dimen: [np.mean, np.var, "count"]})
        means.columns = means.columns.droplevel(0)
        means['mean'] = (means['mean'] - np.mean(x)) / std_dev
        metadata_rows.extend([[dimen, cluster_name, row['mean'], row['var'], row['count']] for cluster_name, row in means.iterrows()])

        # density estimate
        if show_density_estimate:
            # bin starts
            density_bins = np.linspace(np.amin(render_x), np.amax(render_x), density_n_bins)
            bin_width = (density_bins[1]-density_bins[0])
            # add start and end bin for aesthetic reasons
            #density_bins = np.insert(density_bins, 0, density_bins[0]-bin_width)
            #density_bins = np.insert(density_bins, len(density_bins), density_bins[-1]+bin_width)

            # the bin midpoints are where the line is actually plotted
            bin_midpoints = density_bins+bin_width/2
            bin_midpoints = bin_midpoints[:-1]

            xlim_padding = 10 if percentilize else 0.05
            if single_column:
                xlim_padding = 0
            ax.set_xlim(np.amin(bin_midpoints) - xlim_padding, np.amax(bin_midpoints) + xlim_padding)

            density_bg_n_points = 500
            density_bg_x = np.linspace(np.amin(render_x), np.amax(render_x), density_bg_n_points)

            density_bg = cmap.to_rgba(density_bg_x).reshape((1, len(density_bg_x), 4))

            z_order = 0

            groupbys = to_plot.groupby("cluster_id")
            for cluster_id in cluster_ids:
                rows = groupbys.get_group(cluster_id)
                hist, bin_edges = np.histogram(rows[dimen], bins=density_bins, density=True)
                hist_orig = hist

                if smooth_densities:
                    smoother.smooth(hist)
                    hist = smoother.smooth_data[0]

                baseline = cluster_id_to_y[cluster_id]
                hx = bin_midpoints
                hy = baseline + (hist * density_scaling_factor)
                hy_orig = baseline + (hist_orig * density_scaling_factor)

                line,  = ax.plot(hx, hy, c='black', alpha=0.5, linewidth=1,zorder=z_order,solid_capstyle='butt')

                #if smooth_densities:
                #    ax.bar(hx, hist_orig * density_scaling_factor, width=2, color='black', alpha=0.2, linewidth=0.5,zorder=z_order,bottom=baseline)

                xmin, xmax, ymin, ymax = np.amin(bin_midpoints), np.amax(bin_midpoints), baseline, np.amax(hy)
                im = ax.imshow(density_bg, aspect='auto', extent=[xmin, xmax, ymin, ymax],
                         origin='lower', zorder=z_order-1, interpolation='bilinear')

                xy = np.column_stack([hx, hy])
                xy = np.vstack([[xmin, ymin], xy, [xmax, ymin], [xmin, ymin]])
                clip_path = Polygon(xy, facecolor='none', edgecolor='none', closed=True)
                ax.add_patch(clip_path)
                im.set_clip_path(clip_path)
                z_order -= 2

        to_scatter = to_plot
        if scatter_outliers_only:
            to_scatter = to_scatter[np.abs(x - mean) >= render_thresh]
        if not percentilize:
            ax.scatter(to_scatter[dimen], to_scatter["y"], alpha=0.2, c=cmap.to_rgba(to_scatter[dimen]),
                       facecolors='none', marker='|')

        ax.axvline(mean if not percentilize else 50, color='#333333', linestyle='--', linewidth=1)

        if label_outliers:
            minima_indices = to_plot.groupby("cluster_id").apply(lambda x: x.index[np.argsort(x[dimen].values)[-1]])
            maxima_indices = to_plot.groupby("cluster_id").apply(lambda x: x.index[np.argsort(x[dimen].values)[0]])

            for cid in cluster_ids:
                if cid in hide_outliers_for_clusters:
                    continue

                to_plot_for_cluster = to_plot[to_plot["cluster_id"] == cid]

                argsort_idxs = np.argsort(to_plot_for_cluster[dimen].values)
                maxima = to_plot_for_cluster.index[argsort_idxs[-1]]
                minima = to_plot_for_cluster.index[argsort_idxs[0]]

                for outliertype, community in [("minima", minima), ("maxima", maxima)]:
                    if dimen.endswith("_neutral") and outliertype == "minima":
                        continue # Do not show left-side outliers on neutral plots

                    val = x[community]
                    val_for_color = mean + np.sign(val - mean) * np.maximum(std_dev*2, np.abs(val - mean))
                    color = cmap.to_rgba(val_for_color)
                    color_dim_factor = 6
                    fixed_alpha = 0.7
                    color = (color[0]/color_dim_factor, color[1]/color_dim_factor, color[2]/color_dim_factor, fixed_alpha)

                    outlier_x = to_plot_for_cluster[dimen].loc[community]
                    outlier_y = to_plot_for_cluster["y"].loc[community]

                    annotate_outlier(ax, community, outlier_x, outlier_y, cid, outliertype=outliertype, color=color)

        ax.set_xlabel(axis_label(dimen, include_name=False), fontsize=7)
        if percentilize:
            ax.set_xticks([0,100])
            ax.tick_params(axis='x', labelsize=8)

            ticklabels = ax.get_xticklabels()
            ticklabels[0].set_ha("left")
            ticklabels[-1].set_ha("right")

    plt.subplots_adjust(wspace=0.15 if not single_column else 0.15)

    # Add colorbars for axes
    for i, dimen in enumerate(dimens_to_plot):
        ax = axs[i]
        cmap = axis_colormap(dimen)
        cmap = matplotlib.cm.ScalarMappable(cmap=cmap)
        cmap.set_clim(-3,3)

        cbar_width = 0.8
        cbar_height = 0.008
        cbar_y = -0.06 if single_column else -0.05
        x, y, width, height = (0.1, cbar_y - cbar_height, cbar_width, cbar_height)
        t = (ax.transAxes + fig.transFigure.inverted()).transform_point
        x_t, y_t = t((x, y))
        width_t, height_t = t((x + width, y + height))
        width_t -= x_t
        height_t -= y_t
        cax = plt.axes((x_t, y_t, width_t, height_t))

        cbar = fig.colorbar(cmap, cax=cax, orientation='horizontal')
        ctx = [-3, -2, -1, 0, 1, 2, 3]
        cbar.set_ticks(ctx)
        if single_column:
            cbar.set_ticklabels(["-3$\sigma$","","","0","","","3$\sigma$"])
        else:
            cbar.set_ticklabels(["%d$\sigma$" % d for d in ctx])
        cax.tick_params(axis='both', which='major', labelsize=8)

    if export_raw_data:
        assert percentilize
        assert clustering_name == "reddit_all_numbered"
        data_to_export = clustering[['cluster_name']].copy()
        data_to_export[[d + "_pctl" for d in dimens_to_plot]] = scores[dimens_to_plot]
        data_to_export[[d + "_raw" for d in dimens_to_plot]] = scores_nonpctl[dimens_to_plot]
        #data_to_export["over18"] = embedding.metadata["over18"]
        #data_to_export["over18"] = (embedding.metadata["over18"] == "True").astype(int)
        #print(np.sum(~((embedding.metadata["over18"] == "True") | (embedding.metadata["over18"] == "False"))))
        data_to_export = data_to_export.rename(columns={"edginess_pctl":"edgy_pctl", "edginess_raw":"edgy_raw"})
        data_to_export = data_to_export.rename(columns={"partisan_pctl":"partisan_pctl", "partisan_raw":"partisan_raw"})
        data_to_export = data_to_export.rename(columns={"partisan_neutral_pctl":"partisan_neutral_pctl", "partisan_neutral_raw":"partisan_neutral_raw"})
        data_to_export.to_csv("~/research/commembed/source_data/fig_2.csv")

    return fig, glossary

def print_skews_and_variances(pairs=None, clustering_name="reddit_unmerged", bimodal_clusters=[]):
    #
    # Calculate skew and variance figures for text
    #

    clustering = load_clustering(clustering_name)

    single_skew_clusters = [(cluster, dimen) for cluster in clustering["cluster_name"].unique() for dimen in ["age", "gender", "partisan"]]

    percentile_threshold = 20

    scores_pctl = scores_nonpctl.apply(lambda x: 100 * np.argsort(np.argsort(x)) / len(x))
    scores_z = scores_nonpctl.apply(lambda x: (x-np.mean(x)) / np.std(x))
    cscores_pctl = scores_pctl.join(clustering, on='community', how='inner')
    cscores_z = scores_z.join(clustering, on='community', how='inner')

    skew_rows = []

    for cluster, axis in single_skew_clusters:
        rows_pctl = cscores_pctl[cscores_pctl["cluster_name"] == cluster]
        rows_z = cscores_z[cscores_z["cluster_name"] == cluster]
        x_pctl = rows_pctl[axis]
        x_z = rows_z[axis]
        mean_z = np.mean(x_z)

        less, more = axis_less_more(axis)
        modifier_word = less if mean_z < 0 else more

        pctl_text = ""
        if mean_z < 0:
            frac = np.sum(x_pctl < percentile_threshold) * 100.0 / len(x_pctl)
            pctl_text = "%.0f%% are below the %dth percentile" % (frac, percentile_threshold)
            desc = "%s communities skew %s (%s)" % (cluster, modifier_word, pctl_text)
            skew_rows.append([cluster, axis, modifier_word, frac, percentile_threshold, mean_z, desc])
        else:
            frac = np.sum(x_pctl > (100-percentile_threshold)) * 100.0 / len(x_pctl)
            pctl_text = "%.0f%% are above the %dth percentile" % (frac, 100-percentile_threshold)
            desc = "%s communities skew %s (%s)" % (cluster, modifier_word, pctl_text)
            skew_rows.append([cluster, axis, modifier_word, frac, 100-percentile_threshold, mean_z, desc])

    skew_df = pd.DataFrame(skew_rows, columns=["cluster", "axis", "word", "frac", "pctl", "mean_z", "desc"]).set_index(["cluster", "axis"])

    pd.set_option('display.max_colwidth', None)
    if pairs is None:
        display(skew_df.sort_values("frac", ascending=False))
    else:
        display(skew_df.loc[pairs])

    bimodal_window = 20

    for cluster, axis in bimodal_clusters:
        rows_pctl = cscores_pctl[cscores_pctl["cluster_name"] == cluster]
        x_pctl = rows_pctl[axis]

        total_outside = np.sum(x_pctl < bimodal_window) + np.sum(x_pctl > (100 - bimodal_window))
        frac = total_outside * 100.0 / len(x_pctl)

        print("%s communities exhibit a bimodal distribution on the %s axis (%.0f%% are below the %dth percentile or above the %dth percentile)" \
            % (cluster, axis_name(axis), frac, bimodal_window, 100 - bimodal_window))

def render_distribution(fig, gs, selected_dimension, glossary=[], subplot_label=True, sd=None, sd_key=None):

    ax = fig.add_subplot(gs)
    dimen_data = [d[1] for d in dimen_list if d[0] == selected_dimension][0]

    scores_n = scores.apply(lambda x: (x-np.mean(x))/np.std(x), axis=0)
    x = scores_n[selected_dimension]
    norm = Normalize(-3, 3)

    #ks_statistic, p_value = kstest(x, 'norm')
    #print(ks_statistic, p_value)

    # show_gaussian = False
    # if show_gaussian:

    #     x_gauss = np.linspace(-6, 6, 100)
    #     gauss = stats.norm.pdf(x_gauss)
    #     ax.plot(x_gauss, gauss * 100)

    if sd is not None:
        sd[sd_key] = scores_n[[selected_dimension]]

    n, bins, patches = ax.hist(scores_n[selected_dimension], log=False, bins=100)
    #ax.yaxis.tick_right()
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)

    #x = scores[selected_dimension]
    #mean = np.mean(x)
    #max_deviation = np.std(scores[selected_dimension])*3
    #norm = Normalize(mean - max_deviation, mean + max_deviation)

    # set a color for every bar (patch) according
    # to bin value from normalized min-max interval
    cmap = axis_colormap(selected_dimension)
    for bin, patch in zip(bins, patches):
        color = cmap(norm(bin))
        patch.set_facecolor(color)

    ax.set_ylabel("Number of communities")
    ax.set_xlabel(axis_label(selected_dimension))

    def render_nice_line(x, label, y):
        height = y

        ha = 'left'
        x_offset = x + 0.005
        text_offset = 1.5
        if (x > -0.2 and x < 0) or x > 0.2:
            ha = 'right'
            x_offset = x - 0.005
            text_offset = -1.5

        #x = x - (bins[1]-bins[0])*0.5
        color = '#777777'
        textcolor = 'black'
        ax.plot([x, x], [0, height], color=color, linestyle='--', linewidth=1)
        ax.annotate(label, (x_offset, height), xytext=(text_offset, 0), textcoords='offset points',ha=ha, va='top', rotation=0,color=textcolor)

    # render examples
    #sorted_scores = scores.sort_values(selected_dimension)
    #example_x = np.linspace(np.amin(sorted_scores[selected_dimension]), np.amax(sorted_scores[selected_dimension]), 18)
    #for ex in example_x:
    #    indx = np.argmax((sorted_scores[selected_dimension] >= ex).values)
    #    render_nice_line(sorted_scores[selected_dimension].iloc[indx], sorted_scores.index[indx])

    actual_seed = list([d[1] for d in dimen_list if d[0] == selected_dimension][0]['seeds'][0])

    cities = ["hillaryclinton", "The_Donald"] if selected_dimension.startswith("partisan") else []

    #render_nice_line(np.mean(scores.loc[dimen_data["negative_comms"], selected_dimension]), "Average of pairs' negative sides")
    #render_nice_line(np.mean(scores.loc[dimen_data["positive_comms"], selected_dimension]), "Average of pairs' positive sides")

    if selected_dimension.endswith("_neutral"):
        y = [50, 100] if len(cities) == 0 else [50,100,150,200]
    else:
        y = [100, 100] if len(cities) == 0 else [300,250,200,250]

    for i, comm in enumerate(sorted(actual_seed+cities, key=lambda x: scores_n.loc[x,selected_dimension])):
        render_nice_line(scores_n.loc[comm, selected_dimension], comm, y[i])
        glossary.append(comm)

    fig.canvas.draw()

    labels = [item.get_text() + "$\sigma$" if item.get_text() != "0" else item.get_text() for item in ax.get_xticklabels()]
    ax.set_xticklabels(labels)

    if subplot_label:
        plots.add_subplot_label(ax, "c", y=20)

all_word_data = pd.read_parquet(data.DATA_PATH+"/top_words_for_plotting.parquet")
political_word_data = pd.read_parquet(data.DATA_PATH+"/top_words_political_only.parquet")

def render_extreme_subs(fig, gs, selected_dimension, glossary=[], subplot_label=True, political_words=False, sd=None, sd_key=None):

    # cluster colormap
    mean = np.mean(scores[selected_dimension])
    max_deviation = np.std(scores[selected_dimension])*3
    cmap = axis_colormap(selected_dimension)
    cmap = cm.ScalarMappable(cmap=cmap)
    cmap.set_clim(mean - max_deviation, mean + max_deviation)

    ax = fig.add_subplot(gs)

    ax.set_yticks([])
    ax.set_xticks([])
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    if not political_words:
        ax.spines['bottom'].set_visible(False)
    ax.spines['bottom'].set_bounds(-7, 7)

    ax.spines['left'].set_visible(False)

    draw_little_numbers = not political_words

    n = 6 if political_words else 8 # was 10
    to_plot = scores.sort_values(selected_dimension)
    to_plot = pd.concat([to_plot.iloc[0:n][::-1], to_plot[-n:]])
    labelsize=9
    data = []
    colors = []
    norm = Normalize(0, 1)
    ys = list(range(0, int(len(to_plot)/2)))
    ys = ys+ys
    start_y = int(len(to_plot)/2) + 1
    ys = list(range(start_y, start_y+int(len(to_plot)/2))) + list(range(start_y, start_y+int(len(to_plot)/2)))+ys

    sd_rows = []

    for name, row in to_plot.iterrows():
        i = len(data)
        val = row[selected_dimension]
        y = ys[i]
        x = (val - np.mean(scores[selected_dimension]))/np.std(scores[selected_dimension])
        color = cmap.to_rgba(val)

        ha = "left"
        if x < 0:
            ha = "right"

        if draw_little_numbers:
            ax.annotate("%.2f" % val, (np.sign(x) * 0.15, y), c="white", ha=ha, va="center", fontsize=6)
        ax.annotate("r/"+name, (x, y), xytext=(5 * np.sign(x), 0), textcoords='offset points', c="black", ha=ha, va="center",fontsize=labelsize)
        data.append(x)
        colors.append(color)
        glossary.append(name)
        sd_rows.append([name,"community",x])

    # do words
    word_data = political_word_data if political_words else all_word_data
    word_dimen = selected_dimension if selected_dimension in word_data.columns else 'age'
    word_data = word_data.sort_values(word_dimen)

    ban_words = []

    word_data = word_data[~word_data["word"].isin(ban_words)]

    to_plot = pd.concat([word_data.iloc[0:n][::-1], word_data[-n:]]).set_index("word")

    xmax = np.amax(to_plot[word_dimen]-np.mean(scores[word_dimen]))
    anno_x = 17

    if political_words:
        ax.annotate("Communities", (anno_x, ys[len(data)-(n//2)]-(0.5 if (n%2)==0 else 1)), rotation=-90, color='#444444', va='center', annotation_clip=False)

    for name, row in to_plot.iterrows():
        i = len(data)
        val = row[word_dimen]
        y = ys[i]
        x = (val - np.mean(scores[word_dimen])) / np.std(scores[word_dimen])
        color = cmap.to_rgba(val)

        ha = "left"
        if x < 0:
            ha = "right"

        # Redact offensive words
        name = re.sub(r"n.gger", "n-gger", name)
        name = re.sub(r"f.ggot", "f-ggot", name)
        name = re.sub(r"r.tard", "r-tard", name)

        addtl_offset = True#np.abs(x)<0.2

        if draw_little_numbers:
            ax.annotate("%.2f" % val, (np.sign(x) * 0.15, y), c="black", ha=ha, va="center", fontsize=6)
        ax.annotate(name, (x, y), xytext=((10 if addtl_offset else 5) * np.sign(x), 0), textcoords='offset points', c="black", ha=ha, va="center",fontsize=labelsize)
        data.append(x)
        colors.append(color)
        sd_rows.append([name,"word",x])

    if sd is not None:
        sd[sd_key] = pd.DataFrame(sd_rows, columns=["item","type","score"]).set_index("item")

    ax.barh(ys, data, color=colors)

    width = np.amax(data)-np.amin(data)

    ax.set_xlim(np.amin(data)-(width/2), np.amax(data)+(width/2))

    ax.set_xlabel(axis_label(selected_dimension, include_name=False))

    if political_words:
        ax.annotate("Words", (anno_x, ys[-(n//2)]-(0.5 if (n%2)==0 else 1)), rotation=-90, color='#444444', va='center', annotation_clip=False)

        tx = [-6, -3, 0, 3,6]
        ax.set_xticks(tx)
        ax.set_xticklabels(["%d$\sigma$" % d if d != 0 else str(d) for d in tx])

        ax.axvline(0, linewidth=1, color='black', linestyle='--', alpha=0.5)

    if subplot_label:
        plots.add_subplot_label(ax, "d", x=-100, y=-10)

    return ax

def render_scatter(fig, gs):
    ax = fig.add_subplot(gs)

    ax.set_axis_off()

    to_plot = pd.read_csv('fig_1_clustering_projection.csv')
    #print(to_plot.shape)

    color_by_cluster = True
    c = None
    cmap = None
    if color_by_cluster:
        c = to_plot['cluster_id']
        cmap = 'tab20b'
    else:
        c = np.vstack([to_plot["r"],to_plot["g"],to_plot["b"]]).T

    s = 1.5
    ax.scatter(-to_plot["x"], to_plot["y"], c=c, cmap=cmap,
                alpha=0.25,
            #marker=',',
            s=s,
            lw=0)
    ax.axis('equal')
    plots.add_subplot_label(ax, "a", x=-9, y=-16)

    return ax

def render_schematic(fig, gs, selected_dimension, hide_numbers=False, glossary=[]):
    gs = list(gs.subgridspec(1, 3))
    ax = fig.add_subplot(gs[0])
    axs = np.array([ax] + [fig.add_subplot(g,sharey=ax) for g in gs[1:]])
    if not hide_numbers:
        plots.add_subplot_label(axs[0], "b")

    seed_pair = ["r/democrats", "r/Conservative"]
    found_pair = ["r/askhillarysupporters", "r/AskTrumpSupporters"]

    n_random_points = 2

    interesting_pair = np.array([[0, 0.2], [0.6, 0.8]])
    other_interesting_pair = interesting_pair + [0.2, -0.2]

    distort_amt = 0.05
    interesting_pair[0, 0] -= distort_amt
    interesting_pair[0, 1] += distort_amt
    other_interesting_pair[0, 0] += distort_amt
    other_interesting_pair[0, 1] -= distort_amt

    direction = np.array([1,1])
    direction = direction / np.linalg.norm(direction)

    random_points = np.random.rand(n_random_points, 2)
    random_points = np.array([[0.10593992, 0.76416452], [0.04775787, 0.91524662]]) # good looking

    all_points = np.concatenate((random_points, interesting_pair, other_interesting_pair), axis=0)

    color = [sns.color_palette()[7]]

    def draw_arrow_between_pair(ax, pair, subdued=False):
        ax.arrow(pair[0, 0], pair[0, 1], pair[1, 0]-pair[0, 0], pair[1,1]-pair[0, 1],
                 length_includes_head=True, alpha=(0.1 if subdued else 1),
                head_width=0.07, head_length=0.07, color='#666666')

    def draw_all_arrows(ax, pairs):
        for i in range(0, pairs.shape[0]):
            for j in range(0, pairs.shape[0]):
                #if np.random.rand() < 0.9:
                #    continue
                if i != j:
                    draw_arrow_between_pair(ax, pairs[[i, j], :], subdued=True)

    def annotate_community(ax, coords, name, above=False):
        glossary.append(name)
        ax.annotate(name, coords, xytext=(4 if name == "r/Conservative" else 0, 4 if above else -4), textcoords='offset points',
            #arrowprops=dict(arrowstyle="-",shrinkA=0),
            ha='center', va='bottom' if above else 'top')
        #ax.scatter([coords[0]],[coords[1]],c='red' if above else 'blue')

    for i, ax in enumerate(axs.flatten()):
        i = i + 1

        ax.set_xticks([])
        ax.set_yticks([])

        spine_color = '#aaaaaa'
        ax.spines['bottom'].set_color(spine_color)
        ax.spines['top'].set_color(spine_color)
        ax.spines['right'].set_color(spine_color)
        ax.spines['left'].set_color(spine_color)

        ax.margins(0.1)
        ax.axis('equal')
        ax.set_axis_off()

        if i == 3:
            # draw tangent lines for all communities
            for point_i in range(0, all_points.shape[0]):

                point = all_points[point_i, :]
                point_on_line = direction * np.dot(direction, point)
                d = point_on_line - point
                ax.arrow(point[0], point[1], d[0], d[1], alpha=0.1)

        pair1_color = color
        pair2_color = color
        focus_blue = "#3B4CC0"
        focus_red = "#B40426"
        subd_blue = "#84A7FC"
        subd_red = "#F08A6C"

        if i == 1:
            pair1_color = [focus_blue, focus_red]

        if i == 2:
            #pair1_color = [focus_blue, focus_red]#[subd_blue, subd_red]
            pair2_color = [focus_blue, focus_red]

        if i == 3:
            pair1_color = [focus_blue, focus_red]
            pair2_color = [focus_blue, focus_red]

        ax.scatter(random_points[:, 0], random_points[:, 1], c=color)
        ax.scatter(interesting_pair[:, 0], interesting_pair[:, 1], c=pair1_color)
        ax.scatter(other_interesting_pair[:, 0], other_interesting_pair[:, 1], c=pair2_color)

        if not hide_numbers:
            ax.annotate(str(i), [-0.25, 0.9], fontsize=10, bbox=dict(boxstyle='square',fc='#eeeeee',ec='white'))

        if i == 1:
            draw_arrow_between_pair(ax, interesting_pair)
            annotate_community(ax, interesting_pair[0, :], seed_pair[0])
            annotate_community(ax, interesting_pair[1, :], seed_pair[1], above=True)
        if i == 2:
            draw_arrow_between_pair(ax, other_interesting_pair)
            annotate_community(ax, other_interesting_pair[0, :], found_pair[0])
            annotate_community(ax, other_interesting_pair[1, :], found_pair[1], above=False)
        if i == 2:
            draw_all_arrows(ax, all_points)
        if i == 3:
            ax.arrow(0, 0, 0.85, 0.85,
                     length_includes_head=True, alpha=0.5,
                    head_width=0.3, head_length=0.14, color='#666666', width=0.1)
            ax.annotate(axis_name(selected_dimension), (0.5, 0.5), ha='center', va='center', rotation=45)

from matplotlib.transforms import Bbox

def render_fig1_regular(selected_dimension):

    glossary = []
    source_data = {}

    scale_factor = 0.8
    fig = plt.figure(figsize=(9*scale_factor,5.5*scale_factor))

    gs0 = gridspec.GridSpec(6, 9, figure=fig)
    gs0.update(wspace=0.5, hspace=0.5)

    gs_scatter = gs0[0:4, :3]
    gs_schematic = gs0[0:2, 3:]

    gs_word_freqs = gs0[2:, 5:8]
    gs_dist = gs0[3:, :5]

    render_distribution(fig, gs_dist, selected_dimension, glossary=glossary, sd=source_data,sd_key="c")
    ax_subs = render_extreme_subs(fig, gs_word_freqs, selected_dimension, glossary=glossary, political_words=True, sd=source_data,sd_key="d")
    ax_scatter = render_scatter(fig, gs_scatter)
    render_schematic(fig, gs_schematic, selected_dimension, glossary=glossary)

    plots.nudge_axis(ax_scatter, -0.04, 0.06)

    subs_shrink_amt = 0.05
    #plots.nudge_axis(ax_subs, 0, -0.004 - subs_shrink_amt)
    plots.shrink_axis(ax_subs, 0, subs_shrink_amt)

    #plt.savefig('fig-1-composite.pdf', bbox_inches = "tight")
    #plt.show()

    return fig, glossary, source_data

def render_fig1_compact(selected_dimension):

    glossary = []

    fig = plt.figure(figsize=(13,3.5))

    gs0 = gridspec.GridSpec(15, 8, figure=fig)

    gs_word_freqs = gs0[:, 3:]
    gs_dist = gs0[1:-2,:4]

    render_distribution(fig, gs_dist, selected_dimension, glossary=glossary, subplot_label=False)
    ax = render_extreme_subs(fig, gs_word_freqs, selected_dimension, glossary=glossary, subplot_label=False)
    ax.patch.set_alpha(0)

    #plt.savefig('fig-1-composite.pdf', bbox_inches = "tight")
    #plt.show()

    return fig, glossary

def render_fig1_grid(selected_dimensions=["age", "gender", "partisan", "affluence"]):

    glossary = []

    n = len(selected_dimensions)

    fig = plt.figure(figsize=(10,2.8*n))

    gs0 = gridspec.GridSpec(16*n, 8, figure=fig)

    for i, selected_dimension in enumerate(selected_dimensions):
        start = i * 16
        end = (i + 1) * 16

        gs_word_freqs = gs0[start:(end-2), 3:]
        gs_dist = gs0[(start+1):(end-3),:4]

        render_distribution(fig, gs_dist, selected_dimension, glossary=glossary, subplot_label=False)
        ax = render_extreme_subs(fig, gs_word_freqs, selected_dimension, glossary=glossary, subplot_label=False)
        ax.patch.set_alpha(0)

    #plt.savefig('fig-1-composite.pdf', bbox_inches = "tight")
    #plt.show()

    return fig, glossary


# -*- coding: utf-8 -*-

from keybert import KeyBERT
from konlpy.tag import Mecab

def tokenizer(sentence):
    mecab = Mecab(dicpath=r"C:\mecab\mecab-ko-dic")
    pos_lst = ["NNG", "NNP"]
    filtered_lst = [word for word, pos in mecab.pos(sentence) if pos in pos_lst]

    return filtered_lst

def keyword(doc):
    kw_model = KeyBERT()

    keywords = kw_model.extract_keywords(docs=doc)

    candidates_lst = [word for word, _ in keywords]
    candidates_lst = " ".join(candidates_lst)

    fin_keywords = tokenizer(candidates_lst)

    return fin_keywords

if __name__ == "__main__":
    example1 = "강릉여행. 온 가족이 함께 강릉으로 여행을 갔다. 다같이 가서 회도 먹고 물놀이도 하고 너무 즐거웠다"
    example2 = "추석에 온 가족이 모여 찍은 가족사진. 오랜만에 모여 윳놀이도 하고, 고스톱도 치고 재미있었다."
    example3 = "진통이 정말 심했어요. 새벽에 아버지한테 제가 전화해서 너무 아프다고 울었어요. 그래서 그날 가족들 다같이 병원으로 와서 진통하는 걸 봤었지요. 그때 아머지가 계속 손잡고 옆을 지켜주셨어요. 둘째 낳고 아버지가 보양식이며 배냇저고리며 이것저것 엄청 챙겨주셨던 기억이 나요. 낳자마자 아버지가 손수 이름 지어주겠다고 사흘밤낮동안 고민하셨던 것도 눈에 선해요"

    model = KeyBERT()
    test = model.extract_keywords(docs=[example1, example2, example3])
    print(test)