"""Agreement of shipped Core ML packages with native Fast Decisions predictions on bucket-fitting heads."""
import argparse
import json
import re
from pathlib import Path

import coremltools as ct
import numpy as np
from huggingface_hub import snapshot_download

from convert_names import MODEL_ID, MODEL_REVISION
from preprocessing import load_processor, prepare_decision
from runtime import decode


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--native-rows", default="build/fast-decisions-native.rows.jsonl")
    parser.add_argument("--per-domain", type=int, default=20)
    parser.add_argument("--out", default="build/verify-packages.json")
    parser.add_argument("packages", nargs="+")
    args = parser.parse_args()
    data = Path(snapshot_download("fastino/fast-decisions", repo_type="dataset",
                                  revision="1a33070cabf94ce2e29105482dd2ef6c157ad7f2"))
    processor = load_processor(snapshot_download(MODEL_ID, revision=MODEL_REVISION))
    native = {(r["domain"], r["index"], r["task"]): r["predicted"]
              for r in map(json.loads, Path(args.native_rows).open())}
    domains = sorted({key[0] for key in native})
    report = {}
    for package in args.packages:
        length, heads, options = map(int, re.search(r"_L(\d+)_H(\d+)_K(\d+)", package).groups())
        model = ct.models.MLModel(package, compute_units=ct.ComputeUnit.ALL)
        checked = agree = 0
        for domain in domains:
            used = 0
            for index, line in enumerate((data / f"{domain}.jsonl").open()):
                if used >= args.per_domain:
                    break
                row = json.loads(line)
                for head in row["output"]["classifications"]:
                    tasks = {head["task"]: head["labels"]}
                    try:
                        arrays = prepare_decision(processor, row["input"], tasks, length, heads, options)
                    except ValueError:
                        continue
                    got = decode(tasks, np.asarray(model.predict(arrays)["logits"])[0])[head["task"]]
                    checked += 1
                    agree += [got["label"]] == native[(domain, index, head["task"])]
                    used += 1
        report[Path(package).name] = {"checked_heads": checked, "agree_with_native": agree}
        print(Path(package).name, checked, agree, flush=True)
    Path(args.out).write_text(json.dumps(report, indent=2) + "\n")


if __name__ == "__main__":
    main()