causality_simulation.py

import numpy as np
import matplotlib.pyplot as plt
import ipywidgets as wd
import pandas as pd
from IPython.display import display, update_display, Javascript, HTML
from inspect import signature
from graphviz import Digraph
import scipy.stats as sp
import plotly.express as px
import plotly.graph_objects as go
import warnings
import re

'''
Data structure
--------------
CausalNetwork
* complete network of all nodes
* init_data
* no per-experiment data

CausalNode
* causal relations from a single node

Experiment
* group assignment
* intervention
'''

# TODO: assignment(config=) change format config to {'group_name': '1-500'} (but maybe a list of dicts is better for preserving ordering)
# TODO: assignment string must be within the total number of samples
# TODO: assignment mustn't have overlapping groups

display(HTML('''<style>
    [title="Assigned samples:"] { min-width: 150px; }
</style>'''))

def dialog(title, body, button):
    display(Javascript("require(['base/js/dialog'], function(dialog) {dialog.modal({title: '%s', body: '%s', buttons: {'%s': {}}})});" % (title, body, button)))

class CausalNetwork:
    def __init__(self):
        self.nodes = {} # 'node_name': node
        self.init_data = pd.DataFrame()

    def addNode(self, node):
        if node.name in self.nodes.keys(): # check for duplicate nodes
            raise ValueError('A node with the same name %s already exists!' % node['name'])
        self.nodes[node.name] = node

    def addCause(self, effect, cause):
        self.nodes[effect].setCause(func=lambda x: x, causes=[cause]) # func has to be an unlifted function, so don't use identity here!

    # to be run at the beginning of each notebook
    def init(self, data):
        '''
        initialises the network and its constituent nodes
        data: {'some_node': [val, val, ...], ...} all arrays must have same size
        '''
        l = []
        for name, n in self.nodes.items():
            if n.causes is None: # is an init node, populate init_data column, throws error if arrays don't have same size
                self.init_data[name] = data[name]
            n.root_causes = n.traceRoots()
            self.replacePlaceholders(n)                
        # TODO: validate causal network has no single-direction loops (other loops are allowed)

    def generateSingle(self, row):
        '''
        Returns one row of pandas table
        row: single row of DataFrame or dict {'some_node': val, ...}
        '''
        total_nodes = len(self.nodes)
        new_row = {name: row[name] for name in self.nodes.keys() if name in row.keys() and not pd.isnull(row[name])} # populate new_row with values of init nodes or fixed nodes
        while len(new_row) < total_nodes:
            for name, n in self.nodes.items():
                if name not in new_row.keys():
                    ready = True # checks if all its causes have been evaluated or fixed
                    for c in n.root_causes:
                        if c not in new_row.keys():
                            ready = False
                            break
                    if ready: # evaluate this node using values from its .causes
                        cause_vals = {c: new_row[c] for c in n.root_causes}
                        new_row[name] = n.evaluate(cause_vals)
        return new_row

    # replaces all PlaceholderNodes of a node's causes by the actual node searched by name
    def replacePlaceholders(self, node):
        if isinstance(node, PlaceholderNode):
            raise ValueError('A PlaceholderNode cannot be passed directly to replacePlaceholders. Use it on the parent instead.')
        elif node.causes is not None:
            for i, n in enumerate(node.causes):
                if isinstance(n, PlaceholderNode):
                    if n.name not in self.nodes:
                        raise ValueError("Node %s doesn't exist in the causal network!" % n.name)
                    node.causes[i] = self.nodes[n.name] # replace
                else:
                    self.replacePlaceholders(n) # recurse down the tree

    def draw(self, root, intermediate=False):
        g = Digraph(name=root)
        def draw_edges(node, g):
            if node.causes:
                causes = node.causes if intermediate else [self.nodes[name] for name in node.root_causes]
                for c in causes:
                    g.edge(str(c.name), str(node.name))
                    draw_edges(c, g)
        draw_edges(self.nodes[root], g)
        return g

class CausalNode:
    def __init__(self, name=None):
        self.name = name if name else id(self) # given name or use unique id
        self.causes = None # array of node type arguments to f. None means this is an init node
        self.func = None # only takes positional arguments
        self.root_causes = set() # set of names of PlaceholderNodes that this node directly depends on

    def setCause(self, func, causes):
        self.causes = causes
        self.func = func

    def traceRoots(self):
        '''
        Stores a list of names of (string-named) nodes that this node depends on
        '''
        root_causes = set()
        if self.causes is None:
            return root_causes
        for c in self.causes:
            if isinstance(c, PlaceholderNode):
                root_causes.add(c.name)
            else:
                root_causes = root_causes.union(c.traceRoots()) # recurse down the network
        return root_causes

    def evaluate(self, params):
        '''
        Evaluates the value of this node given the values of nodes in root_causes
        kwargs.keys() must match root_causes TODO: validate this somewhere?
        '''
        if self.causes is None:
            raise ValueError('This is an init node. evaluate() should not have been called.')
        cause_vals = [None] * len(self.causes)
        for i, c in enumerate(self.causes):
            if isinstance(c.name, str): # if predefined node (i.e. not intermediate node)
                cause_vals[i] = params[c.name]
            else: # recurse down the tree
                cause_vals[i] = c.evaluate(params)
        return self.func(*cause_vals)

