from typing import Any, Dict, List, Mapping

import dgl
import numpy as np
import torch


def convert_to_single_node_emb(x, offset: int = 128):
    feature_num = x.shape[1] if len(x.shape) > 1 else 1
    feature_offset = 1 + np.arange(0, feature_num * offset, offset, dtype=np.int64)
    x = x + feature_offset
    return x


def convert_to_single_edge_emb(x, offset: int = 8):
    feature_num = x.shape[1] if len(x.shape) > 1 else 1
    feature_offset = 1 + np.arange(0, feature_num * offset, offset, dtype=np.int64)
    x = x + feature_offset
    return x


def preprocess_item(item, keep_features=True):
    if keep_features and "edge_attr" in item.keys():  # edge_attr
        edge_attr = np.asarray(item["edge_attr"], dtype=np.int64)
    else:
        edge_attr = np.ones((len(item["edge_index"][0]), 1), dtype=np.int64)  # same embedding for all

    if keep_features and "node_feat" in item.keys():  # input_nodes
        node_feature = np.asarray(item["node_feat"], dtype=np.int64)
    else:
        node_feature = np.ones((item["num_nodes"], 1), dtype=np.int64)  # same embedding for all

    edge_index = np.asarray(item["edge_index"], dtype=np.int64)

    input_nodes = convert_to_single_node_emb(node_feature)
    num_nodes = item["num_nodes"]

    if len(edge_attr.shape) == 1:
        edge_attr = edge_attr[:, None]
    attn_edge_type = np.zeros([num_nodes, num_nodes, edge_attr.shape[-1]], dtype=np.int64)
    attn_edge_type[edge_index[0], edge_index[1]] = convert_to_single_edge_emb(edge_attr)

    # convert to dgl graph for computing shortest path distance and svd encodings
    g = dgl.graph((edge_index[0], edge_index[1]))
    shortest_path_result = dgl.shortest_dist(g)
    shortest_path_result = torch.where(shortest_path_result == -1, 510, shortest_path_result)
    svd_pe = dgl.svd_pe(g, k=8, padding=True, random_flip=True)

    # combine
    item["input_nodes"] = input_nodes
    item["attn_edge_type"] = attn_edge_type
    item["spatial_pos"] = shortest_path_result
    item["svd_pe"] = svd_pe
    if "labels" not in item:
        item["labels"] = item["y"]

    return item


class EGTDataCollator:
    def __init__(self, on_the_fly_processing=False):
        self.on_the_fly_processing = on_the_fly_processing

    def __call__(self, features: List[dict]) -> Dict[str, Any]:
        if self.on_the_fly_processing:
            features = [preprocess_item(i) for i in features]

        if not isinstance(features[0], Mapping):
            features = [vars(f) for f in features]
        batch = {}

        max_node_num = max(len(i["input_nodes"]) for i in features)
        node_feat_size = len(features[0]["input_nodes"][0])
        edge_feat_size = len(features[0]["attn_edge_type"][0][0])
        svd_pe_size = len(features[0]["svd_pe"][0]) // 2
        batch_size = len(features)

        batch["featm"] = torch.zeros(batch_size, max_node_num, max_node_num, edge_feat_size, dtype=torch.long)
        batch["dm"] = torch.zeros(batch_size, max_node_num, max_node_num, dtype=torch.long)
        batch["node_feat"] = torch.zeros(batch_size, max_node_num, node_feat_size, dtype=torch.long)
        batch["svd_pe"] = torch.zeros(batch_size, max_node_num, svd_pe_size * 2, dtype=torch.float)
        batch["attn_mask"] = torch.zeros(batch_size, max_node_num, dtype=torch.long)

        for ix, f in enumerate(features):
            for k in ["attn_edge_type", "spatial_pos", "input_nodes", "svd_pe"]:
                f[k] = torch.tensor(f[k])

            batch["featm"][ix, : f["attn_edge_type"].shape[0], : f["attn_edge_type"].shape[1], :] = f["attn_edge_type"]
            batch["dm"][ix, : f["spatial_pos"].shape[0], : f["spatial_pos"].shape[1]] = f["spatial_pos"]
            batch["node_feat"][ix, : f["input_nodes"].shape[0], :] = f["input_nodes"]
            batch["svd_pe"][ix, : f["svd_pe"].shape[0], :] = f["svd_pe"]
            batch["attn_mask"][ix, : f["svd_pe"].shape[0]] = 1

        sample = features[0]["labels"]
        if len(sample) == 1:  # one task
            if isinstance(sample[0], float):  # regression
                batch["labels"] = torch.from_numpy(np.concatenate([i["labels"] for i in features]))
            else:  # binary classification
                batch["labels"] = torch.from_numpy(np.concatenate([i["labels"] for i in features]))
        else:  # multi task classification, left to float to keep the NaNs
            batch["labels"] = torch.from_numpy(np.stack([i["labels"] for i in features], dim=0))

        return batch