"""Epoch-wise full TRAIN coverage in 64-prompt, eight-generation groups."""
from __future__ import annotations
import json,math
from pathlib import Path
import torch
from torch.utils.data import Sampler


class ScaleGroupSampler(Sampler[int]):
    def __init__(self, size:int, *, seed:int=42, start_epoch:int=0,
                 prompts_per_step:int=64, generations:int=8,
                 dataloader_repeats:int=512, log_path:Path|None=None):
        if size<64 or prompts_per_step<=0 or generations<=0 or start_epoch<0 or dataloader_repeats<=0:
            raise ValueError('invalid scale sampler configuration')
        self.size=size;self.seed=seed;self.epoch=start_epoch
        self.prompts_per_step=prompts_per_step;self.generations=generations
        self.dataloader_repeats=dataloader_repeats
        self.log_path=log_path

    @property
    def steps_per_epoch(self):
        return math.ceil(self.size/self.prompts_per_step)

    def __len__(self):
        return self.steps_per_epoch*self.prompts_per_step*self.generations*self.dataloader_repeats

    def __iter__(self):
        epoch=self.epoch;self.epoch+=1
        generator=torch.Generator().manual_seed(self.seed+epoch)
        order=torch.randperm(self.size,generator=generator).tolist()
        extra=(-self.size)%self.prompts_per_step
        # Select padding from the epoch head. No prompt is repeated within
        # one64-prompt step, even for the short final group.
        if extra:order.extend(order[:extra])
        if self.log_path:
            self.log_path.parent.mkdir(parents=True,exist_ok=True)
            with self.log_path.open('a') as out:
                out.write(json.dumps(dict(epoch=epoch,shuffle_seed=self.seed+epoch,
                    train_rows=self.size,padding_reused_rows=extra,
                    steps=self.steps_per_epoch,
                    dataloader_repeats_per_optimizer_step=self.dataloader_repeats))+'\n')
        # TRL 0.21 consumes one 512-row dataloader batch for generation and
        # then 511 more dataloader batches while training the 512 buffered
        # microbatches. Those subsequent inputs are ignored. Repeat each
        # complete generation group so the next optimizer step advances to
        # the next group, rather than skipping 511 groups or ending an epoch.
        for start in range(0,len(order),self.prompts_per_step):
            group=order[start:start+self.prompts_per_step]
            for _ in range(self.dataloader_repeats):
                for index in group:
                    for _ in range(self.generations):yield index