# placeholder node that only has a name, used to look up the actual node in the CausalNetwork
class PlaceholderNode(CausalNode):
    pass

class ContinuousNode(CausalNode):
    def __init__(self, name=None, min=0, max=100):
        super().__init__(name=name)
        self.min = min
        self.max = max

class DiscreteNode(CausalNode):
    def __init__(self, name=None, min=0, max=100):
        super().__init__(name=name)
        self.min = min
        self.max = max

class CategoricalNode(CausalNode):
    def __init__(self, name=None, categories=[]):
        super().__init__(name=name)
        self.categories = categories

class Experiment:
    def __init__(self, network):
        self.network = network
        self.N = len(self.network.init_data)
        self.data = self.network.init_data.copy()
        self.assigned = False # has the assignment step been run
        self.done = False # has the experiment been done

    def assignment(self, config=None, hide_random=False):
        '''
        UI for group assignment of samples
        config: list of dicts, each being {'name': group_name, 'samples_str': string}
        samples_str: e.g. '1-25,30,31-34', if all groups have empty string '' then assume randomise
        '''
        self.group_assignment = GroupAssignment(self)
        if config is not None:
            self.group_assignment.setAssignment(config, hide_random)
            self.submitAssignment()

    def setAssignment(self, groups):
        '''
        Populates the 'Group' column of self.data without UI
        groups: list of dicts, each being {'name': group_name, 'samples': [array]}
        '''
        # check for duplicates and populate self.data
        self.data['Group'] = None # eventually the column should all be group names
        seen = []
        for g in groups:
            if g['name'] in seen:
                dialog('Duplicate group names', 'Some of the groups have been given the same name. Please choose a unique name for each group.', 'OK')
                return
            else:
                seen.append(g['name'])
            if g['samples'] is None:
                dialog('Assignment invalid', 'Invalid assignment of samples to groups! Please revise your assignments.', 'OK')
                return
            for i in g['samples']:
                self.data.at[i,'Group'] = g['name']
        self.groups = [g['name'] for g in groups] # list of group names
        self.group_ids = {name: i for i, name in enumerate(self.groups)} # for easier lookup of group ids
        if None in self.data['Group'].unique():
            dialog('Not all samples assigned', 'Not all samples have been assigned to a group! Please revise your assignments.', 'OK')
            return
        self.assigned = True

    def submitAssignment(self, sender=None):
        '''
        Collects the group assignments from UI
        Checks for duplicate group names
        '''
        self.setAssignment(self.group_assignment.getAssignment())
        self.plotAssignment()

    def plotAssignment(self):
        '''
        Can be implemented differently in different scenarios. Extend the Experiment class and override this method
        '''
        pass

    def setting(self, show='all', config=None, disable=[]):
        '''
        Let user design experiment
        disabled: array of names
        show: array of names
        '''
        if not self.assigned:
            dialog('Groups not assigned', 'You have not yet assigned any groups! Click on "Visualise assignment" before running this box.', 'OK')
            return
        disable = self.network.nodes.keys() if disable == 'all' else disable
        self.intervention_setting = InterventionSetting(self, show=show, disable=disable)
        if config is not None:
            self.intervention_setting.setIntervention(config)
            self.doExperiment(config)

    def generate(self, intervention=[]):
        '''
        Generates experimental data for all groups according to intervention. Stores results in .data
        intervention: list of dicts, each being {'name': 'some_group', 'intervention': {'some_node': [args], ...}, ...}
        intervention format:
        ['fixed', val] (val could be number or name of category)
        ['range', start, end]
        ['array', [...]] array size must be n
        '''
        # reset the dataframe to just group assignments
        for key in self.data:
            if key != 'Group':
                self.data.drop(key, axis=1)
        for g in intervention:
            group = g['name']
            inter = g['intervention']
            n = len(self.data.loc[self.data['Group'] == group].index) # number of samples in group
            for name, args in inter.items():
                if args[0] == 'fixed':
                    to_fix = args[1]
                elif args[0] == 'range':
                    to_fix = np.random.permutation(np.linspace(args[1], args[2], n))
                    if isinstance(self, DiscreteNode):
                        to_fix = np.rint(to_fix)
                elif args[0] == 'array':
                    to_fix = args[1]
                # populate only those entries that are intervened on
                self.data.loc[self.data['Group'] == group, name] = to_fix
        # generate the rest of each row using the existing values
        for i in range(len(self.data)):
            new_vals = self.network.generateSingle(self.data.loc[i])
            for name, val in enumerate(new_vals):
                self.data.at[i, name] = val # use loc or iloc?

    def doExperiment(self, intervention, msg=False):
        '''
        Perform experiment under intervention
        intervention: {'group_name': {'node_name': [...]}}
        '''
        self.generate(intervention)
        self.done = True
        if msg:
            display(wd.Label(value='Data from experiment collected!'))

    def plot(self, show='all'):
        '''
        Plots data after doExperiment has been called
        '''
        if not self.done:
            dialog('Experiment not performed', 'You have not yet performed the experiment! Click on "Perform experiment" before running this box.', 'OK')
            return
        p = InteractivePlot(self, show)
        self.p = p
        p.display()

