import itertools
import json
import random
from pathlib import Path

from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer
import torch
from torch.utils.data import Dataset


# official DeepSeek-V3 BPE tokenizer — downloaded once from HF
_tokenizer_path = hf_hub_download('deepseek-ai/DeepSeek-V3', 'tokenizer.json')
_tok = Tokenizer.from_file(_tokenizer_path)
vocab_size = _tok.get_vocab_size()
categories = ('story', 'code', 'math')


def encode(text):
    return _tok.encode(text).ids


def decode(ids):
    return _tok.decode(ids)


def token_text(token_id):
    # the raw BPE token string; Ġ marks a leading space, Ċ a newline
    return _tok.id_to_token(token_id).replace('Ġ', ' ').replace('Ċ', '⏎')


def make_math(n, max_val=1000, seed=0):
    # synthetic arithmetic strings eg '13 + 27 = 40\n'
    rng = random.Random(seed)
    out = []
    for _ in range(n):
        op = rng.choice(['+', '-', '*'])
        a, b = rng.randint(0, max_val), rng.randint(0, max_val)
        c = {'+': a + b, '-': a - b, '*': a * b}[op]
        out.append(f'{a} {op} {b} = {c}\n')
    return out


def _load_or_cache(name, n, text_field, cache_file):
    # download once via HF streaming, then cache as JSON for reproducible reruns
    cache_path = Path('data') / cache_file
    if cache_path.exists():
        return json.loads(cache_path.read_text())

    from datasets import load_dataset
    ds = load_dataset(name, split='train', streaming=True)
    texts = [ex[text_field] for ex in itertools.islice(ds, n)]
    cache_path.parent.mkdir(exist_ok=True)
    cache_path.write_text(json.dumps(texts))
    return texts


def load_stories(n):
    return _load_or_cache('roneneldan/TinyStories', n, 'text', f'stories_{n}.json')


def load_code(n, max_bytes_per_sample=2000):
    texts = _load_or_cache('codeparrot/codeparrot-clean-valid', n, 'content', f'code_{n}.json')
    return [t[:max_bytes_per_sample] for t in texts]


# build one byte stream per category and slide a window over each separately,
# so every (input, target) pair has a single, well-defined category. that
# per-token category is what the heatmap aggregates over later
class MixedDataset(Dataset):
    def __init__(self, samples_per_cat, seq_len, stride):
        self.inputs = []
        self.targets = []
        self.cat_ids = []

        streams = [
            ''.join(load_stories(samples_per_cat)),
            ''.join(load_code(samples_per_cat // 5)),
            ''.join(make_math(samples_per_cat * 10)),
        ]
        for cat_id, text in enumerate(streams):
            tokens = torch.tensor(encode(text), dtype=torch.long)
            for i in range(0, len(tokens) - seq_len, stride):
                self.inputs.append(tokens[i : i + seq_len])
                self.targets.append(tokens[i + 1 : i + seq_len + 1])
                self.cat_ids.append(cat_id)

    def __len__(self):
        return len(self.inputs)

    def __getitem__(self, i):
        return self.inputs[i], self.targets[i], self.cat_ids[i]


# windows over one text stream per category. stride = seq_len gives
# non-overlapping windows, so the same tokens are not seen many times. heavy
# window overlap (small stride) is what let the earlier model memorize
class WindowDataset(Dataset):
    def __init__(self, streams, seq_len, stride):
        self.inputs = []
        self.targets = []
        self.cat_ids = []
        for cat_id, text in enumerate(streams):
            tokens = torch.tensor(encode(text), dtype=torch.long)
            for i in range(0, len(tokens) - seq_len, stride):
                self.inputs.append(tokens[i : i + seq_len])
                self.targets.append(tokens[i + 1 : i + seq_len + 1])
                self.cat_ids.append(cat_id)

    def __len__(self):
        return len(self.inputs)

    def __getitem__(self, i):
        return self.inputs[i], self.targets[i], self.cat_ids[i]


# split whole documents (not windows) into train / val, so no story, code file,
# or equation appears in both. the earlier random split over overlapping windows
# leaked train content into val, which made the val loss look far too good
def build_train_val_datasets(samples_per_cat, seq_len, val_fraction=0.1):
    documents_per_cat = [
        load_stories(samples_per_cat),
        load_code(samples_per_cat // 5),
        make_math(samples_per_cat * 10),
    ]

    train_streams = []
    val_streams = []
    for documents in documents_per_cat:
        num_val = max(1, int(val_fraction * len(documents)))
        train_streams.append(''.join(documents[:-num_val]))
        val_streams.append(''.join(documents[-num_val:]))

    train_dataset = WindowDataset(train_streams, seq_len, stride=seq_len)
    val_dataset   = WindowDataset(val_streams,   seq_len, stride=seq_len)
    return train_dataset, val_dataset


if __name__ == '__main__':
    print(make_math(5))