"""NL→DAG Parser — extract causal DAGs from free-form natural language text.

Four LLM backends:
  1. Claude API (ANTHROPIC_API_KEY) — best accuracy
  2. Ollama local (if running) — free/offline
  3. Local Qwen LoRA (models/nl_dag_adapter) — fine-tuned, no API, offline
  4. Regex fallback — no LLM, pattern matching only

Bridge function converts parsed DAGs to CausalEng input tensors.
"""

import json
import os
import re
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple


# ============================================================
# Data Classes
# ============================================================

@dataclass
class CausalVariable:
    """A variable in a causal DAG."""
    name: str
    index: int
    aliases: List[str] = field(default_factory=list)


@dataclass
class CausalDAG:
    """A causal DAG extracted from text."""
    variables: List[CausalVariable]
    edges: List[Tuple[int, int]]  # (src_idx, dst_idx)
    name_to_idx: Dict[str, int]
    source_text: str
    # Optional enrichments from LLM (e.g. qwen2.5:14b or Claude)
    bidirected_names: List[Tuple[str, str]] = field(default_factory=list)  # hidden confounders
    params: Dict[str, any] = field(default_factory=dict)  # CPT tables if extracted
    suggested_query: Optional[Dict] = None  # {type, treatment, outcome} if LLM found a question

    @property
    def num_nodes(self):
        return len(self.variables)

    @property
    def num_edges(self):
        return len(self.edges)

    def variable_names(self):
        return [v.name for v in self.variables]


@dataclass
class CausalQuery:
    """A causal query parsed against a DAG."""
    query_type: str  # "causal", "independence", "effect"
    x_name: str
    y_name: str
    z_names: List[str] = field(default_factory=list)
    x_idx: int = -1
    y_idx: int = -1
    z_indices: List[int] = field(default_factory=list)


# ============================================================
# Regex Causal Patterns (no LLM needed)
# ============================================================

# Forward patterns: "X causes Y", "X leads to Y", etc.
_FORWARD_PATTERNS = [
    re.compile(r"([A-Za-z][\w\s'-]*?)\s+causes\s+([A-Za-z][\w\s'-]*?)(?:\.|,|;|$)", re.IGNORECASE),
    re.compile(r"([A-Za-z][\w\s'-]*?)\s+leads\s+to\s+([A-Za-z][\w\s'-]*?)(?:\.|,|;|$)", re.IGNORECASE),
    re.compile(r"([A-Za-z][\w\s'-]*?)\s+affects\s+([A-Za-z][\w\s'-]*?)(?:\.|,|;|$)", re.IGNORECASE),
    re.compile(r"([A-Za-z][\w\s'-]*?)\s+influences\s+([A-Za-z][\w\s'-]*?)(?:\.|,|;|$)", re.IGNORECASE),
    re.compile(r"([A-Za-z][\w\s'-]*?)\s+has a direct effect on\s+([A-Za-z][\w\s'-]*?)(?:\.|,|;|$)", re.IGNORECASE),
    re.compile(r"([A-Za-z][\w\s'-]*?)\s+directly affects\s+([A-Za-z][\w\s'-]*?)(?:\.|,|;|$)", re.IGNORECASE),
    re.compile(r"([A-Za-z][\w\s'-]*?)\s+results in\s+([A-Za-z][\w\s'-]*?)(?:\.|,|;|$)", re.IGNORECASE),
    re.compile(r"([A-Za-z][\w\s'-]*?)\s+increases\s+([A-Za-z][\w\s'-]*?)(?:\.|,|;|$)", re.IGNORECASE),
    re.compile(r"([A-Za-z][\w\s'-]*?)\s+decreases\s+([A-Za-z][\w\s'-]*?)(?:\.|,|;|$)", re.IGNORECASE),
    re.compile(r"([A-Za-z][\w\s'-]*?)\s+produces\s+([A-Za-z][\w\s'-]*?)(?:\.|,|;|$)", re.IGNORECASE),
]

