#!/usr/bin/env python3
"""Split an ONNX graph into smaller sub-models driven by sub_config rules.

This script reads a JSON config file (matching the pulsar2 sub_config layout),
extracts the requested subgraphs, and optionally emits any leftover parts of the
model as independent ONNX graphs. A verification utility can run the original
model and the stitched micro-model pipeline to make sure their outputs match.
"""
from __future__ import annotations

import argparse
import json
import logging
from collections import defaultdict, deque
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Sequence, Set

import numpy as np
import onnx
from onnx import utils as onnx_utils

try:
    import onnxruntime as ort
except ImportError:  # pragma: no cover - optional dependency
    ort = None


@dataclass
class SubGraphSpec:
    """Describes a single subgraph to extract from the full model."""

    label: str
    start: List[str]
    end: List[str]
    node_names: Set[str]
    source: str
    output_path: Optional[Path] = None


@dataclass
class GraphIndex:
    """Caches helpful lookups for traversing an ONNX graph."""

    tensor_to_producer: Dict[str, str]
    tensor_to_consumers: Dict[str, List[str]]
    node_inputs: Dict[str, List[str]]
    node_outputs: Dict[str, List[str]]
    graph_inputs: Set[str]
    graph_outputs: Set[str]
    initializer_names: Set[str]
    node_order: List[str]


def sanitize(name: str) -> str:
    keep = [c if c.isalnum() else "_" for c in name] if name else ["anon"]
    sanitized = "".join(keep).strip("_")
    return sanitized or "tensor"


def build_graph_index(model: onnx.ModelProto) -> GraphIndex:
    tensor_to_producer: Dict[str, str] = {}
    tensor_to_consumers: Dict[str, List[str]] = defaultdict(list)
    node_inputs: Dict[str, List[str]] = {}
    node_outputs: Dict[str, List[str]] = {}
    node_order: List[str] = []

    used_names: Set[str] = set()
    for idx, node in enumerate(model.graph.node):
        base = node.name.strip() if node.name else ""
        candidate = base or f"node_{idx}"
        while candidate in used_names:
            candidate = f"{candidate}_{idx}"
        used_names.add(candidate)
        node_name = candidate
        node_order.append(node_name)
        node_inputs[node_name] = [x for x in node.input if x]
        node_outputs[node_name] = [y for y in node.output if y]
        for out_name in node_outputs[node_name]:
            tensor_to_producer[out_name] = node_name
        for inp_name in node_inputs[node_name]:
            tensor_to_consumers[inp_name].append(node_name)

    graph_inputs = {vi.name for vi in model.graph.input}
    graph_outputs = {vi.name for vi in model.graph.output}
    initializer_names = {init.name for init in model.graph.initializer}

    return GraphIndex(
        tensor_to_producer=tensor_to_producer,
        tensor_to_consumers=tensor_to_consumers,
        node_inputs=node_inputs,
        node_outputs=node_outputs,
        graph_inputs=graph_inputs,
        graph_outputs=graph_outputs,
        initializer_names=initializer_names,
        node_order=node_order,
    )


def trace_nodes_between(
    spec: SubGraphSpec,
    index: GraphIndex,
) -> Set[str]:
    boundary = set(spec.start) | index.graph_inputs | index.initializer_names
    visited_tensors: Set[str] = set()
    stack = list(spec.end)
    discovered_nodes: Set[str] = set()

    while stack:
        tensor = stack.pop()
        if tensor in visited_tensors:
            continue
        visited_tensors.add(tensor)
        if tensor in boundary:
            continue
        producer = index.tensor_to_producer.get(tensor)
        if not producer:
            continue
        if producer in discovered_nodes:
            continue
        discovered_nodes.add(producer)
        for upstream in index.node_inputs.get(producer, []):
            if upstream and upstream not in boundary:
                stack.append(upstream)
    return discovered_nodes


def untouched_components(
    all_nodes: Sequence[str],
    covered_nodes: Set[str],
    index: GraphIndex,
) -> List[Set[str]]:
    remaining = [n for n in all_nodes if n not in covered_nodes]
    if not remaining:
        return []
    adjacency: Dict[str, Set[str]] = {name: set() for name in remaining}
    rem_set = set(remaining)
    for node in remaining:
        for out_name in index.node_outputs.get(node, []):
            for consumer in index.tensor_to_consumers.get(out_name, []):
                if consumer in rem_set:
                    adjacency[node].add(consumer)
                    adjacency[consumer].add(node)
        for inp_name in index.node_inputs.get(node, []):
            producer = index.tensor_to_producer.get(inp_name)
            if producer in rem_set:
                adjacency[node].add(producer)
                adjacency[producer].add(node)

    components: List[Set[str]] = []
    visited: Set[str] = set()
    for node in remaining:
        if node in visited:
            continue
        stack = [node]
        comp: Set[str] = set()
        while stack:
            cur = stack.pop()
            if cur in visited:
                continue
            visited.add(cur)
            comp.add(cur)
            stack.extend(adjacency[cur] - visited)
        components.append(comp)
    return components


