import os
from PIL import Image
import numpy as np
import torch
import random
import matplotlib.pyplot as plt
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
from pytorch_grad_cam.utils.image import show_cam_on_image
from scipy.ndimage import rotate
import imageio
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
from scipy.interpolate import interp1d


def one_hot(x, num_classes):
    x = x.squeeze()
    out = np.zeros([x.shape[0], num_classes]).astype(int)
    out[np.arange(x.shape[0]), x] = 1
    return out


def manual_seed(seed=0):
    np.random.seed(seed)
    torch.manual_seed(seed)
    random.seed(seed)

    torch.backends.cudnn.deterministic = True


# @profile
def log_train_val(
    train_loss,
    test_loss=None,
    train_acc=None,
    test_acc=None,
    grad_norm=None,
    log_dir="./log",
):
    plt.figure(figsize=(12, 4))
    plt.plot(np.arange(len(train_loss)), train_loss, label="train_loss")
    if test_loss is not None:
        plt.plot(np.arange(len(test_loss)), test_loss, label="test_loss")
    dest = os.path.join(log_dir, "loss.png")
    plt.legend()
    plt.savefig(dest, dpi=200)
    plt.close("all")

    if train_acc is not None:
        plt.figure(figsize=(12, 4))
        plt.plot(np.arange(len(train_acc)), train_acc, label="train_acc")
        if test_acc is not None:
            plt.plot(np.arange(len(test_acc)), test_acc, label="test_acc")
        dest = os.path.join(log_dir, "acc.png")
        plt.legend()
        plt.savefig(dest, dpi=200)
        plt.close("all")

    if grad_norm is not None:
        plt.figure(figsize=(12, 4))
        plt.plot(np.arange(len(grad_norm)), grad_norm, label="grad_norm")
        dest = os.path.join(log_dir, "grad_norm.png")
        plt.legend()
        plt.savefig(dest, dpi=200)
        plt.close("all")

        plt.figure()
        plt.hist(grad_norm, 10, label="grad_norm_hist")
        percentile_90 = np.percentile(grad_norm, 90)
        dest = os.path.join(log_dir, f"grad_norm_hist_{percentile_90:.4f}.png")
        plt.legend()
        plt.savefig(dest, dpi=200)
        plt.close("all")


def log_img(ep, batch_idx, img, label=None, log_dir="./logs"):
    dest = os.path.join(log_dir, f"ep_{ep}_b_{batch_idx}.png")
    if label is not None:
        dest = dest.replace(".png", f"_label_{label}.png")
    img = (img - img.min()) / (img.max() - img.min())
    if len(img.shape) == 3:
        img = img.transpose((1, 2, 0))
    plt.imsave(dest, img)


def vis_sre_kernel(
    sre_conv, vis_step=1, mode="2D", vmin=None, vmax=None, in_start=0, out_start=0
):
    weight = sre_conv.weight.detach().cpu().numpy()
    kernel = sre_conv._make_weight_matrix(sre_conv.weight).detach().cpu().numpy()
    Cout, Cin, k = weight.shape
    kernel_size = kernel.shape[-1]
    if isinstance(vis_step, int):
        in_vis_step = vis_step
        out_vis_step = vis_step
    else:
        in_vis_step, out_vis_step = vis_step

    in_channels = np.arange(in_start, Cin, in_vis_step)
    out_channels = np.arange(out_start, Cout, out_vis_step)

    if mode == "2D":
        plt.figure(figsize=(2 * len(out_channels), 2 * len(in_channels)))
        idx = 1
        for cin in in_channels:
            for cout in out_channels:
                ax = plt.subplot(len(in_channels), len(out_channels), idx)
                # plt.title(f'K_{cin}_{cout}')
                plt.imshow(kernel[cout, cin, :, :], cmap="Blues", vmin=vmin, vmax=vmax)
                plt.axis("off")
                # cbar = plt.colorbar()

                # if vmin is None:
                #     vmin = kernel[cout, cin, :, :].min()
                # if vmax is None:
                #     vmax = kernel[cout, cin, :, :].max()
                # Apply scientific notation to colorbar
                # # cbar.set_ticks([vmin, (vmin + vmax)/2, vmax])
                # # cbar.set_ticklabels([vmin, (vmin + vmax)/2, vmax])  # Custom labels
                # formatter = ScalarFormatter(useMathText=True)
                # formatter.set_scientific(True)
                # # formatter.set_powerlimits((vmin, vmax))  # You can adjust these limits
                # cbar.ax.yaxis.set_major_formatter(formatter)

                idx += 1
    else:
        plt.figure(figsize=(2 * len(in_channels), 2 * len(out_channels)))
        x = np.arange(k)
        idx = 1
        for cin in in_channels:
            for cout in out_channels:
                ax = plt.subplot(len(in_channels), len(out_channels), idx)
                plt.title(f"K_{cin}_{cout}")
                plt.ylim((vmin, vmax))
                ax.plot(x, weight[cout, cin, :])
                idx += 1
    return plt.gcf()


