from pathlib import Path
import inspect
import numpy as np
import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from numpy import *
import argparse
from PIL import Image
import imageio
import os
from tqdm import tqdm
from utils.metrices import *
import sys
base_dir = Path(__file__).parent.parent
sys.path.append(str(base_dir))
os.chdir(base_dir)


from utils import render
from utils.saver import Saver
from utils.iou import IoU

from data.Imagenet import Imagenet_Segmentation

from sklearn.metrics import precision_recall_curve
import matplotlib.pyplot as plt

import torch.nn.functional as F

plt.switch_backend('agg')

import lovely_tensors as lt
lt.monkey_patch()

# hyperparameters
num_workers = 0
batch_size = 1


# Args
parser = argparse.ArgumentParser(description='Training multi-class classifier')
parser.add_argument('--arc', type=str, default='vgg', metavar='N',
                    help='Model architecture')
parser.add_argument('--train_dataset', type=str, default='imagenet', metavar='N',
                    help='Testing Dataset')
parser.add_argument('--method', type=str,
                    default='grad_rollout',
                    choices=[ 'rollout', 'lrp','transformer_attribution', 'full_lrp', 'lrp_last_layer',
                              'attn_last_layer', 'attn_gradcam', 
                              'attn_last_layer2',
                              'attn_last_layer3',
                              'attn_last_layer4',
                              'attn_last_layer5',
                              'attn_last_layer6',
                              'attn_last_layer7',
                              'attn_last_layer8',
                              'predmap',
                              'predmap2',
                              'predmap3',
                              'predmap4',
                              'predmap5',
                              'predmap6',
                              'predmap7',
                              'predmap8',
                              'predmap9',
                              'predmap10',
                              'predmap11',
                              'predmap12_attn_prev',
                              'predmap13',
                              'predmap14',
                              'predmap15',
                              'predmap16',
                              'predmap_temperature1_1',
                              'predmap_temperature0_9',
                              'predmap_temperature1_2',
                              'predmap_temperature1_3',
                              'predmap_temperature1_35',
                              'predmap_temperature1_4',
                              'predmap_temperature1_45',
                              'predmap_temperature1_5',
                              'predmap_temperature1_6',
                              ],
                    help='')
parser.add_argument('--thr', type=float, default=0.,
                    help='threshold')
parser.add_argument('--K', type=int, default=1,
                    help='new - top K results')
parser.add_argument('--save-img', action='store_true',
                    default=False,
                    help='')
parser.add_argument('--no-ia', action='store_true',
                    default=False,
                    help='')
parser.add_argument('--no-fx', action='store_true',
                    default=False,
                    help='')
parser.add_argument('--no-fgx', action='store_true',
                    default=False,
                    help='')
parser.add_argument('--no-m', action='store_true',
                    default=False,
                    help='')
parser.add_argument('--no-reg', action='store_true',
                    default=False,
                    help='')
parser.add_argument('--is-ablation', type=bool,
                    default=False,
                    help='')
parser.add_argument('--imagenet-seg-path', type=str, required=True)
args = parser.parse_args()



alpha = 2

cuda = torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")

# Data
normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
test_img_trans = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    normalize,
])
test_lbl_trans = transforms.Compose([
    transforms.Resize((224, 224), Image.NEAREST),
])

ds = Imagenet_Segmentation(args.imagenet_seg_path,
                           transform=test_img_trans, target_transform=test_lbl_trans)
dl = DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=False)

cache_dir = Path(__file__).resolve().parents[1] / 'cache'
cache_data = torch.load(cache_dir / f"cache.pt")

iterator = tqdm(dl)

