
import jax 
import jax.numpy as jnp

from functools import partial


@jax.jit
def find_ancestors_jax(mask, node):
    """Find ancestors of a node in a graph.

    Args:
        mask (Array): Adjacency matrix of a directed graph.
        node (int): Node of interest.

    Returns:
        _type_: _description_
    """
    num_nodes = mask.shape[0]
    is_ancestor = jnp.zeros(num_nodes, dtype=jnp.bool_)
    stack = jnp.empty(num_nodes, dtype=jnp.int32)
    stack = stack.at[0].set(node)
    
    def body_fn(carry, i):
        is_ancestor, stack = carry
        current_node = stack[i]
        current_parents = mask[current_node, :]
        
        def inner_body_fn(carry, j):
            is_ancestor, stack = carry
            value = current_parents[j]
            cond = value & (j != current_node) & (~is_ancestor[j])
            
            def true_fn(is_ancestor, stack):
                is_ancestor = is_ancestor.at[j].set(True)
                stack = stack.at[i+1].set(j)
                return is_ancestor, stack
            def false_fn(is_ancestor, stack):
                return is_ancestor, stack
            
            is_ancestor, stack = jax.lax.cond(cond, true_fn, false_fn, is_ancestor, stack)
            return (is_ancestor, stack), None
        
        (is_ancestor, stack), _ = jax.lax.scan(inner_body_fn, (is_ancestor, stack), jnp.arange(num_nodes))
        return (is_ancestor, stack), None
    
    (is_ancestor, stack), _ = jax.lax.scan(body_fn, (is_ancestor, stack), jnp.arange(num_nodes))
    

    return is_ancestor



@partial(jax.jit, static_argnums=(2,))
def faithfull_mask(base_mask, condition_mask, conditioned_nodes="unchanged"):
    """ Faithfull mask update for conditioning"""
    
    graph = base_mask.astype(jnp.bool_).copy()
    base_mask = base_mask.astype(jnp.bool_) # Rows are paraents, columns are children
    condition_mask = condition_mask.astype(jnp.bool_)
    num_nodes = base_mask.shape[0]
    
    def body_fn(carry, i):
        base_mask, condition_mask = carry
        
        def condition_case(base_mask, condition_mask):
            # We need to update all ancestors of i
            is_ancestor = find_ancestors_jax(graph, i)
            is_ancestor = is_ancestor & (~condition_mask)
            all_ancestors = jnp.nonzero(is_ancestor, size=num_nodes, fill_value=i)[0]
            # They will now depend on i
            base_mask = base_mask.at[all_ancestors,i].set(True)
            # They will now depend on each other!
            base_mask = base_mask | (is_ancestor[:,None] & is_ancestor[None,:])
            # The parents of all children of i will now depend on each other
            children_of_i = base_mask[:,i]
            parents_of_children_of_i = base_mask & (children_of_i[:,None] & ~children_of_i[None,:])
            parents_of_children_of_i = jnp.any(parents_of_children_of_i, axis=0)

            return base_mask, condition_mask
        
        def uncondition_case(base_mask, condition_mask):
            return base_mask, condition_mask
        
        base_mask, condition_mask = jax.lax.cond(condition_mask[i], condition_case, uncondition_case, base_mask, condition_mask)
        

        return (base_mask, condition_mask), None

    (base_mask, condition_mask), _ = jax.lax.scan(body_fn, (base_mask, condition_mask), jnp.arange(num_nodes))
        
    
    return base_mask



