"""Run the upstream field classifier with shared-prefix tree attention."""
import argparse
import json
import sys
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from tree_decode import FieldPlan, decode_fields, prefill

ROOT = Path(__file__).resolve().parent
SOURCE = ROOT / 'upstream' if (ROOT / 'upstream').exists() else ROOT
sys.path.insert(0, str(SOURCE))
from core.schema import StructuredSchema


@torch.inference_mode()
def classify(model, tokenizer, context, schema, mode='tree'):
    """Match upstream's one-token candidate scoring, including its limitations.

    Returned probabilities are normalized candidate scores, not empirically
    calibrated confidence. Duplicate candidate tokens are explicitly reported.
    """
    device, dtype = model.device, model.dtype
    meta = schema.compile_parallel_metadata(tokenizer)
    suffixes = [list(map(int, row[:length])) for row, length in zip(meta['suffixes_batch'], meta['suffix_lengths'])]
    prompt = (f'<|im_start|>system\nClassify JSON attributes:\n{schema.to_parallel_schema_str()}<|im_end|>\n'
              f'<|im_start|>user\n{context}<|im_end|>\n<|im_start|>assistant\n{{\n')
    ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
    plan = FieldPlan(suffixes, ids.shape[1], device, dtype, tokenizer.pad_token_id or 0)
    cache = prefill(model, ids)
    logits = decode_fields(model, cache, plan, mode)[0]
    fields = {}
    for i, (name, field) in enumerate(meta['field_items']):
        probs = logits[i, meta['cands_per_field'][i]].float().softmax(-1)
        winner = probs.argmax().item()
        fields[name] = {'value': winner == 0 if field.field_type == 'boolean' else field.choices[winner],
                        'candidate_probability': probs[winner].item()}
    return {'mode': mode, 'fields': fields,
            'candidate_collision_fields': [name for (name, _), collision in zip(meta['field_items'], meta['has_collisions']) if collision]}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--preset', default='fintech_fraud')
    ap.add_argument('--mode', choices=['tree', 'batch'], default='tree')
    ap.add_argument('--device', choices=['auto', 'mps', 'cpu'], default='auto')
    args = ap.parse_args()
    torch.set_num_threads(4)
    name = 'Qwen/Qwen2.5-1.5B-Instruct'
    tokenizer = AutoTokenizer.from_pretrained(name)
    device = args.device
    if device == 'auto':
        device = 'mps' if torch.backends.mps.is_available() else 'cpu'
    model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.float16 if device == 'mps' else torch.float32,
                                                attn_implementation='sdpa').to(device).eval()
    preset = json.loads((SOURCE / 'presets' / f'{args.preset}.json').read_text(encoding='utf-8'))
    result = classify(model, tokenizer, preset['context'], StructuredSchema(preset['schema']), args.mode)
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()