
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import abc
import numpy as np
import sys
import yaml
from tqdm import tqdm
from datetime import datetime
import numpy as np
import cv2
import torch
from torch import tensor
from torchvision import transforms
import torch.nn.functional as F
import matplotlib.pyplot as plt
from matplotlib.ticker import FormatStrFormatter
from PIL import ImageFilter
from sklearn import random_projection
from sklearn.metrics import pairwise_distances
TQDM_PARAMS = {
	"file" : sys.stdout,
	"bar_format" : "   {l_bar}{bar:10}{r_bar}{bar:-10b}",
}

def get_tqdm_params():
    return TQDM_PARAMS

class GaussianBlur:
    def __init__(self, radius : int = 4):
        self.radius = radius
        self.unload = transforms.ToPILImage()
        self.load = transforms.ToTensor()
        self.blur_kernel = ImageFilter.GaussianBlur(radius=4)

    def __call__(self, img):
        map_max = img.max()
        final_map = self.load(
            self.unload(img[0]/map_max).filter(self.blur_kernel)
        )*map_max
        return final_map


def get_coreset_idx_randomp(
    z_lib : tensor, 
    n : int = 1000,
    eps : float = 0.90,
    float16 : bool = True,
    force_cpu : bool = False,
) -> tensor:
    """Returns n coreset idx for given z_lib.
    
    Performance on AMD3700, 32GB RAM, RTX3080 (10GB):
    CPU: 40-60 it/s, GPU: 500+ it/s (float32), 1500+ it/s (float16)

    Args:
        z_lib:      (n, d) tensor of patches.
        n:          Number of patches to select.
        eps:        Agression of the sparse random projection.
        float16:    Cast all to float16, saves memory and is a bit faster (on GPU).
        force_cpu:  Force cpu, useful in case of GPU OOM.

    Returns:
        coreset indices
    """

    print(f"   Fitting random projections. Start dim = {z_lib.shape}.")
    try:
        transformer = random_projection.SparseRandomProjection(eps=eps)
        z_lib = torch.tensor(transformer.fit_transform(z_lib))
        print(f"   DONE.                 Transformed dim = {z_lib.shape}.")
    except ValueError:
        print( "   Error: could not project vectors. Please increase `eps`.")

    select_idx = 0
    last_item = z_lib[select_idx:select_idx+1]
    coreset_idx = [torch.tensor(select_idx)]
    min_distances = torch.linalg.norm(z_lib-last_item, dim=1, keepdims=True)
    # The line below is not faster than linalg.norm, although i'm keeping it in for
    # future reference.
    # min_distances = torch.sum(torch.pow(z_lib-last_item, 2), dim=1, keepdims=True)

    if float16:
        last_item = last_item.half()
        z_lib = z_lib.half()
        min_distances = min_distances.half()
    if torch.cuda.is_available() and not force_cpu:
        last_item = last_item.to("cuda")
        z_lib = z_lib.to("cuda")
        min_distances = min_distances.to("cuda")

    for _ in tqdm(range(n-1), **TQDM_PARAMS):
        distances = torch.linalg.norm(z_lib-last_item, dim=1, keepdims=True) # broadcasting step
        # distances = torch.sum(torch.pow(z_lib-last_item, 2), dim=1, keepdims=True) # broadcasting step
        min_distances = torch.minimum(distances, min_distances) # iterative step
        select_idx = torch.argmax(min_distances) # selection step

        # bookkeeping
        last_item = z_lib[select_idx:select_idx+1]
        min_distances[select_idx] = 0
        coreset_idx.append(select_idx.to("cpu"))

    return torch.stack(coreset_idx)

def print_and_export_results(results : dict, save_file_path: str):
    """Writes results to .yaml and serialized results to .txt."""
    
    print("\n   ╭────────────────────────────╮")
    print(  "   │      Results summary       │")
    print(  "   ┢━━━━━━━━━━━━━━━━━━━━━━━━━━━━┪")
    print( f"   ┃ image rocauc: {results['image rocauc']:.2f} ┃")
    print(  "   ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛\n")

    with open(save_file_path, "w") as outfile:
        yaml.safe_dump(results, outfile, default_flow_style=False)


def serialize_results(results : dict) -> str:
    """Serialize a results dict into something usable in markdown."""
    n_first_col = 20
    ans = []
    for k, v in results.items():
        s = k + " "*(n_first_col-len(k))
        s = s + f"| {v[0]*100:.1f}  |"
        ans.append(s)
    return "\n".join(ans)