# Reverse patterns: "Y depends on X" → edge X→Y
_REVERSE_PATTERNS = [
    re.compile(r"([A-Za-z][\w\s'-]*?)\s+depends on\s+([A-Za-z][\w\s'-]*?)(?:\.|,|;|$)", re.IGNORECASE),
    re.compile(r"([A-Za-z][\w\s'-]*?)\s+is caused by\s+([A-Za-z][\w\s'-]*?)(?:\.|,|;|$)", re.IGNORECASE),
    re.compile(r"([A-Za-z][\w\s'-]*?)\s+is affected by\s+([A-Za-z][\w\s'-]*?)(?:\.|,|;|$)", re.IGNORECASE),
    re.compile(r"([A-Za-z][\w\s'-]*?)\s+is influenced by\s+([A-Za-z][\w\s'-]*?)(?:\.|,|;|$)", re.IGNORECASE),
    re.compile(r"([A-Za-z][\w\s'-]*?)\s+is determined by\s+([A-Za-z][\w\s'-]*?)(?:\.|,|;|$)", re.IGNORECASE),
]

# Query patterns
_QUERY_CAUSAL = re.compile(
    r"does\s+([A-Za-z][\w\s'-]*?)\s+cause\s+([A-Za-z][\w\s'-]*?)\s*\?",
    re.IGNORECASE,
)
_QUERY_EFFECT = re.compile(
    r"what\s+is\s+the\s+effect\s+of\s+([A-Za-z][\w\s'-]*?)\s+on\s+([A-Za-z][\w\s'-]*?)\s*\?",
    re.IGNORECASE,
)
_QUERY_INDEPENDENCE = re.compile(
    r"is\s+([A-Za-z][\w\s'-]*?)\s+independent\s+of\s+([A-Za-z][\w\s'-]*?)\s+given\s+([A-Za-z][\w\s'-]*?)\s*\?",
    re.IGNORECASE,
)
_QUERY_INDEPENDENCE_NO_Z = re.compile(
    r"is\s+([A-Za-z][\w\s'-]*?)\s+independent\s+of\s+([A-Za-z][\w\s'-]*?)\s*\?",
    re.IGNORECASE,
)


# ============================================================
# LLM Prompt
# ============================================================

_LOCAL_SYSTEM_MSG = (
    "You are a causal graph extractor. Convert natural language descriptions of "
    "causal relationships into structured JSON with nodes, edges, and a query. "
    "Respond only with valid JSON."
)

_OLLAMA_SYSTEM_MSG = (
    "You are a causal inference expert. Extract structured causal models from "
    "natural language. Always respond with valid JSON only — no explanation, "
    "no markdown fences, no extra text."
)

# Used for claude backend (plain prompt string)
_LLM_PROMPT = """Extract the causal variables and directed edges from this text.
Return ONLY valid JSON in this exact format:
{{"variables": ["var1", "var2", ...], "edges": [["source", "target"], ...]}}

Rules:
- Variables are nouns/concepts mentioned as causes or effects
- Edges go from cause to effect (source → target)
- Use exact variable names from the text
- Do not add edges that aren't stated or implied

Text:
{text}"""

# Richer prompt for Qwen2.5-14B — extracts full SCM including CPTs and confounders
_OLLAMA_DAG_PROMPT = """Extract a complete causal model from the text below.

Return ONLY a JSON object with these fields:
- "variables": list of variable names (strings)
- "edges": list of [source, target] directed edges (cause → effect)
- "bidirected": list of [A, B] pairs where a hidden common cause exists (omit if none)
- "params": object with conditional probability tables if numeric probabilities are stated,
  e.g. {{"p(X)": 0.3, "p(Y | X)": [0.2, 0.8]}} — omit if no probabilities mentioned
- "query": object with {{"type": "ate"|"ett"|"nde"|"marginal"|"conditional"|"do",
  "treatment": "var", "outcome": "var"}} if the text contains a causal question — omit otherwise

Rules:
- Use short, clean variable names (lowercase, underscores for spaces)
- Edges go from cause to effect only
- Only include bidirected edges when the text explicitly mentions confounding or a hidden common cause
- Only include params when specific probability values appear in the text

Text:
{text}"""

