from sklearn.cluster import KMeans
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axisartist.axislines import SubplotZero
from sklearn import manifold
from matplotlib.colors import LinearSegmentedColormap
import os
import matplotlib.image as imgplt





class FindNiche():

    def __init__(self, output_path):

        """
        FindNiche.

        Parameters
        ----------
        output_path
            Path of the output folder.

        """

        self.name = "FindNiche"
        self.output_path = output_path

        if not os.path.exists(self.output_path + "/Downstream/"):
            os.makedirs(self.output_path + "/Downstream/")



    def Kmeans(self, predicted_raw, predicted_imputed, n_kmeans=8):

        """
        Visualization of loss curve and accuracy curve.

        Parameters
        ----------
        predicted_raw
            Path of predicted raw spots with detailed probabilities.
        predicted_imputed
            Path of predicted imputed spots with detailed probabilities.
        n_kmeans
            Numbers of clusters generated by Kmeans.
        """

        predict = []

        with open(predicted_raw) as f:
            for line in f.readlines():
                temp = line.split()
                predict.append(temp)
        print(len(predict))

        with open(predicted_imputed) as f:
            for line in f.readlines():
                temp = line.split()
                predict.append(temp)
        print(len(predict))

        data = []
        for each in predict:
            data.append(each[1:-2])

        data = np.array(data, dtype='float')

        kmeans = KMeans(n_clusters=n_kmeans).fit(data)
        n_labels = kmeans.labels_

        for i in range(len(predict)):
            predict[i].append(n_labels[i])

        np.savetxt(self.output_path + "/Downstream/kmeans_label.txt", predict, fmt='%s', delimiter='\t')



    def TSNE_Plot(self, path, kmeans_file, Label, colors):


        """
        Visualization of TSNE plot.

        Parameters
        ----------
        path
            Path of predicted files with detailed probabilities.
        kmeans_file
            Path of files with kmeans labels.
        Label
            Dictionary of predicted labels.
        colormap
            Dictionary of colors for labels.
        """

        barcode, predict_data, cell_label, kmeans_label = [], [], [], []

        with open(path + "/Imputedpredict_detail.txt") as f:
            for line in f.readlines():
                temp = line.split()
                barcode.append(temp[0])
                predict_data.append(temp[1:-2])

        with open(path + "/Imputed_predict.txt") as f:
            for line in f.readlines():
                temp = line.split()
                cell_label.append(temp[1])

        with open(path + "/Rawpredict_detail.txt") as f:
            for line in f.readlines():
                temp = line.split()
                barcode.append(temp[0])
                predict_data.append(temp[1:-2])

        with open(path + "/Raw_predict.txt") as f:
            for line in f.readlines():
                temp = line.split()
                cell_label.append(temp[1])

        with open(kmeans_file) as f:
            for line in f.readlines():
                temp = line.split()
                kmeans_label.append(temp[-1])

        data = np.array(predict_data, dtype='float')
        clabel = np.array(cell_label, dtype='int')
        klabel = np.array(kmeans_label, dtype='int')

        tsne = manifold.TSNE(n_components=2, random_state=24)
        embedding = tsne.fit_transform(data)
        print(embedding.shape)

        true_label = np.unique(clabel)
        print(true_label)

        colorslist = []
        for each in Label:
            colorslist.append(colors[Label[each]])
        newcmp = LinearSegmentedColormap.from_list('mylist', colorslist)

        plt.scatter(embedding[:, 0], embedding[:, 1], c=clabel, cmap=newcmp, s=5, vmin=0, vmax=len(Label) - 1)
        plt.gca().set_aspect('equal', 'datalim')
        plt.colorbar(boundaries=np.arange(len(Label) + 1) - 0.5).set_ticks(np.arange(len(Label)))
        plt.title('tSNE ' + str(true_label))
        plt.savefig(self.output_path + "/Downstream/tsne_celltype.pdf")
        plt.clf()

        plt.scatter(embedding[:, 0], embedding[:, 1], c=klabel, cmap="Accent", s=5)
        plt.gca().set_aspect('equal', 'datalim')
        plt.colorbar()
        plt.title('tSNE')
        plt.savefig(self.output_path + "/Downstream/tsne_kmeans.pdf")
        plt.clf()



    def SpatialNiche(self, kmeans_file, image_path, colormap, pointsize=5, back_image=False):

        """
        Spatial plot of cell-type niches.

        Parameters
        ----------
        kmeans_file
            Path of files with kmeans labels.
        image_path
            Path of the original H&E staining image.
        colormap
            Dictionary of colors for labels.
        output_path
            Path of output files.
        pointsize
            Size of scatter points.
        back_image
            Shown of the original H&E staining image.
        """

        Image = imgplt.imread(image_path)
        h, w, _ = Image.shape

        predict = []
        with open(kmeans_file) as f:
            for line in f.readlines():
                temp = line.split()
                predict.append(temp)

        fig = plt.figure(figsize=(10, 10))
        ax = SubplotZero(fig, 1, 1, 1)
        fig.add_subplot(ax)
        ax.set_xlim(left=0, right=int(max(h, w)))
        ax.set_ylim(bottom=int(max(h, w)), top=0)

        if back_image == True:
            plt.imshow(Image)

        for each in predict:
            plt.scatter(float(each[-3]), float(each[-2]), c=colormap[int(each[-1])], s=pointsize, marker='o')
        plt.savefig(self.output_path + "/Downstream/kmeans_spatial.pdf")