def eval_batch(all_attnmap, predmap, labels, method, index):
    # all_attnmap # (blocks, heads, tokens-1)
    n_blocks = all_attnmap.shape[0]
    n_heads = all_attnmap.shape[1]
    all_attnmap = all_attnmap.flatten(0,1) # (blocks*heads, tokens-1)

    res = all_attnmap
    res = res.unflatten(-1, (14,14)) # (blocks*heads, 14, 14)
    # interpolate to full image size (224,224)
    res = res.reshape(-1, 1, 14, 14) # (blocks*heads, 1, 14, 14)
    res = torch.nn.functional.interpolate(res, scale_factor=16, mode='bilinear', align_corners=False).cuda() # (blocks*heads, 1, 224, 224)
    res = res.squeeze(1) # (blocks, 224, 224)

    res = res.flatten(-2) # (blocks, 224*224)

    # threshold between FG and BG is the mean   
    res = (res - res.min(dim=-1, keepdims=True).values) / (res.max(dim=-1, keepdims=True).values - res.min(dim=-1, keepdims=True).values)  # (blocks, pixels)

    ret = res.mean(dim=-1) # (blocks*heads, )
    ret = ret.reshape(-1, 1, 1) # (blocks*heads, 1, 1)

    res = res.unflatten(-1, (224, 224)) # (blocks*heads, 224, 224)


    Res_1 = res.gt(ret).type(res.type())
    Res_0 = res.le(ret).type(res.type())

    Res_1_AP = res
    Res_0_AP = 1-res

    Res_1[Res_1 != Res_1] = 0
    Res_0[Res_0 != Res_0] = 0
    Res_1_AP[Res_1_AP != Res_1_AP] = 0
    Res_0_AP[Res_0_AP != Res_0_AP] = 0


    # TEST
    pred = res.clamp(min=args.thr) / res.amax(dim=(-2,-1), keepdim=True)
    pred = pred.view(-1).data.cpu().numpy()
    target = labels.view(-1).data.cpu().numpy()
    # print("target", target.shape)

    output = torch.stack((Res_0, Res_1), dim=1) # (blocks*heads, 2, 224, 224)
    output_AP = torch.stack((Res_0_AP, Res_1_AP), dim=1) # (blocks*heads, 2, 224, 224)

    # Evaluate Segmentation
    batch_inter, batch_union, batch_correct, batch_label = 0, 0, 0, 0
    batch_ap, batch_f1 = 0, 0

    # Segmentation resutls
    correct_lst, labeled_lst, inter_lst, union_lst = [], [], [], []
    for i in range(output.shape[0]): # foreach block i
            correct, labeled = batch_pix_accuracy(output[i].data.cpu(), labels[0])
            inter, union = batch_intersection_union(output[i].data.cpu(), labels[0], 2)
            correct_lst.append(correct)
            labeled_lst.append(labeled)
            inter_lst.append(inter)
            union_lst.append(union)

    batch_correct += correct
    batch_label += labeled
    batch_inter += inter
    batch_union += union

    # print("output", output.shape)
    # print("ap labels", labels.shape)
    # ap = np.nan_to_num(get_ap_scores(output, labels))
    ap_lst = []
    for i in range(output_AP.shape[0]):
        ap = np.nan_to_num(get_ap_scores(output_AP[i].unsqueeze(0), labels))
        ap_lst.append(ap)
    # f1 = np.nan_to_num(get_f1_scores(output[0, 1].data.cpu(), labels[0]))
    f1 = 0
    batch_ap += ap
    batch_f1 += f1

    inter_lst_np = np.array(inter_lst)  # (blocks*heads, 2)
    union_lst_np = np.array(union_lst)  # (blocks*heads, 2)
    pix_acc = np.float64(1.0) * np.array(correct_lst) / (np.spacing(1, dtype=np.float64) + np.array(labeled_lst))  # (blocks*heads,)
    iou = np.float64(1.0) * inter_lst_np / (np.spacing(1, dtype=np.float64) + union_lst_np)  # (blocks*heads, 2)
    iou = iou.mean(axis=-1)  # (blocks*heads, )
    ap = np.array(ap_lst).squeeze(-1)  # (blocks*heads, )

    # Unflatten arrays
    pix_acc = pix_acc.reshape(-1, n_heads)  # (blocks, heads)
    iou = iou.reshape(-1, n_heads)  # (blocks, heads)
    ap = ap.reshape(-1, n_heads)  # (blocks, heads)
    inter_lst_np = inter_lst_np.reshape(-1, n_heads, 2)  # (blocks, heads, 2)
    union_lst_np = union_lst_np.reshape(-1, n_heads, 2)  # (blocks, heads, 2)

    return pix_acc, iou, inter_lst_np, union_lst_np, ap, pred, target



total_inter, total_union, total_correct, total_label = np.int64(0), np.int64(0), np.int64(0), np.int64(0)
total_ap, total_f1 = [], []

pix_acc_lst, iou_lst, ap_lst = [], [], []
inter_lst, union_lst = [], []
predictions, targets = [], []
for batch_idx, (image, labels) in enumerate(iterator):
    # if batch_idx == 2:
        #  break
    # images = image.cuda()
    # labels = labels.cuda()
    # print("image", image.shape)
    # print("lables", labels.shape)

    all_attnmap = cache_data['all_attnmap'][batch_idx] # (blocks, heads, tokens-1)
    predmap = cache_data['xpredmap'][batch_idx] # (tokens, classes)
    # all_attnmap = all_attnmap.cuda()
    # predmap = predmap.cuda()

    pix_acc, iou, inter, union, ap, pred, target = eval_batch(all_attnmap, predmap, labels, args.method, batch_idx)
    pix_acc_lst.append(pix_acc)
    iou_lst.append(iou)
    ap_lst.append(ap)
    inter_lst.append(inter)
    union_lst.append(union)

    predictions.append(pred)
    targets.append(target)

    # # total_correct += correct.astype('int64')
    # # total_label += labeled.astype('int64')
    # # total_inter += inter.astype('int64')
    # # total_union += union.astype('int64')
    # # total_ap += [ap]
    # # total_f1 += [f1]
    # # pixAcc = np.float64(1.0) * total_correct / (np.spacing(1, dtype=np.float64) + total_label)
    # # IoU = np.float64(1.0) * total_inter / (np.spacing(1, dtype=np.float64) + total_union) # (2, )
    # # mIoU = IoU.mean()
    # # mAp = np.mean(total_ap)
    # # mF1 = np.mean(total_f1)
    # iterator.set_description('pixAcc: %.4f, mIoU: %.4f, mAP: %.4f, mF1: %.4f' % (pixAcc, mIoU, mAp, mF1))

#########################TODO: PROFILE AND DEBUG
pix_acc_np = np.stack(pix_acc_lst, axis=0) # (index, blocks, heads)
iou_np = np.stack(iou_lst, axis=0) # (index, blocks, heads)
ap_np = np.stack(ap_lst, axis=0) # (index, blocks, heads)
inter_np = np.stack(inter_lst, axis=0) # (index, blocks, heads, 2)
union_np = np.stack(union_lst, axis=0) # (index, blocks, heads, 2)


np.save(cache_dir / 'pix_acc.npy', pix_acc_np)
np.save(cache_dir / 'iou.npy', iou_np)
np.save(cache_dir / 'ap.npy', ap_np)
np.save(cache_dir / 'inter.npy', inter_np)
np.save(cache_dir / 'union.npy', union_np)