class TruffulaExperiment(Experiment):
    def __init__(self):
        super().__init__(truffula)

    def plotAssignment(self):
        self.a = AssignmentPlot(self, 'Truffula')

    def plotOrchard(self, gradient=None, show='all'):
        '''
        Takes in the name of the group in the experiment and the name of the 
        variable used to create the color gradient
        '''
        if not self.done:
            dialog('Experiment not performed', 'You have not yet performed the experiment! Click on "Perform experiment" before running this box.', 'OK')
            return
        o = OrchardPlot(self, gradient=gradient, show=show)
        self.o = o
        o.display()

class BasketballExperiment(Experiment):
    def __init__(self):
        super().__init__(basketball)

    def plotAssignment(self):
        pass

class GroupAssignment:
    def __init__(self, experiment):
        '''
        UI for group assignment of samples
        submitAssignment: callback function
        '''
        self.experiment = experiment
        wd.Label(value='Sample size: %d' % self.experiment.N)
        self.randomise_button = wd.Button(description='Randomise assignment', layout=wd.Layout(width='180px'))
        self.group_assignments = [SingleGroupAssignment(1)]
        self.add_group_button = wd.Button(description='Add another group')
        self.submit_button = wd.Button(description='Visualise assignment')
        self.box = wd.VBox([g.box for g in self.group_assignments])
        display(self.randomise_button, self.box, self.add_group_button, self.submit_button)
        self.randomise_button.on_click(self.randomise)
        self.add_group_button.on_click(self.addGroup)
        self.submit_button.on_click(self.experiment.submitAssignment)

    def setAssignment(self, config, hide_random):
        for i in range(len(config)-1):
            self.addGroup()
        self.greyAll()
        for i in range(len(config)):
            self.group_assignments[i].setName(config[i]['name'])
        if ''.join([g['samples_str'] for g in config]) == '':
            self.randomise()
        else:
            for i in range(len(config)):
                self.group_assignments[i].setSamples(config[i]['samples_str'])
        if hide_random:
            self.randomise_button.layout.visibility = 'hidden'

    def addGroup(self, sender=None):
        i = self.group_assignments[-1].i
        self.group_assignments.append(SingleGroupAssignment(i+1))
        self.box.children = [g.box for g in self.group_assignments]

    def getAssignment(self):
        '''
        Reads the settings and returns a list of dictionaries
        '''
        return [g.getAssignment() for g in self.group_assignments]

    def randomise(self, sender=None):
        '''
        Randomly assigns samples to groups and changes settings in UI
        '''
        N = self.experiment.N
        N_group = len(self.group_assignments)
        assigned_ids = randomAssign(N, N_group)
        for i in range(N_group):
            self.group_assignments[i].samples.value = array2Text(assigned_ids[i])

    def greyAll(self):
        self.randomise_button.disabled = True
        self.add_group_button.disabled = True
        self.submit_button.disabled = True
        for g in self.group_assignments:
            g.greyAll()

class SingleGroupAssignment:
    def __init__(self, i):
        '''
        UI for a single line of group assignment
        '''
        self.i = i # Group number
        self.name = 'Group %d' % i
        i_text = wd.Label(value=self.name, layout=wd.Layout(width='70px'))
        self.group_name = wd.Text(description='Name:')
        self.samples = wd.Text(description='Assigned samples:', layout=wd.Layout(width='400px'))
        self.box = wd.HBox([i_text, self.group_name, self.samples])

    def getAssignment(self):
        '''
        Returns dict {'name': group_name, 'samples': [list_of_sample_ids]}
        '''
        assignment = dict()
        self.name = self.name if self.group_name.value == '' else self.group_name.value
        assignment['name'] = self.name
        assignment['samples'] = text2Array(self.samples.value)
        return assignment

    def setName(self, name):
        self.group_name.value = name

    def setSamples(self, samples):
        self.samples.value = samples

    def greyAll(self):
        self.group_name.disabled = True
        self.samples.disabled = True

class InterventionSetting:
    def __init__(self, experiment, show='all', disable=[]):
        self.experiment = experiment
        self.group_settings = [SingleGroupInterventionSetting(self.experiment, g, len(self.experiment.data.loc[self.experiment.data['Group'] == g].index), show=show, disable=disable) for g in self.experiment.groups]
        submit = wd.Button(description='Perform experiment')
        display(submit)
        submit.on_click(self.submit)

    def submit(self, sender=None):
        self.experiment.doExperiment(self.getIntervention(), msg=True)

    def getIntervention(self):
        return [{'name': s.name, 'intervention': s.getIntervention()} for s in self.group_settings]

    def setIntervention(self, config):
        for c in config:
            j = self.experiment.group_ids[c['name']]
            self.group_settings[j].setIntervention(c)

class SingleGroupInterventionSetting:
    def __init__(self, experiment, name, N, show='all', disable=[]):
        '''
        UI settings for a single group
        config: {'name': group_name, 'samples': [sample_ids]}
        '''
        self.experiment = experiment
        self.name = name
        self.N = N
        group_text = wd.Label(value='Group name: %s, %d samples' % (name, N))
        display(group_text)
        to_list = list(self.experiment.network.nodes.keys()) if show == 'all' else show
        to_list = [a for a in to_list if a not in self.experiment.network.init_data.keys()]
        to_list.sort()
        self.node_settings = [SingleNodeInterventionSetting(self.experiment.network.nodes[n], disable=n in disable) for n in to_list]

    def getIntervention(self):
        intervention = dict()
        for s in self.node_settings:
            inter = s.getIntervention()
            if inter is not None:
                intervention[s.name] = inter
        return intervention

    def setIntervention(self, config):
        for s in self.node_settings:
            if s.name in config['intervention'].keys():
                s.setIntervention(config['intervention'][s.name])

