#%%
from PIL import Image
import collections
from pathlib import Path
from matplotlib import ticker
import torchvision.transforms as transforms
import numpy as np
import matplotlib.pyplot as plt
import torch
from lovely_numpy import lo
import lovely_numpy
lovely_numpy.set_config(deeper_width=12)
import lovely_tensors as lt
import cv2
lt.monkey_patch()
lt.set_config(deeper_width=12)
torch.inf = float("Inf")
import sys
import os
base_dir = Path(__file__).parent.parent
sys.path.append(str(base_dir))
os.chdir(base_dir)

from torchvision.datasets import ImageFolder
import torchvision
import tqdm
# %%

from samples.CLS2IDX import CLS2IDX
from baselines.ViT.ViT_LRP import vit_base_patch16_224 as vit_LRP
from baselines.ViT.ViT_new import vit_base_patch16_224 as vit_new
import types

# IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
# IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
# normalize = transforms.Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD)
normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
transform_without_normalize = transforms.Compose([
    # transforms.Resize(256),
    transforms.Resize(224),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
])
transform = transforms.Compose([
    transform_without_normalize, 
    normalize
])

# %%


# initialize ViT pretrained
model = vit_new(pretrained=True).cuda()
model.eval()

model_LRP = vit_LRP(pretrained=True).cuda()
model_LRP.eval()
from baselines.ViT.ViT_explanation_generator import Baselines, LRP
lrp = LRP(model_LRP)


# create heatmap from mask on image
def show_cam_on_image(img, mask):
    heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
    heatmap = np.float32(heatmap) / 255
    cam = heatmap + 1*np.float32(img)
    cam = cam / np.max(cam)
    return cam


# %%
imagenet_r_path = Path(os.environ["IMAGENET_R_PATH"])
imagenet_path = Path(os.environ["IMAGENET_PATH"])
imagenet_val_path = imagenet_path / "val"
data = []

# cls_idx = 943 # cucumber
# cls_wind = "n07718472"
# img_path = imagenet_r_path / f'{cls_wind}/sticker_0.jpg'

cls_idx = 812 # space_shuttle
cls_wind = "n04266014"
img_path = imagenet_r_path / f'{cls_wind}/toy_7.jpg'


data = [
    # (
    #     imagenet_r_path / f'n04266014/toy_7.jpg',
    #     812 # space_shuttle
    # ),
    # (
    #     imagenet_r_path / f'n07718472/sticker_0.jpg',
    #     943 # cucumber
    # ),

    (
        imagenet_r_path / f'n07749582/painting_15.jpg',
        951 # lemon
    ),
    (
        imagenet_r_path / f'n07749582/painting_15.jpg',
        868 # tray
    ),
    (
        imagenet_r_path / f'n04536866/sketch_9.jpg',
        889 # violin
    ),
    (
        imagenet_r_path / f'n04536866/sketch_9.jpg',
        584 # hair slide
    ),
    (
        imagenet_val_path / f'n02123045/ILSVRC2012_val_00039877.JPEG',
        # 281 # tabby cat
        285 # Egyptian cat
    ),
 ]

layer = 11

output_dir = base_dir / Path(f'figures/artifacts/predicatt_decomp/')
output_dir.mkdir(exist_ok=True, parents=True)

