import yaml
import onnx
from pathlib import Path

CONFIG_PATH = Path(__file__).parent.parent / "config.yaml"

# Fixed — graph-specific node, not a runtime input
TARGET = '/model.23/Sigmoid'


def load_config() -> dict:
    with open(CONFIG_PATH, "r") as f:
        return yaml.safe_load(f)["remove_sigmoid"]


def remove_sigmoid(model):
    graph        = model.graph
    sigmoid_node = None

    for node in graph.node:
        if node.name == TARGET:
            sigmoid_node = node
            break

    if sigmoid_node is None:
        available = [n.name for n in graph.node if n.op_type == 'Sigmoid']
        raise RuntimeError(f"Node '{TARGET}' not found.\nAll sigmoids:\n" +
                           '\n'.join(available))

    sig_in  = sigmoid_node.input[0]
    sig_out = sigmoid_node.output[0]

    redirected = 0
    for node in graph.node:
        for i, inp in enumerate(node.input):
            if inp == sig_out:
                node.input[i] = sig_in
                redirected += 1
                print(f"Redirected '{node.name}' input[{i}]: {sig_out} → {sig_in}")

    print(f"Total redirections: {redirected}")
    graph.node.remove(sigmoid_node)
    print(f"Removed: {TARGET}")

    remaining = [n.name for n in graph.node if n.op_type == 'Sigmoid']
    print(f"\nRemaining sigmoids: {len(remaining)}")
    assert all('act/Sigmoid' in n for n in remaining), "WRONG sigmoid was removed!"
    print("Sanity check passed — only output sigmoid was removed")

    return model


def main():
    cfg = load_config()

    print(f"Loading ONNX model from '{cfg['input']}'...")
    model = onnx.load(cfg["input"])
    model = remove_sigmoid(model)

    onnx.checker.check_model(model)
    onnx.save(model, cfg["output"])
    print(f"\nSaved: {cfg['output']}")


if __name__ == "__main__":
    main()