import dataclasses
import logging
import re
from collections import defaultdict
from typing import Tuple, Optional, Any, Dict, List, Union, Mapping

import einops
import seqio
import numpy as np
import tensorflow as tf

from .mm_data import seqio_tokenizer
from .data_utils import pad_to_bounding_box, \
    get_3d_subsegments, _append_to_innermost_axis, resize_and_pad, \
    apply_with_random_selector, get_special_token_ids, make_autoregressive_inputs, \
    trim_and_pad_dataset, assert_not_truncated
from .prompts import apply_keyword_prompt, STYLE_TO_GENERAL_PROMPT, GENERAL_PROMPTS_V1
import .constants as config


def siglip_resize(src, imgsize, truncate):
    """Resize and preprocess for SigLIP ViT in the offical jax implementation"""
    assert src.dtype == tf.uint8
    # SigCLIP removes aspect ratio by default
    resized = tf.image.resize(src, imgsize, method=tf.image.ResizeMethod.BILINEAR, antialias=False)
    dtype = src.dtype
    tf_dtype = tf.type_spec_from_value(src).dtype
    resized = tf.cast(tf.clip_by_value(resized, tf_dtype.min, tf_dtype.max), dtype)

    # Normalize between -1 and 1 without using imagenet standard mean/std
    vmin=-1; vmax=1; in_min=0; in_max=255.0
    in_min_t = tf.constant(in_min, tf.float32)
    in_max_t = tf.constant(in_max, tf.float32)
    image = tf.cast(resized, tf.float32)
    image = (image - in_min_t) / (in_max_t - in_min_t)
    image = vmin + image * (vmax - vmin)
    if truncate:
        image = image[:truncate, :truncate]
    return image


