#!/usr/bin/env python3
"""oracle7-vision on ChartQA (HuggingFaceM4/ChartQA). Development = the VALIDATION split (every rule was written and
tuned there); held-out = the TEST split (2,500 questions: 1,250 human-written, 1,250 machine-generated), scored ONCE.
Metric: ChartQA relaxed accuracy (numbers within 5%, else exact match); abstentions count as wrong. Reports accuracy
overall, per split half, per answering rule, coverage, and accuracy on answered questions.
usage: eval_chartqa.py val|test [limit] [offset]"""
import collections
import io
import json
import os
import sys
import time

import numpy as np
import pyarrow.parquet as pq
from PIL import Image

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
sys.path.insert(0, ROOT)
from vision import chartqa as C  # noqa: E402
from vision import glyphs as G  # noqa: E402


def main():
    split = sys.argv[1]
    limit = int(sys.argv[2]) if len(sys.argv) > 2 else None
    offset = int(sys.argv[3]) if len(sys.argv) > 3 else 0
    t = pq.read_table(os.path.join(ROOT, "data", "chartqa_%s.parquet" % split))
    d = t.slice(offset).to_pydict()
    n = len(d["query"]) if limit is None else min(limit, len(d["query"]))
    fonts = G.discover_fonts([os.path.join(ROOT, "fonts")])
    store = G.GlyphStore([f for f in fonts if not G.is_heldout_family(f[1])])
    cache = {}
    right = answered = 0
    by_rule, by_half = collections.defaultdict(lambda: [0, 0]), collections.defaultdict(lambda: [0, 0])
    samples = []
    t0 = time.time()
    for i in range(n):
        raw = d["image"][i]["bytes"]
        key = hash(raw)
        if key not in cache:
            gray = np.asarray(Image.open(io.BytesIO(raw)).convert("L"))
            cache = {key: G.read_image(store, gray)}           # questions on one chart are adjacent
        ans, rule, ev = C.answer(d["query"][i], cache[key])
        ok = C.relaxed_match(ans, d["label"][i])
        right += ok
        answered += ans is not None
        by_rule[rule][0] += ok
        by_rule[rule][1] += 1
        half = "human" if d["human_or_machine"][i] == 0 else "machine"
        by_half[half][0] += ok
        by_half[half][1] += 1
        if len(samples) < 25:
            samples.append({"q": d["query"][i], "gold": d["label"][i], "answer": ans, "rule": rule, "ok": ok})
        if (i + 1) % 250 == 0:
            print("  %d/%d  acc %.3f  (%.0fs)" % (i + 1, n, right / (i + 1), time.time() - t0), flush=True)
    res = {"split": split, "n": n, "relaxed_accuracy": round(right / n, 4), "coverage": round(answered / n, 4),
           "accuracy_on_answered": round(right / max(1, answered), 4),
           "by_half": {k: {"acc": round(v[0] / v[1], 4), "n": v[1]} for k, v in by_half.items()},
           "by_rule": {k: {"acc": round(v[0] / v[1], 4), "n": v[1]} for k, v in sorted(by_rule.items())},
           "seconds": round(time.time() - t0), "samples": samples}
    print(json.dumps({k: v for k, v in res.items() if k != "samples"}, indent=1))
    os.makedirs(os.path.join(ROOT, "results"), exist_ok=True)
    tag = split if (limit is None and offset == 0) else "%s_%d_%d" % (split, offset, n)
    json.dump(res, open(os.path.join(ROOT, "results", "chartqa_%s.json" % tag), "w"), indent=1)


if __name__ == "__main__":
    main()