import matplotlib.pyplot as plt
from scipy.spatial import Voronoi, voronoi_plot_2d
from .colormap import checkviz_cmap

from sklearn.neighbors import kneighbors_graph


def checkviz(
    scatter_data,
    false_distortion_list,
    missing_distortion_list,
    ax=None,
    point_c="black",
    point_s=1,
    point_alpha=0.5,
    point_marker="o",
):
    vor = Voronoi(scatter_data)
    ## set size
    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 10))

    voronoi_plot_2d(vor, ax=ax, show_vertices=False, show_points=False, line_width=0)

    for idx, region in enumerate(vor.regions[:-1]):
        if not -1 in region:
            polygon = [vor.vertices[i] for i in region]
            ax.fill(
                *zip(*polygon),
                checkviz_cmap(false_distortion_list[idx], missing_distortion_list[idx])
            )

    ax.scatter(
        scatter_data[:, 0],
        scatter_data[:, 1],
        c=point_c,
        zorder=2,
        s=point_s,
        alpha=point_alpha,
        marker=point_marker,
    )

    ax.set_xticks([])
    ax.set_yticks([])


def reliability_map(
    emb,
    false_distortion_list,
    missing_distortion_list,
    k=7,
    ax=None,
    point_c="black",
    point_s=1,
    point_alpha=0.5,
    point_marker="o",
    linewidth=2.5,
    line_alpha=0.8,
):
    ## construct a knn graph
    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 10))

    knn_graph = kneighbors_graph(emb, k, mode="distance", include_self=False)

    ## visualizae points and knn graph
    ax.scatter(
        emb[:, 0],
        emb[:, 1],
        c=point_c,
        zorder=2,
        s=point_s,
        alpha=point_alpha,
        marker=point_marker,
    )
    for i in range(emb.shape[0]):
        for j in knn_graph[i].indices:
            color = checkviz_cmap(
                (false_distortion_list[i] + false_distortion_list[j]) / 2,
                (missing_distortion_list[i] + missing_distortion_list[j]) / 2,
            )
            ax.plot(
                [emb[i, 0], emb[j, 0]],
                [emb[i, 1], emb[j, 1]],
                c=color,
                zorder=1,
                linewidth=linewidth,
                alpha=line_alpha,
            )

    ax.set_xticks([])
    ax.set_yticks([])