def extract_bboxes(text, image_w, image_h):
    points = extract_points(text, image_w, image_h)
    boxes = []
    for i in range(len(points)//2):
        x1, y1 = points[i*2]
        x2, y2 = points[i*2 + 1]
        boxes.append([x1, y1, x2, y2])
    return boxes


def extract_annotated_points(caption, image_w, image_h):
    points = []
    for match in re.finditer("<point x=\"([0-9\\.]*)\" y=\"([0-9\\.]*)\" alt=\"([^\"]*)\">", caption):
        x = float(match.group(1))
        y = float(match.group(2))
        points.append(([[x, y]], match.group(3)))
    for match in re.finditer("<points ([^<]*) alt=\"([^\"]*)\">", caption):
        loc_str = match.group(1)
        locations = defaultdict(dict)
        if loc_str.startswith("points="):
            point_grp = []
            for point_match in re.finditer(r"([0-9]+\.[0-9]),? ([0-9]+\.[0-9])", loc_str):
                try:
                    point = [float(point_match.group(i)) for i in range(1, 3)]
                    point_grp.append(point)
                except ValueError:
                    pass
        else:
            for val in loc_str.split():
                try:
                    key, val = val.split("=")
                    locations[key[1:]][key[:1]] = float(val.strip("\""))
                except ValueError:
                    import pdb; pdb.set_trace()
                    logging.warning(f"Failed to parse {val} from {match.group(0)}")
            point_grp = []
            for key, coords in locations.items():
                if sorted(coords) == ["x", "y"]:
                    point_grp.append([coords["x"], coords["y"]])
        if point_grp:
            points.append((point_grp, match.group(2)))

    normalized = []
    for point_grp, point_text in points:
        normalized.append((
            np.array(point_grp) / 100.0 * np.array([image_w, image_h]),
            point_text,
        ))
    return normalized


def extract_points(text, image_w, image_h):
    all_points = []
    for match in re.finditer(r"Click\(([0-9]+\.[0-9]), ?([0-9]+\.[0-9])\)", text):
        try:
            point = [float(match.group(i)) for i in range(1, 3)]
        except ValueError:
            pass
        else:
            point = np.array(point)
            if np.max(point) > 100:
                # Treat as an invalid output
                continue
            point /= 100.0
            point = point * np.array([image_w, image_h])
            all_points.append(point)

    for match in re.finditer(r"\(([0-9]+\.[0-9]),? ?([0-9]+\.[0-9])\)", text):
        try:
            point = [float(match.group(i)) for i in range(1, 3)]
        except ValueError:
            pass
        else:
            point = np.array(point)
            if np.max(point) > 100:
                # Treat as an invalid output
                continue
            point /= 100.0
            point = point * np.array([image_w, image_h])
            all_points.append(point)
    for match in re.finditer(r'x\d*="\s*([0-9]+(?:\.[0-9]+)?)"\s+y\d*="\s*([0-9]+(?:\.[0-9]+)?)"', text):
        try:
            point = [float(match.group(i)) for i in range(1, 3)]
        except ValueError:
            pass
        else:
            point = np.array(point)
            if np.max(point) > 100:
                # Treat as an invalid output
                continue
            point /= 100.0
            point = point * np.array([image_w, image_h])
            all_points.append(point)
    for match in re.finditer(r'(?:\d+|p)\s*=\s*([0-9]{3})\s*,\s*([0-9]{3})', text):
        try:
            point = [int(match.group(i)) / 10.0 for i in range(1, 3)]
        except ValueError:
            pass
        else:
            point = np.array(point)
            if np.max(point) > 100:
                # Treat as an invalid output
                continue
            point /= 100.0
            point = point * np.array([image_w, image_h])
            all_points.append(point)
    return all_points


def extract_points_from_point_count(text, image_w, image_h):
    all_points = []
    points = re.findall(r"(\d+\.\d+),\s*(\d+\.\d+)", text)
    
    for match in points:
        try:
            point = [float(match[0]), float(match[1])]
        except ValueError:
            pass
        else:
            point = np.array(point)
            if np.max(point) > 100:
                # Treat as an invalid output
                continue
            point = point * np.array([image_w, image_h])
            all_points.append(point)
    return all_points


def select_tiling(h, w, patch_size, max_num_patches):
    """Decide how best to divide in image of size [w, h] in up to max_num_patches of size patch_size"""
    original_size = tf.stack([h, w])  # [1, 2]
    original_res = h * w
    tilings = []
    for i in range(1, max_num_patches+1):
        for j in range(1, max_num_patches+1):
            if i*j <= max_num_patches:
                tilings.append((i, j))
    # sort so argmin and argmax favour smaller tilings in the event of a tie
    tilings.sort(key=lambda x: (x[0]*x[1], x[0]))
    candidate_tilings = tf.constant(tilings, dtype=tf.int32)  # [n_resolutions, 2]
    candidate_resolutions = candidate_tilings * patch_size  # [n_resolutions, 2]

    # How much we would need to scale the image to fit exactly in each tiling
    required_scale_d = tf.cast(candidate_resolutions, tf.float32) / tf.cast(original_size[None, :], tf.float32)
    required_scale = tf.reduce_min(required_scale_d, axis=-1, keepdims=True)  # [n_resolutions, 1]
    if tf.reduce_all(required_scale < 1):
        # We are forced to downscale, so try to minimize the amount of downscaling
        ix = tf.argmax(required_scale)[0]
    else:
        # Pick the resolution that required the least upscaling so that it most closely fits the image
        required_scale = tf.where(required_scale < 1.0, 10e9, required_scale)
        ix = tf.argmin(required_scale)[0]
    return candidate_tilings[ix]


DEMO_STYLES = [
    "point_count",
    "pointing",
    "user_qa",
    "scifi_charts_exp",
    "scifi_charts_exp",
    "scifi_charts_exp",
    "scifi_charts_exp",
    "long_caption",
    "named_entity"
]


@dataclasses.dataclass
class MultiModalPreprocessor:
    """Turns text/image inputs into tensors that can be input to the model"""
    tokenizer: Any

    # How to prompt the model
    prompt_templates: str = "none"  # How to template prompts for examples
    message_format: str = "none"  # How to format messages
    system_prompt: Optional[str] = None  # How to generate system prompts
    prompt_override: Optional[str] = None  # Used for setting prompt manually
    always_start_with_space: bool = False  # Always include a leading space for the first bit of text
    default_inference_len: int = 65  # Inference len for length-conditioned prompting

    # How to crops/resize images
    crop_mode: str = "resize"
    max_crops: int = 6
    overlap_margins: Tuple[int, int] = (4, 4)
    do_random_scale: Optional[bool] = False
    resize: str = "default"
    random_scale_max: float = 1.1
    random_scale_min: float = 0.9
    random_scale_ratio: float = 0.5
    use_col_tokens: bool = True

    # Data about the ViT and connector we need when deciding the crops
    base_image_input_size: Tuple[int, int] = (336, 336)
    image_token_length_w: int = 12
    image_token_length_h: int = 12
    image_patch_size: int = 14
    image_padding_mask: bool = False

    # Other settings
    loss_token_weighting: Optional[str] = None
    unconditioned: Union[bool, float] = False  # Ignore images
    fix_image_input_idx: int = 2  # backwards compatibility fix
    pad_to: Optional[int] = None  # experimental feature

    _special_tokens: Dict[str, int] = None
    split_at: Optional[int] = None

    def get_max_total_crops(self):
        if self.crop_mode == "resize":
            return 1
        elif "resize" in self.crop_mode:
            return 1 + self.max_crops
        else:
            return self.max_crops

    @property
    def image_num_patch(self):
        h, w = self.base_image_input_size
        return h//self.image_patch_size, w//self.image_patch_size

    @property
    def special_token_ids(self):
        if self._special_tokens is None:
            self._special_tokens = get_special_token_ids(self.tokenizer)
        return self._special_tokens

    def image_to_patches_and_tokens(self, image, is_training):
        """Preprocesses an image

        Args:
            image: [h, w, 3] image to preprocessing
        Returns:
            crops: (n_crops, n_patches, patch_dim) individual crops, `n_crops` might
                   change between images but the other dimension are fixed
            tokens: (n_tokens,) tf.int32 tokens, pad tokens indicate where to insert the
                                patch features, might include other special tokens as well
            patch_ordering: (n_crops, n_tokens_per_crop) order image features should be inserted
                            into the `tokens`, negative values indicates patches features to exclude
            padding_mask: (n_crops, h, w) mask of what pixels are padding, can be None
        """
        do_random_scale = self.do_random_scale
        if do_random_scale:
            do_random_scale = is_training

        base_image_input_size = self.base_image_input_size
        if isinstance(base_image_input_size, int):
            base_image_input_size = (base_image_input_size, base_image_input_size)

        image_token_length_w, image_token_length_h = self.image_token_length_w, self.image_token_length_h
        base_image_input_d = self.image_patch_size
        tokens_per_image = image_token_length_w * image_token_length_h
        image_base_patch_w = base_image_input_size[1] // base_image_input_d
        image_base_patch_h = base_image_input_size[0] // base_image_input_d
        extra_image = False
        patch_ordering = None

        if self.resize == "default":
            image = tf.image.convert_image_dtype(image, dtype=tf.float32)
            def _resize(_image, sz):
                return resize_and_pad(
                    _image, sz,
                    do_random_scale=do_random_scale,
                    random_scale_max=self.random_scale_max,
                    random_scale_min=self.random_scale_min,
                    random_scale_ratio=self.random_scale_ratio,
                    return_outputs=False,
                    resize_method='random' if is_training else tf.image.ResizeMethod.BILINEAR)
        elif self.resize == "stretch":
            image = tf.image.convert_image_dtype(image, dtype=tf.float32)
            assert not do_random_scale

            def _resize(_image, sz):
                if not is_training:
                    img = tf.image.resize(_image, sz, antialias=True, method=tf.image.ResizeMethod.BILINEAR)
                else:
                    resize_methods = sorted([k for k in tf.image.ResizeMethod.__dict__.keys() if k.isupper()])
                    img = apply_with_random_selector(
                        _image,
                        lambda x, method_idx: tf.image.resize(x, sz,
                                                              tf.image.ResizeMethod.__dict__[resize_methods[method_idx]],
                                                              antialias=True),
                        num_cases=len(resize_methods))
                return img, tf.ones(tf.shape(img)[:2], dtype=tf.bool)
        elif self.resize in "siglip":
            assert not do_random_scale

            def _resize(_image, sz):
                img = siglip_resize(_image, sz, truncate=None)
                return img, tf.ones(tf.shape(img)[:2], dtype=tf.bool)
        else:
            raise NotImplementedError(self.resize)

        def _img_to_patches(_img, _img_mask, dy=1, dx=1):
            _img = einops.rearrange(
                _img, '(dy h dh) (dx w dw) c -> (dy dx) (h w) (dh dw c)',
                dh=base_image_input_d,
                dw=base_image_input_d,
                dy=dy,
                dx=dx,
                h=image_base_patch_h,
                w=image_base_patch_w
            )
            _img_mask = einops.rearrange(
                _img_mask, '(dy h dh) (dx w dw) -> (dy dx) (h w) (dh dw)',
                dh=base_image_input_d,
                dw=base_image_input_d,
                dy=dy,
                dx=dx,
                h=image_base_patch_h,
                w=image_base_patch_w
            )
            return _img, tf.reduce_mean(tf.cast(_img_mask, tf.float32), -1)

        mode = self.crop_mode
        if mode == "resize":
            patches, img_mask = _resize(image, base_image_input_size)
            patches, img_mask = _img_to_patches(patches, img_mask)
            image_layout_impatch_w = 1
            image_layout_impatch_h = 1
            patch_ordering = tf.range(tokens_per_image)[None, :]

        elif mode in ["overlap", "overlap-and-resize-c2"]:
            original_image_h = tf.shape(image, out_type=tf.int32)[0]
            original_image_w = tf.shape(image, out_type=tf.int32)[1]
            crop_size = base_image_input_size[0]

            # Discard this many patches from the (left/top, right/bottom) of crops
            left_margin, right_margin = self.overlap_margins
            # left_margin, right_margin = 2, 2
            assert left_margin % 2 == 0  # Required for compatibility with 2x2 pooling
            total_margin_pixels = base_image_input_d*(right_margin + left_margin)  # pixels removed per dim
            crop_patches = base_image_input_size[0] // base_image_input_d  # patches per crop dim
            crop_window_patches = crop_patches - (right_margin + left_margin)  # usable patches
            crop_window_size = crop_window_patches * base_image_input_d
            tiling = select_tiling(original_image_h - total_margin_pixels, original_image_w - total_margin_pixels,
                                   crop_window_size, self.max_crops)
            src, img_mask = _resize(
                image, [tiling[0]*crop_window_size+total_margin_pixels, tiling[1]*crop_window_size+total_margin_pixels])

            n_crops = tiling[0]*tiling[1]
            patches_arr = tf.TensorArray(
                tf.float32, n_crops, element_shape=[crop_size, crop_size, 3])
            mask_arr = tf.TensorArray(
                tf.bool, n_crops, element_shape=[crop_size, crop_size])
            # We assume 2x2 pooling, but can allow padding the right/bottom with extra
            # patches if the number of patches per side is not even
            assert (crop_patches+1)//2 == image_token_length_h
            assert (crop_patches+1)//2 == image_token_length_w
            patch_ordering_arr = tf.TensorArray(
                tf.int32, n_crops, element_shape=[image_token_length_h, image_token_length_w])
            on = 0
            on_patch = 0
            for i in range(tiling[0]):
                y0 = i*crop_window_size
                if i == 0:
                    crop_y0 = 0
                else:
                    crop_y0 = left_margin // 2

                crop_h = image_base_patch_h - (right_margin + left_margin)
                if i == 0:
                    crop_h += left_margin
                if i == (tiling[0]-1):
                    crop_h += right_margin
                for j in range(tiling[1]):
                    x0 = j*crop_window_size
                    if j == 0:
                        crop_x0 = 0
                    else:
                        crop_x0 = left_margin // 2

                    crop_w = image_base_patch_w - (right_margin + left_margin)
                    if j == 0:
                        crop_w += left_margin
                    if j == (tiling[1]-1):
                        crop_w += right_margin

                    pooled_w = (crop_w + 1) // 2
                    pooled_h = (crop_h + 1) // 2
                    patch_ordering_arr = patch_ordering_arr.write(
                        on_patch,
                        pad_to_bounding_box(
                            tf.reshape(tf.range(on, on+pooled_h*pooled_w, dtype=tf.int32), (pooled_h, pooled_w, 1)),
                            crop_y0, crop_x0, image_token_length_h, image_token_length_w, value=-1
                        )[:, :, 0]
                    )
                    patches_arr = patches_arr.write(on_patch, src[y0:y0+crop_size, x0:x0+crop_size])
                    mask_arr = mask_arr.write(on_patch, img_mask[y0:y0+crop_size, x0:x0+crop_size])

                    on += pooled_h*pooled_w
                    on_patch += 1
            patches = patches_arr.stack()
            patch_ordering = patch_ordering_arr.stack()
            img_mask = mask_arr.stack()

            image_layout_impatch_w, image_layout_impatch_h = tiling[0], tiling[1]
            patches = einops.rearrange(
                patches, 'p (h dh) (w dw) c -> p (h w) (dh dw c)',
                dh=base_image_input_d,
                dw=base_image_input_d,
                h=image_base_patch_h,
                w=image_base_patch_w
            )
            img_mask = einops.rearrange(
                img_mask, 'p (h dh) (w dw) -> p (h w) (dh dw)',
                dh=base_image_input_d,
                dw=base_image_input_d,
                h=image_base_patch_h,
                w=image_base_patch_w
            )
            img_mask = tf.reduce_mean(tf.cast(img_mask, tf.float32), -1)
            patch_ordering = tf.reshape(patch_ordering, [-1])
            valid = patch_ordering >= 0

            # Transpose, to get left-to-right order
            patch_ordering_rh = tf.reshape(patch_ordering,
                                           [tiling[0], tiling[1], image_token_length_h, image_token_length_w])
            patch_ordering_rh = tf.transpose(patch_ordering_rh, [0, 2, 1, 3])
            patch_ordering_rh = tf.reshape(patch_ordering_rh, [-1])

            # The tranpose will screw up which patches are masked, project the
            # new order into sparse structure of `patch_ordering` to fix this
            patch_ordering = tf.tensor_scatter_nd_update(
                patch_ordering,
                tf.where(valid),
                tf.boolean_mask(patch_ordering_rh, patch_ordering_rh >= 0),
                name="patch_order_transpose_Scatter"
            )

            h = tiling[0]*crop_window_patches + (right_margin+left_margin)
            w = tiling[1]*crop_window_patches + (right_margin+left_margin)
            special_token_ids = self.special_token_ids
            per_row = tf.fill(((w+1)//2,),
                              special_token_ids[config.DEFAULT_IMAGE_PATCH_TOKEN],)
            if self.use_col_tokens:
                per_row = tf.concat([per_row, [special_token_ids[config.DEFAULT_IM_COL_TOKEN]]], 0)

            joint = tf.tile(per_row, [(h+1)//2])
            joint = [
                [special_token_ids[config.DEFAULT_IM_START_TOKEN]],
                joint,
                [special_token_ids[config.DEFAULT_IM_END_TOKEN]]
            ]

            if "resize" in mode:
                resized, resized_mask = _resize(image, base_image_input_size)
                resized, resized_mask = _img_to_patches(resized, resized_mask)
                if 'c2' in mode:
                    patches = tf.concat([resized, patches], 0)
                    image_mask = tf.concat([resized_mask, img_mask], 0)
                else:
                    patches = tf.concat([patches, resized], 0)
                    image_mask = tf.concat([img_mask, resized_mask], 0)

                if patch_ordering is not None:
                    if 'c2' in mode:
                        patch_ordering = tf.where(
                            patch_ordering >= 0,
                            patch_ordering + tokens_per_image,
                            -1
                        )
                        patch_ordering = tf.concat([tf.range(0, tokens_per_image), patch_ordering], 0)
                    else:
                        raise ValueError()
                per_row = tf.fill((image_token_length_w,), special_token_ids[config.DEFAULT_IMAGE_PATCH_TOKEN],)
                if self.use_col_tokens:
                    per_row = tf.concat([per_row, [special_token_ids[config.DEFAULT_IM_COL_TOKEN]]], 0)
                extra_tokens = tf.tile(per_row, [image_token_length_h])
                joint = [
                            [special_token_ids[config.DEFAULT_IM_START_TOKEN]],
                            extra_tokens,
                            [special_token_ids[config.DEFAULT_IM_END_TOKEN]],
                        ] + joint

            joint = tf.concat(joint, 0)
            return patches, joint, patch_ordering, img_mask

        elif mode in ["patchify", "patchify-and-resize", "patchify-v2", "patchify-v2-and-resize", "patchify-v2-and-resize-c2"]:
            original_image_w = tf.shape(image, out_type=tf.int32)[0]
            original_image_h = tf.shape(image, out_type=tf.int32)[1]
            assert base_image_input_size[0] == base_image_input_size[1]
            base_patch_size = base_image_input_size[0]
            tiling = select_tiling(original_image_w, original_image_h, base_patch_size, self.max_crops)

            patches, img_mask = _resize(
                image, [tiling[0]*base_patch_size, tiling[1]*base_patch_size])
            patches, img_mask = _img_to_patches(patches, img_mask, tiling[0], tiling[1])
            if 'v2' in mode:
                # Order patches left-to-right not crop-by-crop
                patch_ordering = tf.reshape(
                    tf.range(tokens_per_image*tiling[0]*tiling[1]),
                    [tiling[0], tiling[1], image_token_length_w, image_token_length_h])
                patch_ordering = tf.transpose(patch_ordering, [0, 2, 1, 3])
                patch_ordering = tf.reshape(patch_ordering, (-1, tokens_per_image))
            else:
                patch_ordering = None

            # given image size, determine the number of patch size.
            image_layout_impatch_w = tiling[0]
            image_layout_impatch_h = tiling[1]

            if "resize" in mode:
                extra_image = True
                resized, resized_mask = _resize(image, base_image_input_size)
                resized, resized_mask = _img_to_patches(resized, resized_mask)
                if 'c2' in mode:
                    patches = tf.concat([resized, patches], 0)
                    image_mask = tf.concat([resized_mask, img_mask], 0)
                else:
                    patches = tf.concat([patches, resized], 0)
                    image_mask = tf.concat([img_mask, resized_mask], 0)

                if patch_ordering is not None:
                    if 'c2' in mode:
                        patch_ordering = tf.concat(
                            [tf.range(0, tokens_per_image)[None, :], patch_ordering+tokens_per_image], 0)
                    else:
                        n = tf.shape(patch_ordering)[0]
                        patch_ordering = tf.concat(patch_ordering, [tf.range(n, n+tokens_per_image)[None, :]], 0)
        else:
            raise NotImplementedError(mode)

        special_token_ids = self.special_token_ids

        per_row = tf.fill((image_token_length_w*image_layout_impatch_w,),
                          special_token_ids[config.DEFAULT_IMAGE_PATCH_TOKEN],)
        if self.use_col_tokens:
            per_row = tf.concat([per_row, [special_token_ids[config.DEFAULT_IM_COL_TOKEN]]], 0)

        joint = tf.tile(per_row, [image_token_length_h * image_layout_impatch_h])
        joint = [
            [special_token_ids[config.DEFAULT_IM_START_TOKEN]],
            joint,
            [special_token_ids[config.DEFAULT_IM_END_TOKEN]]
        ]
        if extra_image:
            assert not self.image_padding_mask
            per_row = tf.fill((image_token_length_w,), special_token_ids[config.DEFAULT_IMAGE_PATCH_TOKEN],)
            if self.use_col_tokens:
                per_row = tf.concat([per_row, [special_token_ids[config.DEFAULT_IM_COL_TOKEN]]], 0)
            extra_tokens = tf.tile(per_row, [image_token_length_h])
            if 'c2' in mode:
                joint = [
                            [special_token_ids[config.DEFAULT_IM_START_TOKEN]],
                            extra_tokens,
                            [special_token_ids[config.DEFAULT_IM_END_TOKEN]],
                        ] + joint
            else:
                joint += [
                    [special_token_ids[config.DEFAULT_IM_START_TOKEN]],
                    extra_tokens,
                    [special_token_ids[config.DEFAULT_IM_END_TOKEN]]
                ]
        if self.pad_to is not None:
            n = [tf.shape(x)[0] for x in joint]
            assert len(joint[-1]) == 1
            to_pad = self.pad_to - tf.reduce_sum(tf.stack(n))
            joint = tf.concat(joint[:-1] + [
                tf.zeros(to_pad, dtype=tf.int32) - 1,
                joint[-1]
            ], axis=0)
        else:
            joint = tf.concat(joint, 0)
        return patches, tf.concat(joint, 0), patch_ordering, img_mask

    def build_image_input_idx(self, input_tokens, patch_order, no_image=None):
        """Builds the index used to insert patch features into `input_tokens`"""
        tokens_per_image = self.image_token_length_w * self.image_token_length_h
        if no_image is not None and no_image:
            return tf.zeros((0, tokens_per_image), tf.int32)

        image_input_idx = input_tokens == self.special_token_ids[config.DEFAULT_IMAGE_PATCH_TOKEN]
        image_input_idx = tf.experimental.numpy.nonzero(image_input_idx)[0]
        image_input_idx = tf.cast(image_input_idx, tf.int32)

        if patch_order is not None:
            n_tokens = tf.shape(image_input_idx)[0]
            # Item N should have the value of image_input_index[where(patch_order == n)] if >= 0 else -1
            patch_order = tf.reshape(patch_order, [-1])
            n_patches = tf.shape(patch_order)[0]
            if n_tokens != n_patches:
                # Most complex case where some patches are dropped
                # First invert the valid tokens
                valid = patch_order >= 0
                sorted_patch_ixs = tf.scatter_nd(
                    tf.boolean_mask(patch_order, valid)[:, None],
                    tf.range(tf.reduce_sum(tf.cast(valid, tf.int32)), dtype=tf.int32),
                    [n_tokens],
                    name="valid_order_scatter"
                )

                # Project the inverted mapping into same sparse structure
                tmp = tf.fill(tf.shape(patch_order), -1)
                sorted_patch_ixs_ex = tf.tensor_scatter_nd_update(
                    tmp,
                    tf.where(valid),
                    sorted_patch_ixs,
                    name="order_with_padding_scatter"
                )

                # Do the gather and then re-masked outputs that were masked in `sorted_patch_ixs`
                valid = tf.cast(sorted_patch_ixs_ex >= 0, tf.int32)
                image_input_idx = tf.gather(image_input_idx, sorted_patch_ixs_ex*valid)
                image_input_idx = image_input_idx*valid - 100*(1 - valid)
            else:
                sorted_patch_ixs = tf.scatter_nd(patch_order[:, None], tf.range(n_patches), [n_patches])
                image_input_idx = tf.gather(tf.reshape(image_input_idx, [-1]), sorted_patch_ixs)
            image_input_idx = tf.reshape(image_input_idx, [-1, tokens_per_image])
        return image_input_idx

    def build_multimodel_features(self, tokens, mask, subsegments, images, is_training):
        """Builds input features by pre-processing `images` and modifying `tokens`
        to include image col/pad/start/end tokens instead image placeholder tokens
        """
        image_token_id = self.special_token_ids[config.IMAGE_PROMPT]
        image_idx = tf.experimental.numpy.nonzero(tokens == image_token_id)[0]
        if images is None or tf.shape(images)[0] == 0:
            tf.debugging.assert_equal(image_idx, tf.cast(0, tf.int64),
                                      "Image placeholders in input, but no images given!")
            tokens_per_image = self.image_token_length_w * self.image_token_length_h
            n_pixels = self.image_patch_size ** 2 * 3
            image_num_patch = np.prod(self.image_num_patch)
            crops = tf.zeros((0, image_num_patch, n_pixels), dtype=tf.float32)
            image_idx = tf.zeros((0, tokens_per_image), tf.int32)
            out = dict(
                target_tokens=tokens,
                images=crops,
                image_input_idx=image_idx,
                loss_masks=mask
            )
            if self.image_padding_mask:
                out["image_masks"] = tf.zeros((0, image_num_patch), dtype=tf.float32)
            if subsegments is not None:
                out["subsegment_ids"] = subsegments
            return out
        elif tf.shape(image_idx)[0] == 0 and tf.shape(images)[0] > 0:
            # As a special case, no image prompt means the images are all at the start
            image_idx = tf.zeros([tf.shape(images)[0]], tf.int64) - 1
        else:
            tf.debugging.assert_equal(
                tf.shape(images)[0], tf.shape(image_idx)[0],
                message="Different number of images and image placeholders")

        # Each image will produce a variable number of crops/tokens, so we aggregate things
        # the results tensor arrays and the concat them
        tokens_per_image = self.image_token_length_w * self.image_token_length_h
        n_pixels = self.image_patch_size*self.image_patch_size*3
        n_patches = self.image_num_patch[0]*self.image_num_patch[1]

        n = tf.shape(images)[0]
        all_crops = tf.TensorArray(dtype=tf.float32, size=n, infer_shape=False,
                                   element_shape=[None, n_patches, n_pixels])
        all_image_idx = tf.TensorArray(dtype=tf.int32, size=n, infer_shape=False,
                                       element_shape=[None, tokens_per_image])
        out_tokens = tf.TensorArray(dtype=tf.int32, size=n, infer_shape=False,
                                    element_shape=[None])
        out_masks = tf.TensorArray(dtype=tf.float32, size=n, infer_shape=False,
                                   element_shape=[None])
        if self.image_padding_mask:
            all_crop_masks = tf.TensorArray(dtype=tf.float32, size=n, infer_shape=False,
                                            element_shape=[None, None])
        else:
            # Dummy array to keep tensorflow's control analysis happy
            all_crop_masks = tf.TensorArray(dtype=tf.float32, size=0, infer_shape=False,
                                            element_shape=[None, None])
        if subsegments is not None:
            out_subsegments = tf.TensorArray(dtype=tf.int32, size=n, element_shape=[None])
        else:
            out_subsegments = tf.TensorArray(dtype=tf.int32, size=0, element_shape=[None])

        image_idx = tf.cast(image_idx, tf.int32)
        for ix in range(tf.shape(image_idx)[0]):
            token_ix = image_idx[ix]
            crops, image_tokens, patch_ordering, img_mask = self.image_to_patches_and_tokens(images[ix], is_training)
            patch_idx = self.build_image_input_idx(image_tokens, patch_ordering)

            if token_ix == -1:  # -1 is an image inserted at the very start
                start = 0
                token_ix = 0
                end = 0
            else:
                start = 0 if ix == 0 else image_idx[ix-1] + 1
                end = token_ix + 1

            all_image_idx = all_image_idx.write(ix, patch_idx + token_ix)
            all_crops = all_crops.write(ix, crops)
            image_token_mask = tf.zeros_like(image_tokens, dtype=tf.float32)

            if ix == (tf.shape(images)[0] - 1):
                tokens_part = tf.concat([tokens[start:token_ix], image_tokens, tokens[end:]], 0)
                mask_part = tf.concat([mask[start:token_ix], image_token_mask, mask[end:]], 0)
            else:
                tokens_part = tf.concat([tokens[start:token_ix], image_tokens], 0)
                mask_part = tf.concat([mask[start:token_ix], image_token_mask], 0)

            out_tokens = out_tokens.write(ix, tokens_part)
            out_masks = out_masks.write(ix, mask_part)
            if self.image_padding_mask:
                all_crop_masks = all_crop_masks.write(ix, img_mask)
            if subsegments is not None:
                parts = tf.fill([tf.shape(image_tokens)[0]], subsegments[token_ix])
                if ix == (tf.shape(images)[0] - 1):
                    seg = tf.concat([subsegments[start:token_ix], parts, subsegments[end:]], 0)
                else:
                    seg = tf.concat([subsegments[start:token_ix], parts], 0)
                out_subsegments = out_subsegments.write(ix, seg)

        out = dict(
            target_tokens=out_tokens.concat(),
            images=all_crops.concat(),
            image_input_idx=all_image_idx.concat(),
            loss_masks=out_masks.concat()
        )
        if self.image_padding_mask:
            out["image_masks"] = all_crop_masks.concat()
        if subsegments is not None:
            out["subsegment_ids"] = out_subsegments.concat()
        return out

    def _format_message(self, args):
        message, ix = args
        return self.format_message(message, ix)

    def format_message(self, message, ix):
        """Applies system formatting to ith message from a sequence of messages"""
        # If the image placeholder text is not preceded by space it will not get tokenized
        # correctly by some tokenizers, so double check it here
        assert config.IMAGE_PROMPT == "<|image|>"
        tf.debugging.assert_equal(
            tf.strings.regex_full_match(message, r".*[^ ]<\|image\|>.*"),
            False,
            message="Image token must always be preceded by a space"
        )
        is_user = ix % 2 == 0
        if self.message_format == "none" or self.message_format is None:
            pass
        elif self.message_format == "role":
            if is_user:
                # We put the "System:" prefix here since it doesn't need a loss
                message = tf.strings.join(["User: ", message, " Assistant:"])
        elif self.message_format == "cleanup":
            if is_user:
                # We put the "System:" prefix here since it doesn't need a loss
                message = tf.strings.join(
                    [
                        "[[User]]: Correct the spelling and punctuation mistakes on the following transcript based on what appears in the image.\n\n{before} ",
                        message,
                        "\n[[Assistant]]: {after}"
                    ]
                )
        elif self.message_format == "mistral":
            if is_user:
                message = tf.strings.join(["[INST] ", message, " [/INST]"])
        else:
            raise NotImplementedError(self.message_format)

        # For now assume a space will be used to separate the messages
        if not self.tokenizer.adds_space:
            if ix != 0 or self.always_start_with_space:
                message = tf.strings.join([" ", message])
        # Else space added automatically by the tokenizer

        return message

    def get_multi_message_token_input(self, conversations, text_weights=None):
        """Build inputs for a ragged tensor of conversations, where each row of the tensor,
        is a different conversation"""
        tf.debugging.assert_equal(tf.reduce_any(tf.strings.regex_full_match(
            conversations.values, re.escape(config.IMAGE_PROMPT))), False, "Segmented prompts must start with the image")

        n_conversation = tf.shape(conversations)[0]
        ar = tf.TensorArray(dtype=tf.int32, infer_shape=False, element_shape=[None],
                            size=n_conversation)
        n_messages_per_conversation = conversations.row_lengths()
        for ix in range(n_conversation):
            ar = ar.write(ix, tf.range(n_messages_per_conversation[ix], dtype=tf.int32))
        message_ix = ar.concat()
        messages = tf.map_fn(
            self._format_message, elems=(conversations.values, message_ix), fn_output_signature=tf.string)
        messages = self.tokenizer.encode_tf(messages)

        # Append EOS
        is_response = message_ix % 2 == 1
        is_response_int = tf.cast(is_response, tf.int32)
        eos = tf.RaggedTensor.from_row_lengths(
            tf.fill([tf.reduce_sum(is_response_int)], self.tokenizer.eos_token_id),
            tf.cast(is_response_int, messages.row_splits.dtype)
        )
        messages = tf.concat([messages, eos], axis=1)

        # Build mask over system responses
        mask = tf.ones_like(messages) * tf.cast(tf.expand_dims(is_response, axis=1), tf.int32)
        decoder_loss_weights = tf.cast(mask.values, tf.float32)

        # Build subsegment ids for each conversation
        tokens_per_message = tf.RaggedTensor.from_row_splits(
            row_splits=conversations.row_splits,
            values=messages.row_lengths()
        )
        token_per_conversation = tf.reduce_sum(tokens_per_message, axis=1)
        subsegment_ids = tf.repeat(tf.range(n_conversation, dtype=tf.int32)+1, token_per_conversation)

        image_ix = self.special_token_ids[config.IMAGE_PROMPT]
        messages = tf.concat([[image_ix], messages.values], axis=0)
        decoder_loss_weights = tf.concat([[0], decoder_loss_weights], axis=0)
        subsegment_ids = tf.concat([[10000], subsegment_ids], axis=0)
        return messages, decoder_loss_weights, subsegment_ids

    def get_multi_response_token_input(self, user_prompt, text, text_weights=None):
        """Build tokens for a multi-response-per-image example"""
        # FIXME this could be relaxed to just having the same prefix
        tf.debugging.assert_equal(tf.reduce_any(tf.strings.regex_full_match(
            user_prompt, re.escape(config.IMAGE_PROMPT))), False, "Segmented prompts must start with the image")
        user_prompt = self.format_message(user_prompt, 0)
        vocab = self.tokenizer
        prompts = vocab.encode_tf(user_prompt)
        response = self.format_message(text, 1)
        responses = vocab.encode_tf(response)
        responses = _append_to_innermost_axis(responses, vocab.eos_token_id)
        response_mask = tf.ones_like(responses, dtype=tf.float32)
        if text_weights is not None:
            response_mask *= text_weights
        image_tokens = tf.constant([self.special_token_ids[config.IMAGE_PROMPT]])

        if len(responses.shape) == 3:
            # Tricky case where we have multiple questions, each of which has multiple answers
            assert len(prompts.shape) == 2

            # Also shift the last tokens to the response segment since that tokens will
            # have multiple possible target tokens to predict
            last_prompt_tokens = prompts[:, -1:]
            last_prompt_tokens = tf.repeat(last_prompt_tokens, responses.row_lengths())
            last_prompt_tokens = tf.RaggedTensor.from_row_splits(
                values=tf.RaggedTensor.from_row_lengths(
                    values=last_prompt_tokens,
                    row_lengths=tf.ones_like(last_prompt_tokens, dtype=responses.row_splits.dtype)
                ),
                row_splits=responses.row_splits
            )
            responses = tf.concat([last_prompt_tokens,  responses], 2)
            prompts = prompts[:, :-1]

            shared_prefix = image_tokens
            segmented_suffix = tf.concat([tf.expand_dims(prompts, 1), responses], 1)
            targets = tf.concat([shared_prefix, segmented_suffix.values.values], 0)

            segmented_mask = tf.concat([
                tf.zeros_like(tf.expand_dims(prompts, 1), dtype=tf.float32),
                tf.concat([
                    tf.zeros_like(last_prompt_tokens, dtype=tf.float32),
                    response_mask
                ], 2)
            ], 1).values.values
            decoder_loss_weights = tf.concat(
                [tf.zeros_like(shared_prefix, dtype=tf.float32), segmented_mask], 0)

            text_segment_ids = get_3d_subsegments(segmented_suffix)
            subsegment_ids = tf.concat([
                tf.zeros_like(shared_prefix) + tf.reduce_max(text_segment_ids)+1,
                text_segment_ids], 0)
            subsegment_ids = tf.cast(subsegment_ids, tf.int32)
        else:
            if len(prompts.shape) == 1:
                # One prompt for all responses, we use the last token of the prompt as the
                # first token of each response segment since there will be multiple targets
                # for that token, the remaining targets are part of the prefix
                shared_prefix = tf.concat([image_tokens, prompts[:-1]], 0)
                prompts = prompts[-1:]
                prompts = tf.tile(tf.expand_dims(prompts, axis=0), [tf.shape(text)[0], 1])
            else:
                shared_prefix = image_tokens

            # Separate prompt for each response
            segmented_suffix = tf.concat([prompts, responses], 1)
            segmented_mask = tf.concat([tf.zeros_like(prompts, dtype=tf.float32), response_mask], 1).values

            targets = tf.concat([shared_prefix, segmented_suffix.values], 0)
            decoder_loss_weights = tf.concat(
                [tf.zeros_like(shared_prefix, dtype=tf.float32), segmented_mask], 0)
            subsegments = tf.ragged.row_splits_to_segment_ids(segmented_suffix.row_splits) + 1
            subsegment_ids = tf.concat([tf.zeros_like(shared_prefix)+10000,
                                        tf.cast(subsegments, tf.int32)], 0)
        return targets, decoder_loss_weights, subsegment_ids

    def get_tokens_input(self, messages, for_inference=False, text_weights=None):
        """Gets the token input for an example, using image placeholder tokens to
        indicate where images features should be inserted

        inputs
        messages: List or tensor users/system text messages, can have image placeholder tokens
        for_inference: bool, if true truncate the messages if it is a system message
        text_weights: Weights per a system message

        returns
        tokens: [n_tokens] tf.int32 token inputs with image placeholder tokens
        loss_mask: [n_tokens] tf.float32 token weights for loss
        subsegment: [n_tokens] tf.int32 or None, subsegment ids used to build more complex
                               attention masks if needed
        """
        if isinstance(messages, tf.RaggedTensor):
            assert not for_inference, "Cannot have multiple target messages for inference"
            return self.get_multi_message_token_input(messages, text_weights)
        elif len(tf.shape(messages[-1])) > 0:
            assert not for_inference, "Cannot have multiple target messages for inference"
            assert len(messages) == 2
            prompt = messages[0]
            response = messages[1]
            return self.get_multi_response_token_input(prompt, response, text_weights)
        else:
            messages = tf.convert_to_tensor(messages)
            if for_inference:
                if tf.shape(messages) % 2 == 0:
                    # Remove the last message since the model should predict it
                    messages = messages[:-1]

        # Apply system formatting
        ix = tf.range(tf.shape(messages)[0])
        is_response = ix % 2 == 1
        messages = tf.map_fn(
            self._format_message, elems=(messages, ix), fn_output_signature=tf.string)

        # Tokenize
        messages = self.tokenizer.encode_tf(messages)

        # Add EOS to system messages
        is_response_int = tf.cast(is_response, tf.int32)
        eos = tf.RaggedTensor.from_row_lengths(
            tf.fill([tf.reduce_sum(is_response_int)], self.tokenizer.eos_token_id),
            tf.cast(is_response_int, messages.row_splits.dtype)
        )
        messages = tf.concat([messages, eos], axis=1)
        targets = messages.values

        # Build mask over system responses
        mask = tf.ones_like(messages) * tf.cast(tf.expand_dims(is_response, axis=1), tf.int32)
        decoder_loss_weights = tf.cast(mask.values, tf.float32)
        if text_weights is not None:
            decoder_loss_weights = decoder_loss_weights * text_weights
        return messages.values, decoder_loss_weights, None

    def preprocess(self, image, input_text, is_training=False,
                   seq_len=None, pad_images=1, style=None, for_inference=True):
        """Get input tensors for the given image/text data

        image: [h, w, 3] numpy uint8 array of image pixels
        input_text: string input text, a list of text for a multi-turn conversation or dictionary
                    of inputs to use to build the prompt from a template
        is_training: allow training-time preprocessing (e.g., image augmentation)
        seq_len: pad input tokens to `seq_len`
        pad_images: pad input images to `self.get_max_total_crops()`
        style: Style to use for prompt templating
        """
        if image is not None and len(tf.shape(image)) == 3:
            image = tf.expand_dims(image, axis=0)

        messages = self.get_messages(input_text, style, is_training, for_inference=for_inference, user_prompt_seed=None, system_prompt_seed=None)
        targets, loss_masks, subsegments = self.get_tokens_input(messages, for_inference=for_inference)
        batch = self.build_multimodel_features(
            targets, loss_masks, subsegments, image, is_training)

        # Optionally padding to get constant sized arrays
        if pad_images:
            max_crops = self.get_max_total_crops() * pad_images
            image = batch["images"]
            n = max_crops - tf.shape(batch["images"])[0]
            batch["images"] = tf.pad(image, [[0, n], [0, 0], [0, 0]], constant_values=-1)
            if self.image_padding_mask:
                m = max_crops - tf.shape(batch["image_masks"])[0]
                batch["image_masks"] = tf.pad(batch["image_masks"], [[0, m], [0, 0]], constant_values=-1)
            batch["image_input_idx"] = tf.pad(batch["image_input_idx"], [[0, n], [0, 0]], constant_values=-1)

        if seq_len is not None:
            targets = batch["target_tokens"]
            if seq_len < len(targets):
                raise ValueError("Sequence length too short")
            n = seq_len - len(targets)
            batch["target_tokens"] = tf.pad(targets, [[0, n]], constant_values=-1)
            batch["loss_masks"] = tf.pad(batch["loss_masks"], [[0, n]], constant_values=-1)

        batch = self.get_post_mixing_preprocessor(pack=False)._convert_example(batch)
        return batch

    def get_user_prompt(self, style, example, is_training=True, for_inference=False, seed=None):
        """Build a list of strings of what a user might type in to the model for the given example,
        and its responses, by applying a prompt template to the fields in `example`

        Can return multiple strings for one message for multi-response examples
        """
        if "style" in example:
            style = example["style"]

        if "prompt" in example:
            # Examples have a complete user prompt pre-specified, usually for eval sets
            prompt = example["prompt"]

        elif self.prompt_templates == "none":
            # Bare-bone prompt with not templating of instructions
            if "prompt" in example:
                prompt = example["prompt"]
            elif "refexp" in example:
                prompt = example["refexp"]
            elif "question" in example and "options" in example:
                prompt = tf.strings.join([example["question"], "\n", example["options"], "\n"])
            elif "question" in example:
                prompt = example["question"]
            else:
                prompt = ""

        elif self.prompt_templates == "uber_model":
            if not isinstance(style, str):
                tf.debugging.assert_equal(tf.logical_or(
                    style == "ai2_diagram_no_letter",
                    style == "ai2_diagram",
                ), True)
                prompt = tf.strings.join([example["question"], "\n", example["options"], "\n"])
            else:
                # We template long captions and pointing since they are "demo" tasks, and use
                # plain text for everything else
                if style == "long_caption":
                    prompt = apply_keyword_prompt(GENERAL_PROMPTS_V1["long_caption"], example, seed)
                elif style == "pointing":
                    prompt = apply_keyword_prompt(GENERAL_PROMPTS_V1["pointing"], example, seed)
                elif style == "point_count":
                    prompt = apply_keyword_prompt(GENERAL_PROMPTS_V1["point_count"], example, seed)
                elif "prompt" in example:
                    prompt = example["prompt"]
                elif "refexp" in example:
                    prompt = example["refexp"]
                elif "question" in example and "options" in example:
                    prompt = tf.strings.join([example["question"], "\n", example["options"], "\n"])
                elif "question" in example:
                    prompt = example["question"]
                else:
                    prompt = ""

        elif self.prompt_templates == "uber_model_pointing":
            if style == "long_caption":
                long_captions = GENERAL_PROMPTS_V1["long_caption_no_pointing"]
                prompt = apply_keyword_prompt(GENERAL_PROMPTS_V1["long_caption"], example, seed)
            elif style == "pointing":
                prompt = apply_keyword_prompt(GENERAL_PROMPTS_V1["pointing"], example, seed)
            elif style in [
                "scifi_charts_explanation",
                "scifi_table_explanation",
                "scifi_document_explanation",
                "scifi_diagram_explanation",
                "user_qa",
                "long_caption",
            ]:
                raise NotImplementedError()
                if style == "long_caption":
                    prompts = GENERAL_PROMPTS_V1["long_caption"]
                elif "prompt" in example:
                    prompts = tf.expand_dims(example["prompt"], axis=0)
                else:
                    prompts = tf.expand_dims(example["question"], axis=0)
                suffixes = []
                for suffix in GENERAL_PROMPTS_V1["no_pointing_suffix"]:
                    if not suffix[0].isspace():
                        suffix = " " + suffix
                    suffixes.append(suffix)
                no_point_prompts = tf.reshape(tf.strings.join([
                    tf.tile(tf.expand_dims(suffixes, 1), [1, tf.shape(prompts)[1]]),
                    tf.tile(prompts, [len(suffixes), 1]),
                ]), [-1])
                # prefixes = []
                # for prefix in GENERAL_PROMPTS_V1["no_pointing_prefix"]:
                #     if not prefix[0].isspace():
                #         prefix = prefix + " "
                #     prefixes.append(prompts + prefix)
                prompt = apply_keyword_prompt(no_point_prompts, example, seed, keywords=[])
            elif "prompt" in example:
                prompt = example["prompt"]
            elif "refexp" in example:
                prompt = example["refexp"]
            elif "question" in example and "options" in example:
                prompt = tf.strings.join([example["question"], "\n", example["options"], "\n"])
            elif "question" in example:
                prompt = example["question"]
            else:
                prompt = ""

        elif self.prompt_templates == "general_instructions_v1":
            if isinstance(style, str):
                prompt = apply_keyword_prompt(GENERAL_PROMPTS_V1[STYLE_TO_GENERAL_PROMPT[style]], example, seed)
            elif isinstance(style, list):
                # This ia bit of hack to allow apply prompts to joint caption/transcript data
                # FIXME ideally we can apply the templating to multiple styles more generally
                def _apply(_style, ix):
                    tmp = dict(example)
                    # prevent apply_keyword_prompt for generating multiple templates
                    tmp["text"] = tmp["text"][0]
                    if _style == "long_caption":
                        return apply_keyword_prompt(GENERAL_PROMPTS_V1["long_caption"], tmp, seed)
                    elif _style == "transcript":
                        return apply_keyword_prompt(GENERAL_PROMPTS_V1["transcript"], tmp, seed)
                    else:
                        raise NotImplementedError(_style)
                prompt = [_apply(x, ix) for ix, x in enumerate(style)]
            else:
                raise NotImplementedError()

        elif self.prompt_templates == "zero_shot_v1":
            assert style is not None
            if not isinstance(style, str):
                # FIXME can we handle tensor style's in a better way?
                if style == "ai2_diagram":
                    prompt = "Question: {question}\nAnswer with correct answer option letter only\nOptions: {options}\nAnswer:"
                    prompt = apply_keyword_prompt([prompt], example, seed)
                elif style == "ai2_diagram_no_letter":
                    prompt = "Question: {question}\nAnswer with correct answer option only\nOptions: {options}\nAnswer:"
                    prompt = apply_keyword_prompt([prompt], example, seed)
                else:
                    prompt = ""
                tf.debugging.assert_equal(prompt != "", True)
            else:
                general_style = STYLE_TO_GENERAL_PROMPT[style]
                if general_style == "short_answer":
                    prompt = apply_keyword_prompt(["Question: {question} Answer with as few words as possible. Answer:"], example, seed)
                elif general_style == "multiple_choice":
                    prompt = apply_keyword_prompt(["Question: {question}\nAnswer with correct answer option letter only\nOptions: {options}\nAnswer:"], example, seed)
                elif general_style == "count_bench":
                    prompt = apply_keyword_prompt(["Question: How many {object} are there?\nRespond with only a number.\nAnswer:"], example, seed)
                else:
                    raise NotImplementedError(general_style)

        elif self.prompt_templates == "zero_shot_v2":
            assert style is not None

            if self.prompt_override:
                prompt = apply_keyword_prompt([self.prompt_override], example, seed)
            elif not isinstance(style, str):
                if style == "ai2_diagram":
                    prompt = "{question} Answer with correct answer option letter only. Options: {options}"
                    prompt = apply_keyword_prompt([prompt], example, seed)
                elif style == "ai2_diagram_no_letter":
                    prompt = "{question} Answer with correct answer option only. Options: {options}"
                    prompt = apply_keyword_prompt([prompt], example, seed)
                else:
                    prompt = ""
                tf.debugging.assert_equal(prompt != "", True)
            else:
                if style in ["vqa2", "gqa", "tally_qa", "okvqa", "a_okvqa_da"]:
                    prompt = "Answer with a single word. {question}"
                elif style in ["text_vqa", "doc_qa", "info_qa", "chart_qa", "st_qa", "ocr_vqa", "dv_qa", "tabwmp_da", "figure_qa", "figure_qa_zero_shot", "plot_qa"]:
                    prompt = "{question}\nRespond as concisely as possible, do not output anything other than the answer."
                elif STYLE_TO_GENERAL_PROMPT[style] == "multiple_choice":
                    prompt = "{question} Answer with correct answer option letter only. Options: {options}"
                elif STYLE_TO_GENERAL_PROMPT[style] == "short_answer":
                    prompt = "{question} Answer with as few words as possible."
                elif style == "vtabfact":
                    prompt = "{question}"
                elif style == "count_bench":
                    prompt = "How many {object} are there?\nRespond with only a number."
                else:
                    raise NotImplementedError(style)
                prompt = apply_keyword_prompt([prompt], example, seed)
        else:
            raise NotImplementedError(self.prompt_templates)

        if for_inference:
            return [prompt]
        else:
            return [prompt, example["text"]]

    def get_system_prompt(self, style, example, for_inference,
                          messages, seed=None):
        if isinstance(style, str) and style == "count_bench":
            style = "ok_vqa"

        if self.system_prompt == "style":
            if isinstance(style, str):
                prefix = style + ":"
            else:
                prefix = tf.strings.join([style, ":"])

        elif self.system_prompt == "demo_or_style":
            if isinstance(style, str):
                if style == "android_control" or style == "demo":
                    # android is a special case since I hacked in prefix in the preprocessor
                    prefix = ""
                elif style in ["scifi_demo", "synthetic_qa"] or style in DEMO_STYLES:
                    if style == "scifi_demo":
                        p_no_prompt = 0.2
                    elif style == "synthetic_qa":
                        p_no_prompt = 0.25
                    else:
                        p_no_prompt = 0.9
                    if len(tf.shape(messages)) > 1:
                        n_messages = tf.shape(messages)[1]
                        style = tf.tile(tf.expand_dims(style, axis=0), [n_messages])
                        r = tf.random.stateless_uniform([n_messages], seed, 0, 1)
                    else:
                        r = tf.random.stateless_uniform((), seed, 0, 1)
                    prefix = tf.where(r < p_no_prompt, "", tf.strings.join([style + ":"]))
                else:
                    prefix = style + ":"
            else:
                if tf.reduce_any(style == tf.constant(DEMO_STYLES + ["scifi_demo", "android_control", "demo"])):
                    prefix = ""
                else:
                    prefix = tf.strings.join([style, ":"])

        elif self.system_prompt in ["long_caption_length_hint", "style_long_caption_length_hint"]:
            if seed is not None:
                raise NotImplementedError("Determinism")
            std = 25
            use_hint = tf.logical_or(
                tf.equal(style, "long_caption"), tf.equal(style, "transcript"))
            if self.system_prompt == "style_long_caption_length_hint":
                default = tf.strings.join([style, ": "])
            else:
                default = ""
            if for_inference:
                assert len(tf.shape(use_hint)) == 0
                if self.default_inference_len and use_hint:
                    prefix = tf.strings.join([style, " ", str(self.default_inference_len), ": "])
                else:
                    prefix = default
            else:
                std = 25
                n = tf.strings.length(messages[-1])
                n += tf.cast(tf.random.normal(n.shape)*std, tf.int32)
                hint = tf.strings.join([style, " ", tf.strings.as_string(n//15), ": "])
                use_hint = tf.logical_and(use_hint, tf.random.uniform(tf.shape(hint)) > 0.1)
                prefix = tf.where(use_hint, hint, default)

        elif for_inference and self.system_prompt in ["style_and_length", "style_and_length_v2"]:
            v2 = self.system_prompt == "style_and_length_v2"
            if example.get("length_cond") is not None:
                # Examples have individual length conditioning
                n = tf.strings.as_string(example["length_cond"])
            else:
                inference_len = self.default_inference_len
                n = None if inference_len is None else str(inference_len)
                logging.warning(f"eval len: {n}")
            if n is not None and tf.strings.length(n) > 0:  # allow empty string to signal unconditioned
                prefix = tf.strings.join([style, " ", n, ":"])
            else:
                prefix = tf.strings.join([style, ":" if v2 else " :"])
        elif self.system_prompt in ["style_and_length", "style_and_length_v2"]:
            v2 = self.system_prompt == "style_and_length_v2"
            std = 25
            logging.info(f"style prompt std={std}, percent=10")
            if seed is not None:
                seeds = tf.random.split(seed)
                p = tf.random.stateless_uniform((), seed=seeds[0])
            else:
                p = tf.random.uniform(())
            if p > 0.10:
                n = tf.strings.length(messages[-1])
                if seed is not None:
                    n += tf.cast(tf.random.stateless_normal(n.shape, seed=seeds[1])*std, tf.int32)
                else:
                    n += tf.cast(tf.random.normal(n.shape)*std, tf.int32)
                n = tf.strings.as_string(n//15)
                prefix = tf.strings.join([style, " ", n, ":"])
            else:
                prefix = tf.strings.join([style, ":" if v2 else " :"])
        else:
            raise NotImplementedError(self.system_prompt)

        return prefix

    def preprend_system_prompt(self, style, example, for_inference, messages, seed=None):
        prefix = self.get_system_prompt(style, example, for_inference, messages, seed=seed)
        separator = tf.where(tf.logical_and(
            tf.strings.length(prefix) > 0, tf.strings.length(messages[0]) > 0), " ", "")
        with_system_prompt = tf.strings.join([prefix, separator, messages[0]])
        if isinstance(messages, list):
            messages = [with_system_prompt] + messages[1:]
        else:
            messages = tf.concat([tf.expand_dims(with_system_prompt, 0), messages[1:]], axis=0)
        return messages

    def get_messages(self, ex, style, is_training, for_inference, user_prompt_seed, system_prompt_seed):
        if isinstance(ex, list):
            messages = ex
        elif isinstance(ex, str):
            messages = [ex]
        elif "messages" in ex:
            messages = ex["messages"]
        else:
            # Apply a prompt template
            messages = self.get_user_prompt(style, ex, is_training, for_inference=for_inference, seed=user_prompt_seed)

        # Maybe add a system prompt. The system prompt gets concatenated with the first user input
        if self.system_prompt and self.system_prompt != "none":
            if isinstance(ex, dict):
                style = ex.get("style", style)

            if isinstance(messages, tf.RaggedTensor):
                n = tf.shape(messages)[0]
                message_arr = tf.TensorArray(dtype=tf.string, size=n, element_shape=(None,))
                seeds = tf.random.split(system_prompt_seed, n)
                for i in range(n):
                    message_arr = message_arr.write(i, self.preprend_system_prompt(style, None, for_inference, messages[i], seed=seeds[i]))
                messages = tf.RaggedTensor.from_row_splits(
                    values=message_arr.concat(), row_splits=messages.row_splits)
            else:
                messages = self.preprend_system_prompt(style, ex, for_inference, messages, seed=system_prompt_seed)

        return messages

    def get_preprocessor(self, is_training, for_inference, style=None, include_metadata=None):
        """Build a preprocessing function that can be applied ot a tf.data.Dataset"""
        vocab = self.tokenizer
        include_response = not for_inference
        if include_metadata is None:
            include_metadata = for_inference

        @seqio.map_over_dataset(num_seeds=2)
        def to_inputs_and_targets(ex, seeds):
            if "unconditioned" in ex:
                raise NotImplementedError()
            if "image" not in ex:
                image = None
            elif ex['image'].dtype == tf.string:
                image = tf.image.decode_image(ex['image'], channels=3, expand_animations=False)
            else:
                image = ex['image']
            raw_image = image
            if image is not None and len(tf.shape(image)) == 3:
                image = tf.expand_dims(image, axis=0)

            unconditioned = self.unconditioned
            if unconditioned and isinstance(unconditioned, float):
                assert image is not None
                if is_training and tf.random.uniform((), 0, 1, dtype=tf.float32) < unconditioned:
                    image = image[:0]
            elif unconditioned:
                image = None

            messages = self.get_messages(ex, style, is_training, for_inference, seeds[0], seeds[1])
            targets, loss_masks, subsegments = self.get_tokens_input(
                messages, for_inference, ex.get("text_weights"))
            # if "scifi" in style and style.endswith("_explanation"):
            #     logging.warning(f"No loss on EOS for {style}")
            #     loss_masks = tf.where(targets == self.tokenizer.eos_token_id, tf.zeros_like(loss_masks), loss_masks)
            out = self.build_multimodel_features(targets, loss_masks, subsegments, image, is_training)

            if include_metadata:
                # FIXME remove these special cases
                if "text" in ex:
                    if len(ex["text"].shape) > 0:
                        # FIXME can this be variable lengths after all?
                        out["metadata/captions"] = tf.strings.reduce_join(
                            tf.strings.regex_replace(ex['text'], "\\s+", " "),
                            separator="\n"
                        )
                    else:
                        out["metadata/captions"] = ex["text"]

                if "image_url" in ex:
                    out["metadata/image_url"] = ex["image_url"]
                elif "url" in ex:
                    out["metadata/image_url"] = ex["url"]
                if "image_id" in ex:
                    out["metadata/image_id"] = ex["image_id"]
                for k, v in ex.items():
                    if k.startswith("metadata"):
                        out[k] = v
                if raw_image is not None and "metadata/image_size" not in out:
                    img_h = tf.shape(raw_image)[0]
                    img_w = tf.shape(raw_image)[1]
                    out["metadata/image_size"] = [img_w, img_h]
                if "metadata/image_url" not in out and raw_image is not None:
                    if len(ex["image"].shape) < 4:
                        # For visualizations FIXME can we make this variable length
                        out["metadata/image"] = tf.io.encode_jpeg(
                            tf.image.convert_image_dtype(raw_image, tf.uint8))
            return out
        return to_inputs_and_targets

    def get_post_mixing_preprocessor(self, pack=False):
        """Build a feature conversion function that can be applied ot a tf.data.Dataset

        This function applies a second stage of pre-processing, but unlike `self.get_preprocessor`
        this stage can be applied after mixing tf.data.Datasets into a mixture
        """
        return MultiModalLMFeatureConverter(
            loss_token_weighting=self.loss_token_weighting,
            bos_id=self.tokenizer.bos_token_id,
            fix_image_input_idx=self.fix_image_input_idx,
            pack=pack,
            special_tokens=list(self.special_token_ids.values()),
        )


class MultiModalLMFeatureConverter:

    def __init__(
        self, pack: bool = False, loss_token_weighting: str=None, bos_id: int = 1,
        special_tokens=None, fix_image_input_idx=2
    ):
        self.pack = pack
        self.bos_id = bos_id
        self.fix_image_input_idx = fix_image_input_idx
        self.special_tokens = tf.constant(special_tokens) if special_tokens else None
        self.loss_token_weighting = loss_token_weighting

    def _convert_example(
        self, features: Mapping[str, tf.Tensor]
    ) -> Mapping[str, tf.Tensor]:
        """Convert an LM example into an example with model features."""
        # targets_segment_id is present only for a packed dataset.
        decoder_input_tokens = make_autoregressive_inputs(
            features["target_tokens"],
            sequence_id=features.get("targets_segment_ids", None),
            bos_id=self.bos_id,
        )

        tf.assert_equal(
            True,
            tf.reduce_all(decoder_input_tokens[-1] != self.special_tokens),
            message="An input ends with an image special token",
        )

        image_input_idx = features["image_input_idx"]
        if self.fix_image_input_idx == 2:
            # plus one sine we have added BOS to the inputs
            image_input_idx = tf.where(image_input_idx < 0,  image_input_idx, image_input_idx + 1)
        else:
            # Some old models trained like this, sometimes image_input_idx will go from -1 -> 0 didn't
            # effect performance but keep this code path for backwards compatiblity with those checkpoints
            image_input_idx = image_input_idx + 1

        d = {
            "target_tokens": features["target_tokens"],
            "input_tokens": decoder_input_tokens,
            "loss_masks": features["loss_masks"],
            "images": features["images"],
            "image_input_idx": image_input_idx
        }
        if "image_masks" in features:
            d["image_masks"] = features["image_masks"]

        has_custom_text_weight = features.get("has_custom_loss_weight", False)

        if "subsegment_ids" in features:
            subsegment_ids = make_autoregressive_inputs(
                features["subsegment_ids"],
                sequence_id=features.get("targets_segment_ids", None),
                bos_id=features["subsegment_ids"][0],
            )

            # Subsegment have a position based on the sum of previous positions they can attend to
            position_ids = tf.zeros_like(subsegment_ids)
            unique_segments = tf.unique(subsegment_ids)[0]
            for i in unique_segments:
                segment_position_ids = tf.cumsum(tf.cast(subsegment_ids >= i, tf.int32)) - 1
                position_ids = tf.where(subsegment_ids == i, segment_position_ids, position_ids)

            # Apply loss weighting, this is done here so it occurs after truncation
            if has_custom_text_weight:
                pass
            elif self.loss_token_weighting in ["subsegments", "root_subsegments"]:
                n_loss_segments = tf.shape(tf.unique(tf.boolean_mask(subsegment_ids, d["loss_masks"] > 0))[0])[0]
                n_loss_segments = tf.maximum(tf.cast(n_loss_segments, tf.float32), 1)
                weight = 1/n_loss_segments if self.loss_token_weighting == "subsegments" else tf.math.rsqrt(n_loss_segments)
                d["loss_masks"] = tf.where(d["loss_masks"] > 0, d["loss_masks"]*weight, d["loss_masks"])
            elif self.loss_token_weighting is not None:
                raise NotImplementedError(self.loss_token_weighting)

            d["subsegment_ids"] = subsegment_ids
            d["position_ids"] = position_ids
        else:
            if self.loss_token_weighting not in [None, "subsegments", "root_subsegments"] and not has_custom_text_weight:
                raise NotImplementedError(self.loss_token_weighting)
        if self.pack:
            d["decoder_segment_ids"] = features["targets_segment_ids"]
            d["decoder_positions"] = features["targets_positions"]

        for k in features:
            if k.startswith("metadata/"):
                d[k] = features[k]
        return d

    def _pack_or_pad(self, ds, task_feature_lengths):
        if self.pack:
            raise NotImplementedError()
        else:
            return trim_and_pad_dataset(ds, task_feature_lengths)

    def __call__(self, ds: tf.data.Dataset, task_feature_lengths: Mapping[str, int]) -> tf.data.Dataset:
        """Convert the dataset to be fed to a language model."""
        task_feature_lengths = dict(task_feature_lengths)

        if "images" in ds.element_spec and "images" in task_feature_lengths:
            # Images should never be truncated
            ds = assert_not_truncated(ds, ["images", "image_input_idx"], task_feature_lengths["images"])

        if any(x.startswith("metadata/") for x in ds.element_spec):
            # Metadata indicates the dataset is being used for inference, inference datasets
            # should not be truncated
            ds = assert_not_truncated(ds, ["target_tokens"], task_feature_lengths["target_tokens"])

        if "image_masks" in ds.element_spec and "images" in task_feature_lengths:
            task_feature_lengths["image_masks"] = task_feature_lengths["images"]
        if "subsegment_ids" in ds.element_spec and "target_tokens" in task_feature_lengths:
            task_feature_lengths["subsegment_ids"] = task_feature_lengths["target_tokens"]
        if "loss_masks" not in task_feature_lengths and "target_tokens" in task_feature_lengths:
            task_feature_lengths["loss_masks"] = task_feature_lengths["target_tokens"]
        ds = self._pack_or_pad(ds, task_feature_lengths)

        return ds.map(
            self._convert_example, num_parallel_calls=tf.data.experimental.AUTOTUNE)