"""
Phase 4C: Counterfactual Engine

Pearl Level 3 structural twin-network logic:
1. Abduction: Infer base probabilities from DAG
2. Action: Mutilate graph (do-calculus), severing parents
3. Prediction: Forward-propagate to see effects of intervention.

Operates on an ephemeral in-memory NetworkX DiGraph extracted from SQLite.
"""
from typing import Dict, Any
import math
import networkx as nx
from pulse.core.brain_db import get_conn

def build_shadow_graph(focal_node: str, radius: int = 3) -> nx.DiGraph:
    """Builds an ephemeral NetworkX DiGraph around a focal node."""
    conn = get_conn()
    G = nx.DiGraph()
    
    # Recursive CTE to get the local causal neighborhood (with path tracking to prevent cycles)
    query = """
    WITH RECURSIVE
      neighborhood(node_id, hops, path) AS (
        SELECT ? as node_id, 0 as hops, ? as path
        UNION ALL
        SELECT e.to_node, w.hops + 1, w.path || ',' || e.to_node
        FROM neighborhood w
        JOIN graph_edges e ON w.node_id = e.from_node
        WHERE w.hops < ? AND e.edge_type IN ('CAUSES', 'PREVENTS')
          AND INSTR(w.path, e.to_node) = 0
        UNION ALL
        SELECT e.from_node, w.hops + 1, w.path || ',' || e.from_node
        FROM neighborhood w
        JOIN graph_edges e ON w.node_id = e.to_node
        WHERE w.hops < ? AND e.edge_type IN ('CAUSES', 'PREVENTS')
          AND INSTR(w.path, e.from_node) = 0
      )
    SELECT DISTINCT node_id FROM neighborhood;
    """
    
    with conn:
        res = conn.execute(query, (focal_node, focal_node, radius, radius)).fetchall()
        node_ids = [r["node_id"] for r in res]
        
        if not node_ids:
            return G
            
        placeholders = ",".join("?" * len(node_ids))
        edges_query = f"""
        SELECT from_node, to_node, edge_type, confidence 
        FROM graph_edges 
        WHERE from_node IN ({placeholders}) AND to_node IN ({placeholders})
          AND edge_type IN ('CAUSES', 'PREVENTS')
        """
        
        edges = conn.execute(edges_query, tuple(node_ids) + tuple(node_ids)).fetchall()
        
        for e in edges:
            weight = e["confidence"] if e["edge_type"] == "CAUSES" else -e["confidence"]
            G.add_edge(e["from_node"], e["to_node"], weight=weight)
            
    return G

def simulate_intervention(target_node: str, do_variable: str, do_value: float, radius: int = 3) -> Dict[str, Any]:
    """
    Evaluates P(target_node | do(do_variable = do_value)).
    Returns a dictionary of expected values for nodes in the downstream path.
    """
    # 1. Abduction (Get the local causal graph)
    G = build_shadow_graph(do_variable, radius=radius)
    if do_variable not in G or target_node not in G:
        return {"error": "Nodes not found or disconnected within radius"}
        
    if not nx.has_path(G, do_variable, target_node):
        return {"error": f"No causal path from {do_variable} to {target_node}"}
        
    # 2. Action (Mutilation - sever all incoming edges to do_variable)
    in_edges = list(G.in_edges(do_variable))
    G.remove_edges_from(in_edges)
    
    # 3. Prediction (Forward Propagation)
    # Start all node values at 0.0, except the intervened variable
    node_values = {n: 0.0 for n in G.nodes()}
    node_values[do_variable] = do_value
    
    # Topologically sort the nodes downstream of do_variable
    descendants = nx.descendants(G, do_variable)
    sub_G = G.subgraph([do_variable] + list(descendants))
    
    try:
        top_order = list(nx.topological_sort(sub_G))
    except nx.NetworkXUnfeasible:
        # Fallback if somehow there's a cycle despite DB checks
        return {"error": "Cycle detected in shadow graph"}
        
    for node in top_order:
        if node == do_variable:
            continue
            
        # Value is sum of (parent_value * edge_weight)
        val = 0.0
        for parent in sub_G.predecessors(node):
            weight = sub_G[parent][node]['weight']
            val += node_values[parent] * weight
            
        # Squash to [-1, 1]
        node_values[node] = max(-1.0, min(1.0, val))
        
    return {
        "intervention": {do_variable: do_value},
        "target_node": target_node,
        "predicted_target_value": node_values[target_node],
        "full_cascade": node_values
    }