class SingleNodeInterventionSetting:
    def __init__(self, node, disable=False):
        '''
        Single line of radio buttons and text boxes for intervening on a single variable in a single group
        '''
        self.name = node.name
        self.disable = disable
        self.is_categorical = isinstance(node, CategoricalNode)
        self.indent = wd.Label(value='', layout=wd.Layout(width='20px'))
        self.text = wd.Label(value=self.name, layout=wd.Layout(width='180px'))
        self.none = wd.RadioButtons(options=['No intervention'], layout=wd.Layout(width='150px'))
        self.fixed = wd.RadioButtons(options=['Fixed'], layout=wd.Layout(width='70px'))
        self.fixed.index = None
        if self.is_categorical:
            fixed_arg = wd.Dropdown(options=node.categories, disabled=True, layout=wd.Layout(width='100px'))
        else:
            fixed_arg = wd.BoundedFloatText(disabled=True, layout=wd.Layout(width='70px'))
        self.fixed_arg = fixed_arg
        self.range_visibility = 'hidden' if self.is_categorical else 'visible'
        self.range = wd.RadioButtons(options=['Range'], layout=wd.Layout(width='70px', visibility=self.range_visibility))
        self.range.index = None
        self.range_arg1_text = wd.Label(value='from', layout=wd.Layout(visibility=self.range_visibility, width='30px'))
        a = 0 if self.is_categorical else node.min
        b = 0 if self.is_categorical else node.max
        self.range_arg1 = wd.BoundedFloatText(min=a, max=b, disabled=True, layout=wd.Layout(width='70px', visibility=self.range_visibility))
        self.range_arg2_text = wd.Label(value='to', layout=wd.Layout(visibility=self.range_visibility, width='15px'))
        self.range_arg2 = wd.BoundedFloatText(min=a, max=b, disabled=True, layout=wd.Layout(width='70px', visibility=self.range_visibility))
        self.none.observe(self.none_observer, names=['value'])
        self.fixed.observe(self.fixed_observer, names=['value'])
        self.range.observe(self.range_observer, names=['value'])
        self.box = wd.HBox([self.indent, self.text, self.none, self.fixed, self.fixed_arg, self.range, self.range_arg1_text, self.range_arg1, self.range_arg2_text, self.range_arg2])
        display(self.box)
        if self.disable:
            self.greyAll()

    def greyAll(self):
        self.none.disabled = True
        self.fixed.disabled = True
        self.fixed_arg.disabled = True
        self.range.disabled = True
        self.range_arg1.disabled = True
        self.range_arg2.disabled = True

    def setIntervention(self, intervention):
        if intervention[0] == 'fixed':
            self.fixed.index = 0
            self.fixed_arg.value = intervention[1]
        elif intervention[0] == 'range':
            self.range.index = 0
            self.range_arg1.value = intervention[1]
            self.range_arg2.value = intervention[2]

    # Radio button .index = None if off, .index = 0 if on
    def none_observer(self, sender):
        if self.none.index == 0:
            self.fixed.index = None
            self.fixed_arg.disabled = True
            self.range.index = None
            self.range_arg1.disabled = True
            self.range_arg2.disabled = True
        if self.disable:
            self.greyAll()

    def fixed_observer(self, sender):
        if self.fixed.index == 0:
            self.none.index = None
            self.fixed_arg.disabled = False
            self.range.index = None
            self.range_arg1.disabled = True
            self.range_arg2.disabled = True
        if self.disable:
            self.greyAll()

    def range_observer(self, sender):
        if self.range.index == 0:
            self.none.index = None
            self.fixed.index = None
            self.fixed_arg.disabled = True
            self.range_arg1.disabled = False
            self.range_arg2.disabled = False
        if self.disable:
            self.greyAll()

    def getIntervention(self):
        '''
        generates intervention from UI settings
        '''
        if self.none.index == 0: # None is deselected, 0 is selected
            return None
        elif self.fixed.index == 0:
            return ['fixed', self.fixed_arg.value]
        elif self.range.index == 0:
            return ['range', self.range_arg1.value, self.range_arg2.value]

