#!/usr/bin/env python3
"""
infer.py — Run a trained Whetstone specialist on text input.

Supports both binary specialists (true/false) and N-way classifiers.

Usage:
    # Single query
    python3 specialists/scripts/infer.py \
        --run-dir /mnt/storage/fabricate_runs/whetstone_verification_type \
        --labels "unit,integration,schema,smoke,docs" \
        --text "Sprint 200 integration summary + regression"

    # Batch from file (one text per line)
    python3 specialists/scripts/infer.py \
        --run-dir /mnt/storage/fabricate_runs/whetstone_verification_type \
        --labels "unit,integration,schema,smoke,docs" \
        --input-file my_tasks.txt

    # Interactive REPL
    python3 specialists/scripts/infer.py \
        --run-dir /mnt/storage/fabricate_runs/whetstone_verification_type \
        --labels "unit,integration,schema,smoke,docs"
"""
import sys
import argparse
import numpy as np
from pathlib import Path

FABRICATE_TOOLS = Path("/home/bill/Documents/WhetstoneAI_Fabricate/tools")


def load_specialist(run_dir: str, labels_str: str | None):
    """Load checkpoint, vocab, and labels. Return inference-ready state."""
    sys.path.insert(0, str(FABRICATE_TOOLS))
    from eval_with_vocab import load_checkpoint, load_vocab, load_layers, build_runtime

    ck  = Path(run_dir) / "checkpoint.bin"
    db  = Path(run_dir) / "fabricate.db"
    if not ck.exists():
        raise FileNotFoundError(f"No checkpoint.bin in {run_dir}")
    if not db.exists():
        raise FileNotFoundError(f"No fabricate.db in {run_dir}")

    weights, _ = load_checkpoint(str(ck))
    _, str_to_id = load_vocab(str(db))
    layers = load_layers(str(db))

    V = len(str_to_id)
    _, emb_cnt = layers['token_embedding']
    D = emb_cnt // V
    _, ff1_cnt = layers['block0.ff1_weight']
    F = ff1_cnt // D
    L = sum(1 for k in layers if k.startswith('block') and k.endswith('.qkv_weight'))
    cfg = dict(seq_len=96, D=D, F=F, V=V, L=L)
    runtime = build_runtime(weights, layers, cfg)

    label_list = None
    if labels_str:
        p = Path(labels_str)
        if p.exists():
            label_list = [l.strip() for l in p.read_text().split(',')]
        else:
            label_list = [l.strip() for l in labels_str.split(',')]

    return weights, layers, cfg, runtime, str_to_id, label_list


def predict(text: str, weights, layers, cfg, runtime, str_to_id, label_list) -> tuple[str, float]:
    """Predict label for a single text. Returns (predicted_label, confidence)."""
    sys.path.insert(0, str(FABRICATE_TOOLS))
    from eval_with_vocab import tokenize_batch, forward_batch

    token_ids = tokenize_batch([text], str_to_id, cfg['seq_len'])
    logits = forward_batch(token_ids, weights, layers, cfg, runtime=runtime)  # (1, V)

    if label_list:
        # N-way: restrict to label vocab IDs, take argmax
        label_ids = np.array([str_to_id.get(l, 0) for l in label_list], dtype=np.int32)
        label_logits = logits[0, label_ids]
        probs = np.exp(label_logits - label_logits.max())
        probs /= probs.sum()
        idx = int(np.argmax(probs))
        return label_list[idx], float(probs[idx])
    else:
        # Binary true/false
        true_id  = str_to_id.get('true',  3)
        false_id = str_to_id.get('false', 4)
        true_logit  = logits[0, true_id]
        false_logit = logits[0, false_id]
        prob_true = float(np.exp(true_logit) / (np.exp(true_logit) + np.exp(false_logit)))
        if prob_true >= 0.5:
            return "true", prob_true
        else:
            return "false", 1.0 - prob_true


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--run-dir", required=True,
                    help="Directory with checkpoint.bin and fabricate.db")
    ap.add_argument("--labels", default=None,
                    help="Comma-sep label list, e.g. 'unit,integration,schema,smoke,docs'")
    ap.add_argument("--text",   default=None, help="Single text to classify")
    ap.add_argument("--input-file", default=None, help="File with one text per line")
    ap.add_argument("--eval-tsv", default=None,
                    help="Evaluate accuracy on a TSV file (label<TAB>hops<TAB>text)")
    args = ap.parse_args()

    print(f"Loading specialist from {args.run_dir}...", flush=True)
    state = load_specialist(args.run_dir, args.labels)
    weights, layers, cfg, runtime, str_to_id, label_list = state
    print(f"Loaded. V={cfg['V']} D={cfg['D']} L={cfg['L']}  labels={label_list}")

    if args.eval_tsv:
        correct = total = 0
        with open(args.eval_tsv) as f:
            for line in f:
                parts = line.strip().split('\t', 2)
                if len(parts) < 3:
                    continue
                true_idx = int(parts[0])
                text = parts[2]
                true_label = label_list[true_idx] if label_list else str(true_idx)
                pred, conf = predict(text, *state[:-1], label_list)
                if pred == true_label:
                    correct += 1
                total += 1
        print(f"Accuracy: {correct}/{total} = {100*correct/max(1,total):.1f}%")
        return

    if args.input_file:
        with open(args.input_file) as f:
            lines = [l.strip() for l in f if l.strip()]
        for text in lines:
            pred, conf = predict(text, weights, layers, cfg, runtime, str_to_id, label_list)
            print(f"{pred:15s}  {conf:.3f}  {text}")
        return

    if args.text:
        pred, conf = predict(args.text, weights, layers, cfg, runtime, str_to_id, label_list)
        print(f"{pred}  ({conf:.3f})")
        return

    # Interactive REPL
    print("Enter task descriptions (Ctrl-D to quit):")
    while True:
        try:
            text = input("> ").strip()
        except (EOFError, KeyboardInterrupt):
            break
        if not text:
            continue
        pred, conf = predict(text, weights, layers, cfg, runtime, str_to_id, label_list)
        print(f"  → {pred}  (conf={conf:.3f})")


if __name__ == "__main__":
    main()