class Colors:
    RED = '\033[91m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    ENDC = '\033[0m'  # Reset color


def print_tok_color(tok, ids, colormask):
    if ids.shape != colormask.shape:
        raise ValueError(f"{ids.shape} != {colormask.shape}")

    # ids and colormask are 2D: (batch, seq_len)
    batch_size, seq_len = ids.shape

    for b in range(batch_size):
        out = []
        for i in range(seq_len):
            token_id = int(ids[b, i])
            masked = bool(colormask[b, i])

            # Convert id → token (use decode if you prefer)
            token = tok.decode(token_id)

            if masked:
                out.append(f"{Colors.RED}{token}{Colors.ENDC}")
            else:
                out.append(token)

        print(" ".join(out))

def print_data_color_in_batch(index, tok, input_ids, data_mask):
    """
    Print decoded tokens for sample `index`, with data-masked tokens in red.

    Args:
        index:     batch index
        tok:       tokenizer
        input_ids: (batch, seq_len) tensor
        data_mask: (batch, seq_len) bool/int tensor — True/1 = data token (will be red)
    """
    RED = "\033[91m"
    RESET = "\033[0m"

    ids = input_ids[index].tolist()
    mask = data_mask[index].bool().tolist()

    parts = []
    cur_text = ""
    cur_is_data = mask[0] if ids else False

    for tid, m in zip(ids, mask):
        decoded = tok.decode([tid])
        if m == cur_is_data:
            cur_text += decoded
        else:
            if cur_text:
                parts.append((cur_is_data, cur_text))
            cur_text = decoded
            cur_is_data = m

    if cur_text:
        parts.append((cur_is_data, cur_text))

    out = ""
    for is_data, text in parts:
        if is_data:
            out += f"{RED}{text}{RESET}"
        else:
            out += text

    print(out)