def derive_interface(
    nodes: Set[str],
    index: GraphIndex,
) -> (List[str], List[str]):
    produced = set()
    for node in nodes:
        produced.update(index.node_outputs.get(node, []))

    start: Set[str] = set()
    for node in nodes:
        for inp in index.node_inputs.get(node, []):
            producer = index.tensor_to_producer.get(inp)
            if producer is None and inp not in index.initializer_names:
                start.add(inp)
            elif producer not in nodes and inp not in index.initializer_names:
                start.add(inp)

    end: Set[str] = set()
    for node in nodes:
        for out in index.node_outputs.get(node, []):
            consumers = index.tensor_to_consumers.get(out, [])
            if not consumers:
                if out in index.graph_outputs:
                    end.add(out)
                continue
            if any(consumer not in nodes for consumer in consumers):
                end.add(out)
    end.update(index.graph_outputs & produced)

    if not end and produced:
        end = produced.copy()

    return sorted(start), sorted(end)


def extract_model_file(
    model_path: Path,
    spec: SubGraphSpec,
    output_dir: Path,
    suffix: str,
) -> Path:
    head = sanitize(spec.start[0]) if spec.start else "const"
    tail = sanitize(spec.end[0]) if spec.end else "out"
    filename = f"{spec.label}_{head}_to_{tail}_{suffix}.onnx"
    destination = output_dir / filename
    onnx_utils.extract_model(
        model_path.as_posix(),
        destination.as_posix(),
        input_names=spec.start,
        output_names=spec.end,
        check_model=False,
    )
    logging.info("Saved %s (start=%s, end=%s)", destination.name, spec.start, spec.end)
    return destination


def ordered_specs(
    specs: Sequence[SubGraphSpec],
    index: GraphIndex,
) -> List[SubGraphSpec]:
    available = set(index.graph_inputs) | index.initializer_names
    pending = list(specs)
    ordered: List[SubGraphSpec] = []
    while pending:
        progressed = False
        for spec in list(pending):
            if set(spec.start).issubset(available):
                ordered.append(spec)
                available.update(spec.end)
                pending.remove(spec)
                progressed = True
        if not progressed:
            missing = {spec.label: sorted(set(spec.start) - available) for spec in pending}
            raise RuntimeError(
                "无法解析子图拓扑，缺少以下张量: %s" % missing
            )
    return ordered


def run_full_model(model_path: Path, feed_dict: Dict[str, np.ndarray], providers: List[str]):
    if ort is None:
        raise RuntimeError("需要 onnxruntime 才能执行验证。")
    session = ort.InferenceSession(model_path.as_posix(), providers=providers)
    outputs = session.run(None, feed_dict)
    names = [meta.name for meta in session.get_outputs()]
    return dict(zip(names, outputs))


def run_split_pipeline(
    ordered_subgraphs: Sequence[SubGraphSpec],
    feed_dict: Dict[str, np.ndarray],
    providers: List[str],
) -> Dict[str, np.ndarray]:
    if ort is None:
        raise RuntimeError("需要 onnxruntime 才能执行验证。")
    tensor_store = dict(feed_dict)
    for spec in ordered_subgraphs:
        if spec.output_path is None:
            raise RuntimeError(f"子图 {spec.label} 尚未生成 ONNX 文件。")
        session = ort.InferenceSession(spec.output_path.as_posix(), providers=providers)
        fetch_inputs = {}
        for name in spec.start:
            if name not in tensor_store:
                raise KeyError(
                    f"子图 {spec.label} 缺少输入张量 {name}，请确认切分顺序。"
                )
            fetch_inputs[name] = tensor_store[name]
        results = session.run(None, fetch_inputs)
        for meta, value in zip(session.get_outputs(), results):
            tensor_store[meta.name] = value
    return tensor_store