class AssignmentPlot:
    def __init__(self, experiment, plot='Truffula'):
        self.experiment = experiment
        self.group_names = experiment.data['Group'].unique()
        self.data = experiment.data
        self.plot = plot
        self.buildTraces()
        if self.plot == 'Truffula':
            self.layout = go.Layout(title=dict(text='Tree Group Assignments'),barmode='overlay', height=650, width=800,
                                  xaxis=dict(title='Longitude', fixedrange=True), yaxis=dict(title='Latitude', fixedrange=True),
                                  hovermode='closest',
                                  margin=dict(b=80, r=200, autoexpand=False),
                                  showlegend=True)
        else:
            self.layout = go.Layout(title=dict(text='Student Group Assignments'),barmode='overlay', height=650, width=800,
                                  xaxis=dict(title='Student', fixedrange=True),
                                  yaxis=dict(title='Height', fixedrange=True, range=(120, 200)),
                                  hovermode='closest',
                                  margin=dict(b=80, r=200, autoexpand=False),
                                  showlegend=True)
        self.plot = go.FigureWidget(data=self.traces, layout=self.layout)
        display(self.plot)
        
    def buildTraces(self):
        self.traces = []
        self.group_names = self.experiment.group_names
        self.data = self.experiment.data
        if self.plot == 'Truffula':
            for i, name in enumerate(self.group_names):
                self.traces += [go.Scatter(x=self.data[name]['Longitude'], y=self.data[name]['Latitude'], mode='markers', hovertemplate='Latitude: %{x} <br>Longitude: %{y} <br>', marker_symbol=i, name=name)]
        else:
            for i, name in enumerate(self.group_names):
                self.traces += [go.Bar(x=self.data[name]['id'], y=self.data[name]['Height (cm)'], hovertemplate='Student: %{x} <br>Height: %{y} cm<br>', name=name)]
        
    def updateAssignments(self):
        self.buildTraces()
        with self.plot.batch_update():
            self.plot.data = []
            for trace in self.traces:
                self.plot.add_traces(trace)
            self.plot.layout = self.layout

class OrchardPlot:
    def __init__(self, experiment, gradient=None, show='all'):
        self.data = experiment.data
        self.experiment = experiment
        self.options = self.data[experiment.group_names[0]].columns.tolist()
        if show != 'all':
            for i in self.options.copy():
                if i not in show:
                    self.options.remove(i)
        for name in experiment.node.nodeDict():
            if experiment.node.nodeDict()[name].vartype == 'categorical' and name in show:
                self.options.remove(name)
        self.options.sort()
        if not gradient:
            gradient = self.options[0]
        self.textbox = wd.Dropdown(
                description='Gradient: ',
                value=gradient,
                options=self.options
        )
        self.textbox.observe(self.response, names="value")
        self.plotOrchard(gradient)
        
    def validate(self):
        return self.textbox.value in self.options

    def response(self, change):
        if self.validate():
            with self.g.batch_update():
                for i, name in enumerate(self.experiment.group_names):
                    self.g.data[i].marker.color = self.data[name][self.textbox.value]
                    self.g.update_layout({'coloraxis':{'colorscale':'Plasma', 'colorbar':{'title':self.textbox.value}}})
                    self.g.data[i].hovertemplate = 'Latitude: %{x} <br>Longitude: %{y} <br>' + self.textbox.value + ': %{marker.color}<br>'

    def plotOrchard(self, gradient):
        """Takes in the name of the group in the experiment and the name of the 
        variable used to create the color gradient"""
        traces = []
        for i, name in enumerate(self.experiment.group_names):
            traces += [go.Scatter(x=self.data[name]['Longitude'], y=self.data[name]['Latitude'],
                                 marker=dict(color=self.data[name][gradient], coloraxis='coloraxis'),
                                 mode='markers',
                                 name=name,
                                 hovertemplate='Latitude: %{x} <br>Longitude: %{y} <br>'+ self.textbox.value + ': %{marker.color}<br>', hoverlabel=dict(namelength=0), marker_symbol=i)]
        width = 700 if (len(self.experiment.group_names) == 1) else 725 + max([len(name) for name in self.experiment.group_names])*6.5
        go_layout = go.Layout(title=dict(text='Orchard Layout'),barmode='overlay', height=650, width=width,
                              xaxis=dict(title='Longitude', fixedrange=True, range=[-50, 1050]), 
                              yaxis=dict(title='Latitude', fixedrange=True, range=[-50, 1050]),
                              hovermode='closest', legend=dict(yanchor="top", y=1, xanchor="left", x=1.25),
                              coloraxis={'colorscale':'Plasma', 'colorbar':{'title':gradient}})
        self.g = go.FigureWidget(data=traces, layout=go_layout)
        
    def display(self):
        container = wd.HBox([self.textbox])
        display(wd.VBox([container, self.g]))

