"""Small JSON files the course widgets draw from (public/llm/data/*.json).

    python course_data.py --work ~/claude-projects/tinyllm-work --out ../../public/llm/data
"""
import argparse
import collections
import glob
import json
import os
import random

import numpy as np

from prepare import SOURCES, clean
from tokenizer import Tokenizer


def corpus(work, out, tok):
    stats = json.load(open(f'{work}/data/corpus_stats.json'))
    rng = random.Random(3)
    samples = {}
    lengths = {}
    for name, read in SOURCES.items():
        docs, n = [], 0
        for text in read(f'{work}/raw'):
            n += 1
            if n % 997 == 0:
                t = clean(text)
                if t:
                    docs.append(t)
            if len(docs) >= 400:
                break
        lens = [len(tok.encode(d, allow_special=False)) for d in docs]
        lengths[name] = np.histogram(lens, bins=20, range=(0, 800))[0].tolist()
        samples[name] = [d[:1400] for d in rng.sample(docs, 24)]
    ids = np.memmap(f'{work}/data/val.bin', dtype=np.uint16, mode='r')
    counts = collections.Counter(ids.tolist())
    top = [{'id': int(i), 'text': tok.decode([int(i)]), 'n': int(c)} for i, c in counts.most_common(400)]
    ranks = sorted(counts.values(), reverse=True)
    zipf = [[r + 1, int(ranks[r])] for r in sorted({int(x) for x in np.unique(np.geomspace(1, len(ranks), 60).astype(int) - 1)})]
    json.dump({'stats': stats, 'samples': samples, 'lengths': lengths, 'length_bins': 40, 'top': top, 'zipf': zipf,
               'val_tokens': int(len(ids))}, open(f'{out}/corpus.json', 'w'), separators=(',', ':'))


def train_log(work, out, run='base', name='train-log'):
    path = f'{work}/runs/{run}/log.jsonl'
    if not os.path.exists(path):
        return
    lines = [json.loads(l) for l in open(path) if l.strip()]
    head = lines[0]
    train = [{k: r[k] for k in ('step', 'loss', 'lr', 'tokens', 'T', 'norm') if k in r} for r in lines[1:] if 'loss' in r]
    evals = [{k: r[k] for k in ('step', 'val', 'tokens', 'time', 'samples') if k in r} for r in lines[1:] if 'val' in r]
    extra = {k: head[k] for k in ('train_rows', 'conversations', 'trainable') if k in head}
    json.dump({'config': head.get('config'), 'args': head.get('args'), 'params': head.get('params'), 'steps': head.get('steps'),
               **extra, 'train': train, 'evals': evals}, open(f'{out}/{name}.json', 'w'), separators=(',', ':'))


def lora_logs(work, out):
    res = {}
    for style in ('pirate', 'poet'):
        path = f'{work}/runs/{style}/log.jsonl'
        if not os.path.exists(path):
            continue
        lines = [json.loads(l) for l in open(path) if l.strip()]
        head = lines[0]
        evals = [r for r in lines[1:] if 'val' in r]
        res[style] = {'rank': head['args']['rank'], 'alpha': head['args']['alpha'], 'trainable': head['trainable'],
                      'params': head['params'], 'steps': head['steps'], 'conversations': head.get('conversations'),
                      'minutes': round((evals[-1]['time'] if evals else 0) / 60, 1),
                      'train': [{'step': r['step'], 'loss': r['loss']} for r in lines[1:] if 'loss' in r],
                      'evals': [{'step': r['step'], 'val': r['val']} for r in evals]}
    if res:
        json.dump(res, open(f'{out}/train-log-lora.json', 'w'), separators=(',', ':'))


def chats(work, out):
    path = f'{work}/data/chats.jsonl'
    if not os.path.exists(path):
        return
    rows = [json.loads(l) for l in open(path)]
    rng = random.Random(5)
    by = collections.defaultdict(list)
    for r in rows:
        by[r['source']].append(r)
    sample = {k: rng.sample(v, min(12, len(v))) for k, v in by.items()}
    json.dump({'counts': {k: len(v) for k, v in by.items()}, 'samples': sample}, open(f'{out}/chats.json', 'w'), separators=(',', ':'))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--work', default='.')
    ap.add_argument('--out', default='../../public/llm/data')
    ap.add_argument('--only', default='')
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)
    tok = Tokenizer.load(f'{args.work}/data/tokenizer.json')
    todo = args.only.split(',') if args.only else ['corpus', 'log', 'chats']
    if 'corpus' in todo:
        corpus(args.work, args.out, tok)
    if 'log' in todo:
        train_log(args.work, args.out)
        for run in ('gpt1', 'gpt4', 'chat'):
            train_log(args.work, args.out, run, f'train-log-{run}')
        lora_logs(args.work, args.out)
    if 'chats' in todo:
        chats(args.work, args.out)


if __name__ == '__main__':
    main()