def verify(
    model_path: Path,
    ordered_subgraphs: Sequence[SubGraphSpec],
    feed_dict: Dict[str, np.ndarray],
    providers: List[str],
    rtol: float,
    atol: float,
) -> None:
    full_outputs = run_full_model(model_path, feed_dict, providers)
    split_store = run_split_pipeline(ordered_subgraphs, feed_dict, providers)
    for name, ref in full_outputs.items():
        cand = split_store.get(name)
        if cand is None:
            raise AssertionError(f"切分流水线未产生模型输出 {name}")
        if not np.allclose(ref, cand, rtol=rtol, atol=atol):
            diff = np.max(np.abs(ref - cand))
            raise AssertionError(
                f"输出 {name} 不匹配，最大偏差 {diff:.3e}"
            )
    logging.info("切分模型与原始模型输出一致 (rtol=%g, atol=%g)。", rtol, atol)


def load_npz_inputs(npz_path: Path) -> Dict[str, np.ndarray]:
    data = np.load(npz_path, allow_pickle=False)
    return {key: data[key] for key in data.files}


def main() -> None:
    parser = argparse.ArgumentParser(description="根据 sub_config 切分 ONNX 模型。")
    parser.add_argument("--model", required=True, type=Path, help="原始 ONNX 路径")
    parser.add_argument("--config", required=True, type=Path, help="pulsar2 配置 JSON")
    parser.add_argument("--output-dir", required=False, default="./split-onnx", type=Path, help="保存子模型的目录")
    parser.add_argument(
        "--verify",
        action="store_true",
        help="生成后立即用 onnxruntime 校验输出是否一致",
    )
    parser.add_argument(
        "--input-npz",
        type=Path,
        help="包含模型所有输入张量的 npz 文件 (verify 模式需要)",
    )
    parser.add_argument(
        "--providers",
        nargs="*",
        default=["CPUExecutionProvider"],
        help="onnxruntime 推理后端顺序",
    )
    parser.add_argument("--rtol", type=float, default=1e-4, help="验证 rtol")
    parser.add_argument("--atol", type=float, default=1e-5, help="验证 atol")
    parser.add_argument("--log", default="INFO", help="日志等级")
    args = parser.parse_args()

    logging.basicConfig(level=getattr(logging, args.log.upper(), logging.INFO))

    model = onnx.load(args.model.as_posix())
    graph_index = build_graph_index(model)

    with args.config.open("r", encoding="utf-8") as f:
        config = json.load(f)

    sub_configs = config.get("compiler", {}).get("sub_configs", [])
    if not sub_configs:
        raise ValueError("配置文件中未找到 compiler.sub_configs。")

    specs: List[SubGraphSpec] = []
    covered_nodes: Set[str] = set()

    for idx, entry in enumerate(sub_configs):
        start = [name for name in entry.get("start_tensor_names", []) if name]
        end = [name for name in entry.get("end_tensor_names", []) if name]
        if not start or not end:
            raise ValueError(f"sub_config[{idx}] 缺少 start/end tensor name。")
        spec = SubGraphSpec(
            label=f"cfg_{idx:02d}",
            start=start,
            end=end,
            node_names=set(),
            source="config",
        )
        nodes = trace_nodes_between(spec, graph_index)
        spec.node_names = nodes
        covered_nodes.update(nodes)
        specs.append(spec)

    leftovers = untouched_components(graph_index.node_order, covered_nodes, graph_index)
    for idx, component in enumerate(leftovers):
        start, end = derive_interface(component, graph_index)
        if not end:
            logging.warning("自动发现的剩余子图 %d 没有输出，跳过。", idx)
            continue
        spec = SubGraphSpec(
            label=f"auto_{idx:02d}",
            start=start,
            end=end,
            node_names=component,
            source="auto",
        )
        specs.append(spec)
        logging.info(
            "自动补充子图 %s: start=%s end=%s (节点数=%d)",
            spec.label,
            spec.start,
            spec.end,
            len(component),
        )

    ordered = ordered_specs(specs, graph_index)

    args.output_dir.mkdir(parents=True, exist_ok=True)
    for spec in ordered:
        spec.output_path = extract_model_file(args.model, spec, args.output_dir, spec.source)

    if args.verify:
        if args.input_npz is None:
            raise ValueError("verify 模式需要 --input-npz 提供输入数据。")
        feed = load_npz_inputs(args.input_npz)
        missing_inputs = graph_index.graph_inputs - feed.keys()
        if missing_inputs:
            raise ValueError(f"npz 中缺少以下模型输入: {sorted(missing_inputs)}")
        verify(args.model, ordered, feed, args.providers, args.rtol, args.atol)


if __name__ == "__main__":
    """
    用法示例:
    python python/VideoX-Fun/scripts/split_onnx_by_subconfigs.py \
        --model /path/to/full.onnx \
        --config python/VideoX-Fun/pulsar2_configs/transformers_subgraph.json \
        --output-dir /tmp/sliced_models \
        --verify \
        --input-npz /path/to/inputs.npz
    """
    main()