from __future__ import annotations
import random
import time
from decon import is_contaminated
STAGE_BOUNDARIES = [0.7, 0.9, 1.0]
STAGE_MIX = [{'web': 0.74, 'math': 0.09, 'code': 0.07, 'arithmetic': 0.05, 'dialogue': 0.05}, {'web': 0.59, 'math': 0.18, 'code': 0.12, 'arithmetic': 0.06, 'dialogue': 0.05}, {'web': 0.5, 'math': 0.24, 'code': 0.11, 'arithmetic': 0.09, 'dialogue': 0.06}]
WEB_MIX = {'fineweb_edu': 0.6, 'dclm': 0.4}
MIX_PRESETS = {'smollm2': {'fineweb_edu': 0.576, 'dclm': 0.384, 'cosmopedia': 0.04}, 'dclm_heavy': {'fineweb_edu': 0.384, 'dclm': 0.576, 'cosmopedia': 0.04}, 'smollm2_math': {'fineweb_edu': 0.55, 'dclm': 0.37, 'cosmopedia': 0.04, 'math': 0.04}, 'dclm_heavy_math': {'fineweb_edu': 0.37, 'dclm': 0.55, 'cosmopedia': 0.04, 'math': 0.04}, 'cooldown': {'fineweb_edu': 0.75, 'cosmopedia': 0.25}}
MUSEKO_MIX = {'web': 0.9, 'arithmetic': 0.1}
EOS = '<|endoftext|>'

def get_stage_mix(progress: float, mix_mode: str='default', arith_frac: float | None=None) -> dict:
    if mix_mode == 'museko':
        mix = dict(MUSEKO_MIX)
    else:
        mix = None
        for boundary, stage in zip(STAGE_BOUNDARIES, STAGE_MIX):
            if progress <= boundary:
                mix = dict(stage)
                break
        if mix is None:
            mix = dict(STAGE_MIX[-1])
    if arith_frac is None:
        return mix
    other = {k: v for k, v in mix.items() if k != 'arithmetic'}
    total_other = sum(other.values()) or 1.0
    scale = (1.0 - arith_frac) / total_other
    out = {k: v * scale for k, v in other.items()}
    out['arithmetic'] = arith_frac
    return out

def _cycling_stream(load_fn, name: str):
    backoff = 5.0
    while True:
        try:
            it = iter(load_fn())
            backoff = 5.0
        except Exception as e:
            msg = str(e)[:140]
            print(f'[data] stream {name} failed to start, retry in {backoff:.0f}s: {msg}', flush=True)
            time.sleep(backoff)
            backoff = min(backoff * 2, 300.0)
            continue
        while True:
            try:
                row = next(it)
            except StopIteration:
                break
            except Exception as e:
                print(f'[data] stream {name} row error, skipping: {e}', flush=True)
                continue
            yield row

def _split(ds, rank: int, world_size: int, name: str=''):
    if world_size <= 1:
        return ds
    n = None
    try:
        n = ds.n_shards
    except Exception:
        pass
    if n is not None and n >= world_size:
        return ds.shard(num_shards=world_size, index=rank)
    print(f'[data] {name}: {n} files < {world_size} slices, striping rows instead', flush=True)
    return _RowStripe(ds, rank, world_size)

class _RowStripe:

    def __init__(self, ds, rank, world_size):
        self.ds, self.rank, self.world_size = (ds, rank, world_size)

    def __iter__(self):
        for i, row in enumerate(self.ds):
            if i % self.world_size == self.rank:
                yield row

def _shard(ds, rank: int, world_size: int):
    if world_size > 1:
        return ds.shard(num_shards=world_size, index=rank)
    return ds

def _rank_filter(gen, rank: int, world_size: int):
    if world_size <= 1:
        yield from gen
        return
    for i, item in enumerate(gen):
        if i % world_size == rank:
            yield item