def create_circle_mask(image_shape, radius):
    # Create grid of coordinates
    x, y = np.ogrid[: image_shape[0], : image_shape[1]]

    center = (image_shape[0] // 2, image_shape[1] // 2)
    # Calculate squared distance from each point to the center
    distance_squared = (x - center[0]) ** 2 + (y - center[1]) ** 2

    # Create mask where True denotes points inside the circle
    mask = distance_squared <= radius**2

    return mask.astype(np.uint8)


def vis_cam(
    nrow,
    ncol,
    model,
    target_layers,
    input_list,
    test_dataset=None,
    img_idx=None,
    save_path=None,
    vis_output=False,
    vis_input=False,
    angle_list=None,
    rot_back=False,
    circle_mask=False,
):
    plt.figure(figsize=(2 * ncol, 2 * nrow))
    plt.margins(x=0)
    for i in range(nrow * ncol):
        if input_list is not None:
            input, label = input_list[i]
        else:
            input, label = test_dataset.__getitem__(img_idx[i])
        if vis_input:
            plt.imsave(
                save_path.replace(".png", f"_{i}_input.png"),
                input.numpy().transpose((1, 2, 0)),
            )
        cam = GradCAM(model=model, target_layers=target_layers)
        input_tensor = input[None]
        targets = [ClassifierOutputTarget(label)]

        grayscale_cam = cam(input_tensor=input_tensor, targets=targets)
        img = input.numpy().transpose((1, 2, 0))
        img = (img - np.min(img)) / (np.max(img) - np.min(img))
        visualization = show_cam_on_image(img, grayscale_cam[0, :], use_rgb=True)

        if angle_list is not None and rot_back:
            angle = angle_list[i]
            h, w, _ = visualization.shape
            visualization = rotate(visualization, -angle, reshape=False)
        if circle_mask:
            h, w, _ = visualization.shape
            mask = create_circle_mask((h, w), 0.5 * h)
            mask = mask == 0
            mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
            visualization = np.ma.masked_array(visualization, mask)
            visualization[mask] = 255

        plt.subplot(nrow, ncol, i + 1)
        plt.imshow(visualization)
        plt.axis("off")
    plt.subplots_adjust(wspace=0, hspace=0)
    plt.tight_layout()
    if save_path is not None:
        plt.savefig(save_path, dpi=300, bbox_inches="tight")


outputs = []


def hook(module, input, output):
    global outputs
    outputs.clear()
    outputs.append(output)


outputs = []


def hook_split(module, input, output):
    global outputs
    outputs.clear()
    final_output, temp_output = module._costume_forward_dw(input[0], return_all=True)
    outputs.append(final_output)
    outputs.append(temp_output)


def vis_feat(
    nrow,
    ncol,
    model,
    target_layer,
    input_list,
    feat_idx=1,
    test_dataset=None,
    img_idx=None,
    save_path=None,
    vis_output=False,
    vis_input=False,
    angle_list=None,
    rot_back=False,
    circle_mask=False,
    create_gif=False,
    frame_duration=0.0005,
    gif_input=False,
    rescale=True,
    split_plot=False,
):
    if create_gif:
        img_list = []
    else:
        plt.figure(figsize=(2 * ncol, 2 * nrow))
        plt.margins(x=0)
    for i in range(nrow * ncol):
        if input_list is not None:
            input, label = input_list[i]
        else:
            input, label = test_dataset.__getitem__(img_idx[i])
        if split_plot:
            handle = getattr(model, target_layer)[1].conv1.register_forward_hook(
                hook_split
            )
        elif "layer" in target_layer:
            handle = getattr(model, target_layer)[1].conv1.register_forward_hook(hook)
        else:
            handle = getattr(model, target_layer).register_forward_hook(hook)
        if vis_input:
            plt.imsave(
                save_path.replace(".png", f"_{i}_input.png"),
                input.numpy().transpose((1, 2, 0)),
            )
        output = model(input[None])
        input_img = input[None].detach().numpy().squeeze().transpose((1, 2, 0))
        if split_plot:
            num_rings = outputs[1].shape[1]
            for j in range(num_rings):
                ring_img = outputs[1][:, j].detach().numpy().squeeze()
                ring_img = ring_img.mean(0)
                ring_img = (ring_img - np.min(ring_img)) / (
                    np.max(ring_img) - np.min(ring_img)
                )
                plt.imsave(
                    save_path.replace(".png", f"_{i}_ring_{j}.png"),
                    ring_img,
                    cmap="viridis",
                )
        feat = outputs[0].detach().numpy().squeeze()
        if feat_idx == "mean":
            img = feat.mean(0)
        else:
            assert isinstance(feat_idx, int)
            img = feat[feat_idx]
        img = (img - np.min(img)) / (np.max(img) - np.min(img))
        if angle_list is not None and rot_back:
            angle = angle_list[i]
            h, w = img.shape
            input_img = rotate(input_img, -angle, reshape=False)
            img = rotate(img, -angle, reshape=False)
        if circle_mask:
            h, w = img.shape
            mask = create_circle_mask(img.shape, 0.5 * h)
            mask = mask == 0
            img = np.ma.masked_array(img, mask)
            h, w = input_img.shape[:2]
            mask = create_circle_mask((h, w), 0.5 * h)
            mask = mask == 0
            mask = mask.repeat(3).reshape(h, w, 3)
            input_img = np.ma.masked_array(input_img, mask)

        if create_gif:
            print(f"Creating gif frame: {i+1}/{nrow*ncol}\r", end="")
            if gif_input:
                img_list.append(input_img)
            else:
                # print(img.shape)
                img_list.append(img)
        else:
            plt.subplot(nrow, ncol, i + 1)
            plt.imshow(img, cmap="viridis")
            plt.axis("off")
            plt.imsave(
                save_path.replace(".png", f"_{i}.png"), img, dpi=300, cmap="viridis"
            )
        # plt.imsave(save_path.replace(".png", f"_{i}.png"), img, dpi=300, cmap='viridis')
        handle.remove()
        if vis_output:
            plt.figure()
            plt.bar(
                np.arange(10), torch.softmax(output[0], 0).detach().numpy().squeeze()
            )
            plt.ylim((0, 1))
            plt.savefig(
                save_path.replace(".png", f"_{i}_output.png"),
                dpi=300,
                bbox_inches="tight",
            )
            plt.close("all")
            plt.cla()
    if create_gif:
        cmap = plt.cm.viridis
        output_path = save_path.replace(".png", ".gif")
        # global rescale
        max_intensity = max([img.max() for img in img_list])
        min_intensity = min([img.min() for img in img_list])
        h, w = img_list[0].shape[:2]
        mask = create_circle_mask((h, w), 0.5 * h)
        mask = mask == 0
        if not gif_input:
            mask = np.repeat(mask[:, :, np.newaxis], 4, axis=2)
        with imageio.get_writer(
            output_path, mode="I", duration=frame_duration
        ) as writer:
            for img in img_list:
                if rescale:
                    img = (img - min_intensity) / (max_intensity - min_intensity)
                if gif_input:
                    image_color = img
                else:
                    image_color = cmap(img)
                image_color = (image_color * 255).astype(np.uint8)
                # fill masked region with white color
                image_color[mask] = 255
                writer.append_data(image_color)
        img_stack = np.stack(img_list, axis=0)  # N, H, W
        rotation_wise_std = np.std(img_stack, axis=0)
        # skip the masked region
        if gif_input:
            rotation_wise_std[mask] = 0
        else:
            rotation_wise_std[mask[:, :, 0]] = 0
        avg_std = np.mean(rotation_wise_std)
        print(f"Avg. rotation-wise signal STD: {avg_std:.4f}\n")
        return rotation_wise_std
    else:
        plt.subplots_adjust(wspace=0, hspace=0)
        plt.tight_layout()
        if save_path is not None:
            plt.savefig(save_path, dpi=300, bbox_inches="tight")
        return None


def acc_at_topk(targets, prob, k=5):
    assert len(targets) == len(prob)
    topk = np.argsort(prob, axis=1)[:, -k:]
    correct = 0
    for i in range(len(targets)):
        if targets[i] in topk[i]:
            correct += 1
    return correct / len(targets)


def force_cudnn_initialization():
    s = 32
    dev = torch.device("cuda")
    torch.nn.functional.conv2d(
        torch.zeros(s, s, s, s, device=dev), torch.zeros(s, s, s, s, device=dev)
    )


def visualize_pca_2d(features, labels, vis=True):
    # Perform PCA
    pca = PCA(n_components=2, random_state=42)
    features = features.astype(np.float64)
    # features = (features - features.mean(axis=0)) / features.std(axis=0)
    features = features / np.linalg.norm(features, axis=1)[:, None]
    pca_result = pca.fit_transform(features)
    if not vis:
        return None, pca_result

    # Create a scatter plot of the PCA results
    plt.figure(figsize=(10, 8))
    classes = np.unique(labels)
    colors = plt.cm.tab10(np.linspace(0, 1, len(classes)))

    for i, class_label in enumerate(classes):
        idx = labels == class_label
        plt.scatter(
            pca_result[idx, 0], pca_result[idx, 1], color=colors[i], label=class_label
        )

    plt.title("PCA Visualization of Data")
    plt.xlabel("Principal Component 1")
    plt.ylabel("Principal Component 2")
    plt.xlim(-1.0, 1.0)
    plt.ylim(-1.0, 1.0)
    plt.grid(True)
    return plt.gcf(), pca_result


def visualize_pca_3d(features, labels, vis=True):
    # Perform PCA
    pca = PCA(n_components=3, random_state=42)
    X_pca = pca.fit_transform(features)
    if not vis:
        return None, X_pca

    # Visualize in 3D space
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection="3d")

    classes = np.unique(labels)
    colors = plt.cm.tab10(
        np.linspace(0, 1, len(classes))
    )  # Assuming there are 3 classes
    for class_label in classes:
        class_indices = np.where(labels == class_label)[0]
        ax.scatter(
            X_pca[class_indices, 0],
            X_pca[class_indices, 1],
            X_pca[class_indices, 2],
            color=colors[class_label],
            label=f"Class {class_label}",
        )

    ax.set_xlabel("Principal Component 1")
    ax.set_ylabel("Principal Component 2")
    ax.set_zlabel("Principal Component 3")
    ax.set_xlim(-1.0, 1.0)
    ax.set_ylim(-1.0, 1.0)
    ax.set_zlim(-1.0, 1.0)
    plt.title("3D PCA Visualization")
    plt.show()


