# %%
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_new import vit_base_patch16_224 as vit_new

import torchvision.datasets

# %%


imagenet_dir = Path("~/ramdisk/datasets/imagenet/").expanduser()

imagenet_ds = torchvision.datasets.ImageNet(imagenet_dir, split='val')


imagenet_indices=[
    440,
4115,
4480,
8628,
# 11341,
13449,
16904,
18517,
# 24681,
25569,
# 29858,
# 36293,
46451,
49031,
49165,
]
# %%

figs_dir = Path("figures/supp/predmap_per_layer/")
for idx in imagenet_indices:
    ds_img_path, cls_idx = imagenet_ds.samples[idx]
    ds_img_path = Path(ds_img_path)

    img_name = Path(ds_img_path).stem
    img_path = figs_dir / f'predmap_{img_name}.png'
    cls_name = CLS2IDX[cls_idx]
    print(f"% {cls_name}")
    print(
r"""
\raisebox{0.01\linewidth}{
    \rotatebox{90}{\makecell{"""+cls_name+"""}}
} &""")
    print(
        rf" \includegraphics[width=0.11\linewidth]{{ {str(img_path)} }} &"
        )
    for layer in range(5,12):
        predmap_path = figs_dir / f'predmap_{img_name}_c{cls_idx}_l{layer+1}.png'
        
        end = "&" if layer != 11 else ""
        print(
            rf"\includegraphics[width=0.11\linewidth]{{ {str(predmap_path)} }} {end}"
            )

    print(r"\\")

# TEMPLATE:
"""
     \\
     \raisebox{0.01\linewidth}{
     \rotatebox{90}{\makecell{Golden \\ Retriever}}
     }
     & 
    \includegraphics[width=0.11\linewidth]{figures/predmap_per_layer/predmap_ILSVRC2012_val_00038781.png} &
    % \includegraphics[width=0.11\linewidth]{figures/predmap_per_layer/predmap_ILSVRC2012_val_00038781_golden retriever_l1.png} &
    % \includegraphics[width=0.11\linewidth]{figures/predmap_per_layer/predmap_ILSVRC2012_val_00038781_golden retriever_l2.png} &
    % \includegraphics[width=0.11\linewidth]{figures/predmap_per_layer/predmap_ILSVRC2012_val_00038781_golden retriever_l3.png} &
    % \includegraphics[width=0.11\linewidth]{figures/predmap_per_layer/predmap_ILSVRC2012_val_00038781_golden retriever_l4.png} &
    % \includegraphics[width=0.11\linewidth]{figures/predmap_per_layer/predmap_ILSVRC2012_val_00038781_golden retriever_l5.png} &
    \includegraphics[width=0.11\linewidth]{figures/predmap_per_layer/predmap_ILSVRC2012_val_00038781_golden retriever_l6.png} &
    \includegraphics[width=0.11\linewidth]{figures/predmap_per_layer/predmap_ILSVRC2012_val_00038781_golden retriever_l7.png} &
    \includegraphics[width=0.11\linewidth]{figures/predmap_per_layer/predmap_ILSVRC2012_val_00038781_golden retriever_l8.png} &
    \includegraphics[width=0.11\linewidth]{figures/predmap_per_layer/predmap_ILSVRC2012_val_00038781_golden retriever_l9.png} &
    \includegraphics[width=0.11\linewidth]{figures/predmap_per_layer/predmap_ILSVRC2012_val_00038781_golden retriever_l10.png} &
    \includegraphics[width=0.11\linewidth]{figures/predmap_per_layer/predmap_ILSVRC2012_val_00038781_golden retriever_l11.png} &
    \includegraphics[width=0.11\linewidth]{figures/predmap_per_layer/predmap_ILSVRC2012_val_00038781_golden retriever_l12.png} 
    \\
"""

# %%