def _fineweb_edu_texts(rank: int=0, world_size: int=1):
    from datasets import load_dataset

    def _load():
        ds = load_dataset('HuggingFaceFW/fineweb-edu', 'sample-100BT', split='train', streaming=True)
        return _split(ds, rank, world_size, 'fineweb-edu')
    for row in _cycling_stream(_load, 'fineweb-edu'):
        yield row['text']

def _dclm_texts(rank: int=0, world_size: int=1):
    from datasets import load_dataset

    def _load():
        ds = load_dataset('mlfoundations/dclm-baseline-1.0', split='train', streaming=True)
        return _split(ds, rank, world_size, 'dclm')
    for row in _cycling_stream(_load, 'dclm'):
        yield row['text']

def _cosmopedia_texts(rank: int=0, world_size: int=1):
    from datasets import load_dataset

    def _load():
        ds = load_dataset('HuggingFaceTB/smollm-corpus', 'cosmopedia-v2', split='train', streaming=True)
        return _split(ds, rank, world_size, 'cosmopedia-v2')
    for row in _cycling_stream(_load, 'cosmopedia-v2'):
        t = row.get('text') or ''
        if t:
            yield t

def _finemath_texts(rank: int=0, world_size: int=1):
    from datasets import load_dataset

    def _load():
        return load_dataset('HuggingFaceTB/finemath', 'finemath-4plus', split='train', streaming=True)
    gen = (row['text'] for row in _cycling_stream(_load, 'finemath'))
    yield from _rank_filter(gen, rank, world_size)

def _code_texts(rank: int=0, world_size: int=1):
    from datasets import load_dataset

    def _load():
        return load_dataset('semeru/Code-Text-Python', split='train', streaming=True)

    def _texts():
        for row in _cycling_stream(_load, 'code-text-python'):
            docstring = row.get('docstring') or ''
            code = row.get('code') or row.get('original_string') or ''
            if not code:
                continue
            yield (f'"""{docstring}"""\n{code}' if docstring else code)
    yield from _rank_filter(_texts(), rank, world_size)

def _smoltalk_texts(rank: int=0, world_size: int=1):
    from datasets import load_dataset

    def _load():
        return load_dataset('HuggingFaceTB/smol-smoltalk', split='train', streaming=True)

    def _texts():
        for row in _cycling_stream(_load, 'smol-smoltalk'):
            parts = []
            for msg in row['messages']:
                role = msg.get('role', 'user')
                parts.append(f"<|{role}|>\n{msg.get('content', '')}")
            yield '\n'.join(parts)
    yield from _rank_filter(_texts(), rank, world_size)

def _realmath_texts(rank: int=0, world_size: int=1):
    import re as _re
    from datasets import load_dataset

    def _orca():
        return load_dataset('microsoft/orca-math-word-problems-200k', split='train', streaming=True)

    def _gsm8k():
        return load_dataset('openai/gsm8k', 'main', split='train', streaming=True)
    calc = _re.compile('<<[^>]*>>')

    def _texts():
        orca = _cycling_stream(_orca, 'orca-math')
        gsm = _cycling_stream(_gsm8k, 'gsm8k')
        while True:
            row = next(orca) if random.random() < 0.85 else next(gsm)
            q = (row.get('question') or '').strip()
            a = calc.sub('', row.get('answer') or '').strip()
            if not q or not a:
                continue
            yield f'{q}\n{a}'
    yield from _rank_filter(_texts(), rank, world_size)

def _arithmetic_texts(rank: int=0, world_size: int=1, source: str='wordmath'):
    import random as _random
    rng = _random.Random(1000 + rank)
    if source == 'drill':
        from gen_arithmetic import gen_example as _gen
        while True:
            yield _gen()
    from gen_wordmath import gen_example as _gen
    if source == 'wordmath':
        while True:
            yield _gen(rng)
    real = _realmath_texts(rank, world_size)
    while True:
        yield (next(real) if rng.random() < 0.75 else _gen(rng))

