import io
import PIL
import torchvision
import matplotlib.pyplot as plt
import umap
from sklearn.preprocessing import StandardScaler
import numpy as np
import pandas as pd


def plot_to_image(fig):
    """Converts the matplotlib plot specified by 'figure' to a PNG image and
    returns it. The supplied figure is closed and inaccessible after
    this call."""
    # Save the plot to a PNG in memory.
    buf = io.BytesIO()
    # image_bites = fig.to_image(format="png")
    # print(image_bites)
    fig.savefig(buf, format="jpeg")
    buf.seek(0)
    # Turn plot into tensor to save on tensorboard
    # buf.seek(0)
    # plot_buf = buf
    image_tensor = PIL.Image.open((buf))
    image = torchvision.transforms.ToTensor()(image_tensor)
    return image


def plot_point_cloud(points):

    fig = plt.figure(figsize=(5, 5))
    ax = plt.axes(projection="3d")

    ax.scatter(points[:, 0], points[:, 1], points[:, 2], marker="o", s=10)
    return fig


def feature_landscape(args):

    np.random.seed(42)
    scalar = StandardScaler()
    scaled_features = scalar.fit_transform(
        np.asarray(args.extracted_features.iloc[:, 0 : args.num_features])
    )
    # UMAP for vizualization
    # n_neighbors=5, min_dist=0.0125
    reducer = umap.UMAP(random_state=42)
    embedding = reducer.fit_transform(scaled_features)
    b = np.zeros((len(embedding), 2))
    b[:, 0] = embedding[:, 0]
    b[:, 1] = embedding[:, 1]

    data = pd.DataFrame(b, columns=["Umap1", "Umap2"])
    data["Treatment"] = args.extracted_features["Treatment"].values