# Used for local backend — matches training format exactly
# Training inputs always had a question appended; output uses nodes/from/to
_LOCAL_DAG_QUERY = "Identify all causal relationships and extract the complete causal graph."

_QUERY_LLM_PROMPT = """Parse this causal query against the given variables.
Return ONLY valid JSON:
{{"query_type": "causal"|"independence"|"effect", "x": "variable_name", "y": "variable_name", "z": ["conditioning_var1", ...]}}

Variables in the DAG: {variables}

Query: {query}"""

_OLLAMA_QUERY_PROMPT = """Identify the causal query type and variables from the question below.

Return ONLY a JSON object:
{{"query_type": "causal"|"independence"|"effect"|"ate"|"ett"|"nde"|"marginal"|"conditional"|"do",
  "x": "treatment_variable", "y": "outcome_variable", "z": ["conditioning_var1", ...]}}

Available variables: {variables}

Question: {query}"""


# ============================================================
# LLM Backends
# ============================================================

def _call_claude(prompt: str, model: str = "claude-sonnet-4-20250514") -> str:
    """Call Claude API. Returns raw text response."""
    import httpx

    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        raise RuntimeError("ANTHROPIC_API_KEY not set")

    resp = httpx.post(
        "https://api.anthropic.com/v1/messages",
        headers={
            "x-api-key": api_key,
            "anthropic-version": "2023-06-01",
            "content-type": "application/json",
        },
        json={
            "model": model,
            "max_tokens": 2048,
            "messages": [{"role": "user", "content": prompt}],
        },
        timeout=60.0,
    )
    resp.raise_for_status()
    return resp.json()["content"][0]["text"]


_LOCAL_MODEL = None  # lazy-loaded singleton
_LOCAL_TOKENIZER = None


def _call_local_qwen(prompt: str, adapter_path: Optional[str] = None) -> str:
    """Run inference with the fine-tuned Qwen2.5-1.5B LoRA adapter.

    Loads model once and caches it. No API key required.
    adapter_path defaults to models/nl_dag_adapter/ relative to this file.
    """
    global _LOCAL_MODEL, _LOCAL_TOKENIZER

    if _LOCAL_MODEL is None:
        import torch
        from transformers import AutoTokenizer, AutoModelForCausalLM
        from peft import PeftModel

        if adapter_path is None:
            here = os.path.dirname(os.path.abspath(__file__))
            adapter_path = os.path.join(here, "..", "models", "nl_dag_adapter")
            adapter_path = os.path.normpath(adapter_path)

        base_id = "Qwen/Qwen2.5-1.5B-Instruct"
        device = "mps" if torch.backends.mps.is_available() else (
            "cuda" if torch.cuda.is_available() else "cpu"
        )

        # Load tokenizer from base model (adapter tokenizer_config may have
        # list-format extra_special_tokens incompatible with older transformers)
        _LOCAL_TOKENIZER = AutoTokenizer.from_pretrained(base_id)
        base = AutoModelForCausalLM.from_pretrained(
            base_id,
            torch_dtype=torch.float16 if device != "cpu" else torch.float32,
            device_map=device,
        )
        _LOCAL_MODEL = PeftModel.from_pretrained(base, adapter_path)
        _LOCAL_MODEL.eval()

    messages = [
        {"role": "system", "content": _LOCAL_SYSTEM_MSG},
        {"role": "user", "content": prompt},
    ]
    text = _LOCAL_TOKENIZER.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = _LOCAL_TOKENIZER(text, return_tensors="pt").to(_LOCAL_MODEL.device)

    import torch
    with torch.no_grad():
        outputs = _LOCAL_MODEL.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.1,
            do_sample=False,
            pad_token_id=_LOCAL_TOKENIZER.eos_token_id,
        )
    generated = outputs[0][inputs["input_ids"].shape[1]:]
    return _LOCAL_TOKENIZER.decode(generated, skip_special_tokens=True)