def visualize_tsne_2d(features, labels, vis=True):
    # Perform t-SNE
    tsne = TSNE(
        n_components=2,
        random_state=42,
        init="pca",
        learning_rate="auto",
        n_iter=1500,
        n_iter_without_progress=300,
        verbose=1,
        perplexity=30,
    )
    # tsne = TSNE(n_components=2, random_state=42, init='random', learning_rate='auto',
    #             n_iter=1500, n_iter_without_progress=300, verbose=1, perplexity=20)
    # features = features.astype(np.float64)
    # features = features / np.linalg.norm(features, axis=1)[:, None]
    tsne_result = tsne.fit_transform(features)
    if not vis:
        return None, tsne_result

    # Create a scatter plot of the t-SNE results
    plt.figure(figsize=(10, 8))
    classes = np.unique(labels)
    classes = np.sort(classes)
    colors = plt.cm.tab10(np.linspace(0, 1, len(classes)))

    for i, class_label in enumerate(classes):
        idx = labels == class_label
        plt.scatter(
            tsne_result[idx, 0],
            tsne_result[idx, 1],
            color=colors[i],
            label=class_label,
            s=5,
        )

    plt.title("t-SNE Visualization of Data")
    plt.xlabel("t-SNE Dimension 1")
    plt.ylabel("t-SNE Dimension 2")
    plt.xlim(-110, 110)
    plt.ylim(-110, 110)
    plt.grid(True)
    return plt.gcf(), tsne_result


