"""Export loading, seeded subsets, steps of similar-length posts and token-budget micro-batches (pure Python).

One optimiser step takes ``posts_per_step`` posts. Steps are drawn from length-sorted pools so their posts have
similar lengths, then shuffled; each step is split into micro-batches whose padded size (rows x longest row)
stays within the token budget. The loss is normalised per step (``step_counts``), so how a step is split into
micro-batches does not change its gradient. No packing: the Gated DeltaNet state does not reset between
packed posts.
"""
from __future__ import annotations

import hashlib
import json
from pathlib import Path
import random


def load_jsonl(path) -> list[dict]:
    with open(path, encoding='utf-8') as fh:
        return [json.loads(line) for line in fh if line.strip()]


def write_jsonl(path, rows) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + '.tmp')
    with open(tmp, 'w', encoding='utf-8') as fh:
        for row in rows:
            fh.write(json.dumps(row, ensure_ascii=False, allow_nan=False) + '\n')
    tmp.replace(path)


def load_split(data_dir, split: str) -> list[dict]:
    return load_jsonl(Path(data_dir) / ('%s.jsonl' % split))


def seeded_order(ids, seed) -> list:
    return sorted(ids, key=lambda i: hashlib.sha256(json.dumps([str(seed), i]).encode()).hexdigest())


def subset(rows: list[dict], n: int, seed) -> list[dict]:
    """The first n rows in seeded hash order of their ids (independent of file order); all rows if n >= len."""
    if not n or n >= len(rows):
        return list(rows)
    keep = set(seeded_order([r['id'] for r in rows], seed)[:n])
    return [r for r in rows if r['id'] in keep]


def steps(lengths: list[int], posts_per_step: int, seed: int, pool: int = 64) -> list[list[int]]:
    """Shuffled steps of similar-length posts: shuffle, sort within pools of posts_per_step*pool, chunk, shuffle."""
    rng = random.Random(seed)
    idx = list(range(len(lengths)))
    rng.shuffle(idx)
    out = []
    for s in range(0, len(idx), posts_per_step * pool):
        chunk = sorted(idx[s:s + posts_per_step * pool], key=lambda i: (lengths[i], i))
        out += [chunk[i:i + posts_per_step] for i in range(0, len(chunk), posts_per_step)]
    rng.shuffle(out)
    return out


def micro_batches(step: list[int], lengths: list[int], token_budget: int) -> list[list[int]]:
    """Split one step so each micro-batch's padded size (rows x longest) is at most token_budget."""
    order = sorted(step, key=lambda i: (-lengths[i], i))
    out, current = [], []
    for i in order:
        width = max([lengths[j] for j in current] + [lengths[i]])
        if current and width * (len(current) + 1) > token_budget:
            out.append(current)
            current = []
        current.append(i)
    if current:
        out.append(current)
    return out


def step_counts(labels: list[list[int]]) -> list[int]:
    """Non-null targets per field over one step's label rows (-1 marks a null); at least 1 per field."""
    if not labels:
        return []
    return [max(1, sum(1 for row in labels if row[f] >= 0)) for f in range(len(labels[0]))]


def length_order(lengths: list[int]) -> list[int]:
    return sorted(range(len(lengths)), key=lambda i: (lengths[i], i))


def inference_batches(lengths: list[int], token_budget: int) -> list[list[int]]:
    """Length-sorted batches for prediction, each within the padded token budget."""
    out, current = [], []
    for i in length_order(lengths):
        if current and lengths[i] * (len(current) + 1) > token_budget:
            out.append(current)
            current = []
        current.append(i)
    if current:
        out.append(current)
    return out