def _call_ollama(
    prompt: str,
    model: str = "qwen2.5:14b",
    system: str = "",
    temperature: float = 0.0,
) -> str:
    """Call local Ollama via the /api/chat endpoint (supports system messages).

    Uses streaming to handle large models without timeout issues.
    Falls back to /api/generate if chat endpoint returns 404.
    """
    import httpx

    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": prompt})

    chunks = []
    try:
        with httpx.stream(
            "POST",
            "http://localhost:11434/api/chat",
            json={
                "model": model,
                "messages": messages,
                "stream": True,
                "options": {"temperature": temperature, "num_predict": 1024},
            },
            timeout=httpx.Timeout(connect=10.0, read=600.0, write=10.0, pool=10.0),
        ) as response:
            response.raise_for_status()
            for line in response.iter_lines():
                if line:
                    try:
                        chunk = json.loads(line)
                        content = chunk.get("message", {}).get("content", "")
                        chunks.append(content)
                        if chunk.get("done"):
                            break
                    except json.JSONDecodeError:
                        continue
    except httpx.HTTPStatusError:
        # Fallback: older Ollama versions may not support /api/chat
        full_prompt = (system + "\n\n" + prompt).strip() if system else prompt
        with httpx.stream(
            "POST",
            "http://localhost:11434/api/generate",
            json={"model": model, "prompt": full_prompt, "stream": True,
                  "options": {"temperature": temperature, "num_predict": 1024}},
            timeout=httpx.Timeout(connect=10.0, read=600.0, write=10.0, pool=10.0),
        ) as response:
            response.raise_for_status()
            for line in response.iter_lines():
                if line:
                    try:
                        chunk = json.loads(line)
                        chunks.append(chunk.get("response", ""))
                        if chunk.get("done"):
                            break
                    except json.JSONDecodeError:
                        continue

    return "".join(chunks)


def _parse_json_response(text: str) -> dict:
    """Extract JSON object from LLM response text."""
    text = text.strip()

    # Try direct parse
    try:
        result = json.loads(text)
        if isinstance(result, dict):
            return result
    except json.JSONDecodeError:
        pass

    # Find JSON object in response
    match = re.search(r'\{[\s\S]*\}', text)
    if match:
        try:
            return json.loads(match.group())
        except json.JSONDecodeError:
            pass

    return {}


# ============================================================
# Core Parser
# ============================================================

