import yaml
import onnx
import kneronnxopt
from pathlib import Path

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


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


def check_onnx_opset(onnx_path):
    model  = onnx.load(str(onnx_path), load_external_data=False)
    opsets = {op.domain: op.version for op in model.opset_import}
    return opsets.get("", opsets.get("ai.onnx", 0))


def main():
    cfg         = load_config()
    input_path  = cfg["input"]
    output_path = cfg["output"]
    input_shape = tuple(int(x) for x in cfg["input_shape"].split(","))
    input_name  = cfg["input_name"]
    dup_weights = cfg["duplicate_shared_weights"]

    print(f"Loading ONNX model from '{input_path}'...")
    model = onnx.load(input_path)

    print("Optimizing model...")
    optimized = kneronnxopt.optimize(
        model,
        duplicate_shared_weights=dup_weights,
        overwrite_input_shapes={input_name: input_shape},
        skip_check=True
    )

    onnx.save(optimized, output_path)
    print(f"✅ Optimized ONNX saved to '{output_path}'")
    print(f"Optimized ONNX opset version: {check_onnx_opset(output_path)}")


if __name__ == "__main__":
    main()