import logging

from math import ceil

from fastprogress.fastprogress import progress_bar
import fire
import numpy as np
import torch
from torch.utils.data import DataLoader
from torchvision import datasets as D, transforms as T

from networks import get_model
from object_discovery import lost


def main(
    data_dir: str,
    arch: str = 'vit_small',
    patch_size: int = 16,
    resnet_dilate: int = 1,
    device: str = 'cuda:0',
    size: int = 384,
    batch_size: int = 8,
    num_workers: int = 8,
    modality: str = 'k',
    k_patches: int = 100,
):
    if 'cuda' in device and not torch.cuda.is_available():
        device = 'cpu'
        logging.warn(f'CUDA not available, overriding device {device} to CPU')
    
    dataset = D.ImageFolder(
        data_dir,
        transform=T.Compose([
            T.Resize((size, size)),
            T.ToTensor(),
            T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
        ]),
    )
    
    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        num_workers=num_workers,
        pin_memory=True
    )
    
    model = get_model(arch, patch_size, resnet_dilate, device)
    
    all_preds = {}
    whwh = np.array([size, size, size, size])
    i = 0
    with torch.inference_mode():
        for image, _ in progress_bar(dataloader):
            initial_image_size = image.shape[-3:]
            if image.size(-1) % patch_size != 0 or image.size(-2) % patch_size != 0:
                padded_w = int(ceil(image.size(-1) / patch_size) * patch_size)
                padded_h = int(ceil(image.size(-2) / patch_size) * patch_size)
                
                padded = torch.zeros((image.size(0), image.size(1), padded_h, padded_w), 
                                     device=device, dtype=image.dtype)
                padded[..., :image.size(-2), :image.size(-1)] = image
                image = padded
            
            h_feature_map = image.size(-2) // patch_size
            w_feature_map = image.size(-1) // patch_size
            
            features_out = {}
            def hook(module, input, output):
                features_out['qkv'] = output
            model._modules['blocks'][-1]._modules['attn']._modules['qkv'].register_forward_hook(hook)
            
            attentions = model.get_last_selfattention(image)
            scales = [patch_size, patch_size]
            
            nb_im, nh, nb_tokens = attentions.shape[:3]
            
            qkv = features_out['qkv'].reshape(nb_im, nb_tokens, 3, nh, -1 // nh).permute(2, 0, 3, 1, 4)
            q, k, v = qkv
            q = q.transpose(1, 2).reshape(nb_im, nb_tokens, -1)
            k = k.transpose(1, 2).reshape(nb_im, nb_tokens, -1)
            v = v.transpose(1, 2).reshape(nb_im, nb_tokens, -1)
            
            if modality == 'k':
                features = k[:, 1:]
            elif modality == 'v':
                features = v[:, 1:]
            elif modality == 'q':
                features = q[:, 1:]
            else:
                raise NotImplemented
            
            for f in features:
                pred, A, scores, seed = lost(
                    f[None, ...], 
                    [h_feature_map, w_feature_map],
                    scales,
                    initial_image_size,
                    k_patches = k_patches
                )
                
                pred = pred / whwh
                key = dataset.samples[i][0]
                all_preds[key] = pred
                i += 1


if __name__ == '__main__':
    fire.Fire(main)