import torch
from PIL import Image
from torchvision import transforms

from vtp.models.vtp_hf import VTPConfig, VTPModel
from vtp.tokenizers import get_tokenizer

model = VTPModel.from_pretrained("/path/to/MiniMaxAI/VTP-Large-f16d64")
model.eval()

# print model parameters
def count_params(m): return sum(p.numel() for p in m.parameters()) / 1e6
print(f"Vision Encoder: {count_params(model.trunk):.1f}M")
print(f"Pixel Decoder:  {count_params(model.pixel_decoder):.1f}M")
print(f"Text Encoder:   {count_params(model.text_transformer):.1f}M")

preprocess = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
image = preprocess(Image.open("figures/dog.png")).unsqueeze(0)

# ---------------------------------------------------------------------------------------
# use it as an auto-encoder; rFID=0.36
# ---------------------------------------------------------------------------------------
denormalize = transforms.Normalize(
    mean=[-0.485/0.229, -0.456/0.224, -0.406/0.225],
    std=[1/0.229, 1/0.224, 1/0.225]
)
with torch.no_grad(), torch.autocast("cuda"):
    latents = model.get_reconstruction_latents(image)  # encode
    recon = model.get_latents_decoded_images(latents)  # decode
recon_image = denormalize(recon[0]).clamp(0, 1).permute(1, 2, 0).cpu().numpy()
Image.fromarray((recon_image * 255).astype("uint8")).save("output/reconstructed.png")


# ---------------------------------------------------------------------------------------
# use it as a clip; zero-shot 78.2
# ---------------------------------------------------------------------------------------
tokenizer = get_tokenizer('ViT-B-32', context_length=model.config.text_context_length)
text = tokenizer(["a diagram", "a dog", "a cat", "a person"])
with torch.no_grad(), torch.autocast("cuda"):
    image_features = model.get_clip_image_feature(image, normalize=True)
    text_features = model.get_clip_text_feature(text, normalize=True)
    text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)
print("Label probs:", [f"{p:.4f}" for p in text_probs[0].tolist()])

# ---------------------------------------------------------------------------------------
# use it as an ssl feature extractor; linear probing 85.7
# ---------------------------------------------------------------------------------------
with torch.no_grad(), torch.autocast("cuda"):
    # get last layer features (cls token + patch tokens)
    features = model.get_last_layer_feature(image)
    cls_token = features['cls_token']      # (B, 1024)
    patch_tokens = features['patch_tokens']  # (B, 256, 1024) for 256x256 image

    # or get intermediate layer features for linear probing
    intermediate = model.get_intermediate_layers_feature(
        image, n=4, return_class_token=True
    )  # returns 4 x (patch_tokens, cls_token), each cls_token is (B, 1024)
    for i in range(1, 5):
        print('Last %d layers:' % i)
        print('Patch tokens shape:', intermediate[-i][0].shape)
        print('Cls token shape:', intermediate[-i][1].shape)