"""Create targeted int8 Jeff Core ML variants without modifying the FP16 source."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import coremltools as ct
import coremltools.optimize.coreml as cto


def package_bytes(path: Path) -> int:
    return sum(file.stat().st_size for file in path.rglob("*") if file.is_file())


def selected_constants(model, scheme: str) -> list[str]:
    metadata = cto.get_weights_metadata(model, weight_threshold=2048)
    selected = []
    for name, weight in metadata.items():
        if not weight.child_ops:
            continue
        consumer = weight.child_ops[0].op_type
        embedding = consumer == "gather" and len(weight.val.shape) == 2
        linear = consumer == "linear" and len(weight.val.shape) == 2
        if (scheme == "e8" and embedding) or (scheme == "w8" and (embedding or linear)):
            selected.append(name)
    if not selected:
        raise ValueError(f"no eligible constants found for {scheme}")
    return selected


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", type=Path, default=Path("build/JeffDecision-L128-FP16.mlpackage"))
    parser.add_argument("--scheme", choices=("e8", "w8"), required=True)
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()
    output = args.output or Path(f"build/JeffDecision-L128-{args.scheme.upper()}.mlpackage")
    if output.exists():
        parser.error(f"output already exists: {output}")
    model = ct.models.MLModel(str(args.source), compute_units=ct.ComputeUnit.CPU_ONLY)
    names = selected_constants(model, args.scheme)
    config = cto.OpLinearQuantizerConfig(mode="linear_symmetric", dtype="int8", granularity="per_channel")
    compressed = cto.linear_quantize_weights(
        model, cto.OptimizationConfig(op_name_configs={name: config for name in names})
    )
    compressed.user_defined_metadata["precision"] = args.scheme
    compressed.save(str(output))
    report = {
        "source": str(args.source.resolve()),
        "source_bytes": package_bytes(args.source),
        "output": str(output.resolve()),
        "output_bytes": package_bytes(output),
        "compressed_constants": len(names),
        "scheme": args.scheme,
    }
    Path(f"build/quantize-{args.scheme}.json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2), flush=True)


if __name__ == "__main__":
    main()