#!/usr/bin/env python3
"""
gen_confidence_tier_data.py

Generate synthetic training data for the confidence_tier specialist.

The Whetstone confidence formula (from TaskitemConfidenceAmbiguity.h):
  base = 85
  - 25 if not queueReady
  - 10 per ambiguous requirement (up to 3)
  - 8 per conflict detected
  - 5 if prereq_count > 2
  - 4 if has_dependencies
  clamped to [0, 100]

Tier labels:
  0 = high    (confidence >= 80)
  1 = medium  (confidence 60–79)
  2 = low     (confidence < 60)

The specialist encodes features as a short token sequence:
  "conflicts=N ambiguity=M deps=D prereqs=P queueready=B"

Usage:
    python3 specialists/scripts/gen_confidence_tier_data.py \
        --train-out specialists/data/generated/confidence_tier_train.tsv \
        --eval-out  specialists/data/generated/confidence_tier_eval.tsv
"""
import random
import argparse
from pathlib import Path
from collections import Counter

LABEL_ORDER = ["high", "medium", "low"]
LABEL_INDEX = {lbl: i for i, lbl in enumerate(LABEL_ORDER)}


def compute_confidence(conflicts, ambiguity, deps, prereqs, queue_ready):
    conf = 85
    if not queue_ready:
        conf -= 25
    conf -= min(ambiguity, 3) * 10
    conf -= conflicts * 8
    if prereqs > 2:
        conf -= 5
    if deps > 0:
        conf -= 4
    return max(0, min(100, conf))


def tier(conf):
    if conf >= 80:
        return "high"
    elif conf >= 60:
        return "medium"
    else:
        return "low"


def encode_input(conflicts, ambiguity, deps, prereqs, queue_ready):
    qr = "yes" if queue_ready else "no"
    return f"conflicts={conflicts} ambiguity={ambiguity} deps={deps} prereqs={prereqs} queueready={qr}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--train-out", default="specialists/data/generated/confidence_tier_train.tsv")
    ap.add_argument("--eval-out",  default="specialists/data/generated/confidence_tier_eval.tsv")
    ap.add_argument("--n-per-class", type=int, default=400)
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    rng = random.Random(args.seed)

    # Generate all unique input combinations and label them
    seen = set()
    by_tier: dict[str, list[str]] = {"high": [], "medium": [], "low": []}

    for conflicts in range(0, 10):
        for ambiguity in range(0, 4):
            for deps in range(0, 3):
                for prereqs in range(1, 5):
                    for queue_ready in (True, False):
                        key = (conflicts, ambiguity, deps, prereqs, queue_ready)
                        if key in seen:
                            continue
                        seen.add(key)
                        conf = compute_confidence(*key)
                        t = tier(conf)
                        text = encode_input(*key)
                        by_tier[t].append(text)

    print("Raw tier distribution:", {k: len(v) for k, v in by_tier.items()})

    # Balance to n_per_class (oversample if needed)
    target = args.n_per_class
    balanced = []
    for lbl, examples in by_tier.items():
        rng.shuffle(examples)
        chosen = examples[:]
        while len(chosen) < target:
            chosen.extend(examples)
        chosen = chosen[:target]
        balanced.extend((lbl, ex) for ex in chosen)

    rng.shuffle(balanced)

    n_eval = max(1, int(len(balanced) * 0.15))
    eval_set  = balanced[:n_eval]
    train_set = balanced[n_eval:]

    Path(args.train_out).parent.mkdir(parents=True, exist_ok=True)
    Path(args.eval_out).parent.mkdir(parents=True, exist_ok=True)

    def write_tsv(path, rows):
        with open(path, "w") as f:
            for lbl, text in rows:
                f.write(f"{LABEL_INDEX[lbl]}\t0\t{text}\n")
        print(f"Wrote {len(rows)} rows → {path}")

    write_tsv(args.train_out, train_set)
    write_tsv(args.eval_out,  eval_set)

    for split_name, split in [("train", train_set), ("eval", eval_set)]:
        c = Counter(lbl for lbl, _ in split)
        print(f"  {split_name}: {dict(sorted(c.items()))}")


if __name__ == "__main__":
    main()