import torch


def ground_caption(captions, n_ground=1, prefix="describe visual inputs:", sort=True):
    """
    For VG

    Args:
        captions
        n_ground

    Returns:
        source_text
        target_text

    Ex) (in vocab ids)
    captions
        ['Yellow banana', 'red crayon', 'black cow', 'blue sky']

    n_ground > 1

    ground_indices
        [1, 0, 2]
    source_text
        describe visual inputs: <vis_extra_id_1> <vis_extra_id_0> <vis_extra_id_2>
    target_text
        <extra_id_0> red crayon <extra_id_1> Yellow banana <extra_id_2> black cow

    n_ground == 1

    source_text
        describe visual inputs: <vis_extra_id_1>
    target_text
        red crayon
    """

    n_boxes = len(captions)

    if sort:
        ground_indices = torch.randperm(n_boxes)[:n_ground].sort().values
    else:
        ground_indices = torch.randperm(n_boxes)[:n_ground]

    ground_indices = ground_indices.tolist()

    source_text = [prefix]
    target_text = []

    if n_ground == 1:
        idx = ground_indices[0]
        source_text.append(f'<vis_extra_id_{idx}>')
        target_text.append(f'{captions[idx]}')
    else:
        for j, idx in enumerate(ground_indices):
            source_text.append(f'<vis_extra_id_{idx}>')

            target_text.append(f'<extra_id_{j}>')
            target_text.append(f'{captions[idx]}')

    # target_text.append('</s>')

    source_text = " ".join(source_text)
    target_text = " ".join(target_text)

    # return ground_indices, source_text, target_text
    return source_text, target_text