import argparse
import shutil
import subprocess
from pathlib import Path

import coremltools as ct
import coremltools.optimize as cto
from coremltools.converters.mil.mil import types


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--src", required=True)
    parser.add_argument("--out-dir", required=True)
    parser.add_argument("--bits", choices=("4", "8"), default="8")
    parser.add_argument("--name", default=None)
    return parser.parse_args()


def main():
    args = parse_args()
    src = Path(args.src)
    out_dir = Path(args.out_dir)
    out_name = args.name or f"model_{args.bits}bit"
    package = out_dir / f"{out_name}.mlpackage"
    compiled = out_dir / f"{out_name}.mlmodelc"
    dtype = types.int4 if args.bits == "4" else types.int8

    out_dir.mkdir(parents=True, exist_ok=True)
    shutil.rmtree(package, ignore_errors=True)
    shutil.rmtree(compiled, ignore_errors=True)

    model = ct.models.MLModel(str(src))
    config = cto.coreml.OptimizationConfig(
        global_config=cto.coreml.OpLinearQuantizerConfig(mode="linear_symmetric", dtype=dtype)
    )
    compressed = cto.coreml.linear_quantize_weights(model, config)
    compressed.save(str(package))

    compile_dir = out_dir / ".compiled"
    shutil.rmtree(compile_dir, ignore_errors=True)
    compile_dir.mkdir(parents=True, exist_ok=True)
    subprocess.run(["xcrun", "coremlcompiler", "compile", str(package), str(compile_dir)], check=True)
    generated = compile_dir / f"{package.stem}.mlmodelc"
    shutil.move(str(generated), str(compiled))
    shutil.rmtree(compile_dir)

    print(package)
    print(compiled)


if __name__ == "__main__":
    main()