NPCArray

class GraphExecutor:
    """
    Executes the lazy computation graph.

    Handles:
    - Topological ordering
    - Parallel execution of independent nodes
    - Caching of intermediate results
    """

    def __init__(
        self,
        parallel: bool = True,
        max_workers: int = 4,
        progress: bool = False
    ):
        self.parallel = parallel
        self.max_workers = max_workers
        self.progress = progress
        self._cache: Dict[int, Any] = {}

    def execute(
        self,
        root: GraphNode,
        specs: List[ModelSpec],
        prompts: Optional[List[str]] = None
    ) -> ResponseTensor:
        """Execute graph starting from root node"""

        ordered = self._topological_sort(root)

        for node in ordered:
            if id(node) in self._cache:
                continue

            parent_results = [self._cache.get(id(p)) for p in node.parents]

            result = self._execute_node(node, specs, prompts, parent_results)
            self._cache[id(node)] = result

        return self._cache[id(root)]

    def _topological_sort(self, root: GraphNode) -> List[GraphNode]:
        """Return nodes in execution order (leaves first)"""
        visited = set()
        ordered = []

        def visit(node):
            if id(node) in visited:
                return
            visited.add(id(node))
            for parent in node.parents:
                visit(parent)
            ordered.append(node)

        visit(root)
        return ordered

    def _execute_node(
        self,
        node: GraphNode,
        specs: List[ModelSpec],
        prompts: Optional[List[str]],
        parent_results: List[Any]
    ) -> ResponseTensor:
        """Execute a single graph node"""

        handlers = {
            OpType.SOURCE: self._exec_source,
            OpType.INFER: self._exec_infer,
            OpType.PREDICT: self._exec_predict,
            OpType.FORWARD: self._exec_forward,
            OpType.FIT: self._exec_fit,
            OpType.MAP: self._exec_map,
            OpType.FILTER: self._exec_filter,
            OpType.REDUCE: self._exec_reduce,
            OpType.CHAIN: self._exec_chain,
            OpType.EVOLVE: self._exec_evolve,
            OpType.JINX: self._exec_jinx,
        }

        handler = handlers.get(node.op_type)
        if handler is None:
            raise ValueError(f"Unknown operation type: {node.op_type}")

        return handler(node, specs, prompts, parent_results)

    def _exec_source(self, node, specs, prompts, parents) -> ResponseTensor:
        """Source node - just returns specs wrapped"""
        return ResponseTensor(
            data=np.array([s.model_ref for s in specs], dtype=object),
            model_specs=specs,
            prompts=prompts
        )

    def _exec_infer(self, node, specs, prompts, parents) -> ResponseTensor:
        """Execute LLM inference across models and prompts"""
        from npcpy.llm_funcs import get_llm_response

        prompts_list = node.params.get("prompts", prompts or [])
        extra_kwargs = {k: v for k, v in node.params.items() if k != "prompts"}

        n_models = len(specs)
        n_prompts = len(prompts_list)

        tasks = []
        for i, spec in enumerate(specs):
            for j, prompt in enumerate(prompts_list):
                tasks.append((i, j, spec, prompt))

        results = np.empty((n_models, n_prompts), dtype=object)

        if self.parallel and len(tasks) > 1:
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                futures = {}
                for i, j, spec, prompt in tasks:
                    future = executor.submit(
                        self._infer_single, spec, prompt, extra_kwargs
                    )
                    futures[future] = (i, j)

                for future in as_completed(futures):
                    i, j = futures[future]
                    try:
                        results[i, j] = future.result()
                    except Exception as e:
                        results[i, j] = f"Error: {e}"
        else:
            for i, j, spec, prompt in tasks:
                try:
                    results[i, j] = self._infer_single(spec, prompt, extra_kwargs)
                except Exception as e:
                    results[i, j] = f"Error: {e}"

        return ResponseTensor(
            data=results,
            model_specs=specs,
            prompts=prompts_list,
            metadata={"operation": "infer", **extra_kwargs}
        )

    def _infer_single(self, spec: ModelSpec, prompt: str, kwargs: Dict) -> str:
        """Single model inference"""
        from npcpy.llm_funcs import get_llm_response

        if spec.model_type == "llm":
            response = get_llm_response(
                prompt,
                model=spec.model_ref,
                provider=spec.provider,
                **{**spec.config, **kwargs}
            )
            return response.get("response", "")

        elif spec.model_type == "npc":
            npc = spec.model_ref
            response = get_llm_response(
                prompt,
                npc=npc,
                **kwargs
            )
            return response.get("response", "")

        else:
            raise ValueError(f"Cannot infer with model type: {spec.model_type}")

    def _exec_predict(self, node, specs, prompts, parents) -> ResponseTensor:
        """Execute sklearn/ML prediction"""
        X = node.params.get("X")

        results = []
        for spec in specs:
            if spec.model_type == "sklearn":
                model = spec.model_ref
                if hasattr(model, 'predict'):
                    pred = model.predict(X)
                    results.append(pred)
                else:
                    results.append(None)
            else:
                results.append(None)

        return ResponseTensor(
            data=np.array(results, dtype=object),
            model_specs=specs,
            metadata={"operation": "predict"}
        )

    def _exec_forward(self, node, specs, prompts, parents) -> ResponseTensor:
        """Execute PyTorch forward pass"""
        inputs = node.params.get("inputs")

        results = []
        for spec in specs:
            if spec.model_type == "torch":
                model = spec.model_ref
                device = spec.config.get("device", "cpu")
                try:
                    import torch
                    model.to(device)
                    model.eval()
                    with torch.no_grad():
                        output = model(inputs.to(device) if hasattr(inputs, 'to') else inputs)
                    results.append(output)
                except Exception as e:
                    results.append(f"Error: {e}")
            else:
                results.append(None)

        return ResponseTensor(
            data=np.array(results, dtype=object),
            model_specs=specs,
            metadata={"operation": "forward"}
        )

    def _exec_fit(self, node, specs, prompts, parents) -> ResponseTensor:
        """Execute model fitting"""
        X = node.params.get("X")
        y = node.params.get("y")

        fitted_specs = []
        for spec in specs:
            if spec.model_type == "sklearn":
                model = copy.deepcopy(spec.model_ref)
                if hasattr(model, 'fit'):
                    model.fit(X, y)
                new_spec = ModelSpec(
                    model_type="sklearn",
                    model_ref=model,
                    config={**spec.config, "fitted": True}
                )
                fitted_specs.append(new_spec)
            elif spec.model_type == "llm":
                fitted_specs.append(spec)
            else:
                fitted_specs.append(spec)

        return ResponseTensor(
            data=np.array([s.model_ref for s in fitted_specs], dtype=object),
            model_specs=fitted_specs,
            metadata={"operation": "fit"}
        )

    def _exec_map(self, node, specs, prompts, parents) -> ResponseTensor:
        """Apply function to each result"""
        fn = node.params.get("fn")
        parent_result = parents[0] if parents else None

        if parent_result is None:
            raise ValueError("Map requires parent result")

        mapped = np.vectorize(fn, otypes=[object])(parent_result.data)

        return ResponseTensor(
            data=mapped,
            model_specs=parent_result.model_specs,
            prompts=parent_result.prompts,
            metadata={**parent_result.metadata, "mapped": True}
        )

    def _exec_filter(self, node, specs, prompts, parents) -> ResponseTensor:
        """Filter results by predicate"""
        predicate = node.params.get("predicate")
        parent_result = parents[0] if parents else None

        if parent_result is None:
            raise ValueError("Filter requires parent result")

        mask = np.vectorize(predicate)(parent_result.data)
        filtered_data = parent_result.data[mask]

        return ResponseTensor(
            data=filtered_data,
            model_specs=parent_result.model_specs,
            prompts=parent_result.prompts,
            metadata={**parent_result.metadata, "filtered": True}
        )

    def _exec_reduce(self, node, specs, prompts, parents) -> ResponseTensor:
        """Reduce results along axis"""
        method = node.params.get("method", "vote")
        axis = node.params.get("axis", 0)
        parent_result = parents[0] if parents else None

        if parent_result is None:
            raise ValueError("Reduce requires parent result")

        data = parent_result.data

        if method == "vote":
            reduced = self._reduce_vote(data, axis)
        elif method == "mean":
            reduced = np.mean(data, axis=axis)
        elif method == "concat":
            reduced = self._reduce_concat(data, axis)
        elif method == "consensus":
            reduced = self._reduce_consensus(data, axis, node.params)
        elif method == "best":
            scores = node.params.get("scores", [])
            reduced = self._reduce_best(data, scores, axis)
        elif callable(method):
            reduced = np.apply_along_axis(method, axis, data)
        else:
            raise ValueError(f"Unknown reduce method: {method}")

        return ResponseTensor(
            data=np.atleast_1d(reduced),
            model_specs=specs if axis != 0 else [specs[0]],
            prompts=prompts,
            metadata={**parent_result.metadata, "reduced": method}
        )

    def _reduce_vote(self, data: np.ndarray, axis: int) -> np.ndarray:
        """Majority voting reduction"""
        from collections import Counter

        def vote_fn(arr):
            counter = Counter(arr)
            return counter.most_common(1)[0][0] if counter else None

        return np.apply_along_axis(vote_fn, axis, data)

    def _reduce_concat(self, data: np.ndarray, axis: int) -> np.ndarray:
        """Concatenate strings"""
        def concat_fn(arr):
            return "\n---\n".join(str(x) for x in arr)

        return np.apply_along_axis(concat_fn, axis, data)

    def _reduce_consensus(self, data: np.ndarray, axis: int, params: Dict) -> np.ndarray:
        """LLM-based consensus"""
        from npcpy.llm_funcs import get_llm_response

        model = params.get("model")
        if not model:
            raise ValueError("No model specified for consensus reduction.")

        def consensus_fn(arr):
            perspectives = "\n".join(f"- {x}" for x in arr)
            prompt = f"Given these different perspectives:\n{perspectives}\n\nProvide a consensus synthesis:"
            response = get_llm_response(prompt, model=model)
            return response.get("response", "")

        return np.apply_along_axis(consensus_fn, axis, data)

    def _reduce_best(self, data: np.ndarray, scores: List[float], axis: int) -> np.ndarray:
        """Select best by score"""
        if axis == 0:
            best_idx = np.argmax(scores)
            return data[best_idx]
        else:
            return data

    def _exec_chain(self, node, specs, prompts, parents) -> ResponseTensor:
        """Chain responses through synthesis function"""
        fn = node.params.get("fn")
        n_rounds = node.params.get("n_rounds", 1)
        parent_result = parents[0] if parents else None

        if parent_result is None:
            raise ValueError("Chain requires parent result")

        current = parent_result.data

        for _ in range(n_rounds):
            new_prompt = fn(current.tolist())

            infer_node = GraphNode(
                op_type=OpType.INFER,
                params={"prompts": [new_prompt]},
                shape=(len(specs), 1)
            )
            current = self._exec_infer(infer_node, specs, [new_prompt], []).data

        return ResponseTensor(
            data=current,
            model_specs=specs,
            prompts=prompts,
            metadata={**parent_result.metadata, "chained": n_rounds}
        )

    def _exec_evolve(self, node, specs, prompts, parents) -> ResponseTensor:
        """Evolve population based on fitness"""
        import random

        fitness_scores = node.params.get("fitness_scores", [])
        mutate_fn = node.params.get("mutate_fn")
        crossover_fn = node.params.get("crossover_fn")
        elite_ratio = node.params.get("elite_ratio", 0.1)

        n = len(specs)
        n_elite = max(1, int(n * elite_ratio))

        sorted_indices = np.argsort(fitness_scores)[::-1]

        new_specs = [specs[i] for i in sorted_indices[:n_elite]]

        while len(new_specs) < n:
            if crossover_fn and random.random() < 0.5:
                parent1 = specs[random.choice(sorted_indices[:n//2])]
                parent2 = specs[random.choice(sorted_indices[:n//2])]
                child_spec = crossover_fn(parent1, parent2)
            else:
                parent = specs[random.choice(sorted_indices[:n//2])]
                child_spec = mutate_fn(parent) if mutate_fn else parent
            new_specs.append(child_spec)

        return ResponseTensor(
            data=np.array([s.model_ref for s in new_specs], dtype=object),
            model_specs=new_specs,
            metadata={"operation": "evolve", "generation": 1}
        )

    def _exec_jinx(self, node, specs, prompts, parents) -> ResponseTensor:
        """Execute a Jinx workflow across models"""
        from npcpy.npc_compiler import NPC, Jinx

        jinx_name = node.params.get("jinx_name")
        inputs = node.params.get("inputs", {})
        extra_kwargs = {k: v for k, v in node.params.items()
                       if k not in ("jinx_name", "inputs")}

        results = []

        def run_jinx_single(spec: ModelSpec) -> str:
            """Run jinx for a single model spec"""
            try:
                if spec.model_type == "npc":
                    npc = spec.model_ref
                else:
                    npc = NPC(
                        name=f"array_npc_{spec.model_ref}",
                        model=spec.model_ref,
                        provider=spec.provider
                    )

                result = npc.execute_jinx(
                    jinx_name=jinx_name,
                    input_values=inputs,
                    **extra_kwargs
                )
                return result.get("output", str(result))
            except Exception as e:
                return f"Error: {e}"

        if self.parallel and len(specs) > 1:
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                futures = {executor.submit(run_jinx_single, spec): i
                          for i, spec in enumerate(specs)}
                results = [None] * len(specs)
                for future in as_completed(futures):
                    idx = futures[future]
                    try:
                        results[idx] = future.result()
                    except Exception as e:
                        results[idx] = f"Error: {e}"
        else:
            results = [run_jinx_single(spec) for spec in specs]

        return ResponseTensor(
            data=np.array(results, dtype=object),
            model_specs=specs,
            metadata={"operation": "jinx", "jinx_name": jinx_name, **inputs}
        )

@dataclass
class GraphNode:
    """A node in the lazy computation graph"""
    op_type: OpType
    params: Dict[str, Any] = field(default_factory=dict)
    parents: List['GraphNode'] = field(default_factory=list)
    result: Any = None
    shape: Optional[Tuple[int, ...]] = None

class LazyResult:
    """
    Lazy result from model operations.

    Builds computation graph without executing until .compute() is called.
    Supports chaining operations like map, filter, reduce.
    """

    def __init__(
        self,
        specs: List[ModelSpec],
        graph: GraphNode,
        prompts: Optional[List[str]] = None
    ):
        self._specs = specs
        self._graph = graph
        self._prompts = prompts
        self._computed = False
        self._result: Optional[ResponseTensor] = None

    @property
    def shape(self) -> Optional[Tuple[int, ...]]:
        """Expected shape of result"""
        return self._graph.shape

    def map(self, fn: Callable[[Any], Any]) -> 'LazyResult':
        """
        Apply function to each response.

        Args:
            fn: Function to apply to each response

        Example:
            >>> result.map(lambda r: len(r))  # Get lengths
            >>> result.map(json.loads)  # Parse JSON
        """
        new_node = GraphNode(
            op_type=OpType.MAP,
            params={"fn": fn},
            parents=[self._graph],
            shape=self._graph.shape
        )

        return LazyResult(self._specs, new_node, self._prompts)

    def filter(self, predicate: Callable[[Any], bool]) -> 'LazyResult':
        """
        Filter responses by predicate.

        Args:
            predicate: Function returning True for responses to keep

        Example:
            >>> result.filter(lambda r: len(r) > 100)
            >>> result.filter(lambda r: 'error' not in r.lower())
        """
        new_node = GraphNode(
            op_type=OpType.FILTER,
            params={"predicate": predicate},
            parents=[self._graph],
            shape=None
        )

        return LazyResult(self._specs, new_node, self._prompts)

    def reduce(
        self,
        method: Union[str, Callable] = "vote",
        axis: int = 0,
        **kwargs
    ) -> 'LazyResult':
        """
        Reduce responses along an axis.

        Args:
            method: Reduction method or custom function
                - 'vote': Majority voting
                - 'mean': Average (for numeric)
                - 'concat': Concatenate strings
                - 'consensus': LLM-based consensus
                - 'best': Select by score
                - callable: Custom reduction
            axis: Axis to reduce (0=models, 1=prompts)
            **kwargs: Additional params for reduction

        Example:
            >>> result.reduce('vote', axis=0)  # Vote across models
            >>> result.reduce('mean', axis=1)  # Average across prompts
        """
        new_node = GraphNode(
            op_type=OpType.REDUCE,
            params={"method": method, "axis": axis, **kwargs},
            parents=[self._graph],
            shape=self._compute_reduced_shape(axis)
        )

        return LazyResult(self._specs, new_node, self._prompts)

    def _compute_reduced_shape(self, axis: int) -> Optional[Tuple[int, ...]]:
        """Compute shape after reduction"""
        if self._graph.shape is None:
            return None
        shape = list(self._graph.shape)
        if axis < len(shape):
            shape.pop(axis)
        return tuple(shape) if shape else (1,)

    def chain(
        self,
        fn: Callable[[List[Any]], str],
        n_rounds: int = 1
    ) -> 'LazyResult':
        """
        Chain outputs through a synthesis function.

        Useful for debate/discussion patterns where outputs
        feed back as context for next round.

        Args:
            fn: Function taking all responses, returning synthesis prompt
            n_rounds: Number of chain rounds

        Example:
            >>> def debate_round(responses):
            ...     return f"Consider these perspectives: {responses}. Synthesize."
            >>> result.chain(debate_round, n_rounds=3)
        """
        new_node = GraphNode(
            op_type=OpType.CHAIN,
            params={"fn": fn, "n_rounds": n_rounds},
            parents=[self._graph],
            shape=self._graph.shape
        )

        return LazyResult(self._specs, new_node, self._prompts)

    def vote(self, axis: int = 0) -> 'LazyResult':
        """Shorthand for reduce('vote', axis)"""
        return self.reduce('vote', axis=axis)

    def consensus(self, axis: int = 0, model: str = None) -> 'LazyResult':
        """Shorthand for reduce('consensus', axis)"""
        return self.reduce('consensus', axis=axis, model=model)

    def variance(self) -> 'LazyResult':
        """Compute variance/disagreement across models"""
        return self.map(_compute_response_variance)

    def argmax(self, scores: List[float]) -> 'LazyResult':
        """Select responses corresponding to max scores"""
        return self.reduce('best', scores=scores)

    def explain(self) -> str:
        """
        Print explanation of the computation graph.

        Returns:
            String representation of the DAG
        """
        lines = ["Computation Graph:"]
        self._explain_node(self._graph, lines, depth=0)
        explanation = "\n".join(lines)
        print(explanation)
        return explanation

    def _explain_node(self, node: GraphNode, lines: List[str], depth: int):
        indent = "  " * depth
        params_str = {k: v for k, v in node.params.items() if k not in ('fn', 'predicate')}
        lines.append(f"{indent}└─ {node.op_type.value}: shape={node.shape}, params={params_str}")
        for parent in node.parents:
            self._explain_node(parent, lines, depth + 1)

    def collect(
        self,
        parallel: bool = True,
        max_workers: int = None,
        progress: bool = False
    ) -> ResponseTensor:
        """
        Execute the computation graph and return results.

        Like Spark's collect(), this materializes the lazy computation.

        Args:
            parallel: Whether to parallelize independent operations
            max_workers: Max parallel workers (default: number of models)
            progress: Show progress bar

        Returns:
            ResponseTensor with materialized results
        """
        if self._computed and self._result is not None:
            return self._result

        executor = GraphExecutor(
            parallel=parallel,
            max_workers=max_workers or len(self._specs),
            progress=progress
        )

        self._result = executor.execute(self._graph, self._specs, self._prompts)
        self._computed = True

        return self._result

    def to_list(self) -> List:
        """Collect and return as Python list"""
        return self.collect().tolist()

    compute = collect

@dataclass
class ModelSpec:
    """Specification for a model in the array"""
    model_type: Literal["llm", "sklearn", "torch", "npc", "custom"]
    model_ref: Any
    provider: Optional[str] = None
    config: Dict[str, Any] = field(default_factory=dict)

    def __hash__(self):
        return hash((self.model_type, str(self.model_ref), self.provider))