from gs_quant.session import Environment, GsSession
from gs_quant.entities.tree_entity import TreeHelper

# external users should substitute their client id and secret; please skip this step if using internal jupyterhub
GsSession.use(Environment.PROD, client_id=None, client_secret=None, scopes=('read_product_data',))

# Initialise the TreeHelper with the MarqueeID of the root Asset and the dataset to be used to build the tree.
# In this example, we are building an STS index tree, so we use STS index datasets.

tree = TreeHelper('MAM9TYYXXXXXXXXX', tree_underlier_dataset='STS_UNDERLIER_WEIGHTS', underlier_column='underlyingAssetId')

# Builds the tree with the underlier dataset passed 
tree.build_tree() 

# Populate weights and attributions of the asset with the appropriate dataset and column name. 
# In this example, we use the datasets for STS indices.

tree.populate_weights(dataset='STS_UNDERLIER_WEIGHTS', weight_column='weight')

tree.populate_attribution(dataset='STS_UNDERLIER_ATTRIBUTION', attribution_column='absoluteAttribution')

# Get tree as a df with the contents stored at this point
tree.to_frame()

# Prints the tree structure formed by the Index for easy visualisation
tree.get_visualisation(visualise_by='asset_name')

# This feature needs treelib to be installed.
tree.get_visualisation(visualise_by='bbid')

# Get the time at which the tree was built
tree.update_time

# This feature needs the notebook package of gs-quant to be installed.
# You can install it using 'pip install gs-quant[notebook]'
tree.get_visualisation(visualise_by='bbid')

# Get the time at which the tree was built
tree.update_time

# Accessing the root node of the tree
root_node = tree.root

# Accessing the child nodes of the root node.
child_nodes_depth_1 = root_node.direct_underlier_assets_as_nodes

# Accessing the child nodes of the first child node of the root.
child_nodes_depth_2 = child_nodes_depth_1[0].direct_underlier_assets_as_nodes