class NLDAGParser:
    """Parse natural language text into causal DAGs.

    Supports three backends:
        - "claude": Claude API (best accuracy)
        - "ollama": Local Ollama (free/offline)
        - "regex": Pattern matching fallback (no LLM)
    """

    def __init__(self, backend: str = "regex", model: Optional[str] = None,
                 adapter_path: Optional[str] = None):
        """
        Args:
            backend: "claude", "ollama", "local", or "regex"
            model: override model name for claude/ollama backends
            adapter_path: path to LoRA adapter dir (local backend only;
                          defaults to models/nl_dag_adapter/)
        """
        self.backend = backend
        self.model = model
        self.adapter_path = adapter_path

    def parse(self, text: str) -> CausalDAG:
        """Extract causal DAG from natural language text.

        Returns:
            CausalDAG with variables, edges, and name mappings
        """
        if self.backend == "regex":
            return self._parse_regex(text)
        elif self.backend == "hybrid":
            # Local model first; fall back to regex if it returns no edges
            dag = self._parse_llm_with_backend("local", text)
            if dag.num_edges == 0:
                dag = self._parse_regex(text)
            return dag
        elif self.backend in ("claude", "ollama", "local"):
            return self._parse_llm(text)
        else:
            raise ValueError(f"Unknown backend: {self.backend}")

    def parse_query(self, query_text: str, dag: CausalDAG) -> CausalQuery:
        """Parse a causal query against a known DAG.

        Args:
            query_text: e.g. "Does A cause C?", "Is A independent of C given B?"
            dag: the CausalDAG to resolve variable names against

        Returns:
            CausalQuery with resolved indices
        """
        if self.backend in ("claude", "ollama", "local", "hybrid"):
            return self._parse_query_llm(query_text, dag)
        return self._parse_query_regex(query_text, dag)

    def parse_with_query(self, text: str, query_text: str) -> Tuple[CausalDAG, CausalQuery]:
        """Parse both DAG and query from text.

        Returns:
            (CausalDAG, CausalQuery)
        """
        dag = self.parse(text)
        query = self.parse_query(query_text, dag)
        return dag, query

    # ── Regex parsing ──────────────────────────────────────────

    def _parse_regex(self, text: str) -> CausalDAG:
        """Extract DAG using regex patterns."""
        raw_edges = []

        # Forward patterns: group(1) → group(2)
        for pattern in _FORWARD_PATTERNS:
            for match in pattern.finditer(text):
                src = match.group(1).strip()
                # Handle multi-target: "X causes B and C"
                targets_str = match.group(2).strip()
                targets = re.split(r'\s+and\s+', targets_str)
                for t in targets:
                    t = t.strip().rstrip('.,;')
                    if t:
                        raw_edges.append((src, t))

        # Reverse patterns: group(1) ← group(2), i.e. edge from group(2) to group(1)
        for pattern in _REVERSE_PATTERNS:
            for match in pattern.finditer(text):
                dst = match.group(1).strip()
                src = match.group(2).strip()
                raw_edges.append((src.rstrip('.,;'), dst.rstrip('.,;')))

        return self._build_dag(raw_edges, text)

    def _build_dag(self, raw_edges: List[Tuple[str, str]], source_text: str) -> CausalDAG:
        """Convert raw name-based edges into a CausalDAG."""
        # Normalize names
        name_map = {}  # normalized -> canonical
        for src, dst in raw_edges:
            for name in (src, dst):
                norm = name.strip().lower()
                if norm not in name_map:
                    name_map[norm] = name.strip()

        # Assign indices
        canonical_names = list(name_map.values())
        # Deduplicate while preserving order
        seen = set()
        unique_names = []
        for n in canonical_names:
            key = n.lower()
            if key not in seen:
                seen.add(key)
                unique_names.append(n)

        name_to_idx = {n.lower(): i for i, n in enumerate(unique_names)}

        variables = [
            CausalVariable(name=n, index=i)
            for i, n in enumerate(unique_names)
        ]

        # Build indexed edges (dedup)
        edges = []
        edge_set = set()
        for src, dst in raw_edges:
            si = name_to_idx.get(src.strip().lower())
            di = name_to_idx.get(dst.strip().lower())
            if si is not None and di is not None and si != di:
                pair = (si, di)
                if pair not in edge_set:
                    edge_set.add(pair)
                    edges.append(pair)

        return CausalDAG(
            variables=variables,
            edges=edges,
            name_to_idx=name_to_idx,
            source_text=source_text,
        )

    def _parse_query_regex(self, query_text: str, dag: CausalDAG) -> CausalQuery:
        """Parse query using regex patterns."""
        # Try independence with conditioning
        m = _QUERY_INDEPENDENCE.match(query_text)
        if m:
            x_name = m.group(1).strip()
            y_name = m.group(2).strip()
            z_str = m.group(3).strip()
            z_names = [z.strip() for z in re.split(r'\s+and\s+|,\s*', z_str)]
            return self._resolve_query("independence", x_name, y_name, z_names, dag)

        # Independence without conditioning
        m = _QUERY_INDEPENDENCE_NO_Z.match(query_text)
        if m:
            x_name = m.group(1).strip()
            y_name = m.group(2).strip()
            return self._resolve_query("independence", x_name, y_name, [], dag)

        # Causal query
        m = _QUERY_CAUSAL.match(query_text)
        if m:
            x_name = m.group(1).strip()
            y_name = m.group(2).strip()
            return self._resolve_query("causal", x_name, y_name, [], dag)

        # Effect query
        m = _QUERY_EFFECT.match(query_text)
        if m:
            x_name = m.group(1).strip()
            y_name = m.group(2).strip()
            return self._resolve_query("effect", x_name, y_name, [], dag)

        raise ValueError(f"Could not parse query: {query_text}")

    def _resolve_query(self, query_type: str, x_name: str, y_name: str,
                       z_names: List[str], dag: CausalDAG) -> CausalQuery:
        """Resolve variable names against DAG using fuzzy matching."""
        x_idx = self._resolve_variable(x_name, dag)
        y_idx = self._resolve_variable(y_name, dag)
        z_indices = [self._resolve_variable(z, dag) for z in z_names]

        return CausalQuery(
            query_type=query_type,
            x_name=x_name,
            y_name=y_name,
            z_names=z_names,
            x_idx=x_idx,
            y_idx=y_idx,
            z_indices=z_indices,
        )

    def _resolve_variable(self, name: str, dag: CausalDAG) -> int:
        """Resolve a variable name against DAG, with fuzzy matching.

        Tries: exact match, case-insensitive, substring, first-letter.
        """
        norm = name.strip().lower()

        # Exact match
        if norm in dag.name_to_idx:
            return dag.name_to_idx[norm]

        # Substring match (name is substring of variable or vice versa)
        for var_name, idx in dag.name_to_idx.items():
            if norm in var_name or var_name in norm:
                return idx

        # First letter match (for single-letter references like "A", "B")
        if len(norm) == 1:
            for var_name, idx in dag.name_to_idx.items():
                if var_name.startswith(norm):
                    return idx

        raise ValueError(f"Cannot resolve variable '{name}' in DAG with variables: "
                         f"{list(dag.name_to_idx.keys())}")

    # ── LLM parsing ──────────────────────────────────────────

    def _parse_llm_with_backend(self, backend: str, text: str) -> CausalDAG:
        """Run _parse_llm using an explicit backend name (used by hybrid)."""
        saved = self.backend
        self.backend = backend
        try:
            return self._parse_llm(text)
        finally:
            self.backend = saved

    def _parse_llm(self, text: str) -> CausalDAG:
        """Extract DAG using LLM backend."""
        if self.backend == "local":
            # Match training format: prose + question, system message added in _call_local_qwen
            prompt = f"{text[:50000].strip()}\n{_LOCAL_DAG_QUERY}"
            raw = _call_local_qwen(prompt, adapter_path=self.adapter_path)
        elif self.backend == "claude":
            prompt = _LLM_PROMPT.format(text=text[:50000])
            model = self.model or "claude-sonnet-4-6"
            raw = _call_claude(prompt, model=model)
        else:
            # Ollama — use richer prompt and system message for instruction models
            prompt = _OLLAMA_DAG_PROMPT.format(text=text[:50000])
            model = self.model or "qwen2.5:14b"
            raw = _call_ollama(prompt, model=model, system=_OLLAMA_SYSTEM_MSG)

        data = _parse_json_response(raw)

        # Handle both "variables" (claude/ollama) and "nodes" (local training format)
        variables = data.get("variables") or data.get("nodes", [])
        raw_edges = data.get("edges", [])

        if not variables:
            # Fallback to regex
            return self._parse_regex(text)

        # Build from LLM output — handle both array ["src","dst"] and
        # dict {"source":"src","target":"dst"} edge formats
        name_pairs = []
        for e in raw_edges:
            if isinstance(e, dict):
                src = e.get("from") or e.get("source") or e.get("src")
                dst = e.get("to") or e.get("target") or e.get("dst")
                if src and dst:
                    name_pairs.append((src, dst))
            elif hasattr(e, "__len__") and len(e) >= 2:
                name_pairs.append((e[0], e[1]))
        # Add any variables mentioned in edges but not in variables list
        var_set = set(v.lower() for v in variables)
        for src, dst in name_pairs:
            if src.lower() not in var_set:
                variables.append(src)
                var_set.add(src.lower())
            if dst.lower() not in var_set:
                variables.append(dst)
                var_set.add(dst.lower())

        dag = self._build_dag(name_pairs, text)

        # Capture enrichments from richer LLM responses (ollama/Claude)
        raw_bidir = data.get("bidirected", [])
        dag.bidirected_names = [
            (e[0], e[1]) if isinstance(e, (list, tuple)) and len(e) >= 2
            else (e.get("from", ""), e.get("to", ""))
            for e in raw_bidir
        ]
        dag.params = data.get("params", {})
        dag.suggested_query = data.get("query")  # {type, treatment, outcome} or None

        return dag

    def _parse_query_llm(self, query_text: str, dag: CausalDAG) -> CausalQuery:
        """Parse query using LLM backend."""
        var_names = dag.variable_names()

        if self.backend == "claude":
            prompt = _QUERY_LLM_PROMPT.format(variables=", ".join(var_names), query=query_text)
            model = self.model or "claude-sonnet-4-6"
            raw = _call_claude(prompt, model=model)
        elif self.backend == "local":
            prompt = _QUERY_LLM_PROMPT.format(variables=", ".join(var_names), query=query_text)
            raw = _call_local_qwen(prompt, adapter_path=self.adapter_path)
        else:
            prompt = _OLLAMA_QUERY_PROMPT.format(variables=", ".join(var_names), query=query_text)
            model = self.model or "qwen2.5:14b"
            raw = _call_ollama(prompt, model=model, system=_OLLAMA_SYSTEM_MSG)

        data = _parse_json_response(raw)

        if not data:
            # Fallback to regex
            return self._parse_query_regex(query_text, dag)

        query_type = data.get("query_type", "causal")
        x_name = data.get("x", "")
        y_name = data.get("y", "")
        z_names = data.get("z", [])

        return self._resolve_query(query_type, x_name, y_name, z_names, dag)