class InteractivePlot:
    def __init__(self, experiment, show='all'):
        self.experiment = experiment
        self.x_options = list(experiment.network.nodes.keys())
        self.y_options = self.x_options.copy()
        if show != 'all':
            for i in self.x_options.copy():
                if i not in show:
                    self.x_options.remove(i)
                    self.y_options.remove(i)
        self.x_options.sort()
        self.y_options.sort()
        self.y_options += ['None (Distributions Only)']
        self.textbox1 = wd.Dropdown(
            description='x-Axis Variable: ',
            value=self.x_options[0],
            options=self.x_options
        )
        self.textbox2 = wd.Dropdown(
            description='y-Axis Variable: ',
            value=self.y_options[0],
            options=self.y_options
        )
        self.button = wd.RadioButtons(
            options=list(experiment.data.keys()) + ['All'],
            layout={'width': 'max-content'},
            description='Group',
            disabled=False
        )
        self.observe()
        self.initTraces()

    def display(self):
        container = wd.HBox([self.textbox1, self.textbox2])
        display(wd.VBox([container, self.g]))
        display(self.button)
        display(Nothing(), display_id='1')
        self.button.layout.display = 'none'
    
    def display_values(self, group):
        text = ""
        xType, yType = self.experiment.node.nodeDict()[self.textbox1.value].vartype, self.experiment.node.nodeDict()[self.textbox2.value].vartype
        if xType != 'categorical' and yType != 'categorical':
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                r = sp.pearsonr(self.experiment.data[group][self.textbox1.value], self.experiment.data[group][self.textbox2.value])
            text += group + ': ' + 'Correlation (r) is ' + '{0:#.3f}, '.format(r[0]) + 'P-value is ' + '{0:#.3g}'.format(r[1])
        return text

    def createTraces(self, x, y):
        traces = []
        annotations = []
        annotation_y = -0.20 - 0.02*len(self.experiment.group_names)
        traceType = self.choose_trace(x, y)
        if traceType == 'histogram':
            for group in self.experiment.group_names:
                data = self.experiment.data[group]
                if self.experiment.node.nodeDict()[x].vartype == 'categorical':
                    opacity = 1
                else:
                    opacity = 0.75
                traces += [go.Histogram(x=data[x], name=group, bingroup=1, opacity=opacity)]
                y = 'Count'
                barmode = 'overlay'
        elif traceType == 'scatter':
            for group in self.experiment.group_names:
                data = self.experiment.data[group]
                traces += [go.Scatter(x=data[x], y=data[y], mode='markers', opacity=0.75, name=group)]
                annotations += [dict(xref='paper',yref='paper',x=0.5, y=annotation_y, showarrow=False, text=self.display_values(group))]
                annotation_y += -0.05
                barmode = 'overlay'
        elif traceType == 'bar':
            for group in self.experiment.group_names:
                avg = self.experiment.data[group].groupby(x).agg('mean')
                std = self.experiment.data[group].groupby(x).agg('std')[y]
                traces += [go.Bar(x=list(avg.index), y=avg[y], name=group, error_y=dict(type='data', array=std))]
                annotations += [dict(xref='paper',yref='paper',x=0.5, y=annotation_y, showarrow=False, text=self.display_values(group))]
                annotation_y += -0.05
                barmode = 'group'
        elif traceType == 'barh':
            for group in self.experiment.group_names:
                avg = self.experiment.data[group].groupby(y).agg('mean')
                std = self.experiment.data[group].groupby(y).agg('std')[x]
                traces += [go.Bar(x=avg[x], y=list(avg.index), name=group, error_x=dict(type='data', array=std), orientation='h')]
                annotations += [dict(xref='paper',yref='paper',x=0.5, y=annotation_y, showarrow=False, text=self.display_values(group))]
                annotation_y += -0.05
                barmode = 'group'
        go_layout = go.Layout(title=dict(text=x if traceType == 'histogram' else x + " vs. " + y ),
                              barmode=barmode,
                              height=500+50,
                              width=800,
                              xaxis=dict(title=x), yaxis=dict(title=y),
                              annotations = annotations,
                              margin=dict(b=80+50, r=200, autoexpand=False))
        return traces, go_layout
    
    def initTraces(self):
        traces, layout = self.createTraces(self.x_options[0], self.y_options[0])
        self.g = go.FigureWidget(layout=layout)
        for t in traces:
            self.g.add_traces(t)
            
    def updateTraces(self):
        self.g.data = []
        traces, layout = self.createTraces(self.textbox1.value, self.textbox2.value)
        for t in traces:
            self.g.add_traces(t)
        self.g.layout.annotations = layout.annotations
        self.g.layout = layout

    def observe(self):
        self.textbox1.observe(self.response, names="value")
        self.textbox2.observe(self.response, names="value")
        self.button.observe(self.update_table, names='value')

    def choose_trace(self, x, y):
        if y == 'None (Distributions Only)':
            return 'histogram'
        xType, yType = self.experiment.node.nodeDict()[x].vartype, self.experiment.node.nodeDict()[y].vartype
        if xType != 'categorical' and yType != 'categorical':
            return 'scatter'
        elif xType == 'categorical' and yType != 'categorical':
            return 'bar'
        elif xType != 'categorical' and yType == 'categorical':
            return 'barh'
        else:
            return 'table'

    def pivot_table(self):
        if self.textbox1.value == self.textbox2.value:
            df = "Cannot create a pivot table with only one variable"
            return df
        if self.button.value == 'All':
            for group in self.experiment.group_names:
                df = pd.DataFrame()
                df = pd.concat([df, self.experiment.data[group]])
            df = df.groupby([self.textbox1.value, self.textbox2.value]).agg('count').reset_index().pivot(self.textbox1.value, self.textbox2.value, self.options[0])
        else:
            df = self.experiment.data[self.button.value].groupby([self.textbox1.value, self.textbox2.value]).agg('count').reset_index().pivot(self.textbox1.value, self.textbox2.value, self.options[0])
        return df

    def update_table(self, change):
        update_display(self.pivot_table(), display_id='1');
        self.button.layout.display = 'flex'

    def validate(self):
        return self.textbox1.value in self.x_options and self.textbox2.value in (self.x_options + ['None (Distributions Only)'])

    def response(self, change):
        if self.validate():
            traceType = self.choose_trace(self.textbox1.value, self.textbox2.value)
            with self.g.batch_update():
                if traceType == 'table':
                    self.g.update_layout({'height':10, 'width':10})
                    self.g.layout.xaxis.title = ""
                    self.g.layout.yaxis.title = ""
                    self.g.layout.title = ""
                    self.button.layout.display = 'flex'
                else:
                    self.updateTraces()
                    update_display(Nothing(), display_id='1')
                    self.button.layout.display = 'none'

