import copy
import inspect
import os.path as osp
import random

# import os,sys,pickle,time,random,glob
import time
from typing import Optional, Tuple, List

import numpy as np
import torch
import torch.nn.functional as F
import torch.utils.data
from sklearn.base import BaseEstimator
from torch_geometric.data import Data
from torch_geometric.nn import GATConv
from torch_sparse import SparseTensor, cat
import torch_geometric.data.data
from numpy import ndarray
from pandas.core.arrays.categorical import Categorical
from scipy.sparse._csr import csr_matrix

# from .utils_func import get_X_y_from_ann


# import pandas as pd


# Seed
seed = 42
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
random.seed(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True


def scipysparse2torchsparse(x: csr_matrix) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Input: scipy csr_matrix
    Returns: torch tensor in experimental sparse format

    REF: Code adatped from [PyTorch discussion forum](https://discuss.pytorch.org/t/better-way-to-forward-sparse-matrix/21915>)
    """
    samples = x.shape[0]
    features = x.shape[1]
    values = x.data
    coo_data = x.tocoo()
    indices = torch.LongTensor(
        [coo_data.row, coo_data.col]
    )  # OR transpose list of index tuples
    t = torch.sparse.FloatTensor(
        indices, torch.from_numpy(values).float(), [samples, features]
    )
    return indices, t


class ClusterData(torch.utils.data.Dataset):
    r"""Clusters/partitions a graph data object into multiple subgraphs, as
    motivated by the `"Cluster-GCN: An Efficient Algorithm for Training Deep
    and Large Graph Convolutional Networks"
    <https://arxiv.org/abs/1905.07953>`_ paper.

    Args:
        data (torch_geometric.data.Data): The graph data object.
        num_parts (int): The number of partitions.
        recursive (bool, optional): If set to :obj:`True`, will use multilevel
            recursive bisection instead of multilevel k-way partitioning.
            (default: :obj:`False`)
        save_dir (string, optional): If set, will save the partitioned data to
            the :obj:`save_dir` directory for faster re-use.
    """

    def __init__(self, data:     torch_geometric.data.data.Data, num_parts: int, recursive: bool=False, save_dir: None=None) -> None:
        assert data.edge_index is not None

        self.num_parts = num_parts
        self.recursive = recursive
        self.save_dir = save_dir

        self.process(data)

    def process(self, data:     torch_geometric.data.data.Data) -> None:
        recursive = "_recursive" if self.recursive else ""
        filename = f"part_data_{self.num_parts}{recursive}.pt"

        path = osp.join(self.save_dir or "", filename)
        if self.save_dir is not None and osp.exists(path):
            data, partptr, perm = torch.load(path)
        else:
            data = copy.copy(data)
            num_nodes = data.num_nodes

            (row, col), edge_attr = data.edge_index, data.edge_attr
            adj = SparseTensor(row=row, col=col, value=edge_attr)
            adj, partptr, perm = adj.partition(self.num_parts, self.recursive)

            for key, item in data:
                if item.size(0) == num_nodes:
                    data[key] = item[perm]

            data.edge_index = None
            data.edge_attr = None
            data.adj = adj

            if self.save_dir is not None:
                torch.save((data, partptr, perm), path)

        self.data = data
        self.perm = perm
        self.partptr = partptr

    def __len__(self) -> int:
        return self.partptr.numel() - 1

    def __getitem__(self, idx):
        start = int(self.partptr[idx])
        length = int(self.partptr[idx + 1]) - start

        data = copy.copy(self.data)
        num_nodes = data.num_nodes

        for key, item in data:
            if item.size(0) == num_nodes:
                data[key] = item.narrow(0, start, length)

        data.adj = data.adj.narrow(1, start, length)

        row, col, value = data.adj.coo()
        data.adj = None
        data.edge_index = torch.stack([row, col], dim=0)
        data.edge_attr = value

        return data

    def __repr__(self):
        return f"{self.__class__.__name__}({self.data}, " f"num_parts={self.num_parts})"


class ClusterLoader(torch.utils.data.DataLoader):
    r"""The data loader scheme from the `"Cluster-GCN: An Efficient Algorithm
    for Training Deep and Large Graph Convolutional Networks"
    <https://arxiv.org/abs/1905.07953>`_ paper which merges partioned subgraphs
    and their between-cluster links from a large-scale graph data object to
    form a mini-batch.

    Args:
        cluster_data (torch_geometric.data.ClusterData): The already
            partioned data object.
        batch_size (int, optional): How many samples per batch to load.
            (default: :obj:`1`)
        shuffle (bool, optional): If set to :obj:`True`, the data will be
            reshuffled at every epoch. (default: :obj:`False`)
    """

    def __init__(self, cluster_data: ClusterData, batch_size: int=1, shuffle: bool=False, **kwargs) -> None:
        class HelperDataset(torch.utils.data.Dataset):
            def __len__(self):
                return len(cluster_data)

            def __getitem__(self, idx):
                start = int(cluster_data.partptr[idx])
                length = int(cluster_data.partptr[idx + 1]) - start

                data = copy.copy(cluster_data.data)
                num_nodes = data.num_nodes
                for key, item in data:
                    if item.size(0) == num_nodes:
                        data[key] = item.narrow(0, start, length)

                return data, idx

        def collate(batch):
            data_list = [data[0] for data in batch]
            parts: List[int] = [data[1] for data in batch]
            partptr = cluster_data.partptr

            adj = cat([data.adj for data in data_list], dim=0)

            adj = adj.t()
            adjs = []
            for part in parts:
                start = partptr[part]
                length = partptr[part + 1] - start
                adjs.append(adj.narrow(0, start, length))
            adj = cat(adjs, dim=0).t()
            row, col, value = adj.coo()

            data = cluster_data.data.__class__()
            data.num_nodes = adj.size(0)
            data.edge_index = torch.stack([row, col], dim=0)
            data.edge_attr = value

            ref = data_list[0]
            keys = list(ref.keys())
            keys.remove("adj")

            for key in keys:
                if ref[key].size(0) != ref.adj.size(0):
                    data[key] = ref[key]
                else:
                    data[key] = torch.cat(
                        [d[key] for d in data_list], dim=ref.__cat_dim__(key, ref[key])
                    )

            return data

        super(ClusterLoader, self).__init__(
            HelperDataset(), batch_size, shuffle, collate_fn=collate, **kwargs
        )


## model
class GAT(torch.nn.Module):  # torch.nn.Module is the base class for all NN modules.
    def __init__(self, n_nodes: int, nFeatures: int, nHiddenUnits: int, nHeads: int, alpha: float, dropout: float) -> None:
        super(GAT, self).__init__()
        # 定义实例属性
        self.n_nodes = n_nodes
        self.nFeatures = nFeatures
        self.nHiddenUnits = nHiddenUnits
        self.nHeads = nHeads
        self.alpha = alpha
        self.dropout = dropout

        self.gat1 = GATConv(
            self.nFeatures,
            out_channels=self.nHiddenUnits,  # 映射到8维
            heads=self.nHeads,
            concat=True,
            negative_slope=self.alpha,
            dropout=self.dropout,
            bias=True,
        )
        self.gat2 = GATConv(
            self.nHiddenUnits * self.nHeads,
            self.n_nodes,  # 最后一层映射到k维度（k=n_class）
            heads=self.nHeads,
            concat=False,
            negative_slope=self.alpha,
            dropout=self.dropout,
            bias=True,
        )

    def forward(self, data:     torch_geometric.data.data.Data) ->     torch.Tensor:
        x, edge_index = data.x, data.edge_index
        x = self.gat1(x, edge_index)  # 第一层输出经过ELU非线性函数
        x = F.elu(x)
        x = self.gat2(x, edge_index)  # 第二层输出经过softmax变成[0, 1]后直接用于分类
        # return F.log_softmax(x, dim=1)
        return x


## sklearn classifier
class GATclassifier(BaseEstimator):
    """A pytorch regressor"""

    def __init__(
        self,
        n_nodes: int=2,
        nFeatures: Optional[int]=None,
        nHiddenUnits: int=8,
        nHeads: int=8,
        alpha: float=0.2,
        dropout: float=0.4,
        clip: None=None,
        rs: int=random.randint(1, 1000000),
        LR: float=0.001,
        WeightDecay: float=5e-4,
        BatchSize: int=256,
        NumParts: int=200,
        nEpochs: int=100,
        fastmode: bool=True,
        verbose: int=0,
        device: str="cpu",
    ) -> None:
        """
        Called when initializing the regressor
        """
        self._history = None
        self._model = None

        args, _, _, values = inspect.getargvalues(inspect.currentframe())
        values.pop("self")

        for arg, val in values.copy().items():
            setattr(self, arg, val)

    def _build_model(self) -> None:

        self._model = GAT(
            self.n_nodes,
            self.nFeatures,
            self.nHiddenUnits,
            self.nHeads,
            self.alpha,
            self.dropout,
        )

    def _train_model(self, X: ndarray, y: Categorical, adj: csr_matrix) -> None:
        # X, y, adj = get_X_y_from_ann(adata_train_final, return_adj=True, n_pc=2, n_neigh=10)

        node_features = torch.from_numpy(X).float()
        labels = torch.LongTensor(y)
        edge_index, _ = scipysparse2torchsparse(adj)

        d = Data(x=node_features, edge_index=edge_index, y=labels)
        cd = ClusterData(d, num_parts=self.NumParts)

        cl = ClusterLoader(cd, batch_size=self.BatchSize, shuffle=True)

        optimizer = torch.optim.Adagrad(
            self._model.parameters(), lr=self.LR, weight_decay=self.WeightDecay
        )

        # Random Seed
        random.seed(self.rs)
        np.random.seed(self.rs)
        torch.manual_seed(self.rs)

        t_total = time.time()
        loss_values = []
        bad_counter = 0
        best = self.nEpochs + 1
        best_epoch = 0

        for epoch in range(self.nEpochs):

            t = time.time()
            epoch_loss = []
            epoch_acc = []
            epoch_acc_val = []
            epoch_loss_val = []

            self._model.train()  # It sets the mode to train

            for batch in cl:  # cl: clusterLoader
                batch = batch.to(self.device)  # move the data to CPU/GPU
                optimizer.zero_grad()  # weight init
                x_output = self._model(batch)  # ncell*2; log_softmax
                output = F.log_softmax(x_output, dim=1)

                loss = F.nll_loss(
                    input=output, target=batch.y
                )  # compute negative log likelihood loss
                # input: ncell*nclass;
                # target: ncell*1, 0 =< value <= nclass-1
                loss.backward()
                if self.clip is not None:
                    torch.nn.utils.clip_grad_norm_(self._model.parameters(), self.clip)
                optimizer.step()
                epoch_loss.append(loss.item())
                # epoch_acc.append(accuracy(output, batch.y).item())

            if not self.fastmode:
                d_val = Data(x=features_val, edge_index=edge_index_val, y=labels_val)
                d_val = d_val.to(self.device)
                self._model.eval()
                x_output = self._model(d_val)
                output = F.log_softmax(x_output, dim=1)

                loss_val = F.nll_loss(output, d_val.y)
                # acc_val = accuracy(output,d_val.y).item() # tensor.item() returns the value of this tensor as a standard Python number.
                if (self.verbose > 0) & ((epoch + 1) % 50 == 0):
                    print(
                        "Epoch {}\t<loss>={:.4f}\tloss_val={:.4f}\tin {:.2f}-s".format(
                            epoch + 1,
                            np.mean(epoch_loss),
                            loss_val.item(),
                            time.time() - t,
                        )
                    )
                loss_values.append(loss_val.item())
            else:
                if (self.verbose > 0) & ((epoch + 1) % 50 == 0):
                    print(
                        "Epoch {}\t<loss>={:.4f}\tin {:.2f}-s".format(
                            epoch + 1, np.mean(epoch_loss), time.time() - t
                        )
                    )
                loss_values.append(np.mean(epoch_loss))

    def fit(self, X: ndarray, y: Categorical, adj: csr_matrix) -> "GATclassifier":
        """
        Trains the pytorch regressor.
        """

        self._build_model()
        self._train_model(X, y, adj)

        return self

    def predict(self, X: ndarray, y: Categorical, adj: csr_matrix) ->     torch.Tensor:
        """
        Makes a prediction using the trained pytorch model
        """

        # X, y, adj = get_X_y_from_ann(adata_test, return_adj=True, n_pc=2, n_neigh=10)

        node_features = torch.from_numpy(X).float()
        labels = torch.LongTensor(y)
        edge_index, _ = scipysparse2torchsparse(adj)

        d_test = Data(x=node_features, edge_index=edge_index, y=labels)

        self._model.eval()  # define the evaluation mode
        d_test = d_test.to(self.device)
        x_output = self._model(d_test)
        output = F.log_softmax(x_output, dim=1)
        preds = output.max(1)[1].type_as(labels)

        return preds

    def predict_proba(self, X: ndarray, y: Categorical, adj: csr_matrix) -> ndarray:

        # X, y, adj = get_X_y_from_ann(adata_test, return_adj=True, n_pc=2, n_neigh=10)

        node_features = torch.from_numpy(X).float()
        labels = torch.LongTensor(y)
        edge_index, _ = scipysparse2torchsparse(adj)

        d_test = Data(x=node_features, edge_index=edge_index, y=labels)

        self._model.eval()  # define the evaluation mode
        d_test = d_test.to(self.device)
        x_output = self._model(d_test)
        output = F.log_softmax(x_output, dim=1)

        probs = torch.exp(output)  # return softmax (output is logsoftmax)
        y_prob = (
            probs.detach().cpu().numpy()
        )  # detach() here prune away the gradients bond with the probs tensor

        return y_prob

    def score(self, X, y, sample_weight=None):
        """
        Scores the data using the trained pytorch model. Under current implementation
        returns negative mae.
        """
        y_pred = self.predict(X, y)
        return F.nll_loss(y_pred, y)