import itertools
import logging
from typing import Any, List, Optional

import torch
from torch import Tensor
from torch_geometric.utils import degree
from torch_geometric.utils import remove_self_loops
from torch_scatter import scatter
from yacs.config import CfgNode



def virtual_node_mask(cfg, batch_idx):
    # get a mask which applies virtual node removal 
    keep = torch.zeros(batch_idx.shape[0], dtype=torch.bool, device=batch_idx.device)

    # Remove virtual node from eigvec preds 
    idx = torch.concat([
        torch.where(batch_idx == i)[0][:-1]
        for i in range(batch_idx.max().item() + 1)
    ])
    keep[idx] = True 

    return keep 

def eigen_mask(cfg, batch_idx):
    # Input: 
    # - batch_idx tensor(N): graph index for every node 
    # Output: 
    # - keep, a binary mask 
    # filters small and large graphs 
    # NOTE: THIS FUNCTION ASSUMES YOU HAVE NOT ALREADY REMOVED VIRTUAL NODE FROM BATCH_IDX
        
    graph_sizes = torch.bincount(batch_idx) 

    keep = torch.ones(batch_idx.shape[0], dtype=torch.bool, device=batch_idx.device)
    
    # want to apply large-graph filtering to EFFECTIVE graph sizes (incl. virtual node), as this is just input dim to concat MLP
    if cfg.posenc_LapPE.MLP_style == "concat":
        too_big_graphs = graph_sizes > cfg.posenc_LapPE.concat_max_nodes
        keep = keep & ~too_big_graphs[batch_idx]

    if cfg.virtual_node: # want to apply small-graph filtering to ACTUAL graph sizes (to avoid nan eigval losses)
        graph_sizes = graph_sizes - 1
        
    too_small_graphs = graph_sizes < cfg.posenc_LapPE.eigen.max_freqs
    keep = keep & ~too_small_graphs[batch_idx]

    return keep



def get_device(device: str, default_device: str) -> str:
    if device == "default":
        device = default_device
    return device


def negate_edge_index(edge_index, batch=None):
    """Negate batched sparse adjacency matrices given by edge indices.

    Returns batched sparse adjacency matrices with exactly those edges that
    are not in the input `edge_index` while ignoring self-loops.

    Implementation inspired by `torch_geometric.utils.to_dense_adj`

    Args:
        edge_index: The edge indices.
        batch: Batch vector, which assigns each node to a specific example.

    Returns:
        Complementary edge index.
    """

    if batch is None:
        batch = edge_index.new_zeros(edge_index.max().item() + 1)

    batch_size = batch.max().item() + 1
    one = batch.new_ones(batch.size(0))
    num_nodes = scatter(one, batch,
                        dim=0, dim_size=batch_size, reduce='add')
    cum_nodes = torch.cat([batch.new_zeros(1), num_nodes.cumsum(dim=0)])

    idx0 = batch[edge_index[0]]
    idx1 = edge_index[0] - cum_nodes[batch][edge_index[0]]
    idx2 = edge_index[1] - cum_nodes[batch][edge_index[1]]

    negative_index_list = []
    for i in range(batch_size):
        n = num_nodes[i].item()
        size = [n, n]
        adj = torch.ones(size, dtype=torch.short,
                         device=edge_index.device)

        # Remove existing edges from the full N x N adjacency matrix
        flattened_size = n * n
        adj = adj.view([flattened_size])
        _idx1 = idx1[idx0 == i]
        _idx2 = idx2[idx0 == i]
        idx = _idx1 * n + _idx2
        zero = torch.zeros(_idx1.numel(), dtype=torch.short,
                           device=edge_index.device)
        scatter(zero, idx, dim=0, out=adj, reduce='mul')

        # Convert to edge index format
        adj = adj.view(size)
        _edge_index = adj.nonzero(as_tuple=False).t().contiguous()
        _edge_index, _ = remove_self_loops(_edge_index)
        negative_index_list.append(_edge_index + cum_nodes[i])

    edge_index_negative = torch.cat(negative_index_list, dim=1).contiguous()
    return edge_index_negative


def flatten_dict(metrics):
    """Flatten a list of train/val/test metrics into one dict to send to wandb.

    Args:
        metrics: List of Dicts with metrics

    Returns:
        A flat dictionary with names prefixed with "train/" , "val/" , "test/"
    """
    prefixes = ['train', 'val', 'test']
    result = {}
    for i in range(len(metrics)):
        # Take the latest metrics.
        stats = metrics[i][-1]
        result.update({f"{prefixes[i]}/{k}": v for k, v in stats.items()})
    return result


