import sys
import yaml
from tqdm import tqdm
from datetime import datetime
import numpy as np

import torch
from torch import tensor
from torchvision import transforms
import matplotlib.pyplot as plt

from PIL import ImageFilter
from sklearn import random_projection

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):
    plt.imshow(sample[0].permute(1,2,0))
    plt.imshow(fmap.permute(1,2,0)*255, cmap='jet', alpha=0.5)

    if predict_class == "good":
        #green
        color = (0., 1., 0.)
    else:
        color = (1., 0., 0.)

    image_height, image_width, _ = sample[0].shape

    x_ratio, y_ratio, w_ratio, h_ratio = (0.1, 0.1, 0.8, 0.8)
    x, y, w, h = int(x_ratio * image_width), int(y_ratio * image_height), int(w_ratio * image_width), int(h_ratio * image_height)
    plt.text(x, y-10, s=predict_class,  color=color)

    plt.gca().set_axis_off()
    plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, 
					hspace = 0, wspace = 0)
    plt.margins(0,0)
    plt.gca().xaxis.set_major_locator(plt.NullLocator())
    plt.gca().yaxis.set_major_locator(plt.NullLocator())
    # print(f_output)
    plt.savefig(f_output, bbox_inches = 'tight',
		pad_inches = 0)	
    plt.clf()