# ============================================================
# Bridge: DAG → CausalEng Input Tensors
# ============================================================

def dag_to_causal_eng_input(dag: CausalDAG, query: CausalQuery,
                            device=None) -> dict:
    """Convert parsed DAG + query into CausalEng input tensors.

    Uses CausalEng.build_node_features() to produce exact [N, 6] features.

    Returns:
        dict with keys: node_features, edge_index, query_x_idx, query_y_idx, z_indices
    """
    import torch
    from models.causal_eng import CausalEng

    if device is None:
        device = torch.device("cpu")

    node_features = CausalEng.build_node_features(
        num_nodes=dag.num_nodes,
        edges=dag.edges,
        query_x=query.x_idx,
        query_y=query.y_idx,
        z_set=query.z_indices,
        device=device,
    )

    # Edge index as [2, E] tensor
    if dag.edges:
        srcs, dsts = zip(*dag.edges)
        edge_index = torch.tensor([list(srcs), list(dsts)], dtype=torch.long, device=device)
    else:
        edge_index = torch.zeros(2, 0, dtype=torch.long, device=device)

    return {
        "node_features": node_features,
        "edge_index": edge_index,
        "query_x_idx": query.x_idx,
        "query_y_idx": query.y_idx,
        "z_indices": query.z_indices,
    }