class Nothing:
    def __init__(self):
        None
    def __repr__(self):
        return ""

def text2Array(text):
    text = text.replace(' ', '')
    if re.fullmatch(r'^((\d+)(|-(\d+)),)*(\d+)(|-(\d+))$', text) is None:
        return None
    matches = re.findall(r'((\d+)-(\d+))|(\d+)', text)
    ids = []
    for m in matches:
        if m[3] != '':
            ids = np.concatenate((ids, [int(m[3])-1])) # Subtract one because text starts at 1, array starts at 0
        else:
            if int(m[2]) < int(m[1]):
                return None
            else:
                ids = np.concatenate((ids, np.arange(int(m[1])-1, int(m[2]))))
    uniq = list(set(ids))
    uniq.sort()
    if len(ids) != len(uniq):
        return None
    return uniq

def array2Text(ids):
    ids.sort()
    ids = np.array(ids)+1 # Add one because text starts at 1, array starts at 0
    segments = []
    start = ids[0]
    end = ids[0]
    for j in range(len(ids)):
        if j == len(ids)-1:
            end = ids[j]
            s = str(start) if start == end else '%d-%d' % (start, end)
            segments.append(s)
        elif ids[j+1] != ids[j]+1:
            end = ids[j]
            s = str(start) if start == end else '%d-%d' % (start, end)
            segments.append(s)
            start = ids[j+1]
    return ','.join(segments)

