#!/usr/bin/env python3
"""Category Extrapolation (Zhao et al., CVPR 2025) for long-tail cuneiform.

For each target ABZ sign, ask an LLM to list visually-similar sibling signs
and provide discriminating feature phrases. Output used as:
  1. Hard-negative pairs for contrastive loss
  2. Sibling-aware confusion graph (cluster hints)
  3. Prompt for VLM re-ranking

Does NOT train or modify models — produces JSON sibling map.
"""
import os, sys, json, argparse, time
from pathlib import Path
from collections import Counter

ROOT = Path("/arf/scratch/stakan/hitit-proje")

def log(m): print(f"[{time.strftime('%H:%M:%S')}] {m}", flush=True)

PROMPT = """You are an expert on Hittite cuneiform signs (Borger ABZ list).
Given the sign **{target}**, list 3–5 other ABZ signs that are most easily
confused with it due to visual similarity in wedge layout (not phonetic).

Respond in this exact JSON format:
{{
  "target": "{target}",
  "siblings": ["SIGN_A", "SIGN_B", "SIGN_C"],
  "distinguishing_features": "<one sentence describing what visually sets {target} apart>"
}}

Only ABZ sign names. No commentary outside JSON. If {target} is ambiguous or unknown, return empty siblings.
"""

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--labels-source', required=True,
                    help='train ckpt with label_to_idx, OR manifest jsonl')
    ap.add_argument('--abz-map', default=str(ROOT / 'hitit_ocr/data/abz_to_sign.json'))
    ap.add_argument('--output', required=True)
    ap.add_argument('--model', default='claude-haiku-4-5-20251001')
    ap.add_argument('--max-labels', type=int, default=0, help='0=all')
    ap.add_argument('--delay-ms', type=int, default=0)
    args = ap.parse_args()

    try: import anthropic
    except ImportError: log("anthropic missing"); sys.exit(1)
    if 'ANTHROPIC_API_KEY' not in os.environ:
        log("ANTHROPIC_API_KEY not set"); sys.exit(1)
    client = anthropic.Anthropic()

    # Load labels
    labels = []
    p = Path(args.labels_source)
    if p.suffix == '.pt':
        import torch
        ck = torch.load(str(p), map_location='cpu', weights_only=False)
        labels = list(ck['label_to_idx'].keys())
    else:
        seen = set()
        for line in open(p):
            r = json.loads(line)
            if r.get('unified_label') and r['unified_label'] not in seen:
                seen.add(r['unified_label']); labels.append(r['unified_label'])
    labels = sorted(set(labels))
    if args.max_labels: labels = labels[:args.max_labels]
    log(f"Labels: {len(labels)}")

    # Resume
    out = {'entries': {}, 'model': args.model}
    if Path(args.output).exists():
        out = json.load(open(args.output))
        out.setdefault('entries', {})
    log(f"Resume: {len(out['entries'])} done")

    for i, lab in enumerate(labels):
        if lab in out['entries']: continue
        try:
            msg = client.messages.create(
                model=args.model, max_tokens=300,
                messages=[{'role': 'user',
                           'content': PROMPT.format(target=lab)}])
            txt = msg.content[0].text.strip()
            # Extract JSON
            jstart = txt.find('{'); jend = txt.rfind('}')
            entry = json.loads(txt[jstart:jend+1])
            out['entries'][lab] = entry
        except Exception as e:
            out['entries'][lab] = {'target': lab, 'siblings': [],
                                    'error': str(e)[:200]}
        if (i + 1) % 20 == 0:
            log(f"  {i+1}/{len(labels)}")
            json.dump(out, open(args.output, 'w'), indent=2)
        if args.delay_ms: time.sleep(args.delay_ms / 1000.0)

    # Build a siblings-only map for easy lookup in training
    sib_map = {k: v.get('siblings', []) for k, v in out['entries'].items()}
    out['sibling_map'] = sib_map

    # Feature matrix: label → indices of siblings within known label set
    label_set = set(labels)
    out['sibling_edges'] = [[a, b] for a, sibs in sib_map.items()
                             for b in sibs if b in label_set and b != a]
    log(f"Sibling edges: {len(out['sibling_edges'])}")
    json.dump(out, open(args.output, 'w'), indent=2)

if __name__ == '__main__':
    main()