# -*- coding: utf-8 -*-
import math
import os
import torch
import random

import anndata
#import cmcrameri
import numpy as np
import scanpy as sc
import networkx as nx
import torch.nn as nn
import matplotlib.pyplot as plt
from scipy.spatial import distance_matrix
from torch_geometric.nn import GCNConv, DeepGraphInfomax
from sklearn.neighbors import kneighbors_graph
from ._util import sparse_mx_to_torch_edge_list, corruption

class SpaceFlow(object):
    """An object for analysis of spatial transcriptomics data.

    :param adata: the `anndata.AnnData` object as input, see `https://anndata.readthedocs.io/en/latest/` for more info about`anndata`.
    :type adata: class:`anndata.AnnData`
    :param count_matrix: count matrix of gene expression, 2D numpy array of size (n_cells, n_genes)
    :type count_matrix: class:`numpy.ndarray`
    :param spatial_locs: spatial locations of cells (or spots) match to rows of the count matrix, 1D numpy array of size (n_cells,)
    :type spatial_locs: class:`numpy.ndarray`
    :param sample_names: list of sample names in 1D numpy str array of size (n_cells,), optional
    :type sample_names: class:`numpy.ndarray` or `list` of `str`
    :param gene_names: list of gene names in 1D numpy str array of size (n_genes,), optional
    :type gene_names: class:`numpy.ndarray` or `list` of `str`

    """

    def __init__(self, adata=None, count_matrix=None, spatial_locs=None, sample_names=None, gene_names=None):
        """
        Inputs
        ------
        adata: an anndata.AnnData type object, optional (either input `adata` or both `count_matrix` and `spatial_locs`)
        count_matrix : count matrix of gene expression, 2D numpy array of size (n_cells, n_genes)
        spatial_locs : spatial locations of cells (or spots) match to rows of the count matrix, 1D numpy array of size (n_cells,)
        sample_names : list of sample names in 1D numpy str array of size (n_cells,), optional
        gene_names : list of gene names in 1D numpy str array of size (n_genes,), optional
        """
        if adata and isinstance(adata, anndata.AnnData):
            self.adata = adata
        elif count_matrix is not None and spatial_locs is not None:
            self.adata = anndata.AnnData(count_matrix.astype(float))
            self.adata.obsm['spatial'] = spatial_locs.astype(float)
            if gene_names:
                self.adata.var_names = np.array(gene_names).astype(str)
            if sample_names:
                self.adata.obs_names = np.array(sample_names).astype(str)
        else:
            print("Please input either an anndata.AnnData or both the count_matrix (count matrix of gene expression, 2D int numpy array of size (n_cells, n_genes)) and spatial_locs (spatial locations of cells (or spots) in 1D float numpy array of size (n_locations,)) to initiate SpaceFlow object.")
            exit(1)

    def plt_setting(self, fig_title_sz=30, font_sz=12, font_weight="bold", axes_title_sz=12, axes_label_sz=12, xtick_sz=10, ytick_sz=10, legend_font_sz=10):
        """
        Setting the plotting configuration
        :param fig_title_sz: fontsize of the figure title, default: 30
        :type fig_title_sz: int, optional
        :param font_sz: controls default text sizes, default: 12
        :type font_sz: int, optional
        :param font_weight: controls default text weights, default: 'bold'
        :type font_weight: str, optional
        :param axes_title_sz: fontsize of the axes title, default: 12
        :type axes_title_sz: int, optional
        :param axes_label_sz: fontsize of the x and y labels, default 12
        :type axes_label_sz: int, optional
        :param xtick_sz: fontsize of the x tick label, default 10
        :type xtick_sz: int, optional
        :param ytick_sz: fontsize of the y tick label, default 10
        :type ytick_sz: int, optional
        :param legend_font_sz: legend fontsize, default 10
        :type legend_font_sz: int, optional
        """
        plt.rc('figure', titlesize=fig_title_sz)  # fontsize of the figure title
        plt.rc('font', size=font_sz, weight=font_weight)  # controls default text sizes
        plt.rc('axes', titlesize=axes_title_sz)  # fontsize of the axes title
        plt.rc('axes', labelsize=axes_label_sz)  # fontsize of the x and y labels
        plt.rc('xtick', labelsize=xtick_sz)  # fontsize of the x tick label
        plt.rc('ytick', labelsize=ytick_sz)  # fontsize of the y tick label
        plt.rc('legend', fontsize=legend_font_sz)  # legend fontsize

    def prepare_figure(self, rsz=4., csz=4., wspace=.4, hspace=.5, left=0.125, right=0.9, bottom=0.1, top=0.9):
        """
        Prepare the figure and axes given the configuration
        :param rsz: row size of the figure in inches, default: 4.0
        :type rsz: float, optional
        :param csz: column size of the figure in inches, default: 4.0
        :type csz: float, optional
        :param wspace: the amount of width reserved for space between subplots, expressed as a fraction of the average axis width, default: 0.4
        :type wspace: float, optional
        :param hspace: the amount of height reserved for space between subplots, expressed as a fraction of the average axis width, default: 0.4
        :type hspace: float, optional
        :param left: the leftmost position of the subplots of the figure in fraction, default: 0.125
        :type left: float, optional
        :param right: the rightmost position of the subplots of the figure in fraction, default: 0.9
        :type right: float, optional
        :param bottom: the bottom position of the subplots of the figure in fraction, default: 0.1
        :type bottom: float, optional
        :param top: the top position of the subplots of the figure in fraction, default: 0.9
        :type top: float, optional
        """
        fig, axs = plt.subplots(1, 1, figsize=(csz, rsz))
        plt.subplots_adjust(wspace=wspace, hspace=hspace, left=left, right=right, bottom=bottom, top=top)
        return fig, axs

    def preprocessing_data(self, n_top_genes=None, n_neighbors=10):
        """
        Preprocessing the spatial transcriptomics data
        Generates:  `self.adata_filtered`: (n_cells, n_locations) `numpy.ndarray`
                    `self.spatial_graph`: (n_cells, n_locations) `numpy.ndarray`
        :param adata: the annData object for spatial transcriptomics data with adata.obsm['spatial'] set to be the spatial locations.
        :type adata: class:`anndata.annData`
        :param n_top_genes: the number of top highly variable genes
        :type n_top_genes: int, optional
        :param n_neighbors: the number of nearest neighbors for building spatial neighbor graph
        :type n_neighbors: int, optional
        :return: a preprocessed annData object of the spatial transcriptomics data
        :rtype: class:`anndata.annData`
        :return: a geometry-aware spatial proximity graph of the spatial spots of cells
        :rtype: class:`scipy.sparse.csr_matrix`
        """
        adata = self.adata
        if not adata:
            print("No annData object found, please run SpaceFlow.SpaceFlow(expr_data, spatial_locs) first!")
            return
        sc.pp.normalize_total(adata, target_sum=1e4)
        sc.pp.log1p(adata)
        sc.pp.highly_variable_genes(adata, n_top_genes=n_top_genes, flavor='cell_ranger', subset=True)
        sc.pp.pca(adata)
        spatial_locs = adata.obsm['spatial']
        spatial_graph = self.graph_alpha(spatial_locs, n_neighbors=n_neighbors)

        self.adata_preprocessed = adata
        self.spatial_graph = spatial_graph

    def graph_alpha(self, spatial_locs, n_neighbors=10):
        """
        Construct a geometry-aware spatial proximity graph of the spatial spots of cells by using alpha complex.
        :param adata: the annData object for spatial transcriptomics data with adata.obsm['spatial'] set to be the spatial locations.
        :type adata: class:`anndata.annData`
        :param n_neighbors: the number of nearest neighbors for building spatial neighbor graph based on Alpha Complex
        :type n_neighbors: int, optional, default: 10
        :return: a spatial neighbor graph
        :rtype: class:`scipy.sparse.csr_matrix`
        """
        import gudhi
        A_knn = kneighbors_graph(spatial_locs, n_neighbors=n_neighbors, mode='distance')
        estimated_graph_cut = A_knn.sum() / float(A_knn.count_nonzero())
        spatial_locs_list = spatial_locs.tolist()
        n_node = len(spatial_locs_list)
        alpha_complex = gudhi.AlphaComplex(points=spatial_locs_list)
        simplex_tree = alpha_complex.create_simplex_tree(max_alpha_square=estimated_graph_cut ** 2)
        skeleton = simplex_tree.get_skeleton(1)
        initial_graph = nx.Graph()
        initial_graph.add_nodes_from([i for i in range(n_node)])
        for s in skeleton:
            if len(s[0]) == 2:
                initial_graph.add_edge(s[0][0], s[0][1])

        extended_graph = nx.Graph()
        extended_graph.add_nodes_from(initial_graph)
        extended_graph.add_edges_from(initial_graph.edges)

        # Remove self edges
        for i in range(n_node):
            try:
                extended_graph.remove_edge(i, i)
            except:
                pass
        
        try:
            G = nx.to_scipy_sparse_matrix(extended_graph, format='csr')
        except AttributeError:
            G = nx.to_scipy_sparse_array(extended_graph, format='csr')
        
        return G

    def train(self, embedding_save_filepath="./embedding.tsv", spatial_regularization_strength=0.1, z_dim=50, lr=1e-3, epochs=1000, max_patience=50, min_stop=100, random_seed=42, gpu=0, regularization_acceleration=True, edge_subset_sz=1000000):
        adata_preprocessed, spatial_graph = self.adata_preprocessed, self.spatial_graph
        """
        Training the Deep GraphInfomax Model
        :param embedding_save_filepath: the default save path for the low-dimensional embeddings
        :type embedding_save_filepath: class:`str`
        :param spatial_regularization_strength: the strength for spatial regularization
        :type spatial_regularization_strength: float, optional, default: 0.1
        :param z_dim: the size of latent dimension
        :type z_dim: int, optional, default: 50
        :param lr: the learning rate for model optimization
        :type lr: float, optional, default: 1e-3
        :param epochs: the max epoch number 
        :type epochs: int, optional, default: 1000
        :param max_patience: the tolerance epoch number without training loss decrease
        :type max_patience: int, optional, default: 50
        :param min_stop: the minimum epoch number for training before any early stop
        :type min_stop: int, optional, default: 100
        :param random_seed: the random seed
        :type random_seed: int, optional, default: 42
        :param gpu: the index for gpu device that will be used for model training, if no gpu detected, cpu will be used.
        :type gpu: int, optional, default: 0
        :param regularization_acceleration: whether or not accelerate the calculation of regularization loss using edge subsetting strategy
        :type regularization_acceleration: bool, optional, default: True
        :param edge_subset_sz: the edge subset size for regularization acceleration
        :type edge_subset_sz: int, optional, default: 1000000
        :return: low dimensional embeddings for the ST data, shape: n_cells x z_dim
        :rtype: class:`numpy.ndarray`
        """
        if not adata_preprocessed:
            print("The data has not been preprocessed, please run preprocessing_data() method first!")
            return
        torch.manual_seed(random_seed)
        random.seed(random_seed)
        np.random.seed(random_seed)

        device = f"cuda:{gpu}" if torch.cuda.is_available() else 'cpu'
        model = DeepGraphInfomax(
            hidden_channels=z_dim, encoder=GraphEncoder(adata_preprocessed.shape[1], z_dim),
            summary=lambda z, *args, **kwargs: torch.sigmoid(z.mean(dim=0)),
            corruption=corruption).to(device)

        expr = adata_preprocessed.X.todense() if type(adata_preprocessed.X).__module__ != np.__name__ else adata_preprocessed.X
        expr = torch.tensor(expr).float().to(device)

        edge_list = sparse_mx_to_torch_edge_list(spatial_graph).to(device)

        model.train()
        min_loss = np.inf
        patience = 0
        optimizer = torch.optim.Adam(model.parameters(), lr=lr)
        best_params = model.state_dict()

        for epoch in range(epochs):
            train_loss = 0.0
            torch.set_grad_enabled(True)
            optimizer.zero_grad()
            z, neg_z, summary = model(expr, edge_list)
            loss = model.loss(z, neg_z, summary)

            coords = torch.tensor(adata_preprocessed.obsm['spatial']).float().to(device)
            if regularization_acceleration or adata_preprocessed.shape[0] > 5000:
                cell_random_subset_1, cell_random_subset_2 = torch.randint(0, z.shape[0], (edge_subset_sz,)).to(
                    device), torch.randint(0, z.shape[0], (edge_subset_sz,)).to(device)
                z1, z2 = torch.index_select(z, 0, cell_random_subset_1), torch.index_select(z, 0, cell_random_subset_2)
                c1, c2 = torch.index_select(coords, 0, cell_random_subset_1), torch.index_select(coords, 0,
                                                                                                 cell_random_subset_1)
                pdist = torch.nn.PairwiseDistance(p=2)

                z_dists = pdist(z1, z2)
                z_dists = z_dists / torch.max(z_dists)

                sp_dists = pdist(c1, c2)
                sp_dists = sp_dists / torch.max(sp_dists)
                n_items = z_dists.size(dim=0)
            else:
                z_dists = torch.cdist(z, z, p=2)
                z_dists = torch.div(z_dists, torch.max(z_dists)).to(device)
                sp_dists = torch.cdist(coords, coords, p=2)
                sp_dists = torch.div(sp_dists, torch.max(sp_dists)).to(device)
                n_items = z.size(dim=0) * z.size(dim=0)

            penalty_1 = torch.div(torch.sum(torch.mul(1.0 - z_dists, sp_dists)), n_items).to(device)
            loss = loss + spatial_regularization_strength * penalty_1

            loss.backward()
            optimizer.step()
            train_loss += loss.item()

            if train_loss > min_loss:
                patience += 1
            else:
                patience = 0
                min_loss = train_loss
                best_params = model.state_dict()
            if epoch % 10 == 1:
                print(f"Epoch {epoch + 1}/{epochs}, Loss: {str(train_loss)}")
            if patience > max_patience and epoch > min_stop:
                break

        model.load_state_dict(best_params)

        z, _, _ = model(expr, edge_list)
        embedding = z.cpu().detach().numpy()
        save_dir = os.path.dirname(embedding_save_filepath)
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        np.savetxt(embedding_save_filepath, embedding[:, :], delimiter="\t")
        print(f"Training complete!\nEmbedding is saved at {embedding_save_filepath}")

        self.embedding = embedding
        return embedding

    def segmentation(self, domain_label_save_filepath="./domains.tsv", n_neighbors=50, resolution=1.0):
        """
        Perform domain segmentation for ST data using Leiden clustering with low-dimensional embeddings as input
        :param domain_label_save_filepath: the default save path for the domain labels
        :type domain_label_save_filepath: class:`str`, optional, default: "./domains.tsv"
        :param n_neighbors: The size of local neighborhood (in terms of number of neighboring data
        points) used for manifold approximation. See `https://scanpy.readthedocs.io/en/stable/generated/scanpy.pp.neighbors.html` for detail
        :type n_neighbors: int, optional, default: 50
        :param resolution: A parameter value controlling the coarseness of the clustering.
        Higher values lead to more clusters. See `https://scanpy.readthedocs.io/en/stable/generated/scanpy.tl.leiden.html` for detail
        :type resolution: float, optional, default: 1.0
        """
        error_message = "No embedding found, please ensure you have run train() method before segmentation!"
        try:
            print("Performing domain segmentation")
            embedding_adata = anndata.AnnData(self.embedding)
            sc.pp.neighbors(embedding_adata, n_neighbors=n_neighbors, use_rep='X')
            sc.tl.leiden(embedding_adata, resolution=float(resolution))
            domains = embedding_adata.obs["leiden"].cat.codes

            save_dir = os.path.dirname(domain_label_save_filepath)
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            np.savetxt(domain_label_save_filepath, domains, fmt='%d', header='', footer='', comments='')
            print(f"Segmentation complete, domain labels of cells or spots saved at {domain_label_save_filepath} !")
            self.domains = domains

        except NameError:
            print(error_message)
        except AttributeError:
            print(error_message)

    def plot_segmentation(self, segmentation_figure_save_filepath="./domain_segmentation.pdf", colormap="tab20", scatter_sz=1., rsz=4., csz=4., wspace=.4, hspace=.5, left=0.125, right=0.9, bottom=0.1, top=0.9):
        """
        Plot the domain segmentation for ST data in spatial
        :param segmentation_figure_save_filepath: the default save path for the figure
        :type segmentation_figure_save_filepath: class:`str`, optional, default: "./domain_segmentation.pdf"
        :param colormap: The colormap to use. See `https://matplotlib.org/stable/tutorials/colors/colormaps.html` for full list of colormaps
        :type colormap: str, optional, default: tab20
        :param scatter_sz: The marker size in points**2
        :type scatter_sz: float, optional, default: 1.0
        :param rsz: row size of the figure in inches, default: 4.0
        :type rsz: float, optional
        :param csz: column size of the figure in inches, default: 4.0
        :type csz: float, optional
        :param wspace: the amount of width reserved for space between subplots, expressed as a fraction of the average axis width, default: 0.4
        :type wspace: float, optional
        :param hspace: the amount of height reserved for space between subplots, expressed as a fraction of the average axis width, default: 0.4
        :type hspace: float, optional
        :param left: the leftmost position of the subplots of the figure in fraction, default: 0.125
        :type left: float, optional
        :param right: the rightmost position of the subplots of the figure in fraction, default: 0.9
        :type right: float, optional
        :param bottom: the bottom position of the subplots of the figure in fraction, default: 0.1
        :type bottom: float, optional
        :param top: the top position of the subplots of the figure in fraction, default: 0.9
        :type top: float, optional
        """
        error_message = "No segmentation data found, please ensure you have run the segmentation() method."
        try:
            fig, ax = self.prepare_figure(rsz=rsz, csz=csz, wspace=wspace, hspace=hspace, left=left, right=right, bottom=bottom, top=top)

            pred_clusters = np.array(self.domains).astype(int)
            uniq_pred = np.unique(pred_clusters)
            n_cluster = len(uniq_pred)
            x, y = self.adata_preprocessed.obsm["spatial"][:, 0], self.adata_preprocessed.obsm["spatial"][:, 1]
            cmap = plt.get_cmap(colormap)
            for cid, cluster in enumerate(uniq_pred):
                color = cmap((cid * (n_cluster / (n_cluster - 1.0))) / n_cluster)
                ind = pred_clusters == cluster
                ax.scatter(x[ind], y[ind], s=scatter_sz, color=color, label=cluster, marker=".")
            ax.set_facecolor("none")
            ax.invert_yaxis()
            ax.set_title("Domain Segmentation", fontsize=14)
            box = ax.get_position()
            height_ratio = .8
            ax.set_position([box.x0, box.y0, box.width * 0.8, box.height * height_ratio])
            lgnd = ax.legend(loc='center left', fontsize=8, bbox_to_anchor=(1, 0.5), scatterpoints=1, handletextpad=0.1,
                             borderaxespad=.1, ncol=int(math.ceil(n_cluster/10)))
            for handle in lgnd.legendHandles:
                handle._sizes = [8]

            save_dir = os.path.dirname(segmentation_figure_save_filepath)
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            plt.savefig(segmentation_figure_save_filepath, dpi=300)
            print(f"Plotting complete, segmentation figure saved at {segmentation_figure_save_filepath} !")
            plt.close('all')
        except NameError:
            print(error_message)
        except AttributeError:
            print(error_message)

    def pseudo_Spatiotemporal_Map(self, pSM_values_save_filepath="./pSM_values.tsv", n_neighbors=20, resolution=1.0):
        """
        Perform pseudo-Spatiotemporal Map for ST data
        :param pSM_values_save_filepath: the default save path for the pSM values
        :type pSM_values_save_filepath: class:`str`, optional, default: "./pSM_values.tsv"
        :param n_neighbors: The size of local neighborhood (in terms of number of neighboring data
        points) used for manifold approximation. See `https://scanpy.readthedocs.io/en/stable/generated/scanpy.pp.neighbors.html` for detail
        :type n_neighbors: int, optional, default: 20
        :param resolution: A parameter value controlling the coarseness of the clustering.
        Higher values lead to more clusters. See `https://scanpy.readthedocs.io/en/stable/generated/scanpy.tl.leiden.html` for detail
        :type resolution: float, optional, default: 1.0
        """
        error_message = "No embedding found, please ensure you have run train() method before calculating pseudo-Spatiotemporal Map!"
        max_cell_for_subsampling = 5000
        try:
            print("Performing pseudo-Spatiotemporal Map")
            adata = anndata.AnnData(self.embedding)
            sc.pp.neighbors(adata, n_neighbors=n_neighbors, use_rep='X')
            sc.tl.umap(adata)
            sc.tl.leiden(adata, resolution=resolution)
            sc.tl.paga(adata)
            if adata.shape[0] < max_cell_for_subsampling:
                sub_adata_x = adata.X
            else:
                indices = np.arange(adata.shape[0])
                selected_ind = np.random.choice(indices, max_cell_for_subsampling, False)
                sub_adata_x = adata.X[selected_ind, :]
            sum_dists = distance_matrix(sub_adata_x, sub_adata_x).sum(axis=1)
            adata.uns['iroot'] = np.argmax(sum_dists)
            sc.tl.diffmap(adata)
            sc.tl.dpt(adata)
            pSM_values = adata.obs['dpt_pseudotime'].to_numpy()
            save_dir = os.path.dirname(pSM_values_save_filepath)
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            np.savetxt(pSM_values_save_filepath, pSM_values, fmt='%.5f', header='', footer='', comments='')
            print(f"pseudo-Spatiotemporal Map(pSM) calculation complete, pSM values of cells or spots saved at {pSM_values_save_filepath}!")
            self.pSM_values = pSM_values
        except NameError:
            print(error_message)
        except AttributeError:
            print(error_message)

    def plot_pSM(self, pSM_figure_save_filepath="./pseudo-Spatiotemporal-Map.pdf", colormap='roma', scatter_sz=1., rsz=4., csz=4., wspace=.4, hspace=.5, left=0.125, right=0.9, bottom=0.1, top=0.9):
        """
        Plot the domain segmentation for ST data in spatial
        :param pSM_figure_save_filepath: the default save path for the figure
        :type pSM_figure_save_filepath: class:`str`, optional, default: "./Spatiotemporal-Map.pdf"
        :param colormap: The colormap to use. See `https://www.fabiocrameri.ch/colourmaps-userguide/` for name list of colormaps
        :type colormap: str, optional, default: roma
        :param scatter_sz: The marker size in points**2
        :type scatter_sz: float, optional, default: 1.0
        :param rsz: row size of the figure in inches, default: 4.0
        :type rsz: float, optional
        :param csz: column size of the figure in inches, default: 4.0
        :type csz: float, optional
        :param wspace: the amount of width reserved for space between subplots, expressed as a fraction of the average axis width, default: 0.4
        :type wspace: float, optional
        :param hspace: the amount of height reserved for space between subplots, expressed as a fraction of the average axis width, default: 0.4
        :type hspace: float, optional
        :param left: the leftmost position of the subplots of the figure in fraction, default: 0.125
        :type left: float, optional
        :param right: the rightmost position of the subplots of the figure in fraction, default: 0.9
        :type right: float, optional
        :param bottom: the bottom position of the subplots of the figure in fraction, default: 0.1
        :type bottom: float, optional
        :param top: the top position of the subplots of the figure in fraction, default: 0.9
        :type top: float, optional
        """
        error_message = "No pseudo Spatiotemporal Map data found, please ensure you have run the pseudo_Spatiotemporal_Map() method."
        try:
            fig, ax = self.prepare_figure(rsz=rsz, csz=csz, wspace=wspace, hspace=hspace, left=left, right=right, bottom=bottom, top=top)
            x, y = self.adata_preprocessed.obsm["spatial"][:, 0], self.adata_preprocessed.obsm["spatial"][:, 1]
            st = ax.scatter(x, y, s=scatter_sz, c=self.pSM_values, cmap=f"cmc.{colormap}", marker=".")
            ax.invert_yaxis()
            clb = fig.colorbar(st)
            clb.ax.set_ylabel("pseudotime", labelpad=10, rotation=270, fontsize=10, weight='bold')
            ax.set_title("pseudo-Spatiotemporal Map", fontsize=14)
            ax.set_facecolor("none")

            save_dir = os.path.dirname(pSM_figure_save_filepath)
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            plt.savefig(pSM_figure_save_filepath, dpi=300)
            print(f"Plotting complete, pseudo-Spatiotemporal Map figure saved at {pSM_figure_save_filepath} !")
            plt.close('all')
        except NameError:
            print(error_message)
        except AttributeError:
            print(error_message)

class GraphEncoder(nn.Module):
    def __init__(self, in_channels, hidden_channels):
        super(GraphEncoder, self).__init__()
        self.conv = GCNConv(in_channels, hidden_channels, cached=False)
        self.prelu = nn.PReLU(hidden_channels)
        self.conv2 = GCNConv(hidden_channels, hidden_channels, cached=False)
        self.prelu2 = nn.PReLU(hidden_channels)

    def forward(self, x, edge_index):
        x = self.conv(x, edge_index)
        x = self.prelu(x)
        x = self.conv2(x, edge_index)
        x = self.prelu2(x)
        return x
