import torch
import transformers

from ..registry import ModelAttribute, model_zoo

# ===============================
# Register single-sentence VIT
# ===============================

config = transformers.ViTConfig(
    num_hidden_layers=4,
    # hidden_size=128,
    # intermediate_size=256,
    num_attention_heads=4)


# define data gen function
def data_gen():
    pixel_values = torch.randn(1, 3, 224, 224)
    return dict(pixel_values=pixel_values)


def data_gen_for_image_classification():
    data = data_gen()
    data['labels'] = torch.tensor([0])
    return data


def data_gen_for_masked_image_modeling():
    data = data_gen()
    num_patches = (config.image_size // config.patch_size)**2
    bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool()
    data['bool_masked_pos'] = bool_masked_pos
    return data


# define output transform function
output_transform_fn = lambda x: x

# function to get the loss
loss_fn_for_vit_model = lambda x: x.pooler_output.mean()
loss_fn_for_image_classification = lambda x: x.logits.mean()
loss_fn_for_masked_image_modeling = lambda x: x.loss

# register the following models
# transformers.ViTModel,
# transformers.ViTForMaskedImageModeling,
# transformers.ViTForImageClassification,
model_zoo.register(name='transformers_vit',
                   model_fn=lambda: transformers.ViTModel(config),
                   data_gen_fn=data_gen,
                   output_transform_fn=output_transform_fn,
                   loss_fn=loss_fn_for_vit_model,
                   model_attribute=ModelAttribute(has_control_flow=True))

model_zoo.register(name='transformers_vit_for_masked_image_modeling',
                   model_fn=lambda: transformers.ViTForMaskedImageModeling(config),
                   data_gen_fn=data_gen_for_masked_image_modeling,
                   output_transform_fn=output_transform_fn,
                   loss_fn=loss_fn_for_masked_image_modeling,
                   model_attribute=ModelAttribute(has_control_flow=True))

model_zoo.register(name='transformers_vit_for_image_classification',
                   model_fn=lambda: transformers.ViTForImageClassification(config),
                   data_gen_fn=data_gen_for_image_classification,
                   output_transform_fn=output_transform_fn,
                   loss_fn=loss_fn_for_image_classification,
                   model_attribute=ModelAttribute(has_control_flow=True))