def build_web_stream(reference_ngrams, decon_enabled=True, rank=0, world_size=1):
    fw = _fineweb_edu_texts(rank, world_size)
    dclm = _dclm_texts(rank, world_size)
    while True:
        src = fw if random.random() < WEB_MIX['fineweb_edu'] else dclm
        text = next(src)
        if decon_enabled and is_contaminated(text, reference_ngrams):
            continue
        yield text

def build_preset_streams(preset: str, reference_ngrams, decon_enabled=True, rank=0, world_size=1):

    def _filtered(gen):
        for text in gen:
            if decon_enabled and is_contaminated(text, reference_ngrams):
                continue
            yield text
    mix = MIX_PRESETS[preset]
    builders = {'fineweb_edu': _fineweb_edu_texts, 'dclm': _dclm_texts, 'cosmopedia': _cosmopedia_texts, 'math': _finemath_texts}
    return {k: _filtered(builders[k](rank, world_size)) for k in mix}

def build_source_streams(reference_ngrams, decon_enabled=True, rank=0, world_size=1, mix_mode='default', arith_source='wordmath'):

    def _filtered(gen, apply_decon):
        for text in gen:
            if apply_decon and decon_enabled and is_contaminated(text, reference_ngrams):
                continue
            yield text
    streams = {'web': build_web_stream(reference_ngrams, decon_enabled, rank, world_size), 'arithmetic': _filtered(_arithmetic_texts(rank, world_size, arith_source), True)}
    if mix_mode != 'museko':
        streams['math'] = _filtered(_finemath_texts(rank, world_size), True)
        streams['code'] = _filtered(_code_texts(rank, world_size), True)
        streams['dialogue'] = _filtered(_smoltalk_texts(rank, world_size), True)
    return streams

class MixedTokenStream:

    def __init__(self, tokenizer, streams: dict, seq_len: int, seed: int=0, mix_mode: str='default', arith_frac: float | None=None, preset: str | None=None):
        self.preset = preset
        self.tokenizer = tokenizer
        self.streams = streams
        self.seq_len = seq_len
        self.eos_id = tokenizer.token_to_id(EOS)
        assert self.eos_id is not None, 'tokenizer is missing <|endoftext|>'
        self.buffer: list[int] = []
        self.rng = random.Random(seed)
        self.progress = 0.0
        self.mix_mode = mix_mode
        self.arith_frac = arith_frac

    def _pull_one_doc_tokens(self) -> list[int]:
        if self.preset is not None:
            mix = MIX_PRESETS[self.preset]
        else:
            mix = get_stage_mix(self.progress, self.mix_mode, self.arith_frac)
        names = list(mix.keys())
        weights = list(mix.values())
        source_name = self.rng.choices(names, weights=weights, k=1)[0]
        text = next(self.streams[source_name])
        ids = self.tokenizer.encode(text).ids
        return ids + [self.eos_id]

    def next_sequence(self):
        while len(self.buffer) < self.seq_len + 1:
            self.buffer.extend(self._pull_one_doc_tokens())
        seq = self.buffer[:self.seq_len + 1]
        self.buffer = self.buffer[self.seq_len:]
        return seq

    def next_batch(self, batch_size: int):
        return [self.next_sequence() for _ in range(batch_size)]
if __name__ == '__main__':
    import time
    from tokenizers import Tokenizer
    tok = Tokenizer.from_file('artifacts/tokenizer.json')
    print('eos id:', tok.token_to_id(EOS), flush=True)
    streams = build_source_streams(reference_ngrams=set(), decon_enabled=False)
    mixer = MixedTokenStream(tok, streams, seq_len=64)
    t0 = time.time()
    for i in range(5):
        seq = mixer.next_sequence()
        print(f'seq {i}: len={len(seq)} sample={seq[:10]}', flush=True)
    print(f'5 sequences in {time.time() - t0:.2f}s', flush=True)