#!/usr/bin/env python
import argparse
import os
import sys
from pathlib import Path
import cv2

import numpy as np
from PIL import Image
from tqdm import tqdm

from sinder import (
    get_tokens,
    load_model,
    load_visual_data,
    pca_array,
)

os.environ['XFORMERS_DISABLED'] = '1'


def parse_args():
    parser = argparse.ArgumentParser(description='Visualize')
    parser.add_argument(
        'imgs', nargs='+', type=str, help='path to image/images'
    )
    parser.add_argument(
        '--model', type=str, default='dinov2_vitg14', help='model name'
    )
    parser.add_argument('--workdir', type=str, default='visualize')
    parser.add_argument(
        '--checkpoint',
        type=str,
        default=None,
        help='Path to checkpoint. Default is None, which loads the official pretrained weights',
    )
    parser.add_argument(
        '--visual_size',
        type=int,
        default=518,
        help='short side size of input image',
    )

    args = parser.parse_args()
    return args


def visualize(args, model, visual_dataset):
    model.eval()

    for d in tqdm(range(len(visual_dataset))):
        visual_image = visual_dataset[d]
        visual_tokens_all = get_tokens(model, visual_image)
        visual_tokens, visual_tokens_cls = zip(*visual_tokens_all)
        filename = Path(visual_dataset.files[d]).stem

        t = visual_tokens[-1].detach().cpu()
        h, w, c = t.shape
        norm = ((t.norm(dim=-1) / t.norm(dim=-1).max()) * 255).byte().numpy()
        norm_img = Image.fromarray(norm).resize((w * 14, h * 14), 0)
        norm = cv2.applyColorMap(np.array(norm_img), cv2.COLORMAP_JET)
        cv2.imwrite(args.folder / f'{filename}_norm.png', norm)

        pca_img = pca_array(visual_tokens[-1])
        pca_img.save(args.folder / f'{filename}_pca.png')


def main():
    args = parse_args()

    args.folder = Path(args.workdir).expanduser()
    os.makedirs(args.folder, exist_ok=True)
    print(args)
    print(' '.join(sys.argv))

    model = load_model(args.model, args.checkpoint)
    visual_dataset = load_visual_data(args, model)
    visualize(args, model, visual_dataset)


if __name__ == '__main__':
    main()