"""Object type mix of the training data (CLAUDE.md section 1, Opus review 2026-10-05).

Target share of the accepted tasks (and of the accepted trajectories) per "kind":
  CLAS/INTF 35 % (CLAS 28, INTF 7), CDS (DDLS) 25 %, FUNC 15 %, PROG 10 %, DDIC 10 % (TABL 8, STRU 2),
  MSAG + exception 5 % (MSAG 2.5, EXC 2.5).
A task has the kind of its main contract object; a CLAS whose reference inherits from CX_ is EXC.
DTEL / DOMA tasks are not made yet (the DDIC share is covered by TABL and STRU).
"""
import json
import os
import re

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
POOL = os.path.join(ROOT, "tasks_gen", "train")
TYPE_SHARE = {"CLAS": 28.0, "INTF": 7.0, "DDLS": 25.0, "FUNC": 15.0, "PROG": 10.0, "TABL": 8.0, "STRU": 2.0,
              "MSAG": 2.5, "EXC": 2.5}
# categories a kind supports (docs/faz1-tasarim.md 8): the generator picks the one with the biggest deficit
KIND_CATEGORIES = {"CLAS": "ACDEFGHI", "FUNC": "ACDEFGHI", "PROG": "BCEGH", "DDLS": "BEFHI", "INTF": "A",
                   "TABL": "B", "STRU": "B", "MSAG": "D", "EXC": "D"}
CATEGORY_SHARE = {"A": 10, "B": 15, "C": 15, "D": 10, "E": 15, "F": 10, "G": 10, "I": 5, "H": 10}
_CACHE = {}


def kind_of_task_dir(task_id):
    """Kind of an accepted training task (cached)."""
    if task_id in _CACHE:
        return _CACHE[task_id]
    d = os.path.join(POOL, task_id)
    try:
        t = json.load(open(os.path.join(d, "task.json")))
    except (OSError, ValueError):
        return None
    otype = t.get("object_type") or "CLAS"
    for c in t.get("contract", []):   # K variants and old tasks: the contract object is the truth
        otype = c.get("type", otype)
        break
    kind = otype
    if otype == "CLAS":
        for o in t.get("reference", []):
            if o.get("file") and o.get("name") == (t.get("contract") or [{}])[0].get("name"):
                try:
                    src = open(os.path.join(d, o["file"])).read()
                except OSError:
                    src = ""
                if re.search(r"INHERITING\s+FROM\s+\S*CX_", src, re.I):
                    kind = "EXC"
    _CACHE[task_id] = kind
    return kind


def deficit_pick(counts, share=TYPE_SHARE, allowed=None):
    """Kind with the largest (target - actual) for the next item. counts: {kind: n}."""
    total = sum(counts.values()) + 1
    best, best_def = None, None
    tot_share = sum(share.values())
    for k, s in share.items():
        if allowed and k not in allowed:
            continue
        d = total * s / tot_share - counts.get(k, 0)
        if best_def is None or d > best_def:
            best, best_def = k, d
    return best


def accepted_task_counts(logs_dir=None):
    """{kind: accepted tasks} from the generation logs of the training pool (K variants count by their kind)."""
    import glob
    logs_dir = logs_dir or os.path.join(POOL, "_logs")
    out = {}
    for f in glob.glob(os.path.join(logs_dir, "G*.json")):
        try:
            l = json.load(open(f))
        except (OSError, ValueError):
            continue
        if l.get("accepted"):
            k = kind_of_task_dir(l["id"])
            if k:
                out[k] = out.get(k, 0) + 1
    return out


def below_target(counts, share=TYPE_SHARE):
    """Kinds whose share of `counts` is below the target for the next item (deficit > 0).
    Since 2026-10-05 (Kral + Opus): only these kinds are generated and run; CLAS and FUNC wait until they are at or
    below their target share."""
    total, ss = sum(counts.values()) + 1, sum(share.values())
    return {k for k, v in share.items() if total * v / ss - counts.get(k, 0) > 0}