# General

import numpy as np

import pandas as pd

#from typing import List, Optional, Tuple, NamedTuple, Union, Callable

# Pytorch

import torch

import torch.nn as nn

from torch_geometric.data import Data

from torch import Tensor

import project_config

def preprocess_graph(args):

    # Read in nodes & edges

    nodes = pd.read_csv(project_config.KG_DIR / args.node_map, sep="\t")

    edges = pd.read_csv(project_config.KG_DIR / args.edgelist, sep="\t")

    # Initialize edge index

    edge_index = torch.LongTensor(edges[['x_idx', 'y_idx']].values.T).contiguous() 

    edge_attr = edges['full_relation']

    # Convert edge attributes to idx

    edge_attr_list = [

                      'effect/phenotype;phenotype_protein;gene/protein',

                      'gene/protein;phenotype_protein;effect/phenotype',

                      'disease;disease_phenotype_negative;effect/phenotype',

                      'effect/phenotype;disease_phenotype_negative;disease',

                      'disease;disease_phenotype_positive;effect/phenotype',

                      'effect/phenotype;disease_phenotype_positive;disease',

                      'gene/protein;protein_pathway;pathway',

                      'pathway;protein_pathway;gene/protein',

                      'disease;disease_protein;gene/protein',

                      'gene/protein;disease_protein;disease',

                      'gene/protein;protein_molfunc;molecular_function',

                      'molecular_function;protein_molfunc;gene/protein',

                      'gene/protein;protein_cellcomp;cellular_component',

                      'cellular_component;protein_cellcomp;gene/protein',

                      'gene/protein;protein_bioprocess;biological_process',

                      'biological_process;protein_bioprocess;gene/protein',

                      'biological_process;bioprocess_bioprocess;biological_process',

                      'biological_process;bioprocess_bioprocess_rev;biological_process',

                      'molecular_function;molfunc_molfunc;molecular_function',

                      'molecular_function;molfunc_molfunc_rev;molecular_function',

                      'cellular_component;cellcomp_cellcomp;cellular_component',

                      'cellular_component;cellcomp_cellcomp_rev;cellular_component',

                      'effect/phenotype;phenotype_phenotype;effect/phenotype',

                      'effect/phenotype;phenotype_phenotype_rev;effect/phenotype',

                      'gene/protein;protein_protein;gene/protein',

                      'gene/protein;protein_protein_rev;gene/protein',

                      'disease;disease_disease;disease',

                      'disease;disease_disease_rev;disease',

                      'pathway;pathway_pathway;pathway',

                      'pathway;pathway_pathway_rev;pathway'

                     ]

    edge_attr_to_idx_dict = {attr:i for i, attr in enumerate(edge_attr_list)}

    edge_attr = torch.LongTensor(np.vectorize(edge_attr_to_idx_dict.get)(edge_attr.values))

    # Get train/val/test masks

    mask = edges["mask"].values

    train_mask = torch.BoolTensor(mask == "train")

    val_mask = torch.BoolTensor(mask == "val")

    test_mask = torch.BoolTensor(mask == "test")

    # Create data object

    data = Data(edge_index = edge_index, edge_attr = edge_attr, train_mask = train_mask, val_mask = val_mask, test_mask = test_mask)

    return data, edge_attr_to_idx_dict, nodes