import torch

def get_tomato_slice(idx):
    if idx == 0:
        result = [0, 1, 2, 3, 463, 464, 465]
    else:
        result = [
            4 + (idx - 1) * 3,
            4 + (idx - 1) * 3 + 1,
            4 + (idx - 1) * 3 + 2,
            157 + (idx - 1) * 6,
            157 + (idx - 1) * 6 + 1,
            157 + (idx - 1) * 6 + 2,
            157 + (idx - 1) * 6 + 3,
            157 + (idx - 1) * 6 + 4,
            157 + (idx - 1) * 6 + 5,
            463 + idx * 3,
            463 + idx * 3 + 1,
            463 + idx * 3 + 2,
        ]
    return result


def get_part_slice(idx_list, func):
    result = []
    for idx in idx_list:
        result.extend(func(idx))
    return result


def expand_mask_to_all(mask, body_scale, hand_scale, face_scale):
    func = get_tomato_slice
    root_slice = get_part_slice([0], func)
    head_slice = get_part_slice([12, 15], func)
    stem_slice = get_part_slice([3, 6, 9], func)
    larm_slice = get_part_slice([14, 17, 19, 21], func)
    rarm_slice = get_part_slice([13, 16, 18, 20], func)
    lleg_slice = get_part_slice([2, 5, 8, 11], func)
    rleg_slice = get_part_slice([1, 4, 7, 10], func)
    lhnd_slice = get_part_slice(range(22, 37), func)
    rhnd_slice = get_part_slice(range(37, 52), func)
    face_slice = range(619, 669)
    B, T = mask.shape[0], mask.shape[1]
    mask = mask.view(B, T, -1)
    all_mask = torch.zeros(B, T, 669).type_as(mask)
    all_mask[:, :, root_slice] = mask[:, :, 0].unsqueeze(-1).repeat(1, 1, len(root_slice))
    all_mask[:, :, head_slice] = mask[:, :, 1].unsqueeze(-1).repeat(1, 1, len(head_slice))
    all_mask[:, :, stem_slice] = mask[:, :, 2].unsqueeze(-1).repeat(1, 1, len(stem_slice))
    all_mask[:, :, larm_slice] = mask[:, :, 3].unsqueeze(-1).repeat(1, 1, len(larm_slice))
    all_mask[:, :, rarm_slice] = mask[:, :, 4].unsqueeze(-1).repeat(1, 1, len(rarm_slice))
    all_mask[:, :, lleg_slice] = mask[:, :, 5].unsqueeze(-1).repeat(1, 1, len(lleg_slice))
    all_mask[:, :, rleg_slice] = mask[:, :, 6].unsqueeze(-1).repeat(1, 1, len(rleg_slice))
    all_mask[:, :, lhnd_slice] = mask[:, :, 7].unsqueeze(-1).repeat(1, 1, len(lhnd_slice))
    all_mask[:, :, rhnd_slice] = mask[:, :, 8].unsqueeze(-1).repeat(1, 1, len(rhnd_slice))
    all_mask[:, :, face_slice] = mask[:, :, 9].unsqueeze(-1).repeat(1, 1, len(face_slice))
    all_mask[:, :, root_slice] *= body_scale
    all_mask[:, :, head_slice] *= body_scale
    all_mask[:, :, stem_slice] *= body_scale
    all_mask[:, :, larm_slice] *= body_scale
    all_mask[:, :, rarm_slice] *= body_scale
    all_mask[:, :, lleg_slice] *= body_scale
    all_mask[:, :, rleg_slice] *= body_scale
    all_mask[:, :, lhnd_slice] *= hand_scale
    all_mask[:, :, rhnd_slice] *= hand_scale
    all_mask[:, :, face_slice] *= face_scale
    return all_mask