from typing import Tuple, List
from tqdm import tqdm
import sys
import os

import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import tensor
from torch.utils.data import DataLoader
import timm
import antialiased_cnns
from antialiased_cnns.blurpool import BlurPool
from torchvision.models.feature_extraction import create_feature_extractor
import numpy as np
from sklearn.metrics import roc_auc_score, roc_curve
from sklearn.metrics import ConfusionMatrixDisplay
from sklearn.metrics import confusion_matrix,classification_report
from sklearn.cluster import KMeans
sys.path.insert(0, os.path.abspath(
    os.path.join(os.path.dirname(__file__), './../../')
))
from ecos_core.indad.utils_indad import GaussianBlur, get_coreset_idx_randomp, get_tqdm_params, save_image, \
	embedding_concat, KNN
from ecos_core.indad.amazon_builder import patchcore, common, sampler, backbones


def get_anti_aliased_model(backbone_name):
	_backbone = {
		"antialiased_wide_resnet50_2": 'antialiased_cnns.wide_resnet50_2(pretrained=True)',
		# "antialiased_wide_resnet101_2": 'antialiased_cnns.wide_resnet101_2(pretrained=True)',
	}
	return eval(_backbone[backbone_name])


class AdaptiveAvgPool2dCustom(nn.Module):
	def __init__(self, output_size):
		super(AdaptiveAvgPool2dCustom, self).__init__()
		self.output_size = np.array(output_size)

	def forward(self, x: torch.Tensor):
		shape_x = x.shape
		if shape_x[-1] < self.output_size[-1]:
			output_shape = (x.shape[0], x.shape[1], self.output_size[0], self.output_size[1])
			padding = [(output_shape[i + 2] - x.shape[i + 2]) // 2 for i in range(2) for _ in range(2)]
			x = F.pad(x, pad=padding)
		stride_size = np.floor(np.array(x.shape[-2:]) / self.output_size).astype(np.int32)
		kernel_size = np.array(x.shape[-2:]) - (self.output_size - 1) * stride_size

		avg = nn.AvgPool2d(kernel_size=list(kernel_size), stride=list(stride_size))
		x = avg(x)
		return x

class KNNExtractor(torch.nn.Module):
	def __init__(
		self,
		backbone_name : str = "resnet50",
		out_indices : Tuple = None,
		pool_last : bool = False,
		layers: List[str] = []
	):
		super().__init__()
		self.layers = layers
		if "antialiased" in backbone_name:
			model = get_anti_aliased_model(backbone_name)

			self.feature_extractor = create_feature_extractor(
				model=model,
				return_nodes={layer: layer for layer in layers},
				tracer_kwargs={"leaf_modules": [BlurPool]},  # for models comes from antialias
        	)
		else:
			self.feature_extractor = timm.create_model(
				backbone_name,
				out_indices=out_indices,
				features_only=True,
				pretrained=True,
			)
		for param in self.feature_extractor.parameters():
			param.requires_grad = False
		self.feature_extractor.eval()

		self.pool = torch.nn.AdaptiveAvgPool2d(1) if pool_last else None
		self.backbone_name = backbone_name # for results metadata
		self.out_indices = out_indices

		self.device = "cuda" if torch.cuda.is_available() else "cpu"
		#print(self.feature_extractor)
		self.feature_extractor = self.feature_extractor.to(self.device)
			
	def call(self, x: tensor):
		with torch.no_grad():
			feature_maps = self.feature_extractor(x.to(self.device))
		# now antialiased implement extract from layer 2, 3
		if "antialiased" in self.backbone_name:
			feature_maps = [feature for layer, feature in feature_maps.items()]
		feature_maps = [fmap.to("cpu") for fmap in feature_maps]
		if self.pool:
			# spit into fmaps and z
			return feature_maps[:-1], self.pool(feature_maps[-1])
		else:
			return feature_maps

	def __call__(self, x: tensor, y: tensor = None):
		if y is not None:
			return self.predict(x,y)
		return self.predict(x)

	def fit(self, _: DataLoader):
		raise NotImplementedError

	def predict(self, _: tensor, y: tensor = None):
		raise NotImplementedError
	
	def evaluate_mask(self, test_dl: DataLoader) -> Tuple[float, float]:
		"""Calls predict step for each test sample."""
		image_preds = []
		image_labels = []
		pixel_preds = []
		pixel_labels = []

		for sample, mask, label in tqdm(test_dl, **get_tqdm_params()):
			z_score, fmap = self.predict(sample)
			
			image_preds.append(z_score.numpy())
			image_labels.append(label)
			
			pixel_preds.extend(fmap.flatten().numpy())
			pixel_labels.extend(mask.flatten().numpy())
			
		image_labels = np.stack(image_labels)
		image_preds = np.stack(image_preds)

		image_rocauc = roc_auc_score(image_labels, image_preds)
		pixel_rocauc = roc_auc_score(pixel_labels, pixel_preds)

		return float(image_rocauc), float(pixel_rocauc)

	def evaluate(self, test_dl: DataLoader) -> float:
		"""Calls predict step for each test sample."""
		image_preds = []
		image_labels = []
		pixel_preds = []

		for sample, label in tqdm(test_dl, **get_tqdm_params()):
			z_score, fmap = self.predict(sample)
			
			image_preds.append(z_score.numpy())
			image_labels.append(label)
			
			pixel_preds.extend(fmap.flatten().numpy())
			
		image_labels = np.stack(image_labels)
		image_preds = np.stack(image_preds)

		image_rocauc = roc_auc_score(image_labels, image_preds)

		return float(image_rocauc)

	def get_parameters(self, extra_params : dict = None) -> dict:
		return {
			"backbone_name": self.backbone_name,
			"out_indices": self.out_indices,
			**extra_params,
		}


class SPADE(KNNExtractor):
	def __init__(
		self,
		k: int = 5,
		size: int = 224,
		backbone_name: str = "resnet18",
	):
		super().__init__(
			backbone_name=backbone_name,
			out_indices=(1,2,3,-1),
			pool_last=True,
		)
		self.k = k
		self.image_size = size
		self.z_lib = []
		self.feature_maps = []
		self.threshold_z = None
		self.threshold_fmaps = None
		self.blur = GaussianBlur(4)

	def fit(self, train_dl):
		for sample, _ in tqdm(train_dl, **get_tqdm_params()):
			feature_maps, z = self.call(sample)

			# z vector
			self.z_lib.append(z)

			# feature maps
			if len(self.feature_maps) == 0:
				for fmap in feature_maps:
					self.feature_maps.append([fmap])
			else:
				for idx, fmap in enumerate(feature_maps):
					self.feature_maps[idx].append(fmap)

		self.z_lib = torch.vstack(self.z_lib)
		
		for idx, fmap in enumerate(self.feature_maps):
			self.feature_maps[idx] = torch.vstack(fmap)

	def predict(self, sample):
		feature_maps, z = self.call(sample)
		distances = torch.linalg.norm(self.z_lib - z, dim=1)

		values, indices = torch.topk(distances.squeeze(), self.k, largest=False)

		z_score = values.mean()

		# Build the feature gallery out of the k nearest neighbours.
		# The authors migh have concatenated all features maps first, then check the minimum norm per pixel.
		# Here, we check for the minimum norm first, then concatenate (sum) in the final layer.
		scaled_s_map = torch.zeros(1,1,self.image_size,self.image_size)
		for idx, fmap in enumerate(feature_maps):

			nearest_fmaps = torch.index_select(self.feature_maps[idx], 0, indices)
			# min() because kappa=1 in the paper
			s_map, _ = torch.min(torch.linalg.norm(nearest_fmaps - fmap, dim=1), 0, keepdims=True)
			scaled_s_map += torch.nn.functional.interpolate(
				s_map.unsqueeze(0), size=(self.image_size, self.image_size), mode='bilinear'
			)

		scaled_s_map = self.blur(scaled_s_map)
		
		return z_score, scaled_s_map

	def get_parameters(self):
		return super().get_parameters({
			"k": self.k,
		})

class PaDiM(KNNExtractor):
	def __init__(
		self,
		size: int = 224,
		d_reduced: int = 100,
		backbone_name: str = "resnet18",
	):
		super().__init__(
			backbone_name=backbone_name,
			out_indices=(1,2,3),
		)
		self.image_size = size
		self.d_reduced = d_reduced # your RAM will thank you
		self.epsilon = 0.04 # cov regularization
		self.patch_lib = []
		self.resize = None
		self.r_indices = None
		self.means_reduced = None
		self.E_inv = None

	def fit(self, train_dl):
		for sample, _ in tqdm(train_dl, **get_tqdm_params()):
			feature_maps = self.call(sample)
			if self.resize is None:
				largest_fmap_size = feature_maps[0].shape[-2:]
				self.resize = torch.nn.AdaptiveAvgPool2d(largest_fmap_size)
			resized_maps = [self.resize(fmap) for fmap in feature_maps]
			self.patch_lib.append(torch.cat(resized_maps, 1))
		self.patch_lib = torch.cat(self.patch_lib, 0)

		# random projection
		if self.patch_lib.shape[1] > self.d_reduced:
			print(f"   PaDiM: (randomly) reducing {self.patch_lib.shape[1]} dimensions to {self.d_reduced}.")
			self.r_indices = torch.randperm(self.patch_lib.shape[1])[:self.d_reduced]
			self.patch_lib_reduced = self.patch_lib[:,self.r_indices,...]
		else:
			print("   PaDiM: d_reduced is higher than the actual number of dimensions, copying self.patch_lib ...")
			self.patch_lib_reduced = self.patch_lib

		# calcs
		self.means = torch.mean(self.patch_lib, dim=0, keepdim=True)
		self.means_reduced = self.means[:,self.r_indices,...]
		x_ = self.patch_lib_reduced - self.means_reduced

		# cov calc
		self.E = torch.einsum(
			'abkl,bckl->ackl',
			x_.permute([1,0,2,3]), # transpose first two dims
			x_,
		) * 1/(self.patch_lib.shape[0]-1)
		self.E += self.epsilon * torch.eye(self.d_reduced).unsqueeze(-1).unsqueeze(-1)
		self.E_inv = torch.linalg.inv(self.E.permute([2,3,0,1])).permute([2,3,0,1])

	def predict(self, sample):
		feature_maps = self.call(sample)
		resized_maps = [self.resize(fmap) for fmap in feature_maps]
		fmap = torch.cat(resized_maps, 1)
		# reduce
		x_ = fmap[:,self.r_indices,...] - self.means_reduced

		left = torch.einsum('abkl,bckl->ackl', x_, self.E_inv)
		s_map = torch.sqrt(torch.einsum('abkl,abkl->akl', left, x_))
		scaled_s_map = torch.nn.functional.interpolate(
			s_map.unsqueeze(0), size=(self.image_size, self.image_size), mode='bilinear'
		)

		return torch.max(s_map), scaled_s_map[0, ...]

	def get_parameters(self):
		return super().get_parameters({
			"d_reduced": self.d_reduced,
			"epsilon": self.epsilon,
		})


class PatchCore(KNNExtractor):
	def __init__(
		self,
		k : int = 3,
		size: int = 224,
		f_coreset: float = 0.01, # fraction the number of training samples
		backbone_name : str = "resnet18",
		coreset_eps: float = 0.90, # sparse projection parameter,
		layers: List[str] = [], # layers to extract
		normalize: bool = False,
	):
		super().__init__(
			backbone_name=backbone_name,
			out_indices=(2,3),
			layers=layers,
		)
		self.f_coreset = f_coreset
		self.coreset_eps = coreset_eps
		self.image_size = size
		self.average = torch.nn.AvgPool2d(3, stride=1)
		self.blur = GaussianBlur(4)
		self.n_reweight = k
		self.normalize = normalize
		self.patch_lib = []
		self.resize = None


	def fit(self, train_dl):
		#Clear patch_lib before starting training
		self.patch_lib = []
		for sample, _ in tqdm(train_dl, **get_tqdm_params()):
			feature_maps = self.call(sample)

			if self.resize is None:
				largest_fmap_size = feature_maps[0].shape[-2:]
				self.resize = torch.nn.AdaptiveAvgPool2d(largest_fmap_size)

			resized_maps = [self.resize(self.average(fmap)) for fmap in feature_maps]
			patch = torch.cat(resized_maps, 1) # torch.Size([1, 1536, 28, 28])
			patch = patch.reshape(patch.shape[1], -1).T # torch.Size([784, 1536])
			self.patch_lib.append(patch)

		self.patch_lib = torch.cat(self.patch_lib, 0)

		if self.f_coreset < 1:
			self.coreset_idx = get_coreset_idx_randomp(
				self.patch_lib,
				n=int(self.f_coreset * self.patch_lib.shape[0]),
				eps=self.coreset_eps,
			)
			self.patch_lib = self.patch_lib[self.coreset_idx]

	def predict(self, sample):		
		feature_maps = self.call(sample)
		resized_maps = [self.resize(self.average(fmap)) for fmap in feature_maps]
		patch = torch.cat(resized_maps, 1)
		patch = patch.reshape(patch.shape[1], -1).T


		dist = torch.cdist(patch, self.patch_lib)
		min_val, min_idx = torch.min(dist, dim=1)
		s_idx = torch.argmax(min_val)
		s_star = torch.max(min_val)

		# reweighting
		m_test = patch[s_idx].unsqueeze(0) # anomalous patch
		m_star = self.patch_lib[min_idx[s_idx]].unsqueeze(0) # closest neighbour
		w_dist = torch.cdist(m_star, self.patch_lib) # find knn to m_star pt.1
		_, nn_idx = torch.topk(w_dist, k=self.n_reweight, largest=False) # pt.2
		# equation 7 from the paper
		m_star_knn = torch.linalg.norm(m_test-self.patch_lib[nn_idx[0,1:]], dim=1)
		# Softmax normalization trick as in transformers.
		# As the patch vectors grow larger, their norm might differ a lot.
		# exp(norm) can give infinities.
		D = torch.sqrt(torch.tensor(patch.shape[1], dtype=torch.float32))
		w = 1-(torch.exp(s_star/D)/(torch.sum(torch.exp(m_star_knn/D))))

		s = w*s_star

		# segmentation map
		s_map = min_val.view(1, 1,*feature_maps[0].shape[-2:])
		s_map = torch.nn.functional.interpolate(
			s_map, size=(self.image_size, self.image_size), mode='bilinear'
		)
		s_map = self.blur(s_map)
		if self.normalize:
			min_s_map = torch.min(s_map)
			max_s_map = torch.max(s_map)
			s_map = (s_map - min_s_map) / (max_s_map - min_s_map)
		return s, s_map


	def get_parameters(self):
		return super().get_parameters({
			"f_coreset": self.f_coreset,
			"n_reweight": self.n_reweight,
		})


class AmazonPatchCore(patchcore.PatchCore):
	def __init__(self,
	      device="cpu",
	      backbone_name:str="wide_resnet50_2",
		  layers_to_extract_from=["layer2", "layer3"],
		  pretrain_embed_dimension:int=1024,
		  target_embed_dimension:int=1024,
		  patchsize:int=3,
		  input_shape=(3, 224, 224),
		  percentage:float=0.1,
		  anomaly_scorer_num_nn:int=1,
		  seed:int=0,
		  num_workers:int=2,
		  is_cuda_used:bool=False
		  ):
		super().__init__(device)
		self.backbone = backbones.load(backbone_name)
		self.backbone.name, self.backbone.seed = backbone_name, seed
		self.nn_method = common.FaissNN(is_cuda_used, num_workers)
		super().load(
			backbone=self.backbone,
			layers_to_extract_from=layers_to_extract_from,
			device=device,
			input_shape=input_shape,
			pretrain_embed_dimension=pretrain_embed_dimension,
			target_embed_dimension=target_embed_dimension,
			patchsize=patchsize,
			featuresampler=sampler.ApproximateGreedyCoresetSampler(percentage, torch.device(device)),
			anomaly_scorer_num_nn=anomaly_scorer_num_nn,
			nn_method=self.nn_method,
			)

	def __call__(self, x:torch.tensor):
		"""Inference

        Args:
            x (torch.tensor): _description_

        Returns:
            tuple(torch.tensor, torch.tensor): image score and pixel score
        """

		return self.predict(x)
	
	def predict(self, x:torch.tensor):
		img_level_score, pxl_level_score = super()._predict(x)
		pxl_score_tensor = torch.tensor(pxl_level_score[0]).unsqueeze(0)
		return torch.tensor(img_level_score), pxl_score_tensor


	@staticmethod
	def get_backbone_from_name(backbone_name):
		return backbones.load(backbone_name)
	
class PatchCoreFAISS(KNNExtractor):
	def __init__(
		self,
		k: int = 3,
		size: int = 224,
		f_coreset: float = 0.01, # fraction the number of training samples
		backbone_name : str = "resnet18",
		coreset_eps: float = 0.90, # sparse projection parameter
		normalize:bool = False,
	):
		super().__init__(
			backbone_name=backbone_name,
			out_indices=(2,3),
		)
		self.f_coreset = f_coreset
		self.coreset_eps = coreset_eps
		self.image_size = size
		self.average = torch.nn.AvgPool2d(3, stride=1)
		self.blur = GaussianBlur(4)
		self.n_reweight = k
		self.normalize = normalize
		self.patch_lib = []
		self.resize = None
		self.nn_method = common.FaissNN(False, 4)
		self.anomaly_scorer = common.NearestNeighbourScorer(
            n_nearest_neighbours=self.n_reweight, nn_method=common.FaissNN(False, 4)
        )


	def fit(self, train_dl):
		#Clear patch_lib before starting training
		self.patch_lib = []
		for sample, _ in tqdm(train_dl, **get_tqdm_params()):
			feature_maps = self.call(sample)

			if self.resize is None:
				largest_fmap_size = feature_maps[0].shape[-2:]
				self.resize = torch.nn.AdaptiveAvgPool2d(largest_fmap_size)

			resized_maps = [self.resize(self.average(fmap)) for fmap in feature_maps]
			patch = torch.cat(resized_maps, 1)
			patch = patch.reshape(patch.shape[1], -1).T

			self.patch_lib.append(patch)
			# fit faiss
			self.anomaly_scorer.fit(detection_features=[patch])

		self.patch_lib = torch.cat(self.patch_lib, 0)

		if self.f_coreset < 1:
			self.coreset_idx = get_coreset_idx_randomp(
				self.patch_lib,
				n=int(self.f_coreset * self.patch_lib.shape[0]),
				eps=self.coreset_eps,
			)
			self.patch_lib = self.patch_lib[self.coreset_idx]

	def predict(self, sample):
		feature_maps = self.call(sample)
		resized_maps = [self.resize(self.average(fmap)) for fmap in feature_maps]
		patch = torch.cat(resized_maps, 1)
		patch = patch.reshape(patch.shape[1], -1).T


		dist = torch.cdist(patch, self.patch_lib)
		min_val, min_idx = torch.min(dist, dim=1)
		s_idx = torch.argmax(min_val)
		s_star = torch.max(min_val)

		# reweighting
		m_test = patch[s_idx].unsqueeze(0) # anomalous patch
		m_star = self.patch_lib[min_idx[s_idx]].unsqueeze(0) # closest neighbour
		w_dist = torch.cdist(m_star, self.patch_lib) # find knn to m_star pt.1

		if self.anomaly_scorer.nn_method.search_index is not None:
			# search with faiss
			s, _, nn_idx = self.anomaly_scorer.predict([patch.numpy()])
		else:
			_, nn_idx = torch.topk(w_dist, k=self.n_reweight, largest=False) # pt.2
		# equation 7 from the paper
		m_star_knn = torch.linalg.norm(m_test-self.patch_lib[nn_idx[0,1:]], dim=1)
		# Softmax normalization trick as in transformers.
		# As the patch vectors grow larger, their norm might differ a lot.
		# exp(norm) can give infinities.
		D = torch.sqrt(torch.tensor(patch.shape[1], dtype=torch.float32))
		w = 1-(torch.exp(s_star/D)/(torch.sum(torch.exp(m_star_knn/D))))

		s = w*s_star

		# segmentation map
		s_map = min_val.view(1, 1,*feature_maps[0].shape[-2:])
		s_map = torch.nn.functional.interpolate(
			s_map, size=(self.image_size, self.image_size), mode='bilinear'
		)
		s_map = self.blur(s_map)
		if self.normalize:
			min_s_map = torch.min(s_map)
			max_s_map = torch.max(s_map)
			s_map = (s_map - min_s_map) / (max_s_map - min_s_map)
		return s, s_map
	
class SAPatchCore(KNNExtractor):
	def __init__(
		self,
		k : int = 3,
		size: int = 224,
		f_coreset: float = 0.01, # fraction the number of training samples
		backbone_name : str = "wide_resnet_50",
		coreset_eps: float = 0.90, # sparse projection parameter,
		layers: List[str] = [], # layers to extract
		normalize: bool = False
	):
		super().__init__(
			backbone_name=backbone_name,
			out_indices=(2,3),
			layers=layers,
		)
		self.f_coreset = f_coreset
		self.coreset_eps = coreset_eps
		self.image_size = size
		self.average = torch.nn.AvgPool2d(3, stride=1, padding=1)
		self.maxp = torch.nn.MaxPool2d(3, 1, 1)
		self.blur = GaussianBlur(4)
		self.n_reweight = k
		self.normalize = normalize
		self.patch_lib = []
		self.resize = None

	def _features_processing(self, features):
		embeddings = []
		for index, fmap in enumerate(features):
			if index != len(features) - 1:
				saconv_out = self.average(fmap)
			else:
				attention_in1 = self.maxp(fmap[:])
				attention_in2 = self.maxp(fmap[:])
				attention_v = self.maxp(fmap[:]) 

				width = attention_in1.shape[3]
		
				attention_in1_flatten = torch.flatten(attention_in1[:], 2, 3)
				attention_in2_flatten = torch.flatten(attention_in1[:], 2, 3)
				attention_v = torch.flatten(attention_v[:], 2, 3)
				attention_in1 = torch.reshape(attention_in1_flatten, (attention_in1_flatten.shape[0], attention_in1_flatten.shape[1], attention_in1_flatten.shape[2], 1))
				attention_in2 = torch.reshape(attention_in2_flatten, (attention_in2_flatten.shape[0], attention_in2_flatten.shape[1], 1, attention_in2_flatten.shape[2]))
				attention_v = torch.reshape(attention_v, (attention_v.shape[0], attention_v.shape[1], attention_v.shape[2], 1))
				attention_out = torch.matmul(attention_in1, attention_in2)
				attention_out = torch.nn.functional.softmax(attention_out, dim=-1)
				attention_out = torch.matmul(attention_out, attention_v)
				saconv_out = torch.reshape(attention_out, (attention_out.shape[0], attention_out.shape[1], width, width)) 
		
			embeddings.append(saconv_out)

		patch = embedding_concat(embeddings[0], embeddings[1])
		patch_tensor = torch.Tensor(patch) # torch.Size([1, 1536, 28, 28])
		patch = patch.reshape(patch_tensor.shape[1], -1).T

		return patch
	
	def fit(self, train_dl):
		#Clear patch_lib before starting training
		self.patch_lib = []
		for sample, _ in tqdm(train_dl, **get_tqdm_params()):
			feature_maps = self.call(sample)

			patch = self._features_processing(feature_maps)
			#print("Patch: ",patch[0].shape, patch[1].shape)
			self.patch_lib.append(patch)

		self.patch_lib = torch.cat(self.patch_lib, 0)
		print("shape: ", self.patch_lib.shape)
		# print("Total embedding shape: ",total_embeddings.shape)
		# print("N: ", int(self.f_coreset * total_embeddings.shape[0]))
		# print("shape patch lib:", self.patch_lib.shape)
		# if self.f_coreset < 1:
		# 	self.coreset_idx = get_coreset_idx_randomp(
		# 		self.patch_lib,
		# 		n=int(self.f_coreset * self.patch_lib.shape[0]),
		# 		eps=self.coreset_eps,
		# 	)
		# 	self.patch_lib = self.patch_lib[self.coreset_idx]
		from sklearn.random_projection import SparseRandomProjection
		from ecos_core.indad.utils_indad import kCenterGreedy
		self.randomprojector = SparseRandomProjection(n_components='auto', eps=0.9) # 'auto' => Johnson-Lindenstrauss lemma  
		self.randomprojector.fit(self.patch_lib)
        # Coreset Subsampling

		selector = kCenterGreedy(self.patch_lib,0,0)
		print("N: ",int(self.patch_lib.shape[0]*self.f_coreset))
		selected_idx = selector.select_batch(model=self.randomprojector, already_selected=[], N=int(self.patch_lib.shape[0]*self.f_coreset))
		self.patch_lib = self.patch_lib[selected_idx]

	def predict(self, sample):		

		feature_maps = self.call(sample)
		
		patch = self._features_processing(feature_maps)
		knn = KNN(self.patch_lib, k=self.n_reweight)
		score_patches = knn(patch)[0].cpu().detach().numpy()
		s_map = score_patches[:,0].reshape(feature_maps[0].shape[-2:])  

		s = max(score_patches[:,0]) # Image-level score
		#gt_np = gt.cpu().numpy()[0,0].astype(int)
		s_map = cv2.resize(s_map, (self.image_size, self.image_size))

		s = torch.tensor(s)
		
		s_map = torch.from_numpy(s_map).unsqueeze(0)
		if self.normalize:
			min_s_map = torch.min(s_map)
			max_s_map = torch.max(s_map)
			s_map = (s_map - min_s_map) / (max_s_map - min_s_map)
		return s, s_map


class PatchCoreCluster(KNNExtractor):
	def __init__(
		self,
		size: int = 256,
		second_size: int = 512,
		max_iter:int = 300,
		n_cluster: int = 100, # number cluster
		backbone_name : str = "resnet18",
		layers: List[str] = [], # layers to extract
		normalize: bool = False,
	):
		super().__init__(
			backbone_name=backbone_name,
			out_indices=(2,3),
			layers=layers,
		)
		self.n_cluster = n_cluster
		self.image_size = size
		self.average = torch.nn.AvgPool2d(3, stride=1)
		self.blur = GaussianBlur(4)

		self.normalize = normalize
		self.patch_lib = []
		self.resize = None
		self.second_resize = None
		self.second_size = second_size
		self.kmeans = KMeans(n_clusters=self.n_cluster, max_iter=max_iter)

	def fit(self, train_dl):
		#Clear patch_lib before starting training
		self.patch_lib = []

		for sample, _ in tqdm(train_dl, **get_tqdm_params()):
			feature_maps = self.call(sample)

			if self.resize is None:
				largest_fmap_size = feature_maps[0].shape[-2:]
				self.resize = torch.nn.AdaptiveAvgPool2d(largest_fmap_size)

			resized_maps = [self.resize(self.average(fmap)) for fmap in feature_maps]
			patch = torch.cat(resized_maps, 1) # torch.Size([1, 1536, 28, 28])
			patch = patch.reshape(patch.shape[1], -1).T # torch.Size([784, 1536])
			# high resolution
			second_feature_maps = self.call(_.get("second_image_trans"))
			if self.second_resize is None:
				largest_sc_fmap_size = second_feature_maps[0].shape[-2:]
				self.second_resize = torch.nn.AdaptiveAvgPool2d(largest_sc_fmap_size)
			second_resized_maps = [self.second_resize(self.average(second_fmap)) for second_fmap in second_feature_maps]
			second_patch = torch.cat(second_resized_maps, 1) 
			second_patch = second_patch.reshape(second_patch.shape[1], -1).T
			# append to bank
			self.patch_lib.append(patch)
			self.patch_lib.append(second_patch)
		
		self.patch_lib = torch.cat(self.patch_lib, 0)

		self.kmeans.fit(torch.tensor(self.patch_lib))
		
		# asign cluster center into patch_lib
		self.patch_lib = torch.from_numpy(self.kmeans.cluster_centers_.astype(np.float32))

	def predict(self, sample, second_sample=None):		
		feature_maps = self.call(sample)
		resized_maps = [self.resize(self.average(fmap)) for fmap in feature_maps]
		patch = torch.cat(resized_maps, 1)
		patch = patch.reshape(patch.shape[1], -1).T
		distances = torch.from_numpy(self.kmeans.transform(patch))
		min_distance, min_cluster = torch.min(distances, dim=1)
		
		second_feature_maps = self.call(second_sample)
		resized_maps = [self.second_resize(self.average(second_fmap)) for second_fmap in second_feature_maps]
		second_patch = torch.cat(resized_maps, 1)
		second_patch = second_patch.reshape(second_patch.shape[1], -1).T
		second_distances = torch.from_numpy(self.kmeans.transform(second_patch))
		second_min_distance, second_min_cluster = torch.min(second_distances, dim=1)
		s_first_patch = torch.mean(min_distance)
		s_second_patch = torch.mean(second_min_distance)
		s = (s_first_patch + s_second_patch)/2
		# segmentation map
		first_s_map = min_distance.view(1, 1,*feature_maps[0].shape[-2:])
		first_s_map = torch.nn.functional.interpolate(
			first_s_map, size=(self.second_size, self.second_size), mode='bilinear'
		)
		second_s_map = second_min_distance.view(1,1,*second_feature_maps[0].shape[-2:])
		second_s_map = torch.nn.functional.interpolate(
			second_s_map, size=(self.second_size, self.second_size), mode='bilinear'
		)
		s_map = (first_s_map + second_s_map)/2
		s_map = self.blur(s_map)
		if self.normalize:
			min_s_map = torch.min(s_map)
			max_s_map = torch.max(s_map)
			s_map = (s_map - min_s_map) / (max_s_map - min_s_map)
		return s, s_map


	def get_parameters(self):
		return super().get_parameters({
			"n_cluster": self.n_cluster,
		})