def save_image(sample, fmap, f_output, predict_class, img_size, display_mode:str="Heatmap", pixel_score=0.0):

    image_width, image_height= img_size
    # resize the figure
    fig, ax = plt.subplots()
    fig.set_size_inches(image_width/100, image_height/100)
    sample_img = F.interpolate(sample[0].unsqueeze(0), size=(image_height, image_width), mode='bilinear', align_corners=False).squeeze(0)
    fmap = F.interpolate(fmap.unsqueeze(0), size=(image_height, image_width), mode='bilinear', align_corners=False).squeeze(0)

    ax.imshow(sample_img.permute(1, 2, 0))
    if display_mode == "HighLight":
        fmap = fmap.permute(1,2,0).numpy()
        if pixel_score == 0:
            pixel_score = np.mean(fmap)

        mask = fmap > pixel_score

        fmap= (fmap/np.max(fmap) * 255).astype(np.uint8)
        fmap = cv2.cvtColor(fmap, cv2.COLOR_RGB2BGR)

        fmap = cv2.applyColorMap(fmap, cv2.COLORMAP_JET)
        mask = mask.astype(np.uint8) * 255
        # draw contours
        contour, _ = cv2.findContours(mask[:,:,0], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if contour:
            cv2.drawContours(fmap, contour, -1, (0, 255, 255), 3)
        # remove bg
        mask = np.where((mask==(0,0,0)).all(axis=2), 0,255).astype(np.uint8)
        fmap = cv2.cvtColor(fmap, cv2.COLOR_BGR2RGBA)
        fmap[:, :, 3] = mask
        ax.imshow(fmap, cmap='jet', alpha=0.5)
    else:
        ax.imshow(fmap.permute(1,2,0) * 255, cmap='jet', alpha=0.5)
    ax.set_axis_off()
    plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, 
					hspace = 0, wspace = 0)
    ax.margins(0,0)
    ax.xaxis.set_major_locator(plt.NullLocator())
    ax.yaxis.set_major_locator(plt.NullLocator())

    fig.savefig(f_output, bbox_inches = 'tight', pad_inches = 0)
    fig.clf()
    plt.close(fig)

def embedding_concat(x, y):
    """Concatenate the given sequence of 2 tensors

    Args:
        x (torch.Tensor): Input feature map with shape (B, C1, H1, W1)
        y (torch.Tensor): Input feature map with shape (B, C2, H2, W2)

    Returns:
        torch.Tensor: Concatenated feature map with shape (B, C1 + C2, H1, W1)
    """
    B, C1, H1, W1 = x.size()
    _, C2, H2, W2 = y.size()
    s = int(H1 / H2)
    x = F.unfold(x, kernel_size=s, dilation=1, stride=s)
    x = x.view(B, C1, -1, H2, W2)
    z = torch.zeros(B, C1 + C2, x.size(2), H2, W2)
    for i in range(x.size(2)):
        z[:, :, i, :, :] = torch.cat((x[:, :, i, :, :], y), 1)
    z = z.view(B, -1, H2 * W2)
    z = F.fold(z, kernel_size=s, output_size=(H1, W1), stride=s)
    return z

def reshape_embedding(embedding):
    embedding_list = []
    for k in range(embedding.shape[0]):
        for i in range(embedding.shape[2]):
            for j in range(embedding.shape[3]):
                embedding_list.append(embedding[k, :, i, j])
    return embedding_list

def distance_matrix(x, y=None, p=2):  # pairwise distance of vectors

    y = x if type(y) == type(None) else y

    n = x.size(0)
    m = y.size(0)
    d = x.size(1)

    x = x.unsqueeze(1).expand(n, m, d)
    y = y.unsqueeze(0).expand(n, m, d)

    #dist = torch.pow(x - y, p).sum(2)
    dist = torch.cdist(x, y, p)

    return dist

class NN():

    def __init__(self, X=None, Y=None, p=2):
        self.p = p
        self.train(X, Y)

    def train(self, X, Y):
        self.train_pts = X
        self.train_label = Y

    def __call__(self, x):
        return self.predict(x)

    def predict(self, x):
        if type(self.train_pts) == type(None) or type(self.train_label) == type(None):
            name = self.__class__.__name__
            raise RuntimeError(f"{name} wasn't trained. Need to execute {name}.train() first")

        dist = distance_matrix(x, self.train_pts, self.p) ** (1 / self.p)
        labels = torch.argmin(dist, dim=1)
        return self.train_label[labels]

class KNN(NN):

	def __init__(self, X=None, Y=None, k=3, p=2):
		self.k = k
		super().__init__(X, Y, p)

	def train(self, X, Y):
		super().train(X, Y)
		if type(Y) != type(None):
			self.unique_labels = self.train_label.unique()

	def predict(self, x):
		# dist = distance_matrix(x, self.train_pts, self.p) ** (1 / self.p)
		dist = torch.cdist(x, self.train_pts, self.p)
		knn = dist.topk(self.k, largest=False)
		return knn

