%matplotlib inline
import matplotlib.pyplot as plt

from sklearn.datasets import load_breast_cancer
import numpy as np
from functools import reduce

# Import our custom utilities
from imp import reload
from utils import irf_jupyter_utils
from utils import irf_utils
reload(irf_jupyter_utils)
reload(irf_utils)

"""
Explanation: Key Requirements for the iRF scikit-learn implementation

The following is a documentation of the main requirements for the iRF implementation

Typical Setup
Import the required dependencies

In particular irf_utils and irf_jupyter_utils
End of explanation
"""


load_breast_cancer = load_breast_cancer()

X_train, X_test, y_train, y_test, rf = irf_jupyter_utils.generate_rf_example(n_estimators=10)

"""
Explanation: Step 1: Fit the Initial Random Forest

Just fit every feature with equal weights per the usual random forest code e.g. DecisionForestClassifier in scikit-learn
End of explanation
"""


print("Training feature dimensions", X_train.shape, sep = ":\n")
print("\n")
print("Training outcome dimensions", y_train.shape, sep = ":\n")
print("\n")
print("Test feature dimensions", X_test.shape, sep = ":\n")
print("\n")
print("Test outcome dimensions", y_test.shape, sep = ":\n")
print("\n")
print("first 5 rows of the training set features", X_train[:2], sep = ":\n")
print("\n")
print("first 5 rows of the training set outcomes", y_train[:2], sep = ":\n")

"""
Explanation: Check out the data
End of explanation
"""


all_rf_tree_data = irf_utils.get_rf_tree_data(rf=rf,
                                              X_train=X_train, y_train=y_train, 
                                              X_test=X_test, y_test=y_test)
#all_rf_tree_data

"""
Explanation: Step 2: Get all Random Forest and Decision Tree Data

Extract in a single dictionary the random forest data and for all of it's decision trees
This is as required for RIT purposes
End of explanation
"""


np.random.seed(12)
gen_random_leaf_paths = irf_utils.generate_rit_samples(all_rf_tree_data=all_rf_tree_data, 
                                                       bin_class_type=1)

#for i in range(7):
#    print(next(gen_random_leaf_paths))

rit0 = irf_utils.build_tree(feature_paths=gen_random_leaf_paths, 
                            max_depth=3, 
                            noisy_split=False, 
                            num_splits=2)

"""
Explanation: STEP 3: Get the RIT data and produce RITs
End of explanation
"""


for node in rit0.traverse_depth_first():
    print(node[0], node[1]._val)

"""
Explanation: Print out all of the intersected nodes and their ids
End of explanation
"""


for idx, node in enumerate(rit0.leaf_nodes()):
    print(idx, node[1]._val)

print("Root:\n", rit0._val)
#print("Some child:\n", tree.children[0].children[1]._val)

len(rit0)

# If noisy split is False, this should pass
#assert(len(rit0) == 1 + 5 + 5**2)
assert(len(rit0) <= 1 + 5 + 5**2)

M = 10
np.random.seed(12)
# Create the weighted randomly sampled paths as a generator
gen_random_leaf_paths = irf_utils.generate_rit_samples(all_rf_tree_data=all_rf_tree_data, 
                                                       bin_class_type=1)
# Create the RIT object
rit = irf_utils.build_tree(feature_paths=gen_random_leaf_paths, 
                            max_depth=3, 
                            noisy_split=False, 
                            num_splits=2)
# Get the intersected node values
rit_intersected_values = [node[1]._val for node in rit.traverse_depth_first()]
# Leaf node values i.e. final intersected features
rit_leaf_node_values = [node[1]._val for node in rit.leaf_nodes()]
rit_leaf_node_union_value = reduce(np.union1d, rit_leaf_node_values)
rit_outputs = {"rit": rit,
                   "rit_intersected_values": rit_intersected_values,
                   "rit_leaf_node_values": rit_leaf_node_values,                   
                   "rit_leaf_node_union_value": rit_leaf_node_union_value}
rit_outputs['rit_intersected_values']

def get_rit_tree_data(all_rf_tree_data,
                      bin_class_type=1,
                      random_state=12, 
                      #M=10, # number of trees (RIT) to build 
                      feature_paths=gen_random_leaf_paths, 
                      max_depth=3, 
                      noisy_split=False, 
                      num_splits=2):
    """
    A wrapper for the Random Intersection Trees (RIT) algorithm
    """
    # Set the random seed for reproducibility
    np.random.seed(12)    
    
    # Create the weighted randomly sampled paths as a generator
    gen_random_leaf_paths = irf_utils.generate_rit_samples(
        all_rf_tree_data=all_rf_tree_data, 
        bin_class_type=bin_class_type)
    
    # Create the RIT object
    rit = irf_utils.build_tree(feature_paths=gen_random_leaf_paths, 
                                max_depth=max_depth, 
                                noisy_split=noisy_split, 
                                num_splits=num_splits)
    
    # Get the intersected node values
    # CHECK remove this for the final value
    rit_intersected_values = [node[1]._val for node in rit.traverse_depth_first()]
    # Leaf node values i.e. final intersected features
    rit_leaf_node_values = [node[1]._val for node in rit.leaf_nodes()]
    rit_leaf_node_union_value = reduce(np.union1d, rit_leaf_node_values)
    rit_outputs = {"rit": rit,
                   "rit_intersected_values": rit_intersected_values,
                   "rit_leaf_node_values": rit_leaf_node_values,                   
                   "rit_leaf_node_union_value": rit_leaf_node_union_value}
    return rit_outputs

a = get_rit_tree_data(all_rf_tree_data=all_rf_tree_data,
                      bin_class_type=1,
                      random_state=12, 
                      #M=10, # number of trees (RIT) to build
                      max_depth=3, 
                      noisy_split=False, 
                      num_splits=2)

a['rit_intersected_values']

b = irf_utils.get_rit_tree_data(all_rf_tree_data=all_rf_tree_data,
                      bin_class_type=1,
                      random_state=12, 
                      M=1,
                      max_depth=3, 
                      noisy_split=False, 
                      num_splits=2)

b['rit0']['rit_intersected_values']

c = irf_utils.get_rit_tree_data(all_rf_tree_data=all_rf_tree_data,
                      bin_class_type=1,
                      random_state=12, 
                      M=10,
                      max_depth=3, 
                      noisy_split=False, 
                      num_splits=2)

c['rit1']

"""
Explanation: Print out all of the leaf node values
End of explanation
"""


# Print the feature ranking
print("Feature ranking:")

feature_importances_rank_idx = all_rf_tree_data['feature_importances_rank_idx']
feature_importances = all_rf_tree_data['feature_importances']

for f in range(X_train.shape[1]):
    print("%d. feature %d (%f)" % (f + 1
                                   , feature_importances_rank_idx[f]
                                   , feature_importances[feature_importances_rank_idx[f]]))

"""
Explanation: Perform Manual CHECKS on the irf_utils

These should be converted to unit tests and checked with nosetests -v test_irf_utils.py

Step 4: Plot some Data
List Ranked Feature Importances
End of explanation
"""


# Plot the feature importances of the forest
feature_importances_std = all_rf_tree_data['feature_importances_std']

plt.figure()
plt.title("Feature importances")
plt.bar(range(X_train.shape[1])
        , feature_importances[feature_importances_rank_idx]
        , color="r"
        , yerr = feature_importances_std[feature_importances_rank_idx], align="center")
plt.xticks(range(X_train.shape[1]), feature_importances_rank_idx)
plt.xlim([-1, X_train.shape[1]])
plt.show()

"""
Explanation: Plot Ranked Feature Importances
End of explanation
"""


uniq_feature_paths = all_rf_tree_data['dtree0']['all_uniq_leaf_paths_features']
leaf_node_classes  = all_rf_tree_data['dtree0']['all_leaf_node_classes']
ones_only = [i for i, j in zip(uniq_feature_paths, leaf_node_classes) 
               if j == 1]
ones_only

"""
Explanation: Manually construct a binary RIT
Filter list unique feature paths where leaf node values == 1

Get the unique feature paths where the leaf node predicted class is just 1
We are just going to get it from the first decision tree for this test case
End of explanation
"""


ones_only_seven = ones_only[-7:]
ones_only_seven

"""
Explanation: Manually extract the last seven values

Just pick the last seven cases, we are going to manually construct
we are going to build a binary RIT of depth 3 i.e. max 2**3 -1 = 7 intersecting nodes
End of explanation
"""


# Construct a binary version of the RIT manually!
# This should come in useful for unit tests!
node0 = ones_only_seven[0]
node1 = np.intersect1d(node0, ones_only_seven[1])
node2 = np.intersect1d(node1, ones_only_seven[2])
node3 = np.intersect1d(node1, ones_only_seven[3])
node4 = np.intersect1d(node0, ones_only_seven[4])
node5 = np.intersect1d(node4, ones_only_seven[5])
node6 = np.intersect1d(node4, ones_only_seven[6])

intersected_nodes_seven = [node0, node1, node2, node3, node4, node5, node6]

for idx, node in enumerate(intersected_nodes_seven):
    print("node" + str(idx), node)

rit_output = reduce(np.union1d, (node2, node3, node5, node6))
rit_output

"""
Explanation: Manually leaf node 1's paths as a generator
Manually create the binary RIT
End of explanation
"""


ones_only_seven_gen = (n for n in ones_only_seven)
#for i in range(len(ones_only_seven)):
#    print(next(ones_only_seven_gen))

rit_man0 = irf_utils.build_tree(
    feature_paths=ones_only_seven_gen,
    max_depth=3,
    noisy_split=False,
    num_splits=2)

print("Root:\n", rit_man0._val)

print("Root:\n", rit_man0._val)
print("node1:\n", rit_man0.children[0]._val)
print("node4:\n", rit_man0.children[1]._val)
print("node2:\n", rit_man0.children[0].children[0]._val)
print("node3:\n", rit_man0.children[0].children[1]._val)
print("node5:\n", rit_man0.children[1].children[0]._val)
print("node6:\n", rit_man0.children[1].children[1]._val)

for node in rit_man0.traverse_depth_first():
    print(node[1]._val)

for node in rit_man0.leaf_nodes():
    print(node[1]._val)

"""
Explanation: Create a binary RIT without the randomized split option
End of explanation
"""


# Now plot the trees individually
irf_jupyter_utils.draw_tree(decision_tree = all_rf_tree_data['rf_obj'].estimators_[0])

"""
Explanation: Check Output of Decision Tree Extracted Data
Decision Tree 0 (First) - Get output
Check the output against the decision tree graph
End of explanation
"""


irf_jupyter_utils.pretty_print_dict(inp_dict = all_rf_tree_data['dtree0'])

# Count the number of samples passing through the leaf nodes
sum(all_rf_tree_data['dtree0']['tot_leaf_node_values'])

"""
Explanation: Compare to our dict of extracted data from the tree
End of explanation
"""


irf_jupyter_utils.pretty_print_dict(inp_dict = all_rf_tree_data['dtree0']['all_leaf_paths_features'])

"""
Explanation: Check output against the diagram
End of explanation
"""


