"""Convert a Fabryka full-vocabulary response to an array indexed by token ID.

Usage: python export_logprobs.py teacher.json --request request.json --out teacher
Requires: numpy. This script makes no network requests.
"""
import argparse
import json
from pathlib import Path

import numpy as np


def export_response(response, request=None):
    lp = response['choices'][0]['logprobs']
    if lp.get('full_vocabulary') is not True:
        raise ValueError('Use logprobs=true and top_logprobs=-1')
    vocab = lp['vocab_size']
    if type(vocab) is not int or vocab <= 0:
        raise ValueError('Invalid vocabulary size')
    positions = lp['content']
    if not positions:
        raise ValueError('No generated positions in response')
    rows = []
    for pos in positions:
        alternatives = pos['top_logprobs']
        ids = [t['id'] for t in alternatives]
        if (len(ids) != vocab or any(type(i) is not int for i in ids)
                or set(ids) != set(range(vocab))):
            raise ValueError('Incomplete vocabulary or duplicate token IDs')
        values = np.array([t['logprob'] for t in alternatives], dtype=np.float32)
        if not np.isfinite(values).all() or (values > 0).any():
            raise ValueError('Invalid log probabilities')
        row = np.empty(vocab, dtype=np.float32)
        row[np.array(ids)] = values
        rows.append(row)
    matrix = np.stack(rows)
    metadata = {
        'model': response.get('model'),
        'response_id': response.get('id'),
        'system_fingerprint': response.get('system_fingerprint'),
        'shape': list(matrix.shape),
        'axis_0': 'generated position',
        'axis_1': 'teacher vocabulary token ID',
        'dtype': 'float32',
        'values': 'natural log probabilities as returned; normalize in training',
        'generated_tokens': [{k: p.get(k) for k in ('id', 'token', 'bytes')} for p in positions],
        'request': request,
        'usage': response.get('usage'),
        'reproducibility_note': 'Also record the exact served checkpoint, tokenizer, chat template and quantization. The public fingerprint is not an immutable revision pin.',
    }
    return matrix, metadata


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('response', type=Path)
    parser.add_argument('--request', type=Path)
    parser.add_argument('--out', default='teacher', type=Path)
    args = parser.parse_args()
    response = json.loads(args.response.read_text())
    request = json.loads(args.request.read_text()) if args.request else None
    matrix, metadata = export_response(response, request)
    np.save(str(args.out) + '.npy', matrix, allow_pickle=False)
    Path(str(args.out) + '.metadata.json').write_text(json.dumps(metadata, indent=2, ensure_ascii=True) + '\n')
    print(f'Saved {matrix.shape[0]} positions x {matrix.shape[1]} tokens to {args.out}.npy')


if __name__ == '__main__':
    main()