def visualize_tsne_3d(features, labels, vis=True):
    # Perform t-SNE
    tsne = TSNE(
        n_components=3,
        random_state=42,
        init="pca",
        learning_rate="auto",
        n_iter=1500,
        n_iter_without_progress=300,
        verbose=1,
        perplexity=30,
    )
    # features = features.astype(np.float64)
    # features = features / np.linalg.norm(features, axis=1)[:, None]
    tsne_result = tsne.fit_transform(features)
    if not vis:
        return None, tsne_result

    # Create a scatter plot of the t-SNE results
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection="3d")
    classes = np.unique(labels)
    classes = np.sort(classes)
    colors = plt.cm.tab10(np.linspace(0, 1, len(classes)))

    for i, class_label in enumerate(classes):
        idx = labels == class_label
        plt.scatter(
            tsne_result[idx, 0],
            tsne_result[idx, 1],
            tsne_result[idx, 2],
            color=colors[i],
            label=class_label,
            s=5,
        )

    plt.title("3D tSNE Visualization")
    ax.set_xlabel("t-SNE Dimension 1")
    ax.set_ylabel("t-SNE Dimension 2")
    ax.set_zlabel("t-SNE Dimension 3")
    ax.set_xlim(-110, 110)
    ax.set_ylim(-110, 110)
    ax.set_zlim(-110, 110)
    return plt.gcf(), tsne_result