@partial(jax.jit, static_argnums=(2,3))
def min_faithfull_mask(mask, condition_mask, top_mode=0, conditioned_nodes="unchanged"):
    """ Minimally faithfull mask update for conditioning"""
    num_nodes = mask.shape[0]
    I = moralize(mask)
    H = jnp.zeros_like(mask, dtype=jnp.bool_)
    # 0 is child, 1 is parent
    UPSTREAM = top_mode
    DOWNSTREAM = 1 - top_mode
    num_parents_or_childs = jnp.sum(mask & (~condition_mask[None, :] & ~condition_mask[:, None]), axis=UPSTREAM)
    #print(num_parents_or_childs)
    S = (num_parents_or_childs == 1) & (~condition_mask) # Frontier set
    M = jnp.zeros((num_nodes), dtype=jnp.bool_) # Marked nodes

    def cond_fn(val):
        S, _, _, _ = val
        return jnp.any(S) 
    
    def body_fn(val):
        S, M, I, H = val
        #print(S)
        # Find the node with the fewest edges added
        v = min_fill_heuristic(mask, I, S,M, top_mode)
        # print("Frontal set: ",S)
        # print("Marked: ", M)
        # print("Selected: ",v)
        # Add edge in I between unmarked neighbours in I 
        neighbours_v = I[v,:] & (~M)
        I = I | (neighbours_v[:,None] & neighbours_v[None,:])
        # Make unmarked neighbours of v, the parents of v in H
        H = H.at[v,:].set(neighbours_v)
        # Remove v from S and mark it
        S = S.at [v].set(False)
        M = M.at[v].set(True)
        
        if top_mode == 1:
            u = mask[:,v] & (~M) # Not marked children
            upstream_u = mask & (u[:, None]  & ~u[None,:]) # Parents of not marked children
            all_upstream_u_marked = ~jnp.any(upstream_u & ~M, axis=1)
        else:
            u = mask[v,:] & (~M) # Not marked parents
            upstream_u = mask & (u[None,:]  & ~u[:, None]) # Children of not marked parents
            all_upstream_u_marked = ~jnp.any(upstream_u & ~M, axis=1)
        
        

        S = S | (u & all_upstream_u_marked)
        S = S & (~condition_mask)

        
        return S, M, I, H
    
    _,_,_, H = jax.lax.while_loop(cond_fn, body_fn, (S, M, I, H))
    H = H | jnp.eye(num_nodes, dtype=jnp.bool_)
    H = jax.lax.cond(jnp.any(condition_mask), lambda x: x, lambda x: mask, H)
    
    # Conditioned nodes will keep the unconditional edges, hence each row of H where condition_mask is true should be equal to "mask"
    if conditioned_nodes == "unchanged":
        H = H & ~condition_mask[:, None] | mask & condition_mask[:, None]
    elif conditioned_nodes == "removed":
        H = H & ~condition_mask[:, None]
    elif conditioned_nodes == "added":
        H = H | condition_mask[:, None]
    
    return H
    
                
        
    
    
@partial(jax.jit, static_argnums=(4,))
def min_fill_heuristic(G, I, S, M, top_mode=0):
    """ Min-fill heuristic for finding a node to eliminate"""
    
    # 0 is child, 1 is parent
    UPSTREAM = top_mode
    DOWNSTREAM = 1 - top_mode
    
    # Find the number of edges that would be added if we eliminated each node
    num_edges_added = I.sum(axis=DOWNSTREAM)
    num_edges_added = S * num_edges_added + (~S) * (I.shape[0] + 1)
    # Find the node that would add the fewest edges
    # Additional constraint: Prefer marked parents
    #print(num_edges_added)
    min_val = jnp.min(num_edges_added)
    marked_parents = -jnp.sum(M[None,:] & G, axis=DOWNSTREAM)
    num_parents= marked_parents * (num_edges_added == min_val) + (I.shape[0] + 1) * (num_edges_added != min_val)
    #node_to_eliminate = jnp.argmin(num_edges_added + num_parents)
    #print(num_parents)
    reversed_array = (num_parents)[::-1]
    index = jnp.argmin(reversed_array)
    node_to_eliminate = len(reversed_array) - 1 - index
    

    return node_to_eliminate 


@jax.jit
def moralize(adj_matrix):
    adj_matrix = adj_matrix.astype(jnp.bool_)
    
    # Make the graph undirected
    undirected_graph = adj_matrix | adj_matrix.T
    
    # Add edges between parents
    undirected_graph = undirected_graph | (adj_matrix.T @ adj_matrix)
    
    return undirected_graph


def minimally_faithfull_mask(mask, condition_mask):
    """ Minimally faithfull mask update for conditioning"""
    I = moralize(mask)
    H = jnp.zeros_like(mask, dtype=jnp.bool_)
    