#  Copyright (c) Meta Platforms, Inc. and affiliates.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
import os

import click
import torch
from aitemplate.testing.benchmark_pt import benchmark_torch_function
from timm.models.vision_transformer import vit_base_patch16_224, vit_large_patch16_384


def create_vit(model_name):
    if model_name == "vit_base_patch16_224":
        model_path = "./vit_base_patch16_224.augreg2_in21k_ft_in1k/pytorch_model.bin"
        model = vit_base_patch16_224(pretrained=True, pretrained_path=model_path).cuda().half()
    elif model_name == "vit_large_patch16_384":
        model_path = "./vit_large_patch16_384.augreg_in21k_ft_in1k/pytorch_model.bin"
        model = vit_large_patch16_384(pretrained=True, pretrained_path=model_path).cuda().half()
    return model


def benchmark(model_name, batch_size, img_size, model):
    if model_name == "vit_base_patch16_224":
        img_size = 224
    elif model_name == "vit_large_patch16_384":
        img_size = 384
    with torch.inference_mode():
        input_shape = (batch_size, 3, img_size, img_size)
        input_data = torch.randn(input_shape).cuda().half()
        # warm up
        benchmark_torch_function(100, model, input_data)
        # benchmark
        t = benchmark_torch_function(100, model, input_data)
        print("batch_size: {}, time: {}".format(batch_size, t))
        dev_flag = os.environ.get("HIP_VISIBLE_DEVICES", "-1")
        dev_flag = dev_flag.replace(",", "_")
        with open(f"{model_name}_pt_benchmark_dev_{dev_flag}.txt", "a") as f:
            f.write("batch_size: {}, latency: {}\n".format(batch_size, t))


@click.command()
@click.option("--model-name", type=str, default="vit_base_patch16_224")
@click.option("--batch-size", default=0, type=int)
def main(model_name, batch_size):
    img_size = 224
    if model_name == "vit_base_patch16_224":
        img_size = 224
    elif model_name == "vit_large_patch16_384":
        img_size = 384
    else:
        raise NotImplementedError
    model = create_vit(model_name)
    if batch_size == 0:
        for batch_size in [1, 2, 4, 8, 16, 32, 64, 128, 256]:
            benchmark(model_name, batch_size, img_size, model)
    else:
        benchmark(model_name, batch_size, img_size, model)


if __name__ == "__main__":
    main()