for img_path, cls_idx in data:
    print(img_path)
    cls_wind = img_path.parent.stem
    image = Image.open(img_path)
    image = transform_without_normalize(image) # (3, 224, 224)
    if image.shape[0] == 1:
        image = image.expand(3, -1, -1)
    if image.shape[0] == 4:
        # remove alpha channel
        image = image[:3] 
    image_norm = normalize(image).unsqueeze(0).cuda() # (1, 3, 224, 224)
    cls_pred_vec = model(image_norm)
    topk = torch.topk(cls_pred_vec, 5)
    print("Top predictions:")
    for ci in topk.indices.squeeze():
        print(f"\t {cls_pred_vec[0, ci].item(): >6.2f} \t {CLS2IDX[ci.item()]}, {ci.item()}, ")

    apply_softmax_classes = True
    xai_map, extras = model.predmap15_batched_layer(image_norm, cls_idx, layer=layer, return_extras=True, apply_softmax_classes=apply_softmax_classes) # (1, tokens-1)
    xai_map = xai_map.detach().cpu().reshape(1, 1, 14,14) # (1, 1, 14, 14)
    xai_map = (xai_map - xai_map.amin())/(xai_map.amax() - xai_map.amin())
    xai_map = torch.nn.functional.interpolate(xai_map, scale_factor=16, mode='bilinear') # (1, 1, 224, 224)
    xai_map = xai_map.squeeze().detach().cpu().numpy() # (224, 224)
    
    ##################
    trans_attr_map = lrp.generate_LRP(image_norm, cls_idx, method="transformer_attribution")
    trans_attr_map = trans_attr_map.detach().cpu().reshape(1, 1, 14,14) # (1, 1, 14, 14)
    trans_attr_map = (trans_attr_map - trans_attr_map.amin())/(trans_attr_map.amax() - trans_attr_map.amin())
    trans_attr_map = torch.nn.functional.interpolate(trans_attr_map, scale_factor=16, mode='bilinear') # (1, 1, 224, 224)
    trans_attr_map = trans_attr_map.squeeze().detach().cpu().numpy() # (224, 224)
    vis_trans_attr_map = show_cam_on_image(image.permute(1, 2, 0).numpy(), trans_attr_map)
    vis_trans_attr_map =  np.uint8(255 * vis_trans_attr_map)
    vis_trans_attr_map = cv2.cvtColor(np.array(vis_trans_attr_map), cv2.COLOR_RGB2BGR)
    ##################

    predmap = extras["predmap"] # (1, tokens-1, classes)
    predmap = predmap[0, :, cls_idx] # (tokens-1)
    predmap = predmap.reshape(1,1,14, 14).detach().cpu() # (1,1,14, 14)
    predmap = (predmap - predmap.amin())/(predmap.amax() - predmap.amin())
    predmap = torch.nn.functional.interpolate(predmap, scale_factor=16, mode='bilinear') # (1, 1, 224, 224)
    predmap = predmap.squeeze().detach().cpu().numpy() # (224, 224)

    weighted_attnmap = extras['weighted_attnmap'] # (1, tokens-1, classes)
    weighted_attnmap = weighted_attnmap[0, :, cls_idx] # (tokens-1)
    weighted_attnmap = weighted_attnmap.reshape(1, 1, 14, 14).detach().cpu() # (1, 1, 14, 14)
    weighted_attnmap = (weighted_attnmap - weighted_attnmap.amin())/(weighted_attnmap.amax() - weighted_attnmap.amin())
    weighted_attnmap = torch.nn.functional.interpolate(weighted_attnmap, scale_factor=16, mode='bilinear') # (1, 1, 224, 224)
    weighted_attnmap = weighted_attnmap.squeeze().detach().cpu().numpy() # (224, 224)

    all_attnmap = extras["all_attnmap"] # (B, blocks, heads, tokens-1)
    all_attnmap = all_attnmap[0] # (blocks, heads, tokens-1)
    all_attnmap = all_attnmap.flatten(0,1) # (blocks*heads, tokens-1)
    all_attnmap = all_attnmap.unflatten(-1, (14,14)) # (blocks*heads, 14, 14)
    all_attnmap = all_attnmap.detach().cpu()
    mean_attnmap = all_attnmap.mean(0) # (14, 14)
    mean_attnmap = mean_attnmap.reshape(1, 1, 14, 14) # (1, 1, 14, 14)
    mean_attnmap = (mean_attnmap - mean_attnmap.amin())/(mean_attnmap.amax() - mean_attnmap.amin())
    mean_attnmap = torch.nn.functional.interpolate(mean_attnmap, scale_factor=16, mode='bilinear') # (1, 1, 224, 224)
    mean_attnmap = mean_attnmap.squeeze().detach().cpu().numpy() # (224, 224)



    mean_last_attnmap = all_attnmap # (blocks*heads, 14, 14)
    mean_last_attnmap = mean_last_attnmap.unflatten(0, (12,12)) # (blocks, heads, 14,14)
    mean_last_attnmap = mean_last_attnmap[-1] # (heads, 14,14)
    mean_last_attnmap = mean_last_attnmap.mean(0) # (14, 14)
    mean_last_attnmap = mean_last_attnmap.reshape(1, 1, 14, 14) # (1, 1, 14, 14)
    mean_last_attnmap = (mean_last_attnmap - mean_last_attnmap.amin())/(mean_last_attnmap.amax() - mean_last_attnmap.amin())
    mean_last_attnmap = torch.nn.functional.interpolate(mean_last_attnmap, scale_factor=16, mode='bilinear') # (1, 1, 224, 224)
    mean_last_attnmap = mean_last_attnmap.squeeze().detach().cpu().numpy() # (224, 224)


    vis_predicatt = show_cam_on_image(image.permute(1, 2, 0).numpy(), xai_map)
    vis_predicatt =  np.uint8(255 * vis_predicatt)
    vis_predicatt = cv2.cvtColor(np.array(vis_predicatt), cv2.COLOR_RGB2BGR)


    vis_predmap = show_cam_on_image(image.permute(1, 2, 0).numpy(), predmap)
    vis_predmap =  np.uint8(255 * vis_predmap)
    vis_predmap = cv2.cvtColor(np.array(vis_predmap), cv2.COLOR_RGB2BGR)

    vis_weighted_attnmap = show_cam_on_image(image.permute(1, 2, 0).numpy(), weighted_attnmap)
    vis_weighted_attnmap =  np.uint8(255 * vis_weighted_attnmap)
    vis_weighted_attnmap = cv2.cvtColor(np.array(vis_weighted_attnmap), cv2.COLOR_RGB2BGR)

    vis_mean_attnmap = show_cam_on_image(image.permute(1, 2, 0).numpy(), mean_attnmap)
    vis_mean_attnmap =  np.uint8(255 * vis_mean_attnmap)
    vis_mean_attnmap = cv2.cvtColor(np.array(vis_mean_attnmap), cv2.COLOR_RGB2BGR)

    vis_mean_last_attnmap = show_cam_on_image(image.permute(1, 2, 0).numpy(), mean_last_attnmap)
    vis_mean_last_attnmap =  np.uint8(255 * vis_mean_last_attnmap)
    vis_mean_last_attnmap = cv2.cvtColor(np.array(vis_mean_last_attnmap), cv2.COLOR_RGB2BGR)


    # Save original image
    path = output_dir / f"decomp_{cls_wind}_{img_path.stem}_original.png"
    plt.imsave(path, image.permute(1, 2, 0).numpy())

    # Save predmap
    path = output_dir / f"decomp_{cls_wind}_{img_path.stem}_p{cls_idx}_predmap_l{layer+1}.png"
    plt.imsave(path, vis_predmap)

    # Save mean attnmap
    path = output_dir / f"decomp_{cls_wind}_{img_path.stem}_p{cls_idx}_mean_attnmap.png"
    plt.imsave(path, vis_mean_attnmap)

    # Save mean attnmap of last layer
    path = output_dir / f"decomp_{cls_wind}_{img_path.stem}_p{cls_idx}_mean_last_attnmap.png"
    plt.imsave(path, vis_mean_last_attnmap)

    fig, axes = plt.subplots(1, 5, figsize=(20, 4))
    ax = axes[0]
    ax.imshow(image.permute(1, 2, 0).numpy())
    pred_str = CLS2IDX[cls_pred_vec[0].argmax().item()]
    ax.set_title(f'Original (pred: {pred_str})')
    ax.axis('off')

    ax = axes[1]
    ax.imshow(vis_trans_attr_map)
    ax.set_title(f'TransAttr\n{CLS2IDX[cls_idx]}')
    ax.axis('off')

    ax = axes[2]
    ax.imshow(vis_predmap)
    ax.set_title(f'PredMap l{layer+1}\n{CLS2IDX[cls_idx]}')
    ax.axis('off')

    ax = axes[3]
    ax.imshow(vis_mean_attnmap)
    ax.set_title(f'mean attnmap')
    ax.axis('off')

    ax = axes[4]
    ax.imshow(vis_mean_last_attnmap)
    ax.set_title(f'mean last attnmap')
    ax.axis('off')

# %%