def interpolate_1d_array(array, K, axis=1, interp_type="linear"):
    """
    Interpolates a 1D array of size N into K evenly spaced points.

    Parameters:
    array (ndarray): Input ND array of size N.
    K (int): Number of points to interpolate.
    axis (int): Axis along which to interpolate.

    Returns:
    ndarray: Interpolated 1D array of size K.
    """
    N = array.shape[axis]
    x = np.linspace(0, N - 1, N)  # Original indices
    xp = np.linspace(0, N - 1, K)  # New indices

    # Create interpolation function
    interpolator = interp1d(x, array, axis=axis, kind=interp_type)

    # Interpolate the array
    interpolated_array = interpolator(xp)

    return interpolated_array


def create_logits_gif(
    predictions, label, degrees, label2cap, color="tab:blue", dest="./tmp/logits.gif"
):
    n_classes = len(label2cap)
    fig_list = []
    for idx, deg in enumerate(degrees):
        cur_pred = predictions[idx]
        plt.figure(figsize=(10, 6))
        plt.bar(np.arange(n_classes), cur_pred, color=color)
        plt.xticks(np.arange(n_classes), label2cap, rotation=15)
        plt.title(f"Predictions at {deg} degrees")
        plt.ylim((0, 1))
        plt.xlim((-0.5, n_classes - 0.5))
        fig = plt.gcf()
        fig.canvas.draw()
        image_array = np.array(fig.canvas.renderer._renderer)
        plt.close("all")
        fig_list.append(image_array)
    with imageio.get_writer(dest, mode="I", duration=0.002) as writer:
        for img in fig_list:
            writer.append_data(img)
    del fig_list