class SamplingMethod(object):
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def __init__(self, X, y, seed, **kwargs):
        self.X = X
        self.y = y
        self.seed = seed

    def flatten_X(self):
        shape = self.X.shape
        flat_X = self.X
        if len(shape) > 2:
            flat_X = np.reshape(self.X, (shape[0],np.product(shape[1:])))
        return flat_X


    @abc.abstractmethod
    def select_batch_(self):
        return

    def select_batch(self, **kwargs):
        return self.select_batch_(**kwargs)

    def to_dict(self):
        return None

class kCenterGreedy(SamplingMethod):

    def __init__(self, X, y, seed, metric='euclidean'):
        self.X = X
        self.y = y
        self.flat_X = self.flatten_X()
        self.name = 'kcenter'
        self.features = self.flat_X
        self.metric = metric
        self.min_distances = None
        self.n_obs = self.X.shape[0]
        self.already_selected = []

    def update_distances(self, cluster_centers, only_new=True, reset_dist=False):
        """Update min distances given cluster centers.

        Args:
            cluster_centers: indices of cluster centers
            only_new: only calculate distance for newly selected points and update
            min_distances.
            rest_dist: whether to reset min_distances.
        """

        if reset_dist:
            self.min_distances = None
        if only_new:
            cluster_centers = [d for d in cluster_centers
                                if d not in self.already_selected]
        if cluster_centers:
            # Update min_distances for all examples given new cluster center.
            x = self.features[cluster_centers]
            dist = pairwise_distances(self.features, x, metric=self.metric)

            if self.min_distances is None:
                self.min_distances = np.min(dist, axis=1).reshape(-1,1)
            else:
                self.min_distances = np.minimum(self.min_distances, dist)

    def select_batch_(self, model, already_selected, N, **kwargs):
        """
        Diversity promoting active learning method that greedily forms a batch
        to minimize the maximum distance to a cluster center among all unlabeled
        datapoints.

        Args:
            model: model with scikit-like API with decision_function implemented
            already_selected: index of datapoints already selected
            N: batch size

        Returns:
            indices of points selected to minimize distance to cluster centers
        """

        try:
            # Assumes that the transform function takes in original data and not
            # flattened data.
            print('Getting transformed features...')
            self.features = model.transform(self.X)
            print('Calculating distances...')
            self.update_distances(already_selected, only_new=False, reset_dist=True)
        except:
            print('Using flat_X as features.')
            self.update_distances(already_selected, only_new=True, reset_dist=False)

        new_batch = []
        for _ in tqdm(range(N)):
            if self.already_selected is None:
            # Initialize centers with a randomly selected datapoint
                ind = np.random.choice(np.arange(self.n_obs))
            else:
                ind = np.argmax(self.min_distances)
            # New examples should not be in already selected since those points
            # should have min_distance of zero to a cluster center.
            assert ind not in already_selected

            self.update_distances([ind], only_new=True, reset_dist=False)
            new_batch.append(ind)
        print('Maximum distance from cluster centers is %0.2f'
                % max(self.min_distances))


        self.already_selected = already_selected

        return new_batch

def visualize_score(scores, labels, output_path, bins=10):
    """Visualize Score

    Args:
        scores (List): list score 
        labels (List): List label, label is 0(good) or 1(ng)
        output_path (str): path to save diagram
        bins (int): the number of equal-width bins in the range. Default 10
    """
    if len(scores) != len(labels):
        return
    fig, ax = plt.subplots(2,1, figsize=(10,8))

    # label: good = 0, ng = 1
    good_index = np.where(np.array(labels) == 0)[0]
    good_scores = np.array(scores)[good_index]
    ng_index = np.where(np.array(labels) == 1)[0]
    ng_scores = np.array(scores)[ng_index]
    # create Score scatter
    # x: index of score, y: score
    ax[0].scatter(good_index, good_scores, alpha=0.5, label='good')
    ax[0].scatter(ng_index, ng_scores, alpha=0.5, label='ng')
    ax[0].set_title("Index vs Score")
    ax[0].set_xlabel("Index")
    ax[0].set_ylabel("Score")
    ax[0].grid()
    ax[0].legend(loc='center left', bbox_to_anchor=(1, 0.5))

    # create space between 2 rows
    fig.subplots_adjust(hspace=0.3)
    # create histogram
    ax[1].hist(good_scores, alpha=0.5, bins=bins, label='good')
    ax[1].hist(ng_scores, alpha=0.5, bins=bins, label='ng')
    ax[1].set_title("Histogram of Data")
    ax[1].set_xlabel("Score")
    ax[1].set_ylabel("Frequency")
    ax[1].grid()
    ax[1].legend(loc='center left', bbox_to_anchor=(1, 0.5))

    fig.savefig(output_path)
    fig.clf()
    plt.close(fig)