from bamt.networks import DiscreteBN
from sklearn.preprocessing import LabelEncoder as le
from pgmpy.estimators import K2

p = Preprocessor([('encoder',le())])
discretized_data, est = p.apply(tdf)
info = p.info

bn = DiscreteBN()
bn.add_nodes(info) # add nodes from info obtained from preprocessing
bn.get_info() # to make sure that the network recognizes the variables as discrete


edges= [('label',x) for x in tdf.columns[1:]]

params = {
            # Defines initial nodes of the network, list of node names
           'init_nodes':['label'],
            # Defines initial edges of the network, list of tuples (node1, node2)
           'init_edges':edges,
            # Strictly set edges where algoritm must learn, list of tuples (node1, node2)
           'white_list':edges,
           # blacklist edges, list of tuples (node1, node2)
           'bl_add': [],
            # Allow algorithm to remove edges defined by user, bool
           'remove_init_edges':False
          }

# Structure learning using K2Score and parameters defined above
bn.add_edges(discretized_data, scoring_function=('K2', K2), params=params)

def create_pairs_itertools(string_array):
    return list(itertools.combinations(string_array, 2))

pos_edges = create_pairs_itertools(tdf.columns[1:])

params = {
            # Defines initial nodes of the network, list of node names
           'init_nodes':['label'],
            # Defines initial edges of the network, list of tuples (node1, node2)
           'init_edges':edges,
            # Strictly set edges where algoritm must learn, list of tuples (node1, node2)
           'white_list':pos_edges,
           # blacklist edges, list of tuples (node1, node2)
           'bl_add': [],
            # Allow algorithm to remove edges defined by user, bool
           'remove_init_edges':False
          }

bn.add_edges(discretized_data, scoring_function=('K2', K2), params=params)

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[54], line 19
      5 params = {
      6             # Defines initial nodes of the network, list of node names
      7            'init_nodes':['label'],
   (...)     15            'remove_init_edges':False
     16           }
     18 # Structure learning using K2Score and parameters defined above
---> 19 bn.add_edges(discretized_data, scoring_function=('K2', K2), params=params)
     21 bn.plot('foo.html') # Plot the network, save it to foo.html, NOT rendered in notebook

File ~/anaconda3/envs/jupyter_env/lib/python3.11/site-packages/bamt/networks/base.py:200, in BaseNetwork.add_edges(self, data, scoring_function, progress_bar, classifier, regressor, params, optimizer, **kwargs)
    196     return None
    198 self.sf_name = scoring_function[0]
--> 200 worker.build(
    201     data=data,
    202     params=params,
    203     classifier=classifier,
    204     regressor=regressor,
    205     progress_bar=progress_bar,
    206     **kwargs,
    207 )
    209 # update family
    210 self.nodes = worker.skeleton["V"]

File ~/anaconda3/envs/jupyter_env/lib/python3.11/site-packages/bamt/builders/hc_builder.py:201, in HCStructureBuilder.build(self, data, progress_bar, classifier, regressor, params, **kwargs)
    199 self.restrict(data, init_nodes, bl_add)
    200 if self.scoring_function[0] == "K2":
--> 201     self.apply_K2(data=data, progress_bar=progress_bar, **self.params)
    202 elif self.scoring_function[0] in ["MI", "LL", "BIC", "AIC"]:
    203     self.apply_group1(data=data, progress_bar=progress_bar, **self.params)

File ~/anaconda3/envs/jupyter_env/lib/python3.11/site-packages/bamt/builders/hc_builder.py:60, in HillClimbDefiner.apply_K2(self, data, init_edges, progress_bar, remove_init_edges, white_list)
     58 if white_list:
     59     all_possible_edges = list(permutations(self.vertices, 2))
---> 60     forbidden_edges |= set(all_possible_edges) - set(white_list)
     62 if not init_edges:
     63     expert_knowledge = ExpertKnowledge(forbidden_edges=forbidden_edges)

TypeError: unhashable type: 'DiscreteNode'

params = {
            # Defines initial nodes of the network, list of node names
           'init_nodes':['label'],
            # Defines initial edges of the network, list of tuples (node1, node2)
           'init_edges':edges,
            # Strictly set edges where algoritm must learn, list of tuples (node1, node2)
           'white_list':None,
           # blacklist edges, list of tuples (node1, node2)
           'bl_add': [],
            # Allow algorithm to remove edges defined by user, bool
           'remove_init_edges':False
          }