# tools/calc_avg_bits_0920.py
import argparse, os
import torch
from pointcept.models.builder import build_model

# 你的项目里怎么读 cfg，就怎么 import：
from pointcept.utils.config import get_cfg_from_file

from pointcept.models.utils.quant_0920 import (
    apply_quantization_0920, print_avg_bits_0920
)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config-file", required=True)
    ap.add_argument("--options", nargs=argparse.REMAINDER, default=None)
    ap.add_argument("--ckpt", default="", help="optional: load checkpoint then wrap quant")
    ap.add_argument("--wbits", type=int, default=2)
    ap.add_argument("--abits", type=int, default=8)
    args = ap.parse_args()

    cfg = get_cfg_from_file(args.config_file)
    if args.options:
        cfg.merge_from_list(args.options)

    model = build_model(cfg.model)
    if args.ckpt and os.path.isfile(args.ckpt):
        sd = torch.load(args.ckpt, map_location="cpu")
        if "state_dict" in sd:
            sd = sd["state_dict"]
        model.load_state_dict(sd, strict=False)

    model = apply_quantization_0920(model, enable=True, w_bits=args.wbits, a_bits=args.abits,
                                    quantize_first_last=False, exclude_name_hints=["cls_head", "embedding.stem"])
    print_avg_bits_0920(model)

if __name__ == "__main__":
    main()