# ============================================================
# End-to-End Pipeline
# ============================================================

def text_to_causal_query(text: str, query: str,
                         backend: str = "regex",
                         model: Optional[str] = None,
                         adapter_path: Optional[str] = None,
                         device=None) -> dict:
    """Full pipeline: text → DAG → tensors → model inference → result.

    Args:
        text: natural language describing causal relationships
        query: causal question (e.g. "Does A cause C?")
        backend: "claude", "ollama", "local", or "regex"
        model: optional model override (claude/ollama backends)
        adapter_path: path to LoRA adapter dir (local backend only)
        device: torch device

    Returns:
        dict with keys: dag, query, tensors, variables, edges
        (model inference is optional — returns tensors ready for CausalEng)
    """
    parser = NLDAGParser(backend=backend, model=model, adapter_path=adapter_path)
    dag, parsed_query = parser.parse_with_query(text, query)

    result = {
        "dag": dag,
        "query": parsed_query,
        "variables": dag.variable_names(),
        "edges": dag.edges,
        "num_nodes": dag.num_nodes,
        "num_edges": dag.num_edges,
    }

    # Generate tensors if torch is available
    try:
        tensors = dag_to_causal_eng_input(dag, parsed_query, device=device)
        result["tensors"] = tensors
    except ImportError:
        result["tensors"] = None

    return result