def randomAssign(N, N_group):
    '''
    Randomly assigns N total items into N_group groups
    Returns a list of lists of ids
    '''
    arr = np.arange(N)
    np.random.shuffle(arr)
    result = []
    for i in range(N_group):
        start = i*N//N_group
        end = min((i+1)*N//N_group, N)
        result.append(arr[start:end])
    return result

# returns a placeholder CausalNode whose name can be used to look up the actual node in the CausalNetwork
def node(name):
    return PlaceholderNode(name=name)

# lifts a function that can takes any non-node inputs into an equivalent function in which the inputs can also be nodes. lift(func) cannot take arrays of nodes as input
def lift(func):
    def liftedFunc(*args, **kwargs): # should have exactly the same format as the unlifted func
        n_args = len(args) # number of positional arguments
        names = list(kwargs.keys())
        vargs = list(args) + [kwargs[name] for name in names] # combine all arguments into positional arguments
        nargs = [v for v in vargs if isinstance(v, CausalNode)] # only node type arguments
        node_inds = [i for i, v in enumerate(vargs) if isinstance(v, CausalNode)] # positions in vargs of all node type arguments. node_inds[i] = position of nargs[i] in vargs
        node_inds_lookup = {j: i for i, j in enumerate(node_inds)} # given index in vargs, gives the index in nargs
        def f(*fnargs): # wrapper of the unlifted func, only given non-node positional arguments that are intended to be lifted to node types. fnargs has the same structure as nargs
            # here reconstruct fargs and fkwargs from fnargs and the non-node type args/kwargs
            fargs = [(fnargs[node_inds_lookup[i]] if isinstance(v, CausalNode) else args[i]) for i, v in enumerate(args)] # replaces the nodes in args by the corresponding values in fnargs
            fkwargs = {n: (fnargs[node_inds_lookup[i+n_args]] if isinstance(kwargs[n], CausalNode) else kwargs[n]) for i, n in enumerate(names)} # replaces the nodes in fwargs by the corresponding values in fnargs
            return func(*fargs, **fkwargs)
        y = CausalNode()
        y.setCause(f, nargs)
        return y
    return liftedFunc

@lift
def identity(x):
    return x

@lift
def toInt(x):
    return round(x)

@lift
def bound(x, floor=-np.inf, ceil=np.inf):
    if floor > ceil:
        raise ValueError('floor has to be less than or equal to ceil.')
    if x<floor:
        return floor
    elif x<ceil:
        return x
    else:
        return ceil

@lift
def sum(*args):
    total = 0
    for a in args:
        total += a
    return total

@lift
def normal(mean, stdev):
    return np.random.normal(mean, stdev)

@lift
def poisson(rate):
    return np.random.poisson(lam=rate)

@lift
def dist(*args, point):
    return np.linalg.norm(np.array(args)-np.array(point))

@lift
def linear(*args, m=1, c=0, points=None):
    '''
    any number of arguments, m is either a number or an array of the same size as args
    points: array of n points, each with n entries, where n = len(args)+1, from which m and c can be solved
    '''
    if points:
        m, c = solveLinear(points)
    elif len(args) == 1:
        m = [m]
    return np.array(m)@np.array(args)+c

@lift
def categoricalLinear(x, category, params):
    m, c = params[category]
    return m*x+c

@lift
def choice(opts, weights=None, replace=True):
    if weights is None:
        chosen = np.random.choice(opts, replace=replace)
    else:
        weights = np.array(weights)
        p = weights/np.sum(weights)
        chosen = np.random.choice(opts, p=p, replace=replace)
    return chosen

# consider m@x+c=y where m,x are arrays and c,y are numbers. returns m,c given an array of points (x1, x2, ..., y)
def solveLinear(points):
    n = len(points)
    A = np.zeros((n, n))
    b = np.zeros(n)
    for i in range(n):
        A[i] = np.append(points[i][0:-1], 1)
        b[i] = points[i][-1]
    sol = np.linalg.solve(A, b)
    return sol[0:-1], sol[-1]

'''
truffula
'''
truffula = CausalNetwork()
truffula.addNode(ContinuousNode(name='Latitude', min=0, max=1000))
truffula.addNode(ContinuousNode(name='Longitude', min=0, max=1000))
truffula.addNode(ContinuousNode(name='Wind speed', min=0, max=40))
truffula.addCause(effect='Wind speed',
    cause=bound(
        normal(
            mean=linear(
                dist(node('Latitude'), node('Longitude'), point=(3000, 2000)), # absolute distance from a special point
                points=[(2236, 20), (3605, 3)] # mean speed=20 if close to point, =3 if far from point
                ),
            stdev=linear(
                dist(node('Latitude'), node('Longitude'), point=(1010, 800)),
                points=[(2236, 10), (3605, 3)] # stdev=10 if close to point, =3 if far from point
                )
            ),
        floor=0)
    )
truffula.addNode(CategoricalNode(name='Supplement', categories=['Water', 'Kombucha', 'Milk', 'Tea']))
truffula.addCause(effect='Supplement',
    cause=identity('Water')
    )
truffula.addNode(ContinuousNode(name='Fertilizer', min=0, max=20))
truffula.addCause(effect='Fertilizer',
    cause=normal(mean=10, stdev=2)
    )
truffula.addNode(ContinuousNode(name='Soil quality', min=0, max=100))
truffula.addCause(effect='Soil quality',
    cause=categoricalLinear(
        sum(
            linear(
                node('Fertilizer'),
                points=[(0, 10), (20, 100)] # quality=10 if 0 fertilizer, =100 if 20 fertilizer
                ),
            normal(mean=0, stdev=5) # formerly "fuzz"
            ),
        category=node('Supplement'),
        params={'Water': (1, 0), 'Kombucha': (0.6, -5), 'Milk': (1.2, 10), 'Tea': (0.7, 0)} # (m, c), m=multiplier, c=additive constant
        )
    )
truffula.addNode(DiscreteNode(name='Number of bees', min=0, max=300))
truffula.addCause(effect='Number of bees',
    cause=toInt(
        categoricalLinear(
            poisson( # rate depends on distance to bee hive and wind speed
                bound(
                    linear(
                        dist(node('Latitude'), node('Longitude'), point=(-50, 30)), # distance from bee hive
                        node('Wind speed'),
                        points=[(30, 0, 250), (1000, 30, 10), (30, 30, 40)] # rate=250 if close to hive with no wind, =10 if far from hive with high wind, =40 if close to hive with high wind
                        ),
                    floor=0
                    )
                ),
            category=node('Supplement'),
            params={'Water': (1, 0), 'Kombucha': (1.3, 0), 'Milk': (1, 0), 'Beer': (0.2, 0)} # effect of supplements on bees
            )
        )
    )
truffula.addNode(DiscreteNode(name='Number of fruits', min=0, max=100))
truffula.addCause(effect='Number of fruits',
    cause=poisson(
        bound(
            linear(
                node('Soil quality'), node('Number of bees'),
                points=[(0, 0, 0), (100, 200, 28), (100, 50, 16)] # rate=0 if poor soil and no bees, =28 if good soil and high bees, =16 if good soil and some bees
                ),
            floor=0
            )
        )
    )

'''
basketball
'''
basketball = CausalNetwork()
basketball.addNode(CategoricalNode(name='Shot type', categories=['Above head', 'Layup', 'Hook shot']))
basketball.addCause(effect='Shot type',
    cause=choice(opts=['Above head', 'Layup', 'Hook shot'], weights=[6, 3, 2])
    )
basketball.addNode(ContinuousNode(name='Hours practised per week', min=0, max=20))
basketball.addCause(effect='Hours practised per week',
    cause=choice(opts=np.linspace(0, 14, 30), weights=1/np.linspace(1, 15, 30))
    )
basketball.addNode(ContinuousNode(name='Height (cm)', min=140, max=200))
basketball.addNode(ContinuousNode(name='Ability', min=0, max=100))
basketball.addCause(effect='Ability',
    cause=linear(
        linear(
            node('Height (cm)'),
            points=[(150, 40), (190, 60)] # 40 if height=150, 60 if height=190
            ),
        m=1,
        c=linear(
            node('Hours practised per week'),
            points=[(0, 0), (10, 20)] # add 0 if hours=0, add 20 if hours=10
            )
        )
    )
basketball.addNode(ContinuousNode(name='Success rate', min=0, max=100))
basketball.addCause(effect='Success rate',
    cause=bound(
        categoricalLinear(
            node('Ability'),
            category=node('Shot type'),
            params={'Above head': (1, 0), 'Layup': (0.6, 0), 'Hook shot': (0.3, 0)}
            ),
        floor=0,
        ceil=100
        )
    )