def cfg_to_dict(cfg_node, key_list=[]):
    """Convert a config node to dictionary.

    Yacs doesn't have a default function to convert the cfg object to plain
    python dict. The following function was taken from
    https://github.com/rbgirshick/yacs/issues/19
    """
    _VALID_TYPES = {tuple, list, str, int, float, bool}

    if not isinstance(cfg_node, CfgNode):
        if type(cfg_node) not in _VALID_TYPES:
            logging.warning(f"Key {'.'.join(key_list)} with "
                            f"value {type(cfg_node)} is not "
                            f"a valid type; valid types: {_VALID_TYPES}")
        return cfg_node
    else:
        cfg_dict = dict(cfg_node)
        for k, v in cfg_dict.items():
            cfg_dict[k] = cfg_to_dict(v, key_list + [k])
        return cfg_dict


def make_wandb_name(cfg):
    # Format dataset name.
    dataset_name = cfg.dataset.format
    if dataset_name.startswith('OGB'):
        dataset_name = dataset_name[3:]
    if dataset_name.startswith('PyG-'):
        dataset_name = dataset_name[4:]
    if dataset_name in ['GNNBenchmarkDataset', 'TUDataset']:
        # Shorten some verbose dataset naming schemes.
        dataset_name = ""
    if cfg.dataset.name != 'none':
        dataset_name += "-" if dataset_name != "" else ""
        if cfg.dataset.name == 'LocalDegreeProfile':
            dataset_name += 'LDP'
        else:
            dataset_name += cfg.dataset.name
    # Format model name.
    model_name = cfg.model.type
    if cfg.model.type in ['gnn', 'custom_gnn']:
        model_name += f".{cfg.gnn.layer_type}"
    elif cfg.model.type == 'GPSModel':
        model_name = f"GPS.{cfg.gt.layer_type}"
    model_name += f".{cfg.name_tag}" if cfg.name_tag else ""
    # Compose wandb run name.
    name = f"{dataset_name}.{model_name}.r{cfg.run_id}"
    return name


def unbatch(src: Tensor, batch: Tensor, dim: int = 0) -> List[Tensor]:
    """
    COPIED FROM NOT YET RELEASED VERSION OF PYG (as of PyG v2.0.4).

    Splits :obj:`src` according to a :obj:`batch` vector along dimension
    :obj:`dim`.

    Args:
        src (Tensor): The source tensor.
        batch (LongTensor): The batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
            entry in :obj:`src` to a specific example. Must be ordered.
        dim (int, optional): The dimension along which to split the :obj:`src`
            tensor. (default: :obj:`0`)
    :rtype: :class:`List[Tensor]`
    """
    sizes = degree(batch, dtype=torch.long).tolist()
    return src.split(sizes, dim)


def unbatch_edge_index(edge_index: Tensor, batch: Tensor) -> List[Tensor]:
    """
    COPIED FROM NOT YET RELEASED VERSION OF PYG (as of PyG v2.0.4).

    Splits the :obj:`edge_index` according to a :obj:`batch` vector.

    Args:
        edge_index (Tensor): The edge_index tensor. Must be ordered.
        batch (LongTensor): The batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
            node to a specific example. Must be ordered.
    :rtype: :class:`List[Tensor]`
    """
    deg = degree(batch, dtype=torch.int64)
    ptr = torch.cat([deg.new_zeros(1), deg.cumsum(dim=0)[:-1]], dim=0)

    edge_batch = batch[edge_index[0]]
    edge_index = edge_index - ptr[edge_batch]
    sizes = degree(edge_batch, dtype=torch.int64).cpu().tolist()
    return edge_index.split(sizes, dim=1)


def grouper(iterable, n: int, *, fillvalue: Optional[Any] = None):
    """Group an interable into chunks of size n.

    Modified from https://docs.python.org/3/library/itertools.html#itertools-recipes

    Args:
        iterable: Iterable to be chunked.
        n: Chunk size.
        fillvalue: Value to fill when the length of the iterable can not be
            devided into exact n sized chunkcs.

    Examples:
        >>> my_list = [1, 2, 3, 4, 5, 6, 7]
        >>> for i in grouper(my_list, 3, fillvalue="x"):
        ...     print(i)
        (1, 2, 3)
        (4, 5, 6)
        (7, 'x', 'x')

    """
    return itertools.zip_longest(*([iter(iterable)] * n), fillvalue=fillvalue)
