# docs/tutorials/model_training/ipu_training_demo.py

# General imports
import os
from os.path import dirname, abspath
import yaml
from copy import deepcopy
from omegaconf import DictConfig
import timeit
from loguru import logger
from pytorch_lightning.utilities.model_summary import ModelSummary

# Current project imports
import graphium
from graphium.config._loader import (
    load_datamodule,
    load_metrics,
    load_architecture,
    load_predictor,
    load_trainer,
)
from graphium.utils.safe_run import SafeRun

# from torch_geometric.nn.aggr import Aggregation
# Aggregation.set_validate_args(False)

# WandB
import wandb

# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(graphium.__file__)))
# CONFIG_FILE= "expts/configs/config_ipu_allsizes.yaml"
CONFIG_FILE = "docs/tutorials/model_training/config_ipu_tutorials.yaml"
os.chdir(MAIN_DIR)

with open(os.path.join(MAIN_DIR, CONFIG_FILE), "r") as f:
    cfg = yaml.safe_load(f)

# Load and initialize the dataset
datamodule = load_datamodule(cfg)
datamodule.prepare_data()

# Initialize the network
model_class, model_kwargs = load_architecture(
    cfg,
    in_dims=datamodule.in_dims,
)

metrics = load_metrics(cfg)
logger.info(metrics)

predictor = load_predictor(cfg, model_class, model_kwargs, metrics)
predictor.set_max_nodes_edges_per_graph(datamodule, stages=["train", "val"])

logger.info(predictor.model)
logger.info(ModelSummary(predictor, max_depth=4))

trainer = load_trainer(cfg, "tutorial-run")

# Run the model training
with SafeRun(name="TRAINING", raise_error=cfg["constants"]["raise_train_error"], verbose=True):
    trainer.fit(model=predictor, datamodule=datamodule)

# Exit WandB
wandb.finish()


# expts/dataset_benchmark.py

import os
from os.path import dirname, abspath
import yaml
from omegaconf import DictConfig
from datetime import datetime
from copy import deepcopy

import graphium
from graphium.config._loader import load_datamodule
import time
from typing import Optional, List, Sequence
import wandb
import statistics
import tqdm
import torch
import numpy as np

# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(graphium.__file__)))
os.chdir(MAIN_DIR)
CONFIG_FILE = "expts/configs/config_pcqmv2_mpnn.yaml"
# CONFIG_FILE = "expts/configs/config_ipu_qm9.yaml"


def benchmark(fn, *args, message="", log2wandb=False, **kwargs):
    start = time.time()
    value = fn(*args, **kwargs)
    duration = time.time() - start
    print(f"{message} {duration:.3f} secs")
    if log2wandb:
        wandb.log({message: duration})
    return value


def benchmark_dataloader(dataloader, name, n_epochs=5, log2wandb=False):
    print(f"length of {name} dataloader: {len(dataloader)}")
    epoch_times = [0] * n_epochs
    tputs = [0] * n_epochs
    n_batches = [0] * n_epochs
    n_graphs = [0] * n_epochs
    n_nodes = [0] * n_epochs
    n_edges = [0] * n_epochs
    for i in range(n_epochs):
        start = time.time()
        for data in tqdm.tqdm(dataloader):
            n_batches[i] += 1
            n_graphs[i] += torch.sum(torch.max(data["features"]["batch"], dim=-1).values).item()
            n_nodes[i] += np.prod(data["features"]["batch"].shape)
            n_edges[i] += np.prod(data["features"]["edge_weight"].shape)
        epoch_times[i] = time.time() - start
        tputs[i] = n_graphs[i] / epoch_times[i]

    average_tput = statistics.mean(tputs)
    print(f"{name} dataloader average tput {average_tput}")
    print(f"{name} dataloader tputs per epoch {tputs}")
    print(f"{name} dataloader epoch times {epoch_times}")
    print(f"{name} dataloader total batches per epoch {n_batches}")
    print(f"{name} dataloader total graphs per epoch {n_graphs}")
    print(f"{name} dataloader total nodes per epoch {n_nodes}")
    print(f"{name} dataloader total edges per epoch {n_edges}")

    if log2wandb:
        wandb.log({"average tput": average_tput})

        for i in range(n_epochs):
            d = {
                "epoch": i,
                "tput per epoch": tputs[i],
                "epoch times ": epoch_times[i],
                "batches per epoch": n_batches[i],
                "graphs per epoch": n_graphs[i],
                "nodes per epoch": n_nodes[i],
                "edges per epoch": n_edges[i],
            }
            print(d)
            wandb.log(d)


def main(
    cfg: DictConfig,
    stages: Optional[Sequence[str]] = None,
    run_name: str = "dataset_benchmark",
    add_date_time: bool = True,
    log2wandb: bool = False,
) -> None:
    if add_date_time:
        run_name += "_" + datetime.now().strftime("%d.%m.%Y_%H.%M.%S")

    if log2wandb:
        wandb.init(project="multitask-gnn", name=run_name, config=cfg)

    cfg = deepcopy(cfg)

    # Load and initialize the dataset
    datamodule = benchmark(load_datamodule, cfg, message="Load duration", log2wandb=log2wandb)

    benchmark(datamodule.prepare_data, message="Prepare duration", log2wandb=log2wandb)

    if False:  # stages is not None:
        for stage in stages:
            benchmark(datamodule.setup, stage, message=f"Setup {stage} duration", log2wandb=log2wandb)
    else:
        benchmark(datamodule.setup, message=f"Setup duration", log2wandb=log2wandb)

    if stages is None or {"train", "fit"}.intersection(stages):
        dataloader = datamodule.train_dataloader()
        benchmark_dataloader(
            dataloader, name="train", n_epochs=cfg["trainer"]["trainer"]["max_epochs"], log2wandb=log2wandb
        )
    if stages is None or {"val", "valid", "validation"}.intersection(stages):
        dataloader = datamodule.val_dataloader()
        benchmark_dataloader(
            dataloader,
            name="validation",
            n_epochs=cfg["trainer"]["trainer"]["max_epochs"],
            log2wandb=log2wandb,
        )
    if stages is None or {"test", "testing"}.intersection(stages):
        dataloader = datamodule.test_dataloader()
        benchmark_dataloader(
            dataloader, name="testing", n_epochs=cfg["trainer"]["trainer"]["max_epochs"], log2wandb=log2wandb
        )


if __name__ == "__main__":
    with open(os.path.join(MAIN_DIR, CONFIG_FILE), "r") as f:
        cfg = yaml.safe_load(f)
    main(cfg, stages=["train"], log2wandb=True)


# expts/main_run_get_fingerprints.py

# General imports
import os
from os.path import dirname, abspath
import yaml
import numpy as np
import pandas as pd
import torch
import fsspec
from pytorch_lightning.utilities.model_summary import ModelSummary

# Current project imports
import graphium
from graphium.config._loader import load_datamodule, load_trainer
from graphium.utils.fs import mkdir
from graphium.trainer.predictor import PredictorModule


# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(graphium.__file__)))
os.chdir(MAIN_DIR)


# MODEL_FILE = "models_checkpoints/micro_ZINC/model.ckpt"
# CONFIG_FILE = "expts/config_micro_ZINC.yaml"


def main() -> None:
    LIST_CONCAT_LAST_LAYERS = [1, 0, [1, 2], [0, 1, 2]]
    DATA_NAME_ALL = ["molbace"]  # , "mollipo", "moltox21", "molHIV"]
    MODEL_PATH = "gs://graphium-public/pretrained-models"
    MODEL_NAME = "graphium-zinc-micro-dummy-test"
    MODEL_FILE = f"{MODEL_PATH}/{MODEL_NAME}/model.ckpt"
    MODEL_CONFIG = f"{MODEL_PATH}/{MODEL_NAME}/configs.yaml"

    predictor = PredictorModule.load_from_checkpoint(MODEL_FILE)

    print(predictor.model)
    print(ModelSummary(predictor, max_depth=4))

    for data_name in DATA_NAME_ALL:
        DATA_CONFIG = f"{MAIN_DIR}/expts/config_{data_name}_pretrained.yaml"

        with fsspec.open(DATA_CONFIG, "r") as f:
            data_cfg = yaml.safe_load(f)
        with fsspec.open(os.path.join(MODEL_CONFIG), "r") as f:
            model_cfg = yaml.safe_load(f)

        # Load and initialize the dataset
        data_cfg["datamodule"]["args"]["featurization"] = model_cfg["datamodule"]["args"]["featurization"]
        datamodule = load_datamodule(data_cfg)
        print("\ndatamodule:\n", datamodule, "\n")

        for concat_last_layers in LIST_CONCAT_LAST_LAYERS:
            export_dir = f"predictions/fingerprints-model-{MODEL_NAME}"
            export_df_path = f"{export_dir}/{data_name}-concatlayers-{concat_last_layers}.csv.gz"

            predictor.model.concat_last_layers = concat_last_layers
            trainer = load_trainer(data_cfg)

            # Run the model prediction
            preds = trainer.predict(model=predictor, datamodule=datamodule)
            if isinstance(preds[0], torch.Tensor):
                preds = [p.detach().cpu().numpy() for p in preds]
            preds = np.concatenate(preds, axis=0)

            # Generate output dataframe
            df = {"SMILES": datamodule.dataset.smiles}

            target = datamodule.dataset.labels
            for ii in range(target.shape[1]):
                df[f"Target-{ii}"] = target[:, ii]

            for ii in range(preds.shape[1]):
                df[f"Preds-{ii}"] = preds[:, ii]
            df = pd.DataFrame(df)
            mkdir(export_dir)
            df.to_csv(export_df_path)
            print(df)
            print(f"file saved to:`{export_df_path}`")


if __name__ == "__main__":
    main()


# expts/main_run_multitask.py

# General imports
import argparse
import os
from os.path import dirname, abspath
import yaml
from copy import deepcopy
from omegaconf import DictConfig
import timeit
from loguru import logger
from datetime import datetime
from pytorch_lightning.utilities.model_summary import ModelSummary

# Current project imports
import graphium
from graphium.config._loader import (
    load_datamodule,
    load_metrics,
    load_architecture,
    load_predictor,
    load_trainer,
    save_params_to_wandb,
    load_accelerator,
)
from graphium.utils.safe_run import SafeRun
from graphium.utils.command_line_utils import update_config, get_anchors_and_aliases


# WandB
import wandb

# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(graphium.__file__)))

# CONFIG_FILE = "expts/configs/config_mpnn_10M_b3lyp.yaml"
# CONFIG_FILE = "expts/configs/config_mpnn_10M_pcqm4m.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_debug.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_large_mpnn.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/debug/config_small_gcn_debug.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_large_gin.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_large_gine.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_small_gcn.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_small_gin.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_small_gine.yaml"

os.chdir(MAIN_DIR)


def main(cfg: dict, run_name: str = "main", add_date_time: bool = True) -> None:
    st = timeit.default_timer()

    date_time_suffix = ""
    if add_date_time:
        date_time_suffix = datetime.now().strftime("%d.%m.%Y_%H.%M.%S")

    cfg = deepcopy(cfg)
    wandb.init(entity="multitask-gnn", project=cfg["constants"]["name"], config=cfg)

    # Initialize the accelerator
    cfg, accelerator_type = load_accelerator(cfg)

    # Load and initialize the dataset
    datamodule = load_datamodule(cfg, accelerator_type)

    # Initialize the network
    model_class, model_kwargs = load_architecture(
        cfg,
        in_dims=datamodule.in_dims,
    )

    datamodule.prepare_data()

    metrics = load_metrics(cfg)
    logger.info(metrics)

    predictor = load_predictor(
        cfg, model_class, model_kwargs, metrics, accelerator_type, datamodule.task_norms
    )

    logger.info(predictor.model)
    logger.info(ModelSummary(predictor, max_depth=4))

    trainer = load_trainer(cfg, run_name, accelerator_type, date_time_suffix)
    save_params_to_wandb(trainer.logger, cfg, predictor, datamodule)

    # Determine the max num nodes and edges in training and validation
    predictor.set_max_nodes_edges_per_graph(datamodule, stages=["train", "val"])

    # Run the model training
    with SafeRun(name="TRAINING", raise_error=cfg["constants"]["raise_train_error"], verbose=True):
        trainer.fit(model=predictor, datamodule=datamodule)

    # Determine the max num nodes and edges in testing
    predictor.set_max_nodes_edges_per_graph(datamodule, stages=["test"])

    # Run the model testing
    with SafeRun(name="TESTING", raise_error=cfg["constants"]["raise_train_error"], verbose=True):
        trainer.test(model=predictor, datamodule=datamodule)  # , ckpt_path=ckpt_path)

    logger.info("--------------------------------------------")
    logger.info("total computation used", timeit.default_timer() - st)
    logger.info("--------------------------------------------")
    wandb.finish()

    return trainer.callback_metrics


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", help="Path to the config file", default=None)

    args, unknown = parser.parse_known_args()
    # Optionally parse the config with the command line
    if args.config is not None:
        CONFIG_FILE = args.config

    with open(os.path.join(MAIN_DIR, CONFIG_FILE), "r") as f:
        cfg = yaml.safe_load(f)
        refs = get_anchors_and_aliases(CONFIG_FILE)
        cfg = update_config(cfg, unknown, refs)
    main(cfg)


# expts/main_run_predict.py

# General imports
import os
from os.path import dirname, abspath
import yaml
from copy import deepcopy
from omegaconf import DictConfig
import numpy as np
import pandas as pd
import torch
from pytorch_lightning.utilities.model_summary import ModelSummary

# Current project imports
import graphium
from graphium.config._loader import load_datamodule, load_trainer
from graphium.utils.fs import mkdir
from graphium.trainer.predictor import PredictorModule


# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(graphium.__file__)))
os.chdir(MAIN_DIR)

DATA_NAME = "molhiv"
MODEL_FILE = "models_checkpoints/ogb-molpcba/model-v2.ckpt"
CONFIG_FILE = f"expts/config_{DATA_NAME}_pretrained.yaml"

# MODEL_FILE = "models_checkpoints/micro_ZINC/model.ckpt"
# CONFIG_FILE = "expts/config_micro_ZINC.yaml"


def main(cfg: DictConfig) -> None:
    cfg = deepcopy(cfg)

    # Load and initialize the dataset
    datamodule = load_datamodule(cfg)
    print("\ndatamodule:\n", datamodule, "\n")

    export_df_path = f"predictions/preds-{DATA_NAME}.csv.gz"

    predictor = PredictorModule.load_from_checkpoint(MODEL_FILE)

    print(predictor.model)
    print(ModelSummary(predictor, max_depth=4))

    trainer = load_trainer(cfg)

    # Run the model prediction
    preds = trainer.predict(model=predictor, datamodule=datamodule)
    if isinstance(preds[0], torch.Tensor):
        preds = [p.detach().cpu().numpy() for p in preds]
    preds = np.concatenate(preds, axis=0)

    # Generate output dataframe
    df = {"SMILES": datamodule.dataset.smiles}

    target = datamodule.dataset.labels
    for ii in range(target.shape[1]):
        df[f"Target-{ii}"] = target[:, ii]

    for ii in range(preds.shape[1]):
        df[f"Preds-{ii}"] = preds[:, ii]
    df = pd.DataFrame(df)
    mkdir("predictions")
    df.to_csv(export_df_path)
    print(df)
    print(f"file saved to:`{export_df_path}`")


if __name__ == "__main__":
    with open(os.path.join(MAIN_DIR, CONFIG_FILE), "r") as f:
        cfg = yaml.safe_load(f)
    main(cfg)


# expts/main_run_test.py

# General imports
import os
from os.path import dirname, abspath
import yaml
from copy import deepcopy
from omegaconf import DictConfig
from pytorch_lightning.utilities.model_summary import ModelSummary

# Current project imports
import graphium
from graphium.config._loader import load_datamodule, load_metrics, load_trainer

from graphium.trainer.predictor import PredictorModule


# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(graphium.__file__)))
os.chdir(MAIN_DIR)

MODEL_FILE = "models_checkpoints/ogb-molpcba/model-v2.ckpt"

CONFIG_FILE = "expts/config_molPCBA.yaml"


def main(cfg: DictConfig) -> None:
    cfg = deepcopy(cfg)

    # Load and initialize the dataset
    datamodule = load_datamodule(cfg)
    print("\ndatamodule:\n", datamodule, "\n")

    metrics = load_metrics(cfg)
    print(metrics)

    predictor = PredictorModule.load_from_checkpoint(MODEL_FILE)
    predictor.metrics = metrics

    print(predictor.model)
    print(ModelSummary(predictor, max_depth=4))

    trainer = load_trainer(cfg)

    # Run the model testing
    trainer.test(model=predictor, datamodule=datamodule, ckpt_path=MODEL_FILE)


if __name__ == "__main__":
    with open(os.path.join(MAIN_DIR, CONFIG_FILE), "r") as f:
        cfg = yaml.safe_load(f)
    main(cfg)


# expts/run_validation.py

# General imports
import argparse
import os
from os.path import dirname, abspath
import yaml
from copy import deepcopy
from omegaconf import DictConfig
import timeit
from loguru import logger
from datetime import datetime
from pytorch_lightning.utilities.model_summary import ModelSummary

# Current project imports
import graphium
from graphium.config._loader import (
    load_datamodule,
    load_metrics,
    load_architecture,
    load_predictor,
    load_trainer,
    save_params_to_wandb,
    load_accelerator,
)
from graphium.utils.safe_run import SafeRun
from graphium.utils.command_line_utils import update_config, get_anchors_and_aliases


# WandB
import wandb

# Set up the working directory
MAIN_DIR = dirname(dirname(abspath(graphium.__file__)))

# CONFIG_FILE = "expts/configs/config_mpnn_10M_b3lyp.yaml"
# CONFIG_FILE = "expts/configs/config_mpnn_10M_pcqm4m.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_debug.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_large_mpnn.yaml"
CONFIG_FILE = "expts/neurips2023_configs/config_large_gcn.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_large_gin.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_large_gine.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_small_gcn.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_small_gin.yaml"
# CONFIG_FILE = "expts/neurips2023_configs/config_small_gine.yaml"
os.chdir(MAIN_DIR)


def main(cfg: DictConfig, run_name: str = "main", add_date_time: bool = True) -> None:
    st = timeit.default_timer()

    date_time_suffix = ""
    if add_date_time:
        date_time_suffix = datetime.now().strftime("%d.%m.%Y_%H.%M.%S")

    cfg = deepcopy(cfg)
    wandb.init(project=cfg["constants"]["name"], config=cfg)

    # Initialize the accelerator
    cfg, accelerator_type = load_accelerator(cfg)

    # Load and initialize the dataset
    datamodule = load_datamodule(cfg, accelerator_type)

    # Initialize the network
    model_class, model_kwargs = load_architecture(
        cfg,
        in_dims=datamodule.in_dims,
    )

    datamodule.prepare_data()

    metrics = load_metrics(cfg)
    logger.info(metrics)

    predictor = load_predictor(
        cfg, model_class, model_kwargs, metrics, accelerator_type, datamodule.task_norms
    )
    logger.info(predictor.model)
    logger.info(ModelSummary(predictor, max_depth=4))

    trainer = load_trainer(cfg, run_name, accelerator_type, date_time_suffix)
    save_params_to_wandb(trainer.logger, cfg, predictor, datamodule)

    # Determine the max num nodes and edges in training and validation
    predictor.set_max_nodes_edges_per_graph(datamodule, stages=["val"])

    # Run the model validation
    with SafeRun(name="VALIDATING", raise_error=cfg["constants"]["raise_train_error"], verbose=True):
        trainer.validate(
            model=predictor,
            ckpt_path=f'{cfg["trainer"]["model_checkpoint"]["dirpath"]}{cfg["trainer"]["model_checkpoint"]["filename"]}.ckpt',
            datamodule=datamodule,
        )

    logger.info("--------------------------------------------")
    logger.info("total computation used", timeit.default_timer() - st)
    logger.info("--------------------------------------------")
    wandb.finish()

    return trainer.callback_metrics


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", help="Path to the config file", default=None)

    args, unknown = parser.parse_known_args()
    # Optionally parse the config with the command line
    if args.config is not None:
        CONFIG_FILE = args.config

    with open(os.path.join(MAIN_DIR, CONFIG_FILE), "r") as f:
        cfg = yaml.safe_load(f)
        refs = get_anchors_and_aliases(CONFIG_FILE)
        cfg = update_config(cfg, unknown, refs)
    main(cfg)


# expts/test_yaml.py

# test_yaml.py


import yaml

CONFIG_FILE = "neurips2023_configs/config_small_mpnn.yaml"

import re
from collections import defaultdict

# def get_anchors_and_aliases(filepath):
#     anchors = defaultdict(list)
#     current_level = {}

#     with open(filepath, "r") as file:
#         for line in file:
#             indent = len(line) - len(line.lstrip(' '))
#             key_match = re.search(r'(\w+):', line)
#             anchor_match = re.search(r'&(\w+)', line)
#             alias_match = re.search(r'\*(\w+)', line)
#             if key_match:
#                 key = key_match.group(1)
#                 # Compute the full path of the current key.
#                 full_path = '.'.join([current_level[i] for i in sorted(current_level.keys()) if i < indent] + [key])
#                 current_level[indent] = key
#                 # Remove any keys that are indented more than the current line.
#                 keys_to_remove = [i for i in current_level if i > indent]
#                 for i in keys_to_remove:
#                     del current_level[i]
#             else:
#                 full_path = '.'.join([current_level[i] for i in sorted(current_level.keys())])
#             if anchor_match:
#                 anchor = anchor_match.group(1)
#                 anchors[anchor].append(full_path)
#             if alias_match:
#                 alias = alias_match.group(1)
#                 anchors[alias].append(full_path)
#     return {k: v for k, v in anchors.items() if len(v) > 1}


def get_anchors_and_aliases(filepath):
    anchors = defaultdict(list)
    current_level = {}
    anchor_to_path = {}

    with open(filepath, "r") as file:
        for line in file:
            indent = len(line) - len(line.lstrip(" "))
            key_match = re.search(r"(\w+):", line)
            anchor_match = re.search(r"&(\w+)", line)
            alias_match = re.search(r"\*(\w+)", line)
            if key_match:
                key = key_match.group(1)
                # Compute the full path of the current key.
                full_path = ".".join(
                    [current_level[i] for i in sorted(current_level.keys()) if i < indent] + [key]
                )
                current_level[indent] = key
                # Remove any keys that are indented more than the current line.
                keys_to_remove = [i for i in current_level if i > indent]
                for i in keys_to_remove:
                    del current_level[i]
            else:
                full_path = ".".join([current_level[i] for i in sorted(current_level.keys())])
            if anchor_match:
                anchor = anchor_match.group(1)
                anchor_to_path[anchor] = full_path
            if alias_match:
                alias = alias_match.group(1)
                if alias in anchor_to_path:
                    anchors[anchor_to_path[alias]].append(full_path)
    return anchors


refs = get_anchors_and_aliases(CONFIG_FILE)
import ipdb

ipdb.set_trace()
pass


# profiling/profile_mol_to_graph.py

from tqdm import tqdm
import datamol as dm
import pickle

from graphium.data.utils import load_micro_zinc
from graphium.features.featurizer import mol_to_pyggraph, mol_to_adj_and_features, mol_to_graph_dict

# Check out this profiling tool: https://kirillstrelkov.medium.com/python-profiling-with-vscode-3a17c0407833


def main():
    df = load_micro_zinc()
    smiles = df["SMILES"].values.tolist()
    smiles = smiles * 10
    print("Num smiles: ", len(smiles))

    pos_enc = {
        "pos_type": "laplacian_eigvec",
        "num_pos": 3,
        "normalization": "none",
        "disconnected_comp": True,
    }
    atom_property_list_float1 = [
        "mass",
        "in-ring",
        "hybridization",
        "chirality",
        "aromatic",
        "degree",
        "formal-charge",
        "single-bond",
        "double-bond",
        "radical-electron",
    ]
    atom_property_list_float2 = ["electronegativity", "vdw-radius", "covalent-radius", "metal"]

    featurizer = {
        "atom_property_list_onehot": ["atomic-number", "valence"],
        "atom_property_list_float": atom_property_list_float1 + atom_property_list_float2,
        "edge_property_list": [
            "bond-type-onehot",
            "bond-type-float",
            "stereo",
            "in-ring",
            "conjugated",
            "estimated-bond-length",
        ],
        "add_self_loop": False,
        "explicit_H": False,
        "use_bonds_weights": False,
        "pos_encoding_as_features": pos_enc,
        # "on_error": "raise",
    }

    graphs = []
    for s in tqdm(smiles):
        mol = dm.to_mol(s)
        graphs.append(mol_to_graph_dict(mol, **featurizer))

    print(graphs[0])


if __name__ == "__main__":
    main()


# profiling/profile_one_of_k_encoding.py

from tqdm import tqdm
from graphium.utils.tensor import one_of_k_encoding


def main():
    CLASSES = ["AA", "BB", "CC", "DD", "EE"]
    CHOICES = CLASSES + ["FF"]

    for ii in tqdm(range(500000)):
        for choice in CHOICES:
            one_of_k_encoding(choice, CLASSES)

    print("DONE :)")


if __name__ == "__main__":
    main()


# profiling/profile_predictor.py

from tqdm import tqdm
import os
from time import time
import yaml
import fsspec

from graphium.config._loader import (
    load_datamodule,
    load_metrics,
    load_trainer,
    load_predictor,
    load_architecture,
)
from pytorch_lightning import Trainer


def main():
    CONFIG_PATH = "expts/config_micro-PCBA.yaml"
    # DATA_PATH = "https://storage.googleapis.com/graphium-public/datasets/graphium-zinc-bench-gnn/smiles_score.csv.gz"

    with fsspec.open(CONFIG_PATH, "r") as f:
        cfg = yaml.safe_load(f)
    cfg["datamodule"]["args"]["cache_data_path"] = "graphium/data/cache/profiling/predictor_data.cache"
    # cfg["datamodule"]["args"]["df_path"] = DATA_PATH
    cfg["trainer"]["trainer"]["max_epochs"] = 5
    cfg["trainer"]["trainer"]["min_epochs"] = 5

    datamodule = load_datamodule(cfg)

    # Initialize the network
    model_class, model_kwargs = load_architecture(
        cfg,
        in_dim_nodes=datamodule.num_node_feats_with_positional_encoding,
        in_dim_edges=datamodule.num_edge_feats,
    )

    metrics = load_metrics(cfg)
    print(metrics)
    predictor = load_predictor(cfg, model_class, model_kwargs, metrics)
    trainer = load_trainer(cfg)
    trainer.fit(model=predictor, datamodule=datamodule)

    print("Done :)")


if __name__ == "__main__":
    main()


# scripts/convert_yml.py

"""
Convert the dependencies from conda's `env.yml` to pip `requirements.txt`
"""

import ruamel.yaml

yaml = ruamel.yaml.YAML()
data = yaml.load(open("env.yml"))

requirements = []
for dep in data["dependencies"]:
    if isinstance(dep, str):
        outputs = dep.split("=")
        if len(outputs) == 1:
            package = outputs[0]
            requirements.append(package)
        elif len(outputs) == 2:
            package, package_version = outputs[0], outputs[1]
            requirements.append(package + "==" + package_version)
        elif len(outputs) == 3:
            package, package_version, python_version = outputs[0], outputs[1], outputs[2]
            requirements.append(package + "==" + package_version)
    elif isinstance(dep, dict):
        for preq in dep.get("pip", []):
            requirements.append(preq)

with open("requirements.txt", "w") as fp:
    for requirement in requirements:
        print(requirement, file=fp)


# tests/conftest.py

import pathlib

import pytest

TEST_DIR_PATH = pathlib.Path(__file__).parent / "data"
DATA_DIR_PATH = TEST_DIR_PATH.absolute()


@pytest.fixture
def datadir(request):
    return DATA_DIR_PATH


# tests/test_architectures.py

"""
Unit tests for the different architectures of graphium/nn/architectures...

The layers are not thoroughly tested due to the difficulty of testing them
"""

import torch
import unittest as ut
from copy import deepcopy
import sys
import traceback

from graphium.nn.architectures import FeedForwardNN, FeedForwardPyg, FullGraphMultiTaskNetwork
from graphium.nn.base_layers import FCLayer
from graphium.nn.residual_connections import (
    ResidualConnectionConcat,
    ResidualConnectionDenseNet,
    ResidualConnectionNone,
    ResidualConnectionSimple,
    ResidualConnectionWeighted,
)
from torch_geometric.data import Data, Batch

from graphium.utils.spaces import LAYERS_DICT


class test_FeedForwardNN(ut.TestCase):
    kwargs = {
        "activation": "relu",
        "last_activation": "none",
        "normalization": "none",
        "dropout": 0.2,
        "name": "LNN",
        "layer_type": FCLayer,
    }

    norms = ["none", "batch_norm", "layer_norm"]

    def test_forward_no_residual(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [5, 6, 7]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="none",
            residual_skip_steps=1,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2])

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_simple_residual_1(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="simple",
            residual_skip_steps=1,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2])

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_norms(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        for normalization in self.norms:
            this_kwargs = deepcopy(self.kwargs)
            this_kwargs["normalization"] = normalization
            err_msg = f"normalization = {normalization}"
            lnn = FeedForwardNN(
                in_dim=in_dim,
                out_dim=out_dim,
                hidden_dims=hidden_dims,
                residual_type="simple",
                residual_skip_steps=1,
                **this_kwargs,
            )

            self.assertEqual(len(lnn.layers), len(hidden_dims) + 1, msg=err_msg)
            self.assertEqual(lnn.layers[0].in_dim, in_dim, msg=err_msg)
            self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0], msg=err_msg)
            self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1], msg=err_msg)
            self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2], msg=err_msg)

            feat = torch.FloatTensor(batch, in_dim)
            feat_out = lnn.forward(feat)

            self.assertListEqual(list(feat_out.shape), [batch, out_dim], msg=err_msg)

    def test_forward_simple_residual_2(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="simple",
            residual_skip_steps=2,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2])
        self.assertEqual(lnn.layers[4].in_dim, hidden_dims[3])
        self.assertEqual(lnn.layers[5].in_dim, hidden_dims[4])

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_concat_residual_1(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="concat",
            residual_skip_steps=1,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, 2 * hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, 2 * hidden_dims[2])
        self.assertEqual(lnn.layers[4].in_dim, 2 * hidden_dims[3])
        self.assertEqual(lnn.layers[5].in_dim, 2 * hidden_dims[4])

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_concat_residual_2(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="concat",
            residual_skip_steps=2,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, 1 * hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, 1 * hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, 2 * hidden_dims[2])
        self.assertEqual(lnn.layers[4].in_dim, 1 * hidden_dims[3])
        self.assertEqual(lnn.layers[5].in_dim, 2 * hidden_dims[4])

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_densenet_residual_1(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="densenet",
            residual_skip_steps=1,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, 2 * hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, 3 * hidden_dims[2])
        self.assertEqual(lnn.layers[4].in_dim, 4 * hidden_dims[3])
        self.assertEqual(lnn.layers[5].in_dim, 5 * hidden_dims[4])

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_densenet_residual_2(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="densenet",
            residual_skip_steps=2,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, 1 * hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, 2 * hidden_dims[2])
        self.assertEqual(lnn.layers[4].in_dim, 1 * hidden_dims[3])
        self.assertEqual(lnn.layers[5].in_dim, 3 * hidden_dims[4])

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_weighted_residual_1(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="weighted",
            residual_skip_steps=1,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2])
        self.assertEqual(lnn.layers[4].in_dim, hidden_dims[3])
        self.assertEqual(lnn.layers[5].in_dim, hidden_dims[4])

        self.assertEqual(len(lnn.residual_layer.residual_list), len(hidden_dims))

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])

    def test_forward_weighted_residual_2(self):
        in_dim = 8
        out_dim = 16
        hidden_dims = [6, 6, 6, 6, 6]
        batch = 2

        lnn = FeedForwardNN(
            in_dim=in_dim,
            out_dim=out_dim,
            hidden_dims=hidden_dims,
            residual_type="weighted",
            residual_skip_steps=2,
            **self.kwargs,
        )

        self.assertEqual(len(lnn.layers), len(hidden_dims) + 1)
        self.assertEqual(lnn.layers[0].in_dim, in_dim)
        self.assertEqual(lnn.layers[1].in_dim, hidden_dims[0])
        self.assertEqual(lnn.layers[2].in_dim, hidden_dims[1])
        self.assertEqual(lnn.layers[3].in_dim, hidden_dims[2])
        self.assertEqual(lnn.layers[4].in_dim, hidden_dims[3])
        self.assertEqual(lnn.layers[5].in_dim, hidden_dims[4])

        self.assertEqual(len(lnn.residual_layer.residual_list), (len(hidden_dims) // 2 + 1))

        feat = torch.FloatTensor(batch, in_dim)
        feat_out = lnn.forward(feat)

        self.assertListEqual(list(feat_out.shape), [batch, out_dim])


class test_FeedForwardGraph(ut.TestCase):
    kwargs = {
        "activation": "relu",
        "last_activation": "none",
        "dropout": 0.2,
        "name": "LNN",
    }

    in_dim = 7
    out_dim = 11
    in_dim_edges = 13
    hidden_dims = [6, 6, 6, 6, 6]

    edge_idx1 = (torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3]))
    edge_idx2 = (torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0]))
    num_edges1 = len(edge_idx1[0])
    num_nodes1 = max(edge_idx1[0].max(), edge_idx1[1].max()) + 1
    num_edges2 = len(edge_idx2[0])
    num_nodes2 = max(edge_idx2[0].max(), edge_idx2[1].max()) + 1
    h1 = torch.zeros(num_nodes1, in_dim, dtype=torch.float32)
    e1 = torch.ones(num_edges1, in_dim_edges, dtype=torch.float32)
    h2 = torch.ones(num_nodes2, in_dim, dtype=torch.float32)
    e2 = torch.zeros(num_edges2, in_dim_edges, dtype=torch.float32)

    g1 = Data(feat=h1, edge_index=torch.stack(edge_idx1), edge_feat=e1)
    g2 = Data(feat=h2, edge_index=torch.stack(edge_idx2), edge_feat=e2)
    data_list = [g1, g2, deepcopy(g1), deepcopy(g2)]
    batch_pyg = Batch.from_data_list(data_list)
    num_nodes = batch_pyg.num_nodes
    num_edges = batch_pyg.num_edges
    batch_size = len(data_list)

    virtual_nodes = ["none", "mean", "sum"]
    norms = ["none", "batch_norm", "layer_norm"]
    pna_kwargs = {"aggregators": ["mean", "max", "sum"], "scalers": ["identity", "amplification"]}

    layers_kwargs = {
        "pyg:gin": {},
        "pyg:gine": {"in_dim_edges": in_dim_edges},
        "pyg:gated-gcn": {"in_dim_edges": in_dim_edges, "hidden_dims_edges": hidden_dims},
        "pyg:pna-msgpass#1": {"layer_kwargs": pna_kwargs, "in_dim_edges": 0},
        "pyg:pna-msgpass#2": {"layer_kwargs": pna_kwargs, "in_dim_edges": in_dim_edges},
    }

    def test_forward_no_residual(self):
        for residual_skip_steps in [1, 2]:
            for virtual_node in self.virtual_nodes:
                for normalization in self.norms:
                    for layer_name, this_kwargs in self.layers_kwargs.items():
                        err_msg = f"virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                        layer_type = layer_name.split("#")[0]

                        # PYG
                        if layer_type.startswith("pyg:"):
                            layer_class = FeedForwardPyg
                            bg = deepcopy(self.batch_pyg)

                        gnn = layer_class(
                            in_dim=self.in_dim,
                            out_dim=self.out_dim,
                            hidden_dims=self.hidden_dims,
                            residual_type="none",
                            residual_skip_steps=residual_skip_steps,
                            layer_type=layer_type,
                            normalization=normalization,
                            **this_kwargs,
                            **self.kwargs,
                        )
                        # gnn.to(torch.float32)

                        self.assertIsInstance(gnn.residual_layer, ResidualConnectionNone)
                        self.assertEqual(len(gnn.layers), len(self.hidden_dims) + 1, msg=err_msg)
                        self.assertEqual(gnn.layers[0].out_dim, self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[1].out_dim, self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[2].out_dim, self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[3].out_dim, self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[4].out_dim, self.hidden_dims[4], msg=err_msg)
                        self.assertEqual(gnn.layers[5].out_dim, self.out_dim, msg=err_msg)

                        f = gnn.layers[0].out_dim_factor
                        self.assertEqual(gnn.layers[0].in_dim, self.in_dim, msg=err_msg)
                        self.assertEqual(gnn.layers[1].in_dim, f * self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[2].in_dim, f * self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[3].in_dim, f * self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[4].in_dim, f * self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[5].in_dim, f * self.hidden_dims[4], msg=err_msg)

                        out_g = gnn.forward(bg)
                        feat_out = out_g["feat"]
                        edge_feat_out = out_g["edge_feat"]

                        out_dim_edges = (
                            gnn.full_dims_edges[-1] if gnn.full_dims_edges is not None else self.in_dim_edges
                        )

                        self.assertListEqual(
                            list(feat_out.shape), [self.num_nodes, self.out_dim], msg=err_msg
                        )
                        self.assertListEqual(
                            list(edge_feat_out.shape), [self.num_edges, out_dim_edges], msg=err_msg
                        )

    def test_forward_simple_residual(self):
        for residual_skip_steps in [1, 2]:
            for virtual_node in self.virtual_nodes:
                for normalization in self.norms:
                    for layer_name, this_kwargs in self.layers_kwargs.items():
                        err_msg = f"virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                        layer_type = layer_name.split("#")[0]

                        # PYG
                        if layer_type.startswith("pyg:"):
                            layer_class = FeedForwardPyg
                            bg = deepcopy(self.batch_pyg)

                        gnn = layer_class(
                            in_dim=self.in_dim,
                            out_dim=self.out_dim,
                            hidden_dims=self.hidden_dims,
                            residual_type="simple",
                            residual_skip_steps=1,
                            layer_type=layer_type,
                            **this_kwargs,
                            **self.kwargs,
                        )
                        # gnn.to(torch.float32)

                        self.assertIsInstance(gnn.residual_layer, ResidualConnectionSimple)
                        self.assertEqual(len(gnn.layers), len(self.hidden_dims) + 1, msg=err_msg)
                        self.assertEqual(gnn.layers[0].out_dim, self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[1].out_dim, self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[2].out_dim, self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[3].out_dim, self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[4].out_dim, self.hidden_dims[4], msg=err_msg)
                        self.assertEqual(gnn.layers[5].out_dim, self.out_dim, msg=err_msg)

                        f = gnn.layers[0].out_dim_factor
                        self.assertEqual(gnn.layers[0].in_dim, self.in_dim, msg=err_msg)
                        self.assertEqual(gnn.layers[1].in_dim, f * self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[2].in_dim, f * self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[3].in_dim, f * self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[4].in_dim, f * self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[5].in_dim, f * self.hidden_dims[4], msg=err_msg)

                        out_g = gnn.forward(bg)
                        feat_out = out_g["feat"]
                        edge_feat_out = out_g["edge_feat"]

                        out_dim_edges = (
                            gnn.full_dims_edges[-1] if gnn.full_dims_edges is not None else self.in_dim_edges
                        )

                        self.assertListEqual(
                            list(feat_out.shape), [self.num_nodes, self.out_dim], msg=err_msg
                        )
                        self.assertListEqual(
                            list(edge_feat_out.shape), [self.num_edges, out_dim_edges], msg=err_msg
                        )

    def test_forward_weighted_residual(self):
        for residual_skip_steps in [1, 2]:
            for virtual_node in self.virtual_nodes:
                for normalization in self.norms:
                    for layer_name, this_kwargs in self.layers_kwargs.items():
                        err_msg = f"virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                        layer_type = layer_name.split("#")[0]

                        # PYG
                        if layer_type.startswith("pyg:"):
                            layer_class = FeedForwardPyg
                            bg = deepcopy(self.batch_pyg)

                        gnn = layer_class(
                            in_dim=self.in_dim,
                            out_dim=self.out_dim,
                            hidden_dims=self.hidden_dims,
                            residual_type="weighted",
                            residual_skip_steps=residual_skip_steps,
                            layer_type=layer_type,
                            **this_kwargs,
                            **self.kwargs,
                        )
                        # gnn.to(torch.float32)

                        self.assertIsInstance(gnn.residual_layer, ResidualConnectionWeighted)
                        self.assertEqual(len(gnn.layers), len(self.hidden_dims) + 1, msg=err_msg)
                        self.assertEqual(gnn.layers[0].out_dim, self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[1].out_dim, self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[2].out_dim, self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[3].out_dim, self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[4].out_dim, self.hidden_dims[4], msg=err_msg)
                        self.assertEqual(gnn.layers[5].out_dim, self.out_dim, msg=err_msg)

                        f = gnn.layers[0].out_dim_factor
                        self.assertEqual(gnn.layers[0].in_dim, self.in_dim, msg=err_msg)
                        self.assertEqual(gnn.layers[1].in_dim, f * self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[2].in_dim, f * self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[3].in_dim, f * self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[4].in_dim, f * self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[5].in_dim, f * self.hidden_dims[4], msg=err_msg)

                        out_g = gnn.forward(bg)
                        feat_out = out_g["feat"]
                        edge_feat_out = out_g["edge_feat"]

                        out_dim_edges = (
                            gnn.full_dims_edges[-1] if gnn.full_dims_edges is not None else self.in_dim_edges
                        )

                        self.assertListEqual(
                            list(feat_out.shape), [self.num_nodes, self.out_dim], msg=err_msg
                        )
                        self.assertListEqual(
                            list(edge_feat_out.shape), [self.num_edges, out_dim_edges], msg=err_msg
                        )

    def test_forward_concat_residual(self):
        for residual_skip_steps in [1, 2]:
            for virtual_node in self.virtual_nodes:
                for normalization in self.norms:
                    for layer_name, this_kwargs in self.layers_kwargs.items():
                        err_msg = f"virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                        layer_type = layer_name.split("#")[0]

                        # PYG
                        if layer_type.startswith("pyg:"):
                            layer_class = FeedForwardPyg
                            bg = deepcopy(self.batch_pyg)

                        gnn = layer_class(
                            in_dim=self.in_dim,
                            out_dim=self.out_dim,
                            hidden_dims=self.hidden_dims,
                            residual_type="concat",
                            residual_skip_steps=residual_skip_steps,
                            layer_type=layer_type,
                            **this_kwargs,
                            **self.kwargs,
                        )
                        # gnn.to(torch.float32)

                        self.assertIsInstance(gnn.residual_layer, ResidualConnectionConcat)
                        self.assertEqual(len(gnn.layers), len(self.hidden_dims) + 1, msg=err_msg)
                        self.assertEqual(gnn.layers[0].out_dim, self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[1].out_dim, self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[2].out_dim, self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[3].out_dim, self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[4].out_dim, self.hidden_dims[4], msg=err_msg)
                        self.assertEqual(gnn.layers[5].out_dim, self.out_dim, msg=err_msg)

                        f = gnn.layers[0].out_dim_factor
                        f2 = [2 * f if ((ii % residual_skip_steps) == 0 and ii > 0) else f for ii in range(6)]
                        self.assertEqual(gnn.layers[0].in_dim, self.in_dim, msg=err_msg)
                        self.assertEqual(gnn.layers[1].in_dim, f2[0] * self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[2].in_dim, f2[1] * self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[3].in_dim, f2[2] * self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[4].in_dim, f2[3] * self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[5].in_dim, f2[4] * self.hidden_dims[4], msg=err_msg)

                        out_g = gnn.forward(bg)
                        feat_out = out_g["feat"]
                        edge_feat_out = out_g["edge_feat"]

                        out_dim_edges = (
                            gnn.full_dims_edges[-1] if gnn.full_dims_edges is not None else self.in_dim_edges
                        )

                        self.assertListEqual(
                            list(feat_out.shape), [self.num_nodes, self.out_dim], msg=err_msg
                        )
                        self.assertListEqual(
                            list(edge_feat_out.shape), [self.num_edges, out_dim_edges], msg=err_msg
                        )

    def test_forward_densenet_residual(self):
        for residual_skip_steps in [1, 2]:
            for virtual_node in self.virtual_nodes:
                for normalization in self.norms:
                    for layer_name, this_kwargs in self.layers_kwargs.items():
                        err_msg = f"virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}"
                        layer_type = layer_name.split("#")[0]

                        # PYG
                        if layer_type.startswith("pyg:"):
                            layer_class = FeedForwardPyg
                            bg = deepcopy(self.batch_pyg)

                        gnn = layer_class(
                            in_dim=self.in_dim,
                            out_dim=self.out_dim,
                            hidden_dims=self.hidden_dims,
                            residual_type="densenet",
                            residual_skip_steps=residual_skip_steps,
                            layer_type=layer_type,
                            **this_kwargs,
                            **self.kwargs,
                        )
                        # gnn.to(torch.float32)

                        self.assertIsInstance(gnn.residual_layer, ResidualConnectionDenseNet)
                        self.assertEqual(len(gnn.layers), len(self.hidden_dims) + 1, msg=err_msg)
                        self.assertEqual(gnn.layers[0].out_dim, self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[1].out_dim, self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[2].out_dim, self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[3].out_dim, self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[4].out_dim, self.hidden_dims[4], msg=err_msg)
                        self.assertEqual(gnn.layers[5].out_dim, self.out_dim, msg=err_msg)

                        f = gnn.layers[0].out_dim_factor
                        f2 = [
                            ((ii // residual_skip_steps) + 1) * f
                            if ((ii % residual_skip_steps) == 0 and ii > 0)
                            else f
                            for ii in range(6)
                        ]
                        self.assertEqual(gnn.layers[0].in_dim, self.in_dim, msg=err_msg)
                        self.assertEqual(gnn.layers[1].in_dim, f2[0] * self.hidden_dims[0], msg=err_msg)
                        self.assertEqual(gnn.layers[2].in_dim, f2[1] * self.hidden_dims[1], msg=err_msg)
                        self.assertEqual(gnn.layers[3].in_dim, f2[2] * self.hidden_dims[2], msg=err_msg)
                        self.assertEqual(gnn.layers[4].in_dim, f2[3] * self.hidden_dims[3], msg=err_msg)
                        self.assertEqual(gnn.layers[5].in_dim, f2[4] * self.hidden_dims[4], msg=err_msg)

                        out_g = gnn.forward(bg)
                        feat_out = out_g["feat"]
                        edge_feat_out = out_g["edge_feat"]

                        out_dim_edges = (
                            gnn.full_dims_edges[-1] if gnn.full_dims_edges is not None else self.in_dim_edges
                        )

                        self.assertListEqual(
                            list(feat_out.shape), [self.num_nodes, self.out_dim], msg=err_msg
                        )
                        self.assertListEqual(
                            list(edge_feat_out.shape), [self.num_edges, out_dim_edges], msg=err_msg
                        )


if __name__ == "__main__":
    ut.main()


# tests/test_attention.py

"""
Unit tests for the attention layer
"""

from ast import Assert
import numpy as np
import torch
import unittest as ut
from torch_geometric.data import Data, Batch
from copy import deepcopy
from graphium.ipu.to_dense_batch import to_dense_batch

from graphium.nn.base_layers import MultiheadAttentionMup


def seed_everything(seed: int):
    import random, os
    import numpy as np
    import torch

    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


class test_MultiHeadAttention(ut.TestCase):
    seed_everything(42)
    in_dim = 12
    out_dim = 12
    in_dim_edges = 10
    out_dim_edges = 10
    edge_idx1 = torch.stack([torch.tensor([0, 1, 2, 3, 2]), torch.tensor([1, 2, 3, 0, 0])])
    edge_idx2 = torch.stack([torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])])
    x1 = torch.randn(edge_idx1.max() + 1, in_dim, dtype=torch.float32)
    e1 = torch.randn(edge_idx1.shape[-1], in_dim_edges, dtype=torch.float32)
    x2 = torch.randn(edge_idx2.max() + 1, in_dim, dtype=torch.float32)
    e2 = torch.randn(edge_idx2.shape[-1], in_dim_edges, dtype=torch.float32)
    g1 = Data(feat=x1, edge_index=edge_idx1, edge_feat=e1)
    g2 = Data(feat=x2, edge_index=edge_idx2, edge_feat=e2)
    bg = Batch.from_data_list([g1, g2])

    attn_kwargs = {"embed_dim": in_dim, "num_heads": 2, "batch_first": True}

    def test_attention_class(self):
        bg = deepcopy(self.bg)
        seed_everything(42)
        attention_layer = MultiheadAttentionMup(biased_attention=False, **self.attn_kwargs)
        attention_layer.eval()
        seed_everything(42)
        attention_layer_bias = MultiheadAttentionMup(biased_attention=True, **self.attn_kwargs)
        attention_layer_bias.eval()

        h_dense, mask, _ = to_dense_batch(
            bg.feat,
            batch=bg.batch,
            batch_size=None,
            max_num_nodes_per_graph=None,
            drop_nodes_last_graph=False,
        )
        # attn_bias [batch, num_heads, nodes, nodes]
        nodes = h_dense.size()[1]
        attn_bias_3d = torch.zeros(2, 2, nodes, nodes)
        h_dense.nodepair_gaussian_bias_3d = attn_bias_3d
        # Apply attention layer and attention layer with bias.
        h_attn_output = attention_layer(
            h_dense,
            h_dense,
            h_dense,
            attn_bias=None,
            attn_mask=None,
            key_padding_mask=~mask,
            need_weights=False,
        )[0]
        h_attn_output_biased = attention_layer_bias(
            h_dense,
            h_dense,
            h_dense,
            attn_bias=attn_bias_3d,
            attn_mask=None,
            key_padding_mask=~mask,
            need_weights=False,
        )[0]
        self.assertEqual(h_attn_output.size(), h_attn_output_biased.size())
        np.testing.assert_array_almost_equal(
            h_attn_output.detach().numpy(), h_attn_output_biased.detach().numpy(), decimal=0
        )


if __name__ == "__main__":
    ut.main()


# tests/test_base_layers.py

"""
Unit tests for the different layers of graphium/nn/base_layers
"""

import torch
import unittest as ut
from torch_geometric.data import Data, Batch
from copy import deepcopy

from graphium.nn.base_layers import DropPath, TransformerEncoderLayerMup
from graphium.ipu.to_dense_batch import to_dense_batch, to_sparse_batch


class test_Base_Layers(ut.TestCase):
    in_dim = 21
    out_dim = 11
    in_dim_edges = 13
    out_dim_edges = 17

    edge_idx1 = torch.stack([torch.tensor([0, 1, 2, 3, 2]), torch.tensor([1, 2, 3, 0, 0])])
    edge_idx2 = torch.stack([torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])])
    x1 = torch.randn(edge_idx1.max() + 1, in_dim, dtype=torch.float32)
    e1 = torch.randn(edge_idx1.shape[-1], in_dim_edges, dtype=torch.float32)
    x2 = torch.randn(edge_idx2.max() + 1, in_dim, dtype=torch.float32)
    e2 = torch.randn(edge_idx2.shape[-1], in_dim_edges, dtype=torch.float32)

    g1 = Data(feat=x1, edge_index=edge_idx1, edge_feat=e1)
    g2 = Data(feat=x2, edge_index=edge_idx2, edge_feat=e2)
    bg = Batch.from_data_list([g1, g2])

    batch_size = 2
    max_num_nodes_per_graph = max(x1.shape[0], x2.shape[0])

    # for drop_rate=0.5, test if the output shape is correct
    def test_droppath_layer_0p5(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        layer = DropPath(drop_rate=0.5)
        h_out = layer.forward(feat_in, bg.batch)
        self.assertEqual(h_out.shape, feat_in.shape)

    # for drop_rate=1.0, test if the output are all zeros
    def test_droppath_layer_1p0(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        zero_tesor = torch.zeros(feat_in.shape)
        layer = DropPath(drop_rate=1.0)
        h_out = layer.forward(feat_in, bg.batch)
        self.assertTrue(torch.allclose(zero_tesor, h_out.detach()))

    # for drop_rate=0.0, test if the output matches the original output
    def test_droppath_layer_0p0(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        layer = DropPath(drop_rate=0.0)
        h_out = layer.forward(feat_in, bg.batch)
        self.assertTrue(torch.allclose(feat_in.detach(), h_out.detach()))

    # test output shape is correct for TransformerEncoderLayerMup
    def test_transformer_encoder_layer_mup(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        layer = TransformerEncoderLayerMup(
            biased_attention=False, d_model=self.in_dim, nhead=1, dim_feedforward=4 * self.in_dim
        )

        feat_dense, key_padding_mask, idx = to_dense_batch(
            feat_in,
            batch=bg.batch,
            batch_size=self.batch_size,
            max_num_nodes_per_graph=self.max_num_nodes_per_graph,
            drop_nodes_last_graph=False,
        )
        attn_mask = None
        key_padding_mask = ~key_padding_mask

        h_out_dense = layer.forward(feat_dense)

        h_out = to_sparse_batch(h_out_dense, mask_idx=idx)

        self.assertEqual(h_out.shape, feat_in.shape)


# tests/test_collate.py

"""
Unit tests for the collate
"""

import unittest as ut
from copy import deepcopy
import torch
import numpy as np
from torch_geometric.data import Data

from graphium.data.collate import collate_labels, graphium_collate_fn


class test_Collate(ut.TestCase):
    def test_collate_labels(self):
        # Create fake labels
        labels_size_dict = {
            "graph_label1": [1],
            "graph_label2": [3],
            "node_label2": [5],
            "edge_label3": [5, 2],
            "node_label4": [5, 1],
        }
        fake_label = {
            "graph_label1": torch.FloatTensor([1]),
            "graph_label2": torch.FloatTensor([1, 2, 3]),
            "node_label2": torch.FloatTensor([1, 2, 3, 4, 5]),
            "edge_label3": torch.FloatTensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
            "node_label4": torch.FloatTensor([[1], [2], [3], [4], [5]]),
        }
        fake_labels = []
        num_labels = 10
        for i in range(num_labels):
            pyg_labels = Data(x=torch.empty(5, 0), edge_index=torch.empty(2, 5))
            for key, val in fake_label.items():
                pyg_labels[key] = val + 17 * 2
            fake_labels.append(pyg_labels)

        # Collate labels and check for the right shapes
        collated_labels = collate_labels(deepcopy(fake_labels), deepcopy(labels_size_dict))
        self.assertEqual(collated_labels["graph_label1"].shape, torch.Size([num_labels, 1]))  # , 1
        self.assertEqual(collated_labels["graph_label2"].shape, torch.Size([num_labels, 3]))  # , 1
        self.assertEqual(collated_labels["node_label2"].shape, torch.Size([num_labels * 5, 1]))  # , 5
        self.assertEqual(collated_labels["edge_label3"].shape, torch.Size([num_labels * 5, 2]))  # , 5, 2
        self.assertEqual(collated_labels["node_label4"].shape, torch.Size([num_labels * 5, 1]))  # , 5, 1

        # Check that the values are correct
        graph_label1_true = deepcopy(torch.stack([this_label["graph_label1"] for this_label in fake_labels]))
        graph_label2_true = deepcopy(torch.stack([this_label["graph_label2"] for this_label in fake_labels]))
        label2_true = deepcopy(torch.stack([this_label["node_label2"] for this_label in fake_labels]))
        label3_true = deepcopy(torch.stack([this_label["edge_label3"] for this_label in fake_labels]))
        label4_true = deepcopy(torch.stack([this_label["node_label4"] for this_label in fake_labels]))

        # NOTE: Flatten due to the way Data objects are collated (concat along first dim, instead of stacked)
        np.testing.assert_array_equal(collated_labels["graph_label1"].numpy(), graph_label1_true.numpy())
        np.testing.assert_array_equal(collated_labels["graph_label2"].numpy(), graph_label2_true.numpy())
        np.testing.assert_array_equal(
            collated_labels["node_label2"].numpy(), label2_true.flatten(0, 1).unsqueeze(1).numpy()
        )
        np.testing.assert_array_equal(
            collated_labels["edge_label3"].numpy(), label3_true.flatten(0, 1).numpy()
        )
        np.testing.assert_array_equal(
            collated_labels["node_label4"].numpy(), label4_true.flatten(0, 1).numpy()
        )

        # Remove some labels and check that the collation still works and puts `nan` in the right places
        missing_labels = {
            "graph_label1": [1, 3, 5, 7, 9],
            "graph_label2": [0, 4, 3, 1, 7],
            "node_label2": [0, 2, 4, 6, 8],
            "edge_label3": [0, 1, 2, 3, 4],
            "node_label4": [5, 1, 4, 9, 6],
        }
        for key, missing_idx in missing_labels.items():
            for idx in missing_idx:
                fake_labels[idx].pop(key)
        graph_label1_true[missing_labels["graph_label1"]] = float("nan")
        graph_label2_true[missing_labels["graph_label2"]] = float("nan")
        label2_true[missing_labels["node_label2"]] = float("nan")
        label3_true[missing_labels["edge_label3"]] = float("nan")
        label4_true[missing_labels["node_label4"]] = float("nan")

        # Collate labels and check for the right shapes
        labels_size_dict = {
            "graph_label1": [1],
            "graph_label2": [3],
            "node_label2": [5],
            "edge_label3": [5, 2],
            "node_label4": [5, 1],
        }
        collated_labels = collate_labels(deepcopy(fake_labels), deepcopy(labels_size_dict))
        self.assertEqual(collated_labels["graph_label1"].shape, torch.Size([num_labels, 1]))  # , 1
        self.assertEqual(collated_labels["graph_label2"].shape, torch.Size([num_labels, 3]))  # , 1
        self.assertEqual(collated_labels["node_label2"].shape, torch.Size([num_labels * 5, 1]))  # , 5
        self.assertEqual(collated_labels["edge_label3"].shape, torch.Size([num_labels * 5, 2]))  # , 5, 2
        self.assertEqual(collated_labels["node_label4"].shape, torch.Size([num_labels * 5, 1]))  # , 5, 1

        # Check that the values are correct when some labels are missing
        # NOTE: Flatten due to the way Data objects are collated (concat along first dim, instead of stacked)
        np.testing.assert_array_equal(collated_labels["graph_label1"].numpy(), graph_label1_true.numpy())
        np.testing.assert_array_equal(collated_labels["graph_label2"].numpy(), graph_label2_true.numpy())
        np.testing.assert_array_equal(
            collated_labels["node_label2"].numpy(), label2_true.flatten(0, 1).unsqueeze(1).numpy()
        )
        np.testing.assert_array_equal(
            collated_labels["edge_label3"].numpy(), label3_true.flatten(0, 1).numpy()
        )
        np.testing.assert_array_equal(
            collated_labels["node_label4"].numpy(), label4_true.flatten(0, 1).numpy()
        )
        # Now test the `graphium_collate_fn` function when only labels are given
        fake_labels2 = [{"labels": this_label} for this_label in fake_labels]
        collated_labels = graphium_collate_fn(deepcopy(fake_labels2), labels_size_dict=labels_size_dict)[
            "labels"
        ]
        self.assertEqual(collated_labels["graph_label1"].shape, torch.Size([num_labels, 1]))
        self.assertEqual(collated_labels["graph_label2"].shape, torch.Size([num_labels, 3]))
        self.assertEqual(collated_labels["node_label2"].shape, torch.Size([num_labels * 5, 1]))  # , 5
        self.assertEqual(collated_labels["edge_label3"].shape, torch.Size([num_labels * 5, 2]))  # , 5, 2
        self.assertEqual(collated_labels["node_label4"].shape, torch.Size([num_labels * 5, 1]))  # , 5, 1

        # Check that the values are correct when some labels are missing
        # NOTE: Flatten due to the way Data objects are collated (concat along first dim, instead of stacked)
        np.testing.assert_array_equal(collated_labels["graph_label1"].numpy(), graph_label1_true.numpy())
        np.testing.assert_array_equal(collated_labels["graph_label2"].numpy(), graph_label2_true.numpy())
        np.testing.assert_array_equal(
            collated_labels["node_label2"].numpy(), label2_true.flatten(0, 1).unsqueeze(1).numpy()
        )
        np.testing.assert_array_equal(
            collated_labels["edge_label3"].numpy(), label3_true.flatten(0, 1).numpy()
        )
        np.testing.assert_array_equal(
            collated_labels["node_label4"].numpy(), label4_true.flatten(0, 1).numpy()
        )


if __name__ == "__main__":
    ut.main()


# tests/test_data_utils.py

import pandas as pd
import unittest as ut
import graphium
import tempfile


class TestDataUtils(ut.TestCase):
    def test_list_datasets(
        self,
    ):
        datasets = graphium.data.utils.list_graphium_datasets()
        assert isinstance(datasets, set)
        assert len(datasets) > 0

    def test_download_datasets(self):
        dataset_dir = tempfile.TemporaryDirectory()
        data_path = graphium.data.utils.download_graphium_dataset(
            "graphium-zinc-micro", output_path=dataset_dir.name
        )

        fpath = graphium.utils.fs.join(data_path, "ZINC-micro.csv")
        df = pd.read_csv(fpath)
        assert df.shape == (1000, 4)  # type: ignore
        assert df.columns.tolist() == ["SMILES", "SA", "logp", "score"]  # type: ignore


if __name__ == "__main__":
    ut.main()


# tests/test_datamodule.py

import unittest as ut
import numpy as np
import torch
import pandas as pd
import datamol as dm

import graphium
from graphium.utils.fs import rm, exists, get_size
from graphium.data import GraphOGBDataModule, MultitaskFromSmilesDataModule

TEMP_CACHE_DATA_PATH = "tests/temp_cache_0000"


class Test_DataModule(ut.TestCase):
    def test_ogb_datamodule(self):
        # other datasets are too large to be tested
        dataset_names = ["ogbg-molhiv", "ogbg-molpcba", "ogbg-moltox21", "ogbg-molfreesolv"]
        dataset_name = dataset_names[3]

        # Setup the featurization
        featurization_args = {}
        featurization_args["atom_property_list_float"] = []  # ["weight", "valence"]
        featurization_args["atom_property_list_onehot"] = ["atomic-number", "degree"]
        # featurization_args["conformer_property_list"] = ["positions_3d"]
        featurization_args["edge_property_list"] = ["bond-type-onehot"]
        featurization_args["add_self_loop"] = False
        featurization_args["use_bonds_weights"] = False
        featurization_args["explicit_H"] = False

        # Config for datamodule
        task_specific_args = {}
        task_specific_args["task_1"] = {"task_level": "graph", "dataset_name": dataset_name}
        dm_args = {}
        dm_args["cache_data_path"] = None
        dm_args["featurization"] = featurization_args
        dm_args["batch_size_training"] = 16
        dm_args["batch_size_inference"] = 16
        dm_args["num_workers"] = 0
        dm_args["pin_memory"] = True
        dm_args["featurization_n_jobs"] = 2
        dm_args["featurization_progress"] = True
        dm_args["featurization_backend"] = "loky"
        dm_args["featurization_batch_size"] = 50

        ds = GraphOGBDataModule(task_specific_args, **dm_args)

        ds.prepare_data()

        # Check the keys in the dataset
        ds.setup(save_smiles_and_ids=False)
        assert set(ds.train_ds[0].keys()) == {"features", "labels"}

        ds.setup(save_smiles_and_ids=True)
        assert set(ds.train_ds[0].keys()) == {"smiles", "mol_ids", "features", "labels"}

        # test module
        assert ds.num_edge_feats == 5
        assert ds.num_node_feats == 50
        assert len(ds) == 642

        # test batch loader
        batch = next(iter(ds.train_dataloader()))
        assert len(batch["smiles"]) == 16
        assert len(batch["labels"]["graph_task_1"]) == 16
        assert len(batch["mol_ids"]) == 16

    def test_none_filtering(self):
        # Create the objects to filter
        list_of_num = [ii for ii in range(100)]
        list_of_str = [str(ii) for ii in list_of_num]
        tuple_of_num = tuple(list_of_num)
        array_of_num = np.asarray(list_of_num)
        array_of_str = np.asarray(list_of_str)
        tensor_of_num = torch.as_tensor(array_of_num)
        arrays_of_num = np.stack([list_of_num, list_of_num, list_of_num], axis=1)
        arrays_of_str = np.stack([list_of_str, list_of_str, list_of_str], axis=1)
        tensors_of_num = torch.as_tensor(arrays_of_num)
        dic = {"str": list_of_str, "num": list_of_num}
        df = pd.DataFrame(dic)
        df_shuffled = df.sample(frac=1)
        series_num = df["num"]
        series_num_shuffled = df_shuffled["num"]

        # Create different indexes to use for filtering
        all_idx_none = [[3, 17, 88], [22, 33, 44, 55, 66, 77, 88], [], np.arange(len(list_of_num))]

        # Loop all the indexes and filter the objects.
        for ii, idx_none in enumerate(all_idx_none):
            msg = f"Failed for ii={ii}"

            # Create the true filtered sequences
            filtered_num = [ii for ii in range(100) if ii not in idx_none]
            filtered_str = [str(ii) for ii in filtered_num]
            assert len(filtered_num) == len(list_of_num) - len(idx_none)
            assert len(filtered_str) == len(list_of_str) - len(idx_none)

            # Filter the sequences from the Datamodule function
            (
                list_of_num_2,
                list_of_str_2,
                tuple_of_num_2,
                array_of_num_2,
                array_of_str_2,
                tensor_of_num_2,
                df_2,
                df_shuffled_2,
                dic_2,
                arrays_of_num_2,
                arrays_of_str_2,
                tensors_of_num_2,
                series_num_2,
                series_num_shuffled_2,
            ) = graphium.data.MultitaskFromSmilesDataModule._filter_none_molecules(
                idx_none,
                list_of_num,
                list_of_str,
                tuple_of_num,
                array_of_num,
                array_of_str,
                tensor_of_num,
                df,
                df_shuffled,
                dic,
                arrays_of_num,
                arrays_of_str,
                tensors_of_num,
                series_num,
                series_num_shuffled,
            )

            df_shuffled_2 = df_shuffled_2.sort_values(by="num", axis=0)
            series_num_shuffled_2 = series_num_shuffled_2.sort_values(axis=0)

            # Assert the filtering is done correctly
            self.assertListEqual(list_of_num_2, filtered_num, msg=msg)
            self.assertListEqual(list_of_str_2, filtered_str, msg=msg)
            self.assertListEqual(list(tuple_of_num_2), filtered_num, msg=msg)
            self.assertListEqual(array_of_num_2.tolist(), filtered_num, msg=msg)
            self.assertListEqual(array_of_str_2.tolist(), filtered_str, msg=msg)
            self.assertListEqual(tensor_of_num_2.tolist(), filtered_num, msg=msg)
            for jj in range(arrays_of_num.shape[1]):
                self.assertListEqual(arrays_of_num_2[:, jj].tolist(), filtered_num, msg=msg)
                self.assertListEqual(arrays_of_str_2[:, jj].tolist(), filtered_str, msg=msg)
                self.assertListEqual(tensors_of_num_2[:, jj].tolist(), filtered_num, msg=msg)
            self.assertListEqual(dic_2["num"], filtered_num, msg=msg)
            self.assertListEqual(dic_2["str"], filtered_str, msg=msg)
            self.assertListEqual(df_2["num"].tolist(), filtered_num, msg=msg)
            self.assertListEqual(df_2["str"].tolist(), filtered_str, msg=msg)
            self.assertListEqual(series_num_2.tolist(), filtered_num, msg=msg)

            # When the dataframe is shuffled, the lists are different because the filtering
            # is done on the row indexes, not the dataframe indexes.
            bool_to_check = (len(idx_none) == 0) or (len(idx_none) == len(df_shuffled))
            self.assertIs(df_shuffled_2["num"].tolist() == filtered_num, bool_to_check, msg=msg)
            self.assertIs(df_shuffled_2["str"].tolist() == filtered_str, bool_to_check, msg=msg)
            self.assertIs(series_num_shuffled_2.tolist() == filtered_num, bool_to_check, msg=msg)

    def test_caching(self):
        # other datasets are too large to be tested
        dataset_name = "ogbg-molfreesolv"

        # Setup the featurization
        featurization_args = {}
        featurization_args["atom_property_list_float"] = []  # ["weight", "valence"]
        featurization_args["atom_property_list_onehot"] = ["atomic-number", "degree"]
        # featurization_args["conformer_property_list"] = ["positions_3d"]
        featurization_args["edge_property_list"] = ["bond-type-onehot"]
        featurization_args["add_self_loop"] = False
        featurization_args["use_bonds_weights"] = False
        featurization_args["explicit_H"] = False

        # Config for datamodule
        task_specific_args = {}
        task_specific_args["task_1"] = {"task_level": "graph", "dataset_name": dataset_name}
        dm_args = {}
        dm_args["featurization"] = featurization_args
        dm_args["batch_size_training"] = 16
        dm_args["batch_size_inference"] = 16
        dm_args["num_workers"] = 0
        dm_args["pin_memory"] = True
        dm_args["featurization_n_jobs"] = 2
        dm_args["featurization_progress"] = True
        dm_args["featurization_backend"] = "loky"
        dm_args["featurization_batch_size"] = 50

        # Delete the cache if already exist
        if exists(TEMP_CACHE_DATA_PATH):
            rm(TEMP_CACHE_DATA_PATH, recursive=True)

        # Prepare the data. It should create the cache there
        assert not exists(TEMP_CACHE_DATA_PATH)
        ds = GraphOGBDataModule(task_specific_args, cache_data_path=TEMP_CACHE_DATA_PATH, **dm_args)
        assert not ds.load_data_from_cache(verbose=False)
        ds.prepare_data()

        # Check the keys in the dataset
        ds.setup(save_smiles_and_ids=False)
        assert set(ds.train_ds[0].keys()) == {"features", "labels"}

        ds.setup(save_smiles_and_ids=True)
        assert set(ds.train_ds[0].keys()) == {"smiles", "mol_ids", "features", "labels"}

        # Make sure that the cache is created
        full_cache_path = ds.get_data_cache_fullname(compress=False)
        assert exists(full_cache_path)
        assert get_size(full_cache_path) > 10000

        # Check that the data is loaded correctly from cache
        assert ds.load_data_from_cache(verbose=False)

        # test module
        assert ds.num_edge_feats == 5
        assert ds.num_node_feats == 50
        assert len(ds) == 642

        # test batch loader
        batch = next(iter(ds.train_dataloader()))
        assert len(batch["smiles"]) == 16
        assert len(batch["labels"]["graph_task_1"]) == 16
        assert len(batch["mol_ids"]) == 16

    def test_datamodule_with_none_molecules(self):
        # Setup the featurization
        featurization_args = {}
        featurization_args["atom_property_list_float"] = []  # ["weight", "valence"]
        featurization_args["atom_property_list_onehot"] = ["atomic-number", "degree"]
        featurization_args["edge_property_list"] = ["bond-type-onehot"]

        # Config for datamodule
        bad_csv = "tests/data/micro_ZINC_corrupt.csv"
        task_specific_args = {}
        task_kwargs = {"df_path": bad_csv, "split_val": 0.0, "split_test": 0.0}
        task_specific_args["task_1"] = {
            "task_level": "graph",
            "label_cols": "SA",
            "smiles_col": "SMILES1",
            **task_kwargs,
        }
        task_specific_args["task_2"] = {
            "task_level": "graph",
            "label_cols": "logp",
            "smiles_col": "SMILES2",
            **task_kwargs,
        }
        task_specific_args["task_3"] = {
            "task_level": "graph",
            "label_cols": "score",
            "smiles_col": "SMILES3",
            **task_kwargs,
        }

        # Read the corrupted dataset and get stats
        df = pd.read_csv(bad_csv)
        bad_smiles = (df["SMILES1"] == "XXX") & (df["SMILES2"] == "XXX") & (df["SMILES3"] == "XXX")
        num_bad_smiles = sum(bad_smiles)

        # Test the datamodule
        datamodule = MultitaskFromSmilesDataModule(
            task_specific_args=task_specific_args,
            featurization_args=featurization_args,
            featurization_n_jobs=0,
            featurization_batch_size=1,
        )
        datamodule.prepare_data()
        datamodule.setup(save_smiles_and_ids=True)

        # Check that the number of molecules is correct
        smiles = df["SMILES1"].tolist() + df["SMILES2"].tolist() + df["SMILES3"].tolist()
        num_unique_smiles = len(set(smiles)) - 1  # -1 because of the XXX
        # self.assertEqual(len(datamodule.train_ds), num_unique_smiles - num_bad_smiles)

        # Change the index of the dataframe
        index_smiles = []
        for ii in range(len(df)):
            if df["SMILES1"][ii] != "XXX":
                smiles = df["SMILES1"][ii]
            elif df["SMILES2"][ii] != "XXX":
                smiles = df["SMILES2"][ii]
            elif df["SMILES3"][ii] != "XXX":
                smiles = df["SMILES3"][ii]
            else:
                smiles = "XXX"
            index_smiles.append(smiles)
        df["idx_smiles"] = index_smiles
        df = df.set_index("idx_smiles")

        # Convert the smilies from the train_ds to a list, and check the content
        train_smiles = [d["smiles"] for d in datamodule.train_ds]

        # Check that the set of smiles are the same
        train_smiles_flat = list(set([item for sublist in train_smiles for item in sublist]))
        train_smiles_flat.sort()
        index_smiles_filt = list(set([smiles for smiles in index_smiles if smiles != "XXX"]))
        index_smiles_filt.sort()
        self.assertListEqual(train_smiles_flat, index_smiles_filt)

        # Check that the smiles are correct for each datapoint in the dataset
        for smiles in train_smiles:
            self.assertEqual(len(set(smiles)), 1)  # Check that all smiles are the same
            this_smiles = smiles[0]
            true_smiles = df.loc[this_smiles][["SMILES1", "SMILES2", "SMILES3"]]
            num_true_smiles = sum(true_smiles != "XXX")
            self.assertEqual(len(smiles), num_true_smiles)  # Check that the number of smiles is correct
            self.assertEqual(
                this_smiles, true_smiles[true_smiles != "XXX"].values[0]
            )  # Check that the smiles are correct

        # Convert the labels from the train_ds to a dataframe
        train_labels = [{task: val[0] for task, val in d["labels"].items()} for d in datamodule.train_ds]
        train_labels_df = pd.DataFrame(train_labels)
        train_labels_df = train_labels_df.rename(
            columns={"graph_task_1": "graph_SA", "graph_task_2": "graph_logp", "graph_task_3": "graph_score"}
        )
        train_labels_df["smiles"] = [s[0] for s in datamodule.train_ds.smiles]
        train_labels_df = train_labels_df.set_index("smiles")
        train_labels_df = train_labels_df.sort_index()

        # Check that the labels are correct
        df2 = df.reset_index()[~bad_smiles].set_index("idx_smiles").sort_index()
        labels = train_labels_df[["graph_SA", "graph_logp", "graph_score"]].values
        nans = np.isnan(labels)
        true_nans = df2[["SMILES1", "SMILES2", "SMILES3"]].values == "XXX"
        true_labels = df2[["SA", "logp", "score"]].values
        true_labels[true_nans] = np.nan
        np.testing.assert_array_equal(nans, true_nans)  # Check that the nans are correct
        np.testing.assert_array_almost_equal(
            labels, true_labels, decimal=5
        )  # Check that the label values are correct

    def test_datamodule_multiple_data_files(self):
        # Test single CSV files
        csv_file = "tests/data/micro_ZINC_shard_1.csv"
        task_kwargs = {"df_path": csv_file, "split_val": 0.0, "split_test": 0.0}
        task_specific_args = {
            "task": {"task_level": "graph", "label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}
        }

        ds = MultitaskFromSmilesDataModule(task_specific_args)
        ds.prepare_data()
        ds.setup()

        self.assertEqual(len(ds.train_ds), 10)

        # Test multi CSV files
        csv_file = "tests/data/micro_ZINC_shard_*.csv"
        task_kwargs = {"df_path": csv_file, "split_val": 0.0, "split_test": 0.0}
        task_specific_args = {
            "task": {"task_level": "graph", "label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}
        }

        ds = MultitaskFromSmilesDataModule(task_specific_args)
        ds.prepare_data()
        ds.setup()

        self.assertEqual(len(ds.train_ds), 20)

        # Test single Parquet files
        parquet_file = "tests/data/micro_ZINC_shard_1.parquet"
        task_kwargs = {"df_path": parquet_file, "split_val": 0.0, "split_test": 0.0}
        task_specific_args = {
            "task": {"task_level": "graph", "label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}
        }

        ds = MultitaskFromSmilesDataModule(task_specific_args)
        ds.prepare_data()
        ds.setup()

        self.assertEqual(len(ds.train_ds), 10)

        # Test multi Parquet files
        parquet_file = "tests/data/micro_ZINC_shard_*.parquet"
        task_kwargs = {"df_path": parquet_file, "split_val": 0.0, "split_test": 0.0}
        task_specific_args = {
            "task": {"task_level": "graph", "label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}
        }

        ds = MultitaskFromSmilesDataModule(task_specific_args)
        ds.prepare_data()
        ds.setup()

        self.assertEqual(len(ds.train_ds), 20)


if __name__ == "__main__":
    ut.main()

    # Delete the cache
    if exists(TEMP_CACHE_DATA_PATH):
        rm(TEMP_CACHE_DATA_PATH, recursive=True)


# tests/test_dataset.py

import unittest as ut

from graphium.data import load_micro_zinc
from graphium.data.dataset import SingleTaskDataset, MultitaskDataset
from graphium.data.smiles_transform import smiles_to_unique_mol_ids


class Test_Multitask_Dataset(ut.TestCase):
    # Then we can choose different rows and columns for the tests as we see fit.
    # Remember tests are supposed to be FAST, and reading from the file system multiple times slows things down.

    # Make sure that the inputs to single task datasets are always lists!
    # Do not pass a data frame itself, but turn it into a list to satisfy the type required.

    def test_multitask_dataset_case_1(self):
        """Case: different tasks, all with the same smiles set.
        - Check that for each task, all smiles are received from the initial DF.
        - Check that for each task, you have the same label values as the initial DF.
        """

        df_micro_zinc = load_micro_zinc()  # Has about 1000 molecules
        df = df_micro_zinc.iloc[0:4]
        num_unique_mols = 4

        # Here we take the microzinc dataset and split the labels up into 'SA', 'logp' and 'score' in order to simulate having multiple single-task datasets
        df_micro_zinc_SA = df[["SMILES", "SA"]]
        df_micro_zinc_logp = df[["SMILES", "logp"]]
        df_micro_zinc_score = df[["SMILES", "score"]]

        # We need to turn these dataframes into single-task datasets.
        # We don't need to do featurization yet.
        ds_micro_zinc_SA = SingleTaskDataset(
            smiles=df_micro_zinc_SA.loc[:, "SMILES"].tolist(), labels=df_micro_zinc_SA.loc[:, "SA"].tolist()
        )

        ds_micro_zinc_logp = SingleTaskDataset(
            smiles=df_micro_zinc_logp.loc[:, "SMILES"].tolist(),
            labels=df_micro_zinc_logp.loc[:, "logp"].tolist(),
        )
        ds_micro_zinc_score = SingleTaskDataset(
            smiles=df_micro_zinc_score.loc[:, "SMILES"].tolist(),
            labels=df_micro_zinc_score.loc[:, "score"].tolist(),
        )

        # Create the multitask dataset
        datasets_dict = {"SA": ds_micro_zinc_SA, "logp": ds_micro_zinc_logp, "score": ds_micro_zinc_score}
        multitask_microzinc = MultitaskDataset(
            datasets_dict, save_smiles_and_ids=True
        )  # Can optionally have features

        # Check: The number of unique molecules equals the number of datapoints in the multitask dataset.
        self.assertEqual(num_unique_mols, multitask_microzinc.__len__())

        # Check that for each task, you have the same label values as the initial DF.
        for idx in range(multitask_microzinc.__len__()):
            smiles = df[["SMILES"]].iloc[idx].values[0]
            # label = df[['SA']].iloc[idx]
            label_SA = ds_micro_zinc_SA.labels[idx]
            label_logp = ds_micro_zinc_logp.labels[idx]
            label_score = ds_micro_zinc_score.labels[idx]

            # Search for the mol id in the multitask dataset
            mol_ids = smiles_to_unique_mol_ids([smiles])
            mol_id = mol_ids[0]
            found_idx = -1
            for i, id in enumerate(multitask_microzinc.mol_ids):
                if mol_id == id:
                    found_idx = i

            # Compare labels
            self.assertEqual(label_SA, multitask_microzinc.labels[found_idx]["SA"])
            self.assertEqual(label_logp, multitask_microzinc.labels[found_idx]["logp"])
            self.assertEqual(label_score, multitask_microzinc.labels[found_idx]["score"])

    def test_multitask_dataset_case_2(self):
        """Case: Different tasks, but with no intersection in the smiles (each task has a unique set of smiles)
        - Check that the total dataset has as much smiles as all tasks together
        - Check that, for each task, only the smiles related to that task have values, and ensure the value is what's expected from the initial DF
        """
        df = load_micro_zinc()  # Has about 1000 molecules

        # Choose non-overlapping smiles by choosing specific rows from the original dataframe.
        df_rows_SA = df.iloc[0:200]  # 200 data points
        df_rows_logp = df.iloc[200:400]  # 200 data points
        df_rows_score = df.iloc[400:750]  # 350 data points
        total_data_points = 750

        # Here we split the data according to the task we care about.
        df_micro_zinc_SA = df_rows_SA[["SMILES", "SA"]]
        df_micro_zinc_logp = df_rows_logp[["SMILES", "logp"]]
        df_micro_zinc_score = df_rows_score[["SMILES", "score"]]

        # We need to turn these dataframes into single-task datasets.
        # We don't need to do featurization yet.
        ds_micro_zinc_SA = SingleTaskDataset(
            smiles=df_micro_zinc_SA.loc[:, "SMILES"].tolist(), labels=df_micro_zinc_SA.loc[:, "SA"].tolist()
        )
        ds_micro_zinc_logp = SingleTaskDataset(
            smiles=df_micro_zinc_logp.loc[:, "SMILES"].tolist(),
            labels=df_micro_zinc_logp.loc[:, "logp"].tolist(),
        )
        ds_micro_zinc_score = SingleTaskDataset(
            smiles=df_micro_zinc_score.loc[:, "SMILES"].tolist(),
            labels=df_micro_zinc_score.loc[:, "score"].tolist(),
        )

        # Create the multitask dataset
        datasets_dict = {"SA": ds_micro_zinc_SA, "logp": ds_micro_zinc_logp, "score": ds_micro_zinc_score}
        multitask_microzinc = MultitaskDataset(
            datasets_dict, save_smiles_and_ids=True
        )  # Can optionally have features

        # The total dataset has as many molecules as there are smiles in all tasks put together
        self.assertEqual(total_data_points, multitask_microzinc.__len__())

        # For each task, only the smiles related to that task have values, and the value is what's expected from the initial DF.
        for idx in range(len(ds_micro_zinc_SA)):
            smiles = df[["SMILES"]].iloc[idx].values[0]

            task = "task"
            if idx in range(0, 200):
                task = "SA"
            elif idx in range(200, 400):
                task = "logp"
            elif idx in range(400, 750):
                task = "score"

            # Labels of that molecule
            label_SA = df[["SA"]].iloc[idx].values[0]
            label_logp = df[["logp"]].iloc[idx].values[0]
            label_score = df[["score"]].iloc[idx].values[0]

            # Search for that molecule in the multitask dataset
            mol_ids = smiles_to_unique_mol_ids([smiles])
            mol_id = mol_ids[0]
            found_idx = -1
            for i, id in enumerate(multitask_microzinc.mol_ids):
                if mol_id == id:
                    found_idx = i

            if task == "SA":
                self.assertEqual(label_SA, multitask_microzinc.labels[found_idx]["SA"])
                self.assertFalse("score" in multitask_microzinc.labels[found_idx].keys)
                self.assertFalse("logp" in multitask_microzinc.labels[found_idx].keys)
            elif task == "logp":
                self.assertEqual(label_logp, multitask_microzinc.labels[found_idx]["logp"])
                self.assertFalse("score" in multitask_microzinc.labels[found_idx].keys)
                self.assertFalse("SA" in multitask_microzinc.labels[found_idx].keys)
            elif task == "score":
                self.assertEqual(label_score, multitask_microzinc.labels[found_idx]["score"])
                self.assertFalse("SA" in multitask_microzinc.labels[found_idx].keys)
                self.assertFalse("logp" in multitask_microzinc.labels[found_idx].keys)

    def test_multitask_dataset_case_3(self):
        """Case: Different tasks, but with semi-intersection (some smiles unique per task, some intersect)
        - Check that the total dataset has as much smiles as the unique number of smiles.
        - Check that for each task, you retrieve the same smiles as expected from the initial DF
        """
        df_micro_zinc = load_micro_zinc()  # Has about 1000 molecules
        df = df_micro_zinc.iloc[0:5]

        # Choose OVERLAPPING smiles by choosing specific rows from the original dataframe. The tasks will not necessarily have unique smiles.
        df_rows_SA = df.iloc[0:3]
        df_rows_logp = df.iloc[1:4]
        df_rows_score = df.iloc[3:5]
        total_data_points = 5

        # Here we split the data according to the task we care about.
        df_micro_zinc_SA = df_rows_SA[["SMILES", "SA"]]
        df_micro_zinc_logp = df_rows_logp[["SMILES", "logp"]]
        df_micro_zinc_score = df_rows_score[["SMILES", "score"]]

        # We need to turn these dataframes into single-task datasets.
        # We don't need to do featurization yet.
        ds_micro_zinc_SA = SingleTaskDataset(
            smiles=df_micro_zinc_SA.loc[:, "SMILES"].tolist(), labels=df_micro_zinc_SA.loc[:, "SA"].tolist()
        )
        ds_micro_zinc_logp = SingleTaskDataset(
            smiles=df_micro_zinc_logp.loc[:, "SMILES"].tolist(),
            labels=df_micro_zinc_logp.loc[:, "logp"].tolist(),
        )
        ds_micro_zinc_score = SingleTaskDataset(
            smiles=df_micro_zinc_score.loc[:, "SMILES"].tolist(),
            labels=df_micro_zinc_score.loc[:, "score"].tolist(),
        )

        # Create the multitask dataset
        datasets_dict = {"SA": ds_micro_zinc_SA, "logp": ds_micro_zinc_logp, "score": ds_micro_zinc_score}
        multitask_microzinc = MultitaskDataset(
            datasets_dict, save_smiles_and_ids=True
        )  # Can optionally have features

        # The multitask dataset has as many molecules as there are unique smiles across the single task datasets.
        self.assertEqual(total_data_points, multitask_microzinc.__len__())


if __name__ == "__main__":
    ut.main()


# tests/test_dict_tensor.py

"""
Unit tests for the metrics and wrappers of graphium/utils/...
"""

from typing import Dict
from graphium.utils.dict_tensor import DictTensor
import torch
from torch import Tensor

import unittest as ut
from copy import deepcopy


class test_dict_tensor(ut.TestCase):
    def _assert_dict_tensor_equal(self, dict_ten1, dict_ten2, msg=""):
        self.assertSetEqual(set(dict_ten1.keys()), set(dict_ten2.keys()), msg=msg)
        for key, tensor1 in dict_ten1.items():
            tensor2 = dict_ten2[key]
            if isinstance(tensor1, Tensor):
                msg2 = msg + f"\ntensor1: \n{tensor1}\n\ntensor2: \n{tensor2}\n\n"
                self.assertTrue(torch.all(tensor1 == tensor2), msg=msg2)
            else:
                for ii in range(len(tensor1)):
                    msg2 = msg + f"\ntensor1: \n{tensor1[ii]}\n\ntensor2: \n{tensor2[ii]}\n\n"
                    self.assertTrue(torch.all(tensor1[ii] == tensor2[ii]), msg=msg2)

    def test_tensor_funcs(self):
        dict_ten = DictTensor(
            {
                "a": torch.randn(5, 6),
                "b": torch.randn(2, 3, 4),
                "c": torch.randn(1, 2, 3, 5),
            }
        )

        # Check `DictTensor.func` returns right value
        self.assertDictEqual(dict_ten.shape, {k: v.shape for k, v in dict_ten.items()})
        self._assert_dict_tensor_equal(
            dict_ten.pinverse(), {k: v.pinverse() for k, v in dict_ten.items()}, msg="dict_ten.pinverse()"
        )
        self._assert_dict_tensor_equal(
            dict_ten.sum(), {k: v.sum() for k, v in dict_ten.items()}, msg="dict_ten.sum()"
        )
        self._assert_dict_tensor_equal(
            dict_ten.sum(axis=0), {k: v.sum(axis=0) for k, v in dict_ten.items()}, msg="dict_ten.sum(axis=0)"
        )
        self._assert_dict_tensor_equal(
            dict_ten.sum(0), {k: v.sum(0) for k, v in dict_ten.items()}, msg="dict_ten.sum(0)"
        )
        self._assert_dict_tensor_equal(
            dict_ten.sum(axis=1), {k: v.sum(axis=1) for k, v in dict_ten.items()}, msg="dict_ten.sum(axis=1)"
        )
        self._assert_dict_tensor_equal(
            dict_ten.to(dtype=int),
            {k: v.to(dtype=int) for k, v in dict_ten.items()},
            msg="dict_ten.to(dtype=int)",
        )
        self._assert_dict_tensor_equal(
            dict_ten.max(axis=0),
            {k: v.max(axis=0) for k, v in dict_ten.items()},
            msg="dict_ten.max(axis=0)[0]",
        )
        self._assert_dict_tensor_equal(
            dict_ten.transpose(0, 1),
            {k: v.transpose(0, 1) for k, v in dict_ten.items()},
            msg="dict_ten.max(axis=0)[0]",
        )

        # Check `DictTensor.func` returns right type
        self.assertIsInstance(dict_ten.shape, dict)
        self.assertIsInstance(dict_ten.shape["a"], torch.Size)
        self.assertIsInstance(dict_ten.pinverse(), DictTensor)
        self.assertIsInstance(dict_ten.sum(), DictTensor)
        self.assertIsInstance(dict_ten.sum(axis=0), DictTensor)
        self.assertIsInstance(dict_ten.sum(0), DictTensor)
        self.assertIsInstance(dict_ten.sum(axis=1), DictTensor)
        self.assertIsInstance(dict_ten.to(dtype=int), DictTensor)
        self.assertIsInstance(dict_ten.max(axis=0), dict)
        self.assertIsInstance(dict_ten.max(axis=0)["a"][0], Tensor)

    def test_torch_funcs(self):
        dict_ten = DictTensor(
            {
                "a": torch.randn(5, 6),
                "b": torch.randn(2, 3, 4),
                "c": torch.randn(1, 2, 3, 5),
            }
        )

        # Check `torch.func(DictTensor, *)` returns right value
        self._assert_dict_tensor_equal(
            torch.pinverse(
                dict_ten,
            ),
            {
                k: torch.pinverse(
                    v,
                )
                for k, v in dict_ten.items()
            },
            msg="torch.pinverse(dict_ten, )",
        )
        self._assert_dict_tensor_equal(
            torch.sum(
                dict_ten,
            ),
            {
                k: torch.sum(
                    v,
                )
                for k, v in dict_ten.items()
            },
            msg="torch.sum(dict_ten, )",
        )
        self._assert_dict_tensor_equal(
            torch.sum(dict_ten, 0),
            {k: torch.sum(v, 0) for k, v in dict_ten.items()},
            msg="torch.sum(dict_ten, 0)",
        )
        self._assert_dict_tensor_equal(
            torch.sum(dict_ten, axis=0),
            {k: torch.sum(v, axis=0) for k, v in dict_ten.items()},
            msg="torch.sum(dict_ten, axis=0)",
        )
        self._assert_dict_tensor_equal(
            torch.sum(dict_ten, axis=1),
            {k: torch.sum(v, axis=1) for k, v in dict_ten.items()},
            msg="torch.sum(dict_ten, axis=1)",
        )
        self._assert_dict_tensor_equal(
            torch.max(dict_ten, axis=0),
            {k: torch.max(v, axis=0) for k, v in dict_ten.items()},
            msg="torch.max(dict_ten, axis=0)[0]",
        )

        # Check `torch.func(DictTensor, *)` returns right type
        self.assertIsInstance(
            torch.pinverse(
                dict_ten,
            ),
            DictTensor,
        )
        self.assertIsInstance(
            torch.sum(
                dict_ten,
            ),
            DictTensor,
        )
        self.assertIsInstance(torch.sum(dict_ten, axis=0), DictTensor)
        self.assertIsInstance(torch.sum(dict_ten, 0), DictTensor)
        self.assertIsInstance(torch.sum(dict_ten, axis=1), DictTensor)
        self.assertIsInstance(torch.max(dict_ten, axis=0), dict)
        self.assertIsInstance(torch.max(dict_ten, axis=0)["a"][0], Tensor)

    def test_apply(self):
        dict_ten = DictTensor(
            {
                "a": torch.randn(5, 6),
                "b": torch.randn(2, 3, 4),
                "c": torch.randn(1, 2, 3, 5),
            }
        )

        # Check `DictTensor.apply` function returns right value
        self._assert_dict_tensor_equal(
            dict_ten.apply(
                torch.pinverse,
            ),
            {
                k: torch.pinverse(
                    v,
                )
                for k, v in dict_ten.items()
            },
            msg="torch.pinverse(dict_ten, )",
        )
        self._assert_dict_tensor_equal(
            dict_ten.apply(
                torch.sum,
            ),
            {
                k: torch.sum(
                    v,
                )
                for k, v in dict_ten.items()
            },
            msg="torch.sum(dict_ten, )",
        )
        self._assert_dict_tensor_equal(
            dict_ten.apply(torch.sum, 0),
            {k: torch.sum(v, 0) for k, v in dict_ten.items()},
            msg="torch.sum(dict_ten, 0)",
        )
        self._assert_dict_tensor_equal(
            dict_ten.apply(torch.sum, axis=0),
            {k: torch.sum(v, axis=0) for k, v in dict_ten.items()},
            msg="torch.sum(dict_ten, axis=0)",
        )
        self._assert_dict_tensor_equal(
            dict_ten.apply(torch.sum, axis=1),
            {k: torch.sum(v, axis=1) for k, v in dict_ten.items()},
            msg="torch.sum(dict_ten, axis=1)",
        )
        self._assert_dict_tensor_equal(
            dict_ten.apply(torch.max, axis=0),
            {k: torch.max(v, axis=0) for k, v in dict_ten.items()},
            msg="torch.max(dict_ten, axis=0)[0]",
        )

        # Check `DictTensor.apply` function returns right type
        self.assertIsInstance(
            dict_ten.apply(
                torch.pinverse,
            ),
            DictTensor,
        )
        self.assertIsInstance(
            dict_ten.apply(
                torch.sum,
            ),
            DictTensor,
        )
        self.assertIsInstance(dict_ten.apply(torch.sum, axis=0), DictTensor)
        self.assertIsInstance(dict_ten.apply(torch.sum, 0), DictTensor)
        self.assertIsInstance(dict_ten.apply(torch.sum, axis=1), DictTensor)
        self.assertIsInstance(dict_ten.apply(torch.max, axis=0), dict)
        self.assertIsInstance(dict_ten.apply(torch.max, axis=0)["a"][0], Tensor)

    def test_tensor_operators(self):
        dict_ten = DictTensor(
            {
                "a": torch.randn(5, 6),
                "b": torch.randn(4, 5, 6),
                "c": torch.randn(1, 2, 5, 6),
            }
        )

        # Check product, division, sum, subtraction by a constant
        FACTOR = 0.5
        self._assert_dict_tensor_equal(
            dict_ten * FACTOR, {k: FACTOR * v for k, v in dict_ten.items()}, msg="DT * factor"
        )
        self._assert_dict_tensor_equal(
            FACTOR * dict_ten, {k: FACTOR * v for k, v in dict_ten.items()}, msg="factor * DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten + FACTOR, {k: FACTOR + v for k, v in dict_ten.items()}, msg="DT + factor"
        )
        self._assert_dict_tensor_equal(
            FACTOR + dict_ten, {k: FACTOR + v for k, v in dict_ten.items()}, msg="factor + DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten - FACTOR, {k: v - FACTOR for k, v in dict_ten.items()}, msg="DT - factor"
        )
        self._assert_dict_tensor_equal(
            FACTOR - dict_ten, {k: FACTOR - v for k, v in dict_ten.items()}, msg="factor - DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten / FACTOR, {k: v / FACTOR for k, v in dict_ten.items()}, msg="DT / factor"
        )
        self._assert_dict_tensor_equal(
            dict_ten // FACTOR, {k: v // FACTOR for k, v in dict_ten.items()}, msg="DT // factor"
        )

        # Check product, division, sum, subtraction by a Tensor of shape [6]
        FACTOR = torch.rand(6)
        self._assert_dict_tensor_equal(
            dict_ten * FACTOR, {k: FACTOR * v for k, v in dict_ten.items()}, msg="DT * tensor"
        )
        self._assert_dict_tensor_equal(
            FACTOR * dict_ten, {k: FACTOR * v for k, v in dict_ten.items()}, msg="tensor * DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten + FACTOR, {k: FACTOR + v for k, v in dict_ten.items()}, msg="DT + tensor"
        )
        self._assert_dict_tensor_equal(
            FACTOR + dict_ten, {k: FACTOR + v for k, v in dict_ten.items()}, msg="tensor + DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten - FACTOR, {k: v - FACTOR for k, v in dict_ten.items()}, msg="DT - tensor"
        )
        self._assert_dict_tensor_equal(
            FACTOR - dict_ten, {k: FACTOR - v for k, v in dict_ten.items()}, msg="tensor - DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten / FACTOR, {k: v / FACTOR for k, v in dict_ten.items()}, msg="DT / tensor"
        )
        self._assert_dict_tensor_equal(
            dict_ten // FACTOR, {k: v // FACTOR for k, v in dict_ten.items()}, msg="DT // tensor"
        )

        # Check product, division, sum, subtraction by a Tensor of shape [5, 6]
        FACTOR = torch.rand(5, 6)
        self._assert_dict_tensor_equal(
            dict_ten * FACTOR, {k: FACTOR * v for k, v in dict_ten.items()}, msg="DT * tensor"
        )
        self._assert_dict_tensor_equal(
            FACTOR * dict_ten, {k: FACTOR * v for k, v in dict_ten.items()}, msg="tensor * DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten + FACTOR, {k: FACTOR + v for k, v in dict_ten.items()}, msg="DT + tensor"
        )
        self._assert_dict_tensor_equal(
            FACTOR + dict_ten, {k: FACTOR + v for k, v in dict_ten.items()}, msg="tensor + DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten - FACTOR, {k: v - FACTOR for k, v in dict_ten.items()}, msg="DT - tensor"
        )
        self._assert_dict_tensor_equal(
            FACTOR - dict_ten, {k: FACTOR - v for k, v in dict_ten.items()}, msg="tensor - DT"
        )
        self._assert_dict_tensor_equal(
            dict_ten / FACTOR, {k: v / FACTOR for k, v in dict_ten.items()}, msg="DT / tensor"
        )
        self._assert_dict_tensor_equal(
            dict_ten // FACTOR, {k: v // FACTOR for k, v in dict_ten.items()}, msg="DT // tensor"
        )

    def test_comparison_operators(self):
        dict_ten = DictTensor(
            {
                "a": torch.randn(5, 6),
                "b": torch.randn(4, 5, 6),
                "c": torch.randn(1, 2, 5, 6),
            }
        )
        dict_ten2 = deepcopy(dict_ten).abs()

        # Comparison operators with float
        self._assert_dict_tensor_equal(
            dict_ten > 0.2, {k: v > 0.2 for k, v in dict_ten.items()}, msg="DT > 0.2"
        )
        self._assert_dict_tensor_equal(
            dict_ten < 0.2, {k: v < 0.2 for k, v in dict_ten.items()}, msg="DT < 0.2"
        )
        self._assert_dict_tensor_equal(
            dict_ten >= 0.2, {k: v >= 0.2 for k, v in dict_ten.items()}, msg="DT >= 0.2"
        )
        self._assert_dict_tensor_equal(
            dict_ten <= 0.2, {k: v <= 0.2 for k, v in dict_ten.items()}, msg="DT <= 0.2"
        )
        self._assert_dict_tensor_equal(
            dict_ten.to(int) == 0, {k: v.to(int) == 0 for k, v in dict_ten.items()}, msg="DT.int == 0"
        )
        self._assert_dict_tensor_equal(
            dict_ten.to(int) != 0, {k: v.to(int) != 0 for k, v in dict_ten.items()}, msg="DT.int != 0"
        )

        # Comparison operators with Tensor
        tensor = torch.rand(5, 6)
        self._assert_dict_tensor_equal(
            dict_ten > tensor, {k: v > tensor for k, v in dict_ten.items()}, msg="DT > DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten < tensor, {k: v < tensor for k, v in dict_ten.items()}, msg="DT < DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten >= tensor, {k: v >= tensor for k, v in dict_ten.items()}, msg="DT >= DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten <= tensor, {k: v <= tensor for k, v in dict_ten.items()}, msg="DT <= DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten.to(int) == tensor,
            {k: v.to(int) == tensor for k, v in dict_ten.items()},
            msg="DT.int == DT2",
        )
        self._assert_dict_tensor_equal(
            dict_ten.to(int) != tensor,
            {k: v.to(int) != tensor for k, v in dict_ten.items()},
            msg="DT.int != DT2",
        )

        # Comparison operators with DictTensor
        self._assert_dict_tensor_equal(
            dict_ten > dict_ten2, {k: v > dict_ten2[k] for k, v in dict_ten.items()}, msg="DT > DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten < dict_ten2, {k: v < dict_ten2[k] for k, v in dict_ten.items()}, msg="DT < DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten >= dict_ten2, {k: v >= dict_ten2[k] for k, v in dict_ten.items()}, msg="DT >= DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten <= dict_ten2, {k: v <= dict_ten2[k] for k, v in dict_ten.items()}, msg="DT <= DT2"
        )
        self._assert_dict_tensor_equal(
            dict_ten.to(int) == dict_ten2,
            {k: v.to(int) == dict_ten2[k] for k, v in dict_ten.items()},
            msg="DT.int == DT2",
        )
        self._assert_dict_tensor_equal(
            dict_ten.to(int) != dict_ten2,
            {k: v.to(int) != dict_ten2[k] for k, v in dict_ten.items()},
            msg="DT.int != DT2",
        )

    def test_dict_functions(self):
        dict1 = {
            "a": torch.randn(5, 6),
            "b": torch.randn(4, 5, 6),
            "c": torch.randn(1, 2, 5, 6),
        }
        dict_ten1 = DictTensor(deepcopy(dict1))
        dict2 = {
            "c": torch.randn(1, 2, 6),
            "d": torch.randn(1, 2, 3, 4),
        }
        dict_ten2 = DictTensor(deepcopy(dict2))

        # Check update
        dict1_temp = deepcopy(dict1)
        dict_ten1_temp = deepcopy(dict_ten1)
        dict1_temp.update(dict2)
        truth = DictTensor(dict1_temp)
        dict_ten1_temp.update(dict_ten2)
        self._assert_dict_tensor_equal(dict_ten1_temp, truth)

        # Check pop
        dict1_temp = deepcopy(dict1)
        dict_ten1_temp = deepcopy(dict_ten1)
        truth_a = dict1_temp.pop("a")
        truth = DictTensor(dict1_temp)
        val_a = dict_ten1_temp.pop("a")
        self._assert_dict_tensor_equal(dict_ten1_temp, truth)
        self.assertTrue(torch.all(truth_a == val_a))

        # Check indexing
        dict1_temp = deepcopy(dict1)
        dict_ten1_temp = deepcopy(dict_ten1)
        truth_b = dict1_temp["b"]
        val_b = dict_ten1_temp["b"]
        self.assertTrue(torch.all(truth_b == val_b))

        # Check adding element
        dict1_temp = deepcopy(dict1)
        dict_ten1_temp = deepcopy(dict_ten1)
        dict1_temp["new"] = dict2["d"]
        truth = DictTensor(dict1_temp)
        dict_ten1_temp["new"] = dict_ten2["d"]
        self._assert_dict_tensor_equal(dict_ten1_temp, truth)


if __name__ == "__main__":
    ut.main()


# tests/test_featurizer.py

"""
Unit tests for the different datasets of graphium/features/featurizer.py
"""

import numpy as np
import unittest as ut
from copy import deepcopy
from rdkit import Chem
import datamol as dm

from graphium.features.featurizer import (
    get_mol_atomic_features_onehot,
    get_mol_atomic_features_float,
    get_mol_edge_features,
    mol_to_adj_and_features,
    mol_to_pyggraph,
)


class test_featurizer(ut.TestCase):
    smiles = [
        "C",
        "CC",
        "C1(C[N]CCC1)=O",
        "CC(C)CC1=CC=C(C=C1)C(C)C(=O)O",
        "OCCc1c(C)[n+](cs1)Cc2cnc(C)nc2N",
        "O1C=C[C@H]([C@H]1O2)c3c2cc(OC)c4c3OC(=O)C5=C4CCC(=O)5",
        "CC1(C2C1C(N(C2)C(=O)C(C(C)=B)NC(=O)C(F)(F)F)C(=O)NC(C(C3CCNC3=O)[Cl])C#N)C",
    ]

    smiles_noble = ["[He].[He]", "[He][He]", "[Kr][Kr]"]

    atomic_onehot_props = [
        "atomic-number",
        "valence",
        "degree",
        "implicit-valence",
        "hybridization",
        "chirality",
        "phase",
        "type",
        "group",
        "period",
    ]

    atomic_float_props = [
        "atomic-number",
        "mass",
        "valence",
        "implicit-valence",
        "hybridization",
        "chirality",
        "aromatic",
        "in-ring",
        "min-ring",
        "max-ring",
        "num-ring",
        "degree",
        "radical-electron",
        "formal-charge",
        "vdw-radius",
        "covalent-radius",
        "electronegativity",
        "ionization",
        "melting-point",
        "metal",
        "single-bond",
        "aromatic-bond",
        "double-bond",
        "triple-bond",
        "is-carbon",
        "group",
        "period",
    ]

    edge_props = [
        "bond-type-onehot",
        "bond-type-float",
        "stereo",
        "in-ring",
        "conjugated",
        "estimated-bond-length",
        "conformer-bond-length",
    ]

    def test_get_mol_atomic_features_onehot(self):
        props = deepcopy(self.atomic_onehot_props)
        bad_props = ["bob"]

        all_smiles = self.smiles + self.smiles_noble

        for s in all_smiles:
            err_msg = f"\n\tError for params:\n\t\tSMILES: {s}"
            mol = dm.to_mol(s)

            for ii in range(len(props)):
                this_props = props[:ii]
                err_msg2 = err_msg + f"\n\t\tprops: {this_props}"
                prop_dict = get_mol_atomic_features_onehot(mol, property_list=this_props)
                self.assertListEqual(list(prop_dict.keys()), this_props, msg=err_msg)
                for key, val in prop_dict.items():
                    err_msg3 = err_msg2 + f"\n\t\tkey: {key}"
                    self.assertEqual(val.shape[0], mol.GetNumAtoms(), msg=err_msg3)
                    self.assertGreater(val.shape[1], 1, msg=err_msg3)
                    self.assertTrue(np.all((val == 0) | (val == 1)), msg=err_msg3)

            with self.assertRaises(ValueError, msg=err_msg):
                get_mol_atomic_features_onehot(mol, property_list=bad_props)

    def test_get_mol_atomic_features_float(self):
        props = deepcopy(self.atomic_float_props)

        bad_props = ["bob"]

        all_smiles = self.smiles + self.smiles_noble
        for s in all_smiles:
            err_msg = f"\n\tError for params:\n\t\tSMILES: {s}"
            mol = dm.to_mol(s)

            for ii in range(len(props)):
                this_props = props[:ii]
                err_msg2 = err_msg + f"\n\t\tprops: {this_props}"
                prop_dict = get_mol_atomic_features_float(mol, property_list=this_props, mask_nan=None)
                self.assertListEqual(list(prop_dict.keys()), this_props, msg=err_msg)
                for key, val in prop_dict.items():
                    err_msg3 = err_msg2 + f"\n\t\tkey: {key}"
                    self.assertListEqual(list(val.shape), [mol.GetNumAtoms()], msg=err_msg3)

            with self.assertRaises(ValueError, msg=err_msg):
                get_mol_atomic_features_float(mol, property_list=bad_props)

    def test_get_mol_atomic_features_float_nan_mask(self):
        for s in self.smiles_noble:
            mol = dm.to_mol(s)

            # Nothing happens when `mask_nan = None`, nans are still in the property array
            prop_dict = get_mol_atomic_features_float(
                mol, property_list=self.atomic_float_props, mask_nan=None
            )
            prop_array = np.concatenate(list(prop_dict.values()), axis=0)
            nans = np.isnan(prop_array)

            # Capture a raised error when `mask_nan = "raise"`
            with self.assertRaises(ValueError):
                prop_dict = get_mol_atomic_features_float(
                    mol, property_list=self.atomic_float_props, mask_nan="raise"
                )

            # Not sure how to Capture a logged warning when `mask_nan = "warn"`
            # Here, I'm testing a behaviour similar to `mask_nan = None`
            prop_dict = get_mol_atomic_features_float(
                mol, property_list=self.atomic_float_props, mask_nan="warn"
            )
            prop_array = np.concatenate(list(prop_dict.values()), axis=0)
            self.assertEqual(len(self.atomic_float_props), len(prop_dict))
            self.assertTrue(any(np.isnan(prop_array)))

            # NaNs are replaced by `42` when `mask_nan=42`
            prop_dict = get_mol_atomic_features_float(mol, property_list=self.atomic_float_props, mask_nan=42)
            prop_array = np.concatenate(list(prop_dict.values()), axis=0)
            self.assertEqual(len(self.atomic_float_props), len(prop_dict))
            self.assertFalse(any(np.isnan(prop_array)))
            self.assertTrue(all(prop_array[nans] == 42))

    def test_get_mol_edge_features(self):
        props = deepcopy(self.edge_props)
        bad_props = ["bob"]

        all_smiles = self.smiles + self.smiles_noble
        for s in all_smiles:
            err_msg = f"\n\tError for params:\n\t\tSMILES: {s}"
            mol = dm.to_mol(s)
            for ii in range(len(props)):
                this_props = props[: ii + 1]
                err_msg2 = err_msg + f"\n\t\tprops: {this_props}"
                prop_dict = get_mol_edge_features(mol, property_list=this_props)
                self.assertListEqual(list(prop_dict.keys()), this_props, msg=err_msg)
                for key, val in prop_dict.items():
                    err_msg3 = err_msg2 + f"\n\t\tkey: {key}"
                    self.assertEqual(val.shape[0], mol.GetNumBonds(), msg=err_msg3)

            if mol.GetNumBonds() > 0:
                with self.assertRaises(ValueError, msg=err_msg):
                    get_mol_edge_features(mol, property_list=bad_props)

    def test_mol_to_adj_and_features(self):
        np.random.seed(42)

        for s in self.smiles:
            err_msg = f"\n\tError for params:\n\t\tSMILES: {s}"
            mol = dm.to_mol(s)
            mol_Hs = Chem.AddHs(mol)  # type: ignore
            mol_No_Hs = Chem.RemoveHs(mol)  # type: ignore

            for explicit_H in [True, False]:
                this_mol = mol_Hs if explicit_H else mol_No_Hs
                for ii in np.arange(0, 5, 0.2):
                    num_props = int(round(ii))
                    err_msg2 = err_msg + f"\n\t\texplicit_H: {explicit_H}\n\t\tii: {ii}"

                    adj, ndata, edata, _, _ = mol_to_adj_and_features(
                        mol=mol,
                        atom_property_list_onehot=np.random.choice(
                            self.atomic_onehot_props, size=num_props, replace=False
                        ),
                        atom_property_list_float=np.random.choice(
                            self.atomic_float_props, size=num_props, replace=False
                        ),
                        edge_property_list=np.random.choice(self.edge_props, size=num_props, replace=False),
                        add_self_loop=False,
                        explicit_H=explicit_H,
                        use_bonds_weights=False,
                    )

                    self.assertEqual(adj.shape[0], this_mol.GetNumAtoms(), msg=err_msg2)
                    if num_props > 0:
                        self.assertEqual(ndata.shape[0], this_mol.GetNumAtoms(), msg=err_msg2)
                        if this_mol.GetNumBonds() > 0:
                            self.assertEqual(edata.shape[0], this_mol.GetNumBonds(), msg=err_msg2)
                            self.assertGreaterEqual(edata.shape[1], num_props, msg=err_msg2)
                        self.assertGreaterEqual(ndata.shape[1], num_props, msg=err_msg2)

    def test_mol_to_pyggraph(self):
        np.random.seed(42)

        for s in self.smiles:
            err_msg = f"\n\tError for params:\n\t\tSMILES: {s}"
            mol = dm.to_mol(s)
            mol_Hs = Chem.AddHs(mol)  # type: ignore
            mol_No_Hs = Chem.RemoveHs(mol)  # type: ignore

            graph = mol_to_pyggraph(
                mol=mol,
                atom_property_list_onehot=[],
                atom_property_list_float=["atomic-number"],
                edge_property_list=["bond-type-float"],
                add_self_loop=False,
                explicit_H=False,
                use_bonds_weights=False,
                on_error="raise",
            )

            # Check the number of nodes and edges
            self.assertListEqual(list(graph["feat"].shape), [mol.GetNumAtoms(), 1], msg=err_msg)
            self.assertListEqual(list(graph["edge_feat"].shape), [2 * mol.GetNumBonds(), 1], msg=err_msg)

            # Check the node features
            feat = graph["feat"].to_dense().numpy() * 5 + 6  # Undo the scaling
            atom_nums = np.asarray([atom.GetAtomicNum() for atom in mol.GetAtoms()])
            np.testing.assert_array_almost_equal(feat[:, 0], atom_nums, decimal=5, err_msg=err_msg)

            # Check the edge features
            edge_feat = graph["edge_feat"].to_dense().numpy()
            bond_types = np.asarray([bond.GetBondTypeAsDouble() for bond in mol.GetBonds()]).repeat(2)
            np.testing.assert_array_almost_equal(edge_feat[:, 0], bond_types, decimal=5, err_msg=err_msg)

            # Check the edge indices
            if mol.GetNumBonds() > 0:
                edge_index = graph["edge_index"].to_dense().numpy()
                true_edge_index = []
                for bond in mol.GetBonds():
                    true_edge_index.append([bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()])
                    true_edge_index.append([bond.GetEndAtomIdx(), bond.GetBeginAtomIdx()])
                true_edge_index = np.asarray(true_edge_index).T
                np.testing.assert_array_equal(edge_index, true_edge_index, err_msg=err_msg)

            # Loop over many possible combinations of properties
            for explicit_H in [True, False]:
                this_mol = mol_Hs if explicit_H else mol_No_Hs
                for ii in np.arange(0, 5, 0.2):
                    num_props = int(round(ii))
                    err_msg2 = err_msg + f"\n\t\texplicit_H: {explicit_H}\n\t\tii: {ii}"

                    graph = mol_to_pyggraph(
                        mol=mol,
                        atom_property_list_onehot=np.random.choice(
                            self.atomic_onehot_props, size=num_props, replace=False
                        ),
                        atom_property_list_float=np.random.choice(
                            self.atomic_float_props, size=num_props, replace=False
                        ),
                        edge_property_list=np.random.choice(self.edge_props, size=num_props, replace=False),
                        add_self_loop=False,
                        explicit_H=explicit_H,
                        use_bonds_weights=False,
                        on_error="raise",
                    )

                    self.assertEqual(graph.num_nodes, this_mol.GetNumAtoms(), msg=err_msg2)
                    self.assertEqual(graph.num_edges, 2 * this_mol.GetNumBonds(), msg=err_msg2)
                    if num_props > 0:
                        ndata = graph["feat"]
                        edata = graph["edge_feat"]
                        self.assertEqual(ndata.shape[0], this_mol.GetNumAtoms(), msg=err_msg2)
                        self.assertEqual(edata.shape[0], 2 * this_mol.GetNumBonds(), msg=err_msg2)
                        self.assertGreaterEqual(ndata.shape[1], num_props, msg=err_msg2)
                        self.assertGreaterEqual(edata.shape[1], num_props, msg=err_msg2)


if __name__ == "__main__":
    ut.main()


# tests/test_ipu_dataloader.py

# General imports
import yaml
import unittest as ut
import numpy as np
from copy import deepcopy
from warnings import warn
from pytorch_lightning import Trainer, LightningModule
from pytorch_lightning.strategies import IPUStrategy
from functools import partial

import torch
from torch.utils.data.dataloader import default_collate

# Current library imports
from graphium.config._loader import load_datamodule, load_metrics, load_architecture, load_accelerator


def random_packing(num_nodes, batch_size):
    ipu_batch_size = int(len(num_nodes) / batch_size)
    indices = np.arange(len(num_nodes))
    np.random.shuffle(indices)
    indices = np.reshape(indices, (ipu_batch_size, batch_size)).tolist()
    return indices


def global_batch_collator(batch_size, batches):
    packs = []
    for pack_idx in range(0, len(batches), batch_size):
        packs.append(default_collate(batches[pack_idx : pack_idx + batch_size]))
    global_batch = default_collate(packs)
    global_batch = (global_batch[0], tuple(global_batch[1]))
    return global_batch


class test_DataLoading(ut.TestCase):
    class TestSimpleLightning(LightningModule):
        # Create a basic Ligthning for testing the batch sizes
        def __init__(self, batch_size, node_feat_size, edge_feat_size, num_batch) -> None:
            super().__init__()
            self.batch_size = batch_size
            self.node_feat_size = node_feat_size
            self.edge_feat_size = edge_feat_size
            self.layer = torch.nn.Linear(node_feat_size, 1)
            self.loss_fn = torch.nn.L1Loss()
            self.num_batch = num_batch

        def validation_step(self, batch, batch_idx):
            self.assert_shapes(batch, batch_idx, "val")
            loss = self.forward(batch)
            return loss

        def training_step(self, batch, batch_idx):
            self.assert_shapes(batch, batch_idx, "train")
            loss = self.forward(batch)
            return loss

        def forward(self, batch):
            out = self.layer(batch[1][0]).squeeze(-1)
            loss = self.loss_fn(out, batch[0])
            return loss

        def assert_shapes(self, batch, batch_idx, step):
            # Test the shape of the labels
            this_shape = list(batch[0].shape)
            true_shape = [1, self.batch_size]
            assert (
                this_shape == true_shape
            ), f"Shape of the labels is `{this_shape}` but should be {true_shape}"

            # Test the shape of the first feature
            this_shape = list(batch[1][0].shape)
            true_shape = [1, self.batch_size, self.node_feat_size]
            assert (
                this_shape == true_shape
            ), f"Shape of the feature 0 is `{this_shape}` but should be {true_shape}"

            # Test the shape of the second feature
            this_shape = list(batch[1][1].shape)
            true_shape = [1, self.batch_size, self.edge_feat_size]
            assert (
                this_shape == true_shape
            ), f"Shape of the feature 0 is `{this_shape}` but should be {true_shape}"

        def configure_optimizers(self):
            return torch.optim.Adam(self.parameters(), lr=1e-3)

    class TestDataset(torch.utils.data.Dataset):
        # Create a simple dataset for testing the Lightning integration
        def __init__(self, labels, node_features, edge_features):
            self.labels = labels
            self.node_features = node_features
            self.edge_features = edge_features

        def __len__(self):
            return len(self.labels)

        def __getitem__(self, idx):
            # [label, [feat1, feat2]]
            return [self.labels[idx], [self.node_features[idx], self.edge_features[idx]]]

    def test_poptorch_simple_deviceiterations_gradient_accumulation(self):
        """
        Test a simple version of the device-iterations and gradient accumulation
        to make sure that the dataloader and models handle them correcly.
        """

        # Run this test only if poptorch is available
        try:
            import poptorch
        except Exception as e:
            warn(f"Skipping this test because poptorch is not available.\n{e}")
            return

        # Initialize constants
        gradient_accumulation = 2
        device_iterations = 3
        batch_size = 5
        num_replicate = 7
        node_feat_size = 11
        edge_feat_size = 13

        # Initialize the batch info and poptorch options
        opts = poptorch.Options()
        opts.deviceIterations(device_iterations)
        training_opts = deepcopy(opts)
        training_opts.Training.gradientAccumulation(gradient_accumulation)
        inference_opts = deepcopy(opts)

        # Initialize the dataset
        num_batch = device_iterations * gradient_accumulation * num_replicate
        data_size = num_batch * batch_size
        dataset = self.TestDataset(
            labels=np.random.rand(data_size).astype(np.float32),
            node_features=[np.random.rand(node_feat_size).astype(np.float32) for ii in range(data_size)],
            edge_features=[np.random.rand(edge_feat_size).astype(np.float32) for ii in range(data_size)],
        )

        # Initialize the dataloader
        train_dataloader = poptorch.DataLoader(
            options=training_opts,
            dataset=deepcopy(dataset),
            batch_size=batch_size,
            collate_fn=partial(global_batch_collator, batch_size),
        )

        val_dataloader = poptorch.DataLoader(
            options=inference_opts,
            dataset=deepcopy(dataset),
            batch_size=batch_size,
            collate_fn=partial(global_batch_collator, batch_size),
        )

        # Build the model, and run it on IPU
        model = self.TestSimpleLightning(batch_size, node_feat_size, edge_feat_size, num_batch)
        strategy = IPUStrategy(training_opts=training_opts, inference_opts=inference_opts)
        trainer = Trainer(
            logger=False,
            enable_checkpointing=False,
            max_epochs=2,
            strategy=strategy,
            num_sanity_val_steps=0,
            ipus=1,
        )
        trainer.fit(model=model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)

    def test_poptorch_graphium_deviceiterations_gradient_accumulation(self):
        """
        Test the device-iterations and gradient accumulation in a way
        that is very similar to the Graphium code
        to make sure that the dataloader and models handle them correcly.
        """

        try:
            import poptorch
        except Exception as e:
            warn(f"Skipping this test because poptorch is not available.\n{e}")
            return

        from graphium.ipu.ipu_wrapper import PredictorModuleIPU

        class TestPredictor(PredictorModuleIPU):
            # Create a basic Ligthning for testing the batch sizes
            def __init__(self, batch_size, node_batch_size, edge_batch_size, in_dims, **kwargs) -> None:
                super().__init__(**kwargs)
                self.in_dims = in_dims
                self.batch_size = batch_size
                self.node_batch_size = node_batch_size
                self.edge_batch_size = edge_batch_size

                import poptorch

                self.pop = poptorch

            def validation_step(self, features, labels):
                features, labels = self.squeeze_input_dims(features, labels)
                dict_input = {"features": features, "labels": labels}
                self.assert_shapes(dict_input, "val")
                preds = self.forward(dict_input)["preds"]
                loss = self.compute_loss(
                    preds, dict_input["labels"], weights=None, loss_fun=self.loss_fun, target_nan_mask=None
                )
                loss = loss[0]
                return loss

            def training_step(self, features, labels):
                features, labels = self.squeeze_input_dims(features, labels)
                dict_input = {"features": features, "labels": labels}
                preds = self.forward(dict_input)["preds"]
                loss = self.compute_loss(
                    preds, dict_input["labels"], weights=None, loss_fun=self.loss_fun, target_nan_mask=None
                )
                loss = self.pop.identity_loss(loss[0], reduction="mean")
                return loss

            def assert_shapes(self, dict_input, step):
                msg = f", step=`{step}`"

                # Test the shape of the labels
                this_shape = list(dict_input["labels"]["homo"].shape)
                true_shape = [self.batch_size, 2]
                assert (
                    this_shape == true_shape
                ), f"Shape of the labels is `{this_shape}` but should be {true_shape}. {msg}"

                # Test the shape of the node feature
                this_shape = list(dict_input["features"]["feat"].shape)
                true_shape = [self.node_batch_size, self.in_dims["feat"]]
                assert (
                    this_shape == true_shape
                ), f"Shape of the feature 0 is `{this_shape}` but should be {true_shape}. {msg}"

                # Test the shape of the node feature
                this_shape = list(dict_input["features"]["edge_feat"].shape)
                true_shape = [self.edge_batch_size, self.in_dims["edge_feat"]]
                assert (
                    this_shape == true_shape
                ), f"Shape of the feature 0 is `{this_shape}` but should be {true_shape}. {msg}"

            def get_progress_bar_dict(self):
                return {}

            def on_train_batch_end(self, *args, **kwargs):
                return

            def on_validation_batch_end(self, *args, **kwargs):
                return

            def validation_epoch_end(self, *args, **kwargs):
                return

            def on_train_epoch_end(self) -> None:
                return

            def configure_optimizers(self):
                return torch.optim.Adam(self.parameters(), lr=1e-3)

            def squeeze_input_dims(self, features, labels):
                for key, tensor in features:
                    if isinstance(tensor, torch.Tensor):
                        features[key] = features[key].squeeze(0)

                for key in labels:
                    labels[key] = labels[key].squeeze(0)

                return features, labels

        from graphium.ipu.ipu_wrapper import DictIPUStrategy

        gradient_accumulation = 3
        device_iterations = 5
        batch_size = 7

        # Initialize the batch info and poptorch options
        opts = poptorch.Options()
        opts.deviceIterations(device_iterations)
        training_opts = deepcopy(opts)
        training_opts.Training.gradientAccumulation(gradient_accumulation)
        inference_opts = deepcopy(opts)

        # Load the configuration file for the model
        CONFIG_FILE = "tests/config_test_ipu_dataloader.yaml"
        with open(CONFIG_FILE, "r") as f:
            cfg = yaml.safe_load(f)
        cfg, accelerator = load_accelerator(cfg)
        cfg["datamodule"]["args"]["batch_size_training"] = batch_size
        cfg["datamodule"]["args"]["batch_size_inference"] = batch_size
        node_factor = cfg["datamodule"]["args"]["ipu_dataloader_training_opts"]["max_num_nodes_per_graph"]
        edge_factor = cfg["datamodule"]["args"]["ipu_dataloader_training_opts"]["max_num_edges_per_graph"]

        # Load the datamodule, and prepare the data
        datamodule = load_datamodule(cfg, accelerator_type="ipu")
        datamodule.prepare_data()
        datamodule.setup()
        datamodule.ipu_training_opts = training_opts
        datamodule.ipu_inference_opts = inference_opts

        model_class, model_kwargs = load_architecture(cfg, in_dims=datamodule.in_dims)
        metrics = load_metrics(cfg)
        predictor = TestPredictor(
            batch_size=batch_size,
            node_batch_size=node_factor * batch_size,
            edge_batch_size=edge_factor * batch_size,
            in_dims=datamodule.in_dims,
            model_class=model_class,
            model_kwargs=model_kwargs,
            metrics=metrics,
            **cfg["predictor"],
        )
        strategy = DictIPUStrategy(training_opts=training_opts, inference_opts=inference_opts)
        trainer = Trainer(
            logger=False, enable_checkpointing=False, max_epochs=2, strategy=strategy, num_sanity_val_steps=0
        )
        trainer.fit(model=predictor, datamodule=datamodule)


if __name__ == "__main__":
    ut.main()


# tests/test_ipu_losses.py

import unittest as ut
import torch
from torch.nn import BCELoss, MSELoss, L1Loss
from copy import deepcopy

from graphium.ipu.ipu_losses import BCELossIPU, MSELossIPU, L1LossIPU


class test_Losses(ut.TestCase):
    torch.manual_seed(42)
    preds = torch.rand((100, 10), dtype=torch.float32)
    target = torch.rand((100, 10), dtype=torch.float32)

    th = 0.7
    nan_th = 0.2
    preds_greater = preds > th
    target_greater = (target > th).to(torch.float32)
    target_greater_nan = deepcopy(target_greater)
    is_nan = target < nan_th
    target_greater_nan[target < nan_th] = torch.nan
    target_nan = deepcopy(target)
    target_nan[target < nan_th] = torch.nan

    def test_bce(self):
        preds = deepcopy(self.preds)
        target = deepcopy(self.target_greater)
        target_nan = deepcopy(self.target_greater_nan)

        # Regular loss
        loss_true = BCELoss()(preds, target)
        loss_ipu = BCELossIPU()(preds, target)
        self.assertFalse(loss_true.isnan(), "Regular BCELoss is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular BCELoss is different"
        )

        # Weighted loss
        weight = torch.rand(preds.shape[1], dtype=torch.float32)
        loss_true = BCELoss(weight=weight)(preds, target)
        loss_ipu = BCELossIPU(weight=weight)(preds, target)
        self.assertFalse(loss_true.isnan(), "Regular BCELoss is NaN")
        self.assertAlmostEqual(loss_true.item(), loss_ipu.item(), msg="Weighted BCELoss is different")

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        loss_true = BCELoss()(preds[not_nan], target[not_nan])
        loss_ipu = BCELossIPU()(preds, target_nan)
        self.assertFalse(loss_true.isnan(), "Regular BCELoss with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Regular BCELossIPU with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular BCELoss with NaN is different"
        )

        # Weighted loss with NaNs in target
        not_nan = ~target_nan.isnan()
        weight = torch.rand(preds.shape, dtype=torch.float32)
        loss_true = BCELoss(weight=weight[not_nan])(preds[not_nan], target_nan[not_nan])
        loss_ipu = BCELossIPU(weight=weight)(preds, target_nan)
        self.assertFalse(loss_true.isnan(), "Weighted BCELoss with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Weighted BCELossIPU with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Weighted BCELoss with NaN is different"
        )

    def test_mse(self):
        preds = deepcopy(self.preds)
        target = deepcopy(self.target)
        target_nan = deepcopy(self.target_nan)

        # Regular loss
        loss_true = MSELoss()(preds, target)
        loss_ipu = MSELossIPU()(preds, target)
        self.assertFalse(loss_true.isnan(), "Regular MSELoss is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular MSELoss is different"
        )

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        loss_true = MSELoss()(preds[not_nan], target[not_nan])
        loss_ipu = MSELossIPU()(preds, target_nan)
        self.assertFalse(loss_true.isnan(), "Regular MSELoss with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Regular MSELossIPU with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular MSELoss with NaN is different"
        )

    def test_l1(self):
        preds = deepcopy(self.preds)
        target = deepcopy(self.target)
        target_nan = deepcopy(self.target_nan)

        # Regular loss
        loss_true = L1Loss()(preds, target)
        loss_ipu = L1LossIPU()(preds, target)
        self.assertFalse(loss_true.isnan(), "Regular MAELoss is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular MAELoss is different"
        )

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        loss_true = L1Loss()(preds[not_nan], target[not_nan])
        loss_ipu = L1LossIPU()(preds, target_nan)
        self.assertFalse(loss_true.isnan(), "Regular MAELoss with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Regular MAELossIPU with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular MAELoss with NaN is different"
        )


# tests/test_ipu_metrics.py

import unittest as ut
import torch
from torchmetrics.functional import (
    auroc,
    average_precision,
    precision,
    accuracy,
    recall,
    pearson_corrcoef,
    spearman_corrcoef,
    r2_score,
    f1_score,
    fbeta_score,
    mean_squared_error,
    mean_absolute_error,
)
from copy import deepcopy

from graphium.ipu.ipu_metrics import (
    auroc_ipu,
    average_precision_ipu,
    precision_ipu,
    accuracy_ipu,
    recall_ipu,
    pearson_ipu,
    spearman_ipu,
    r2_score_ipu,
    f1_score_ipu,
    fbeta_score_ipu,
    mean_squared_error_ipu,
    mean_absolute_error_ipu,
)


class test_Metrics(ut.TestCase):
    torch.manual_seed(42)
    preds = torch.rand((100, 10), dtype=torch.float32)
    target = torch.rand((100, 10), dtype=torch.float32)

    th = 0.7
    nan_th = 0.2
    preds_greater = preds > th
    target_greater = (target > th).to(torch.float32)
    target_greater_nan = deepcopy(target_greater)
    is_nan = target < nan_th
    target_greater_nan[target < nan_th] = torch.nan
    target_nan = deepcopy(target)
    target_nan[target < nan_th] = torch.nan

    def test_auroc(self):
        preds = deepcopy(self.preds)[:, 0]
        target = deepcopy(self.target)[:, 0]
        target_nan = deepcopy(self.target_nan)[:, 0]

        target[target < 0.5] = 0
        target[target >= 0.5] = 1

        target_nan[target_nan < 0.5] = 0
        target_nan[target_nan >= 0.5] = 1

        # Regular loss
        score_true = auroc(preds, target.to(int))
        score_ipu = auroc_ipu(preds, target)
        self.assertFalse(score_true.isnan(), "Regular AUROC score is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Regular AUROC score is different"
        )

        # Weighted loss (As in BCE)
        sample_weights = torch.rand(preds.shape[0], dtype=torch.float32)
        score_true = auroc(preds, target.to(int), sample_weights=sample_weights)
        score_ipu = auroc_ipu(preds, target, sample_weights=sample_weights)
        self.assertFalse(score_true.isnan(), "Regular AUROC score is NaN")
        self.assertAlmostEqual(score_true.item(), score_ipu.item(), msg="Weighted AUROC score is different")

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = auroc(preds[not_nan], target[not_nan].to(int))
        score_ipu = auroc_ipu(preds, target_nan)
        self.assertFalse(score_true.isnan(), "Regular AUROC score with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Regular AUROCIPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Regular AUROC score with NaN is different"
        )

        # Weighted loss with NaNs in target (As in BCE)
        not_nan = ~target_nan.isnan()
        sample_weights = torch.rand(preds.shape, dtype=torch.float32)
        loss_true = auroc(preds[not_nan], target_nan[not_nan].to(int), sample_weights=sample_weights[not_nan])
        loss_ipu = auroc_ipu(preds, target_nan, sample_weights=sample_weights)
        self.assertFalse(loss_true.isnan(), "Weighted AUROC score with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Weighted AUROC IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            # AssertionError: 0.6603766679763794 != 0.6234951615333557 within 2 places
            loss_true.item(),
            loss_ipu.item(),
            places=6,
            msg="Weighted AUROC with NaN is different",
        )

    def test_average_precision(self):  # TODO: Make work with multi-class
        preds = deepcopy(self.preds)[:, 0]
        target = deepcopy(self.target)[:, 0]
        target_nan = deepcopy(self.target_nan)[:, 0]

        target[target < 0.5] = 0
        target[target >= 0.5] = 1

        target_nan[target_nan < 0.5] = 0
        target_nan[target_nan >= 0.5] = 1

        # Regular loss
        score_true = average_precision(preds, target.to(int), task="binary")
        score_ipu = average_precision_ipu(preds, target.to(int), task="binary")
        self.assertFalse(score_true.isnan(), "Regular Average Precision is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Regular Average Precision is different"
        )

        # Regular average precision with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = average_precision(preds[not_nan], target[not_nan].to(int), task="binary")
        score_ipu = average_precision_ipu(preds, target_nan, task="binary")
        self.assertFalse(score_true.isnan(), "Regular Average Precision with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Regular Average Precision IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Regular Average Precision with NaN is different",
        )

    def test_precision(self):
        preds = deepcopy(self.preds)[:, :4]
        target = deepcopy(self.target)[:, 0]
        t = deepcopy(target)

        target[t < 0.4] = 0
        target[(t >= 0.4) & (t < 0.6)] = 1
        target[(t >= 0.6) & (t < 0.8)] = 2
        target[(t >= 0.8)] = 3

        target_nan = deepcopy(target)
        target_nan[self.is_nan[:, 0]] = float("nan")
        target_nan_bin = deepcopy(target_nan)
        target_nan_bin[target_nan > 0] = 1

        # Micro precision binary
        score_true = precision(preds[:, 0], target.to(int) > 0, average="micro")
        score_ipu = precision_ipu(preds[:, 0], target > 0, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Precision binary is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Precision binary is different"
        )

        # Micro precision binary with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = precision(preds[:, 0][not_nan], target_nan_bin[not_nan].to(int), average="micro")
        score_ipu = precision_ipu(preds[:, 0], target_nan_bin, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Precision binary with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro Precision binary IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Precision with NaN is different"
        )

        # Micro precision
        score_true = precision(preds, target.to(int), average="micro")
        score_ipu = precision_ipu(preds, target, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Precision is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Precision is different"
        )

        # Micro precision with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = precision(preds[not_nan], target[not_nan].to(int), average="micro")
        score_ipu = precision_ipu(preds, target_nan, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Precision with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro Precision IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Precision with NaN is different"
        )

        # Macro precision
        score_true = precision(preds, target.to(int), average="macro", num_classes=4)
        score_ipu = precision_ipu(preds, target, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro Precision is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Macro Precision is different"
        )

        # Macro precision multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = precision(preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4)
        score_ipu = precision_ipu(preds, target_nan, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro Precision multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Macro Precision multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Macro Precision multiclass with NaN is different",
        )

        # Macro precision multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = precision(preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4)
        score_ipu = precision_ipu(preds, target_nan, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro Precision multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Macro Precision multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Macro Precision multiclass with NaN is different",
        )

        # Weighted precision multiclass
        score_true = precision(preds, target.to(int), average="weighted", num_classes=4)
        score_ipu = precision_ipu(preds, target, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted Precision multiclass is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Weighted Precision multiclass is different"
        )

        # Weighted precision multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = precision(preds[not_nan], target[not_nan].to(int), average="weighted", num_classes=4)
        score_ipu = precision_ipu(preds, target_nan, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted Precision multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Weighted Precision multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Regular Average Precision multiclass with NaN is different",
        )

    def test_accuracy(self):
        preds = deepcopy(self.preds)[:, :4]
        target = deepcopy(self.target)[:, 0]
        t = deepcopy(target)

        target[t < 0.4] = 0
        target[(t >= 0.4) & (t < 0.6)] = 1
        target[(t >= 0.6) & (t < 0.8)] = 2
        target[(t >= 0.8)] = 3

        target_nan = deepcopy(target)
        target_nan[self.is_nan[:, 0]] = float("nan")
        target_nan_bin = deepcopy(target_nan)
        target_nan_bin[target_nan > 0] = 1

        # Micro accuracy binary
        score_true = accuracy(preds[:, 0], target.to(int) > 0, average="micro")
        score_ipu = accuracy_ipu(preds[:, 0], target > 0, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Accuracy binary is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Accuracy binary is different"
        )

        # Micro accuracy binary with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = accuracy(preds[:, 0][not_nan], target_nan_bin[not_nan].to(int), average="micro")
        score_ipu = accuracy_ipu(preds[:, 0], target_nan_bin, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Accuracy binary with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro Accuracy binary IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Accuracy with NaN is different"
        )

        # Micro accuracy
        score_true = accuracy(preds, target.to(int), average="micro")
        score_ipu = accuracy_ipu(preds, target, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Accuracy is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Accuracy is different"
        )

        # Micro accuracy with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = accuracy(preds[not_nan], target[not_nan].to(int), average="micro")
        score_ipu = accuracy_ipu(preds, target_nan, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Accuracy with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro Accuracy IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Accuracy with NaN is different"
        )

        # Macro accuracy
        score_true = accuracy(preds, target.to(int), average="macro", num_classes=4)
        score_ipu = accuracy_ipu(preds, target, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro Accuracy is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Macro Accuracy is different"
        )

        # Macro accuracy with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = accuracy(preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4)
        score_ipu = accuracy_ipu(preds, target_nan, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro Accuracy with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Macro Accuracy IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Macro Accuracy with NaN is different"
        )

        # Weighted accuracy
        score_true = accuracy(preds, target.to(int), average="weighted", num_classes=4)
        score_ipu = accuracy_ipu(preds, target, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted Accuracy is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Weighted Accuracy is different"
        )

        # Weighted accuracy with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = accuracy(preds[not_nan], target[not_nan].to(int), average="weighted", num_classes=4)
        score_ipu = accuracy_ipu(preds, target_nan, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted Accuracy with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Weighted Accuracy IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Regular Accuracy with NaN is different"
        )

    def test_recall(self):
        preds = deepcopy(self.preds)[:, :4]
        target = deepcopy(self.target)[:, 0]
        t = deepcopy(target)

        target[t < 0.4] = 0
        target[(t >= 0.4) & (t < 0.6)] = 1
        target[(t >= 0.6) & (t < 0.8)] = 2
        target[(t >= 0.8)] = 3

        target_nan = deepcopy(target)
        target_nan[self.is_nan[:, 0]] = float("nan")
        target_nan_bin = deepcopy(target_nan)
        target_nan_bin[target_nan > 0] = 1

        # Micro recall binary
        score_true = recall(preds[:, 0], target.to(int) > 0, average="micro")
        score_ipu = recall_ipu(preds[:, 0], target > 0, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Recall binary is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Recall binary is different"
        )

        # Micro recall binary with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = recall(preds[:, 0][not_nan], target_nan_bin[not_nan].to(int), average="micro")
        score_ipu = recall_ipu(preds[:, 0], target_nan_bin, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Recall binary with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro Recall binary IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Recall binary with NaN is different"
        )

        # Micro recall
        score_true = recall(preds, target.to(int), average="micro")
        score_ipu = recall_ipu(preds, target, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Recall is NaN")
        self.assertAlmostEqual(score_true.item(), score_ipu.item(), places=6, msg="Micro Recall is different")

        # Micro recall with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = recall(preds[not_nan], target[not_nan].to(int), average="micro")
        score_ipu = recall_ipu(preds, target_nan, average="micro")
        self.assertFalse(score_true.isnan(), "Micro Recall with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro Recall IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro Recall with NaN is different"
        )

        # Macro recall multiclass
        score_true = recall(preds, target.to(int), average="macro", num_classes=4)
        score_ipu = recall_ipu(preds, target, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro Recall is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Macro Recall multiclass is different"
        )

        # Macro recall multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = recall(preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4)
        score_ipu = recall_ipu(preds, target_nan, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro Recall multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Macro Recall multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Macro Recall multiclass with NaN is different"
        )

        # Weighted recallmulticlass
        score_true = recall(preds, target.to(int), average="weighted", num_classes=4)
        score_ipu = recall_ipu(preds, target, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted Recall is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Weighted Recall is different"
        )

        # Weighted recall multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = recall(preds[not_nan], target[not_nan].to(int), average="weighted", num_classes=4)
        score_ipu = recall_ipu(preds, target_nan, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted Recall multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Weighted Recall multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Regular Recall multiclass with NaN is different",
        )

    def test_pearsonr(self):
        preds = deepcopy(self.preds)[:, 0]
        target = deepcopy(self.target)[:, 0] + preds
        target_nan = deepcopy(target)
        target_nan[self.is_nan[:, 0]] = float("nan")

        # Regular loss
        score_true = pearson_corrcoef(preds, target)
        score_ipu = pearson_ipu(preds, target)
        self.assertFalse(score_true.isnan(), "Pearson is NaN")
        self.assertAlmostEqual(score_true.item(), score_ipu.item(), places=4, msg="Pearson is different")

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = pearson_corrcoef(preds[not_nan], target[not_nan])
        score_ipu = pearson_ipu(preds, target_nan)
        self.assertFalse(score_true.isnan(), "Regular PearsonR with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "IPU PearsonR score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=4, msg="Pearson with NaN is different"
        )

    def test_spearmanr(self):
        preds = deepcopy(self.preds)[:, 0]
        target = deepcopy(self.target)[:, 0] + preds
        target_nan = deepcopy(target)
        target_nan[self.is_nan[:, 0]] = float("nan")

        # Regular loss
        score_true = spearman_corrcoef(preds, target)
        score_ipu = spearman_ipu(preds, target)
        self.assertFalse(score_true.isnan(), "Spearman is NaN")
        self.assertAlmostEqual(score_true.item(), score_ipu.item(), places=4, msg="Spearman is different")

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = spearman_corrcoef(preds[not_nan], target[not_nan])
        score_ipu = spearman_ipu(preds, target_nan)
        self.assertFalse(score_true.isnan(), "Regular Spearman with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "IPU Spearman score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=4, msg="Spearman with NaN is different"
        )

    def test_r2_score(self):
        preds = deepcopy(self.preds)
        target = deepcopy(self.target) + preds
        target_nan = deepcopy(target)
        target_nan[self.is_nan] = float("nan")

        # Regular loss
        score_true = r2_score(preds, target)
        score_ipu = r2_score_ipu(preds, target)
        self.assertFalse(score_true.isnan(), "r2_score is NaN")
        self.assertAlmostEqual(score_true.item(), score_ipu.item(), places=4, msg="r2_score is different")

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        score_ipu = r2_score_ipu(preds, target_nan, multioutput="raw_values")
        for ii in range(preds.shape[1]):
            score_true = r2_score(
                preds[:, ii][not_nan[:, ii]], target_nan[:, ii][not_nan[:, ii]], multioutput="raw_values"
            )
            self.assertFalse(score_true.isnan().any(), f"{ii}: r2_score with target_nan is NaN")
            self.assertFalse(score_ipu[ii].isnan().any(), f"{ii}: IPU r2_score with target_nan is NaN")
            self.assertAlmostEqual(
                score_true.item(), score_ipu[ii].item(), places=4, msg=f"{ii}: r2_score with NaN is different"
            )

    def test_fbeta_score(self):
        preds = deepcopy(self.preds)[:, :4]
        target = deepcopy(self.target)[:, 0]
        t = deepcopy(target)

        target[t < 0.4] = 0
        target[(t >= 0.4) & (t < 0.6)] = 1
        target[(t >= 0.6) & (t < 0.8)] = 2
        target[(t >= 0.8)] = 3

        target_nan = deepcopy(target)
        target_nan[self.is_nan[:, 0]] = float("nan")
        target_nan_bin = deepcopy(target_nan)
        target_nan_bin[target_nan > 0] = 1

        # Micro fbeta_score binary
        score_true = fbeta_score(preds[:, 0], target.to(int) > 0, average="micro", beta=0.5)
        score_ipu = fbeta_score_ipu(preds[:, 0], target > 0, average="micro", beta=0.5)
        self.assertFalse(score_true.isnan(), "Micro FBETA_score binary is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro FBETA_score binary is different"
        )

        # Micro fbeta_score binary with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = fbeta_score(
            preds[:, 0][not_nan], target_nan_bin[not_nan].to(int), average="micro", beta=0.5
        )
        score_ipu = fbeta_score_ipu(preds[:, 0], target_nan_bin, average="micro", beta=0.5)
        self.assertFalse(score_true.isnan(), "Micro FBETA_score binary with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro FBETA_score binary IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Micro FBETA_score binary with NaN is different",
        )

        # Micro fbeta_score
        score_true = fbeta_score(preds, target.to(int), average="micro", beta=0.5)
        score_ipu = fbeta_score_ipu(preds, target, average="micro", beta=0.5)
        self.assertFalse(score_true.isnan(), "Micro FBETA_score is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro FBETA_score is different"
        )

        # Micro fbeta_score with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = fbeta_score(preds[not_nan], target[not_nan].to(int), average="micro", beta=0.5)
        score_ipu = fbeta_score_ipu(preds, target_nan, average="micro", beta=0.5)
        self.assertFalse(score_true.isnan(), "Micro FBETA_score with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro FBETA_score IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro FBETA_score with NaN is different"
        )

        # Macro fbeta_score multiclass
        score_true = fbeta_score(preds, target.to(int), average="macro", num_classes=4, beta=0.5)
        score_ipu = fbeta_score_ipu(preds, target, average="macro", num_classes=4, beta=0.5)
        self.assertFalse(score_true.isnan(), "Macro FBETA_score is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Macro FBETA_score multiclass is different"
        )

        # Macro fbeta_score multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = fbeta_score(
            preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4, beta=0.5
        )
        score_ipu = fbeta_score_ipu(preds, target_nan, average="macro", num_classes=4, beta=0.5)
        self.assertFalse(score_true.isnan(), "Macro FBETA_score multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Macro FBETA_score multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Macro FBETA_score multiclass with NaN is different",
        )

        # Weighted fbeta_scoremulticlass
        score_true = fbeta_score(preds, target.to(int), average="weighted", num_classes=4, beta=0.5)
        score_ipu = fbeta_score_ipu(preds, target, average="weighted", num_classes=4, beta=0.5)
        self.assertFalse(score_true.isnan(), "Weighted FBETA_score is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Weighted FBETA_score is different"
        )

        # Weighted fbeta_score multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = fbeta_score(
            preds[not_nan], target[not_nan].to(int), average="weighted", num_classes=4, beta=0.5
        )
        score_ipu = fbeta_score_ipu(preds, target_nan, average="weighted", num_classes=4, beta=0.5)
        self.assertFalse(score_true.isnan(), "Weighted FBETA_score multiclass with target_nan is NaN")
        self.assertFalse(
            score_ipu.isnan(), "Weighted FBETA_score multiclass IPU score with target_nan is NaN"
        )
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Regular FBETA_score multiclass with NaN is different",
        )

    def test_f1_score(self):
        preds = deepcopy(self.preds)[:, :4]
        target = deepcopy(self.target)[:, 0]
        t = deepcopy(target)

        target[t < 0.4] = 0
        target[(t >= 0.4) & (t < 0.6)] = 1
        target[(t >= 0.6) & (t < 0.8)] = 2
        target[(t >= 0.8)] = 3

        target_nan = deepcopy(target)
        target_nan[self.is_nan[:, 0]] = float("nan")
        target_nan_bin = deepcopy(target_nan)
        target_nan_bin[target_nan > 0] = 1

        # Micro f1_score binary
        score_true = f1_score(preds[:, 0], target.to(int) > 0, average="micro")
        score_ipu = f1_score_ipu(preds[:, 0], target > 0, average="micro")
        self.assertFalse(score_true.isnan(), "Micro F1_score binary is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro F1_score binary is different"
        )

        # Micro f1_score binary with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = f1_score(preds[:, 0][not_nan], target_nan_bin[not_nan].to(int), average="micro")
        score_ipu = f1_score_ipu(preds[:, 0], target_nan_bin, average="micro")
        self.assertFalse(score_true.isnan(), "Micro F1_score binary with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro F1_score binary IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro F1_score binary with NaN is different"
        )

        # Micro f1_score
        score_true = f1_score(preds, target.to(int), average="micro")
        score_ipu = f1_score_ipu(preds, target, average="micro")
        self.assertFalse(score_true.isnan(), "Micro F1_score is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro F1_score is different"
        )

        # Micro f1_score with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = f1_score(preds[not_nan], target[not_nan].to(int), average="micro")
        score_ipu = f1_score_ipu(preds, target_nan, average="micro")
        self.assertFalse(score_true.isnan(), "Micro F1_score with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Micro F1_score IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Micro F1_score with NaN is different"
        )

        # Macro f1_score multiclass
        score_true = f1_score(preds, target.to(int), average="macro", num_classes=4)
        score_ipu = f1_score_ipu(preds, target, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro F1_score is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Macro F1_score multiclass is different"
        )

        # Macro f1_score multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = f1_score(preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4)
        score_ipu = f1_score_ipu(preds, target_nan, average="macro", num_classes=4)
        self.assertFalse(score_true.isnan(), "Macro F1_score multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Macro F1_score multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Macro F1_score multiclass with NaN is different",
        )

        # Weighted f1_scoremulticlass
        score_true = f1_score(preds, target.to(int), average="weighted", num_classes=4)
        score_ipu = f1_score_ipu(preds, target, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted F1_score is NaN")
        self.assertAlmostEqual(
            score_true.item(), score_ipu.item(), places=6, msg="Weighted F1_score is different"
        )

        # Weighted f1_score multiclass with NaNs in target
        not_nan = ~target_nan.isnan()
        score_true = f1_score(preds[not_nan], target[not_nan].to(int), average="weighted", num_classes=4)
        score_ipu = f1_score_ipu(preds, target_nan, average="weighted", num_classes=4)
        self.assertFalse(score_true.isnan(), "Weighted F1_score multiclass with target_nan is NaN")
        self.assertFalse(score_ipu.isnan(), "Weighted F1_score multiclass IPU score with target_nan is NaN")
        self.assertAlmostEqual(
            score_true.item(),
            score_ipu.item(),
            places=6,
            msg="Regular F1_score multiclass with NaN is different",
        )

    def test_mse(self):
        preds = deepcopy(self.preds)
        target = deepcopy(self.target)
        target_nan = deepcopy(self.target_nan)
        squared = True

        # Regular loss
        loss_true = mean_squared_error(preds, target, squared)
        loss_ipu = mean_squared_error_ipu(preds=preds, target=target, squared=squared)
        self.assertFalse(loss_true.isnan(), "Regular Mean Squared Error is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular Mean Squared Error is different"
        )

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        loss_true = mean_squared_error(preds[not_nan], target[not_nan], squared)
        loss_ipu = mean_squared_error_ipu(preds=preds, target=target_nan, squared=squared)
        self.assertFalse(loss_true.isnan(), "Regular Mean Squared Error with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Regular Mean Squared Error IPU with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(),
            loss_ipu.item(),
            places=6,
            msg="Regular Mean Squared Error with NaN is different",
        )

        squared = False

        # Regular loss
        loss_true = mean_squared_error(preds, target, squared)
        loss_ipu = mean_squared_error_ipu(preds=preds, target=target, squared=squared)
        self.assertFalse(loss_true.isnan(), "Regular Mean Squared Error is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular Mean Squared Error is different"
        )

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        loss_true = mean_squared_error(preds[not_nan], target[not_nan], squared)
        loss_ipu = mean_squared_error_ipu(preds=preds, target=target_nan, squared=squared)
        self.assertFalse(loss_true.isnan(), "Regular Mean Squared Error with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Regular Mean Squared Error IPU with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(),
            loss_ipu.item(),
            places=6,
            msg="Regular Mean Squared Error with NaN is different",
        )

    def test_mae(self):
        preds = deepcopy(self.preds)
        target = deepcopy(self.target)
        target_nan = deepcopy(self.target_nan)

        # Regular loss
        loss_true = mean_absolute_error(preds, target)
        loss_ipu = mean_absolute_error_ipu(preds=preds, target=target)
        self.assertFalse(loss_true.isnan(), "Regular Mean Absolute Error is NaN")
        self.assertAlmostEqual(
            loss_true.item(), loss_ipu.item(), places=6, msg="Regular Mean Absolute Error is different"
        )

        # Regular loss with NaNs in target
        not_nan = ~target_nan.isnan()
        loss_true = mean_absolute_error(preds[not_nan], target[not_nan])
        loss_ipu = mean_absolute_error_ipu(preds=preds, target=target_nan)
        self.assertFalse(loss_true.isnan(), "Regular Mean Absolute Error with target_nan is NaN")
        self.assertFalse(loss_ipu.isnan(), "Regular Mean Absolute Error IPU with target_nan is NaN")
        self.assertAlmostEqual(
            loss_true.item(),
            loss_ipu.item(),
            places=6,
            msg="Regular Mean Absolute Error with NaN is different",
        )


if __name__ == "__main__":
    ut.main()


# tests/test_loaders.py

from graphium.config._loader import merge_dicts
from copy import deepcopy
import unittest as ut


class TestLoader(ut.TestCase):
    def test_merge_dicts(self):
        dict_a = {"a": {"b": {"c": 1, "d": 2}, "e": 3}, "f": 4}

        dict_b = {"a": {"b": {"g": 6}}, "h": 7}

        dict_c = {
            "a": {
                "b": {
                    "c": 1,
                },
            }
        }

        # Check that the new keys are added correctly
        merge_dicts(dict_a, dict_b)
        self.assertEqual(dict_a["a"]["b"]["g"], dict_b["a"]["b"]["g"])
        self.assertEqual(dict_a["h"], dict_b["h"])

        # Check that no error is thrown if a key exists, but the value is identical
        merge_dicts(dict_a, dict_c)
        self.assertEqual(dict_a["a"]["b"]["c"], dict_c["a"]["b"]["c"])

        # Check that an error is thrown if a key exists, but the value is different
        dict_d = deepcopy(dict_c)
        dict_d["a"]["b"]["c"] = 2
        with self.assertRaises(ValueError):
            merge_dicts(dict_a, dict_d)


# Main
if __name__ == "__main__":
    ut.main()


# tests/test_losses.py

"""
Unit tests for the metrics and wrappers of graphium/trainer/metrics/...
"""

import torch
import unittest as ut
from torch.nn import functional as F

from graphium.trainer.losses import HybridCELoss
from graphium.trainer.predictor_options import EvalOptions


def _parse(loss_fun):
    eval_options = EvalOptions(loss_fun=loss_fun, metrics_on_progress_bar=None)
    return eval_options.parse_loss_fun(loss_fun)


class test_HybridCELoss(ut.TestCase):
    input = torch.Tensor([[0.1, 0.1, 0.3, 0.5, 0.0], [0.1, 0.0, 0.7, 0.2, 0.0]])
    target = torch.Tensor([3, 0]).long()
    brackets = torch.Tensor([0, 1, 2, 3, 4])
    regression_input = torch.Tensor([2.2, 2.0])  # inner product of input and brackets
    regression_target = torch.Tensor([3, 0]).float()

    def test_pure_ce_loss(self):
        loss = HybridCELoss(n_brackets=len(self.brackets), alpha=1.0, reduction="none")

        assert torch.allclose(
            loss(self.input, self.target),
            F.cross_entropy(self.input, self.target, reduction="none"),
        )
        assert loss(self.input, self.target).shape == (2,)

    def test_pure_mae_loss(self):
        loss = HybridCELoss(
            n_brackets=len(self.brackets),
            alpha=0.0,
            regression_loss="mae",
            reduction="none",
        )

        assert torch.allclose(
            loss(self.input, self.target),
            F.l1_loss(self.regression_input, self.regression_target, reduction="none"),
        )
        assert loss(self.input, self.target).shape == (2,)

    def test_pure_mse_loss(self):
        loss = HybridCELoss(
            n_brackets=len(self.brackets),
            alpha=0.0,
            regression_loss="mse",
            reduction="none",
        )

        assert torch.allclose(
            loss(self.input, self.target),
            F.mse_loss(self.regression_input, self.regression_target, reduction="none"),
        )
        assert loss(self.input, self.target).shape == (2,)

    def test_hybrid_loss(self):
        loss = HybridCELoss(n_brackets=len(self.brackets), alpha=0.5, regression_loss="mse")

        ce_loss = F.cross_entropy(self.input, self.target)
        mse_loss = F.mse_loss(self.regression_input, self.regression_target)

        assert torch.allclose(loss(self.input, self.target), 0.5 * ce_loss + 0.5 * mse_loss)
        assert loss(self.input, self.target).shape == torch.Size([])

    def test_loss_parser(self):
        # HybridCE cannot be parsed from a string because it requires specifying n_brackets
        loss_fun = "hybrid_ce"
        self.assertRaises(TypeError, _parse, loss_fun=loss_fun)

        # HybridCE requires n_brackets to be specified
        loss_fun = {"name": "hybrid_ce"}
        self.assertRaises(TypeError, _parse, loss_fun=loss_fun)

        loss_fun = {"name": "hybrid_ce", "n_brackets": 3}
        parsed = _parse(loss_fun)

        assert isinstance(parsed, HybridCELoss)
        assert len(parsed.brackets) == 3
        assert parsed.regression_loss == F.mse_loss

        loss_fun = {"name": "hybrid_ce", "n_brackets": 5, "regression_loss": "mae"}
        parsed = _parse(loss_fun)

        assert isinstance(parsed, HybridCELoss)
        assert len(parsed.brackets) == 5
        assert parsed.regression_loss == F.l1_loss


class test_BCELoss(ut.TestCase):
    def test_loss_parser(self):
        loss_fun = "bce"
        parsed = _parse(loss_fun)

        assert isinstance(parsed, torch.nn.BCELoss)

        loss_fun = {"name": "bce"}
        parsed = _parse(loss_fun)

        assert isinstance(parsed, torch.nn.BCELoss)
        assert parsed.reduction != "sum"

        loss_fun = {"name": "bce", "reduction": "sum"}
        parsed = _parse(loss_fun)

        assert isinstance(parsed, torch.nn.BCELoss)
        assert parsed.reduction == "sum"


if __name__ == "__main__":
    ut.main()


# tests/test_metrics.py

"""
Unit tests for the metrics and wrappers of graphium/trainer/metrics/...
"""

import torch
import unittest as ut
import tempfile
import os
import operator as op

from graphium.trainer.metrics import (
    MetricWrapper,
    Thresholder,
)

from torchmetrics.functional import mean_squared_error


class test_Metrics(ut.TestCase):
    def test_thresholder(self):
        torch.manual_seed(42)
        preds = torch.rand(100, dtype=torch.float32)
        target = torch.rand(100, dtype=torch.float32)

        th = 0.7
        preds_greater = preds > th
        target_greater = target > th

        # Test thresholder greater
        for th_on_preds in [True, False]:
            for th_on_target in [True, False]:
                thresholder = Thresholder(
                    threshold=th, operator="greater", th_on_target=th_on_target, th_on_preds=th_on_preds
                )
                preds2, target2 = thresholder(preds, target)
                if th_on_preds:
                    self.assertListEqual(preds2.tolist(), preds_greater.tolist())
                else:
                    self.assertListEqual(preds2.tolist(), preds.tolist())
                if th_on_target:
                    self.assertListEqual(target2.tolist(), target_greater.tolist())
                else:
                    self.assertListEqual(target2.tolist(), target.tolist())

        # Test thresholder lower
        for th_on_preds in [True, False]:
            for th_on_target in [True, False]:
                thresholder = Thresholder(
                    threshold=th, operator="lower", th_on_target=th_on_target, th_on_preds=th_on_preds
                )
                preds2, target2 = thresholder(preds, target)
                if th_on_preds:
                    self.assertListEqual(preds2.tolist(), (~preds_greater).tolist())
                else:
                    self.assertListEqual(preds2.tolist(), preds.tolist())
                if th_on_target:
                    self.assertListEqual(target2.tolist(), (~target_greater).tolist())
                else:
                    self.assertListEqual(target2.tolist(), target.tolist())


class test_MetricWrapper(ut.TestCase):
    def test_target_nan_mask(self):
        torch.random.manual_seed(42)

        for sz in [(100,), (100, 1), (100, 10)]:
            err_msg = f"Error for `sz = {sz}`"

            # Generate prediction and target matrices
            target = torch.rand(sz, dtype=torch.float32)
            preds = (0.5 * target) + (0.5 * torch.rand(sz, dtype=torch.float32))
            is_nan = torch.rand(sz) > 0.3
            target = (target > 0.5).to(torch.float32)
            target[is_nan] = float("nan")

            # Compute score with different ways of ignoring NaNs
            metric = MetricWrapper(metric="mse", target_nan_mask=None)
            score1 = metric(preds, target)
            self.assertTrue(torch.isnan(score1), msg=err_msg)

            # Replace NaNs by 0
            metric = MetricWrapper(metric="mse", target_nan_mask=0)
            score2 = metric(preds, target)

            this_target = target.clone()
            this_target[is_nan] = 0
            this_preds = preds.clone()
            self.assertAlmostEqual(score2, mean_squared_error(this_preds, this_target), msg=err_msg)

            # Replace NaNs by 1.5
            metric = MetricWrapper(metric="mse", target_nan_mask=1.5)
            score3 = metric(preds, target)

            this_target = target.clone()
            this_target[is_nan] = 1.5
            this_preds = preds.clone()
            self.assertAlmostEqual(score3, mean_squared_error(this_preds, this_target), msg=err_msg)

            # Flatten matrix and ignore NaNs
            metric = MetricWrapper(metric="mse", target_nan_mask="ignore", multitask_handling="flatten")
            score4 = metric(preds, target)

            this_target = target.clone()[~is_nan]
            this_preds = preds.clone()[~is_nan]
            self.assertAlmostEqual(score4, mean_squared_error(this_preds, this_target), msg=err_msg)

            # Ignore NaNs in each column and average the score
            metric = MetricWrapper(
                metric="mse", target_nan_mask="ignore", multitask_handling="mean-per-label"
            )
            score5 = metric(preds, target)

            this_target = target.clone()
            this_preds = preds.clone()
            this_is_nan = is_nan.clone()
            if len(sz) == 1:
                this_target = target.unsqueeze(-1)
                this_preds = preds.unsqueeze(-1)
                this_is_nan = is_nan.unsqueeze(-1)

            this_target = [this_target[:, ii][~this_is_nan[:, ii]] for ii in range(this_target.shape[1])]
            this_preds = [this_preds[:, ii][~this_is_nan[:, ii]] for ii in range(this_preds.shape[1])]
            mse = []
            for ii in range(len(this_preds)):
                mse.append(mean_squared_error(this_preds[ii], this_target[ii]))
            mse = torch.mean(torch.stack(mse))
            self.assertAlmostEqual(score5.tolist(), mse.tolist(), msg=err_msg)

    def test_pickling(self):
        pickle_file = os.path.join(tempfile.gettempdir(), "test_metric_pickled.pkl")
        metrics = ["mae", "mse", mean_squared_error]
        target_nan_masks = [None, 2, "ignore"]
        multitask_handlings = [None, "flatten", "mean-per-label"]
        squeeze_targets = [True, False]
        target_to_ints = [True, False]
        other_kwargs = [{}, {"squared": False}]
        thresholds = [
            None,
            {"threshold": 0.2, "operator": "greater"},
            {"threshold": 0.3, "operator": op.lt},
            {"threshold": 0.4, "operator": "lower"},
            {"threshold": 0.5, "operator": "lower", "th_on_preds": False, "th_on_target": True},
            {"threshold": 0.6, "operator": "lower"},
        ]

        for metric in metrics:
            for target_nan_mask in target_nan_masks:
                for kwargs in other_kwargs:
                    for threshold_kwargs in thresholds:
                        for multitask_handling in multitask_handlings:
                            for squeeze_target in squeeze_targets:
                                for target_to_int in target_to_ints:
                                    err_msg = f"{metric} - {target_nan_mask} - {kwargs} - {threshold_kwargs}"

                                    if (multitask_handling is None) and (target_nan_mask == "ignore"):
                                        # Raise with incompatible options
                                        with self.assertRaises(ValueError):
                                            MetricWrapper(
                                                metric=metric,
                                                threshold_kwargs=threshold_kwargs,
                                                target_nan_mask=target_nan_mask,
                                                multitask_handling=multitask_handling,
                                                squeeze_target=squeeze_target,
                                                target_to_int=target_to_int,
                                                **kwargs,
                                            )

                                    else:
                                        metric_wrapper = MetricWrapper(
                                            metric=metric,
                                            threshold_kwargs=threshold_kwargs,
                                            target_nan_mask=target_nan_mask,
                                            multitask_handling=multitask_handling,
                                            squeeze_target=squeeze_target,
                                            target_to_int=target_to_int,
                                            **kwargs,
                                        )

                                        # Check that the metric can be saved and re-loaded without error
                                        torch.save(metric_wrapper, pickle_file)
                                        metric_wrapper2 = torch.load(pickle_file)
                                        self.assertTrue(metric_wrapper == metric_wrapper2, msg=err_msg)

                                        # Check that the metric only contains primitive types
                                        state = metric_wrapper.__getstate__()
                                        if state["threshold_kwargs"] is not None:
                                            self.assertIsInstance(
                                                state["threshold_kwargs"], dict, msg=err_msg
                                            )
                                        if isinstance(metric, str):
                                            self.assertIsInstance(state["metric"], str, msg=err_msg)

    def test_classifigression_target_squeezing(self):
        preds = torch.Tensor([[0.1, 0.1, 0.3, 0.5, 0.0, 0.1, 0.0, 0.7, 0.2, 0.0]])
        target = torch.Tensor([3, 0])
        expected_scores = [0.5, 0.75]
        n_brackets = 5
        metrics = ["accuracy", "averageprecision"]
        other_kwargs = [
            {"task": "multiclass", "num_classes": n_brackets, "top_k": 1},
            {"task": "multiclass", "num_classes": n_brackets},
        ]

        for metric, kwargs, expected_score in zip(metrics, other_kwargs, expected_scores):
            metric_wrapper = MetricWrapper(
                metric=metric,
                multitask_handling="mean-per-label",
                squeeze_targets=True,
                target_to_int=True,
                **kwargs,
            )
            score = metric_wrapper(preds, target)

            assert score == expected_score


if __name__ == "__main__":
    ut.main()


# tests/test_mtl_architecture.py

"""
Unit tests for the different architectures of graphium/nn/architectures...

The layers are not thoroughly tested due to the difficulty of testing them
"""

import torch
import unittest as ut
from torch_geometric.data import Data, Batch
from copy import deepcopy
from itertools import product
import sys
import traceback

import graphium
from graphium.config._loader import load_architecture
from graphium.nn.architectures import TaskHeads, FullGraphMultiTaskNetwork, GraphOutputNN
from graphium.nn.base_layers import FCLayer

from graphium.utils.spaces import LAYERS_DICT

kwargs = {
    "activation": "relu",
    "last_activation": "none",
    "normalization": "none",
    "dropout": 0.2,
    "name": "LNN",
    "layer_type": FCLayer,
    "residual_type": "none",
    "residual_skip_steps": 1,
}
# task kwargs
task_1_kwargs = {
    "out_dim": 5,
    "task_level": "node",
    "hidden_dims": [5, 6, 7],
}
task_2_kwargs = {
    "out_dim": 3,
    "task_level": "edge",
    "hidden_dims": [8, 9, 10],
}
task_3_kwargs = {
    "out_dim": 4,
    "task_level": "graph",
    "hidden_dims": [2, 2, 2],
}
task_4_kwargs = {
    "out_dim": 4,
    "task_level": "nodepair",
    "hidden_dims": [2, 2, 2],
}

# level-wise kwargs
node_level_kwargs = {
    "out_dim": 8,
    "hidden_dims": [8, 9, 10],
    "activation": "relu",
    "last_activation": "none",
    "normalization": "none",
    "dropout": 0.2,
    "name": "LNN",
    "layer_type": FCLayer,
    "residual_type": "none",
    "residual_skip_steps": 1,
}
graph_level_kwargs = {
    "pooling": ["sum", "max"],
    "out_dim": 8,
    "hidden_dims": [8, 9, 10],
    "activation": "relu",
    "last_activation": "none",
    "normalization": "none",
    "dropout": 0.2,
    "name": "LNN",
    "layer_type": FCLayer,
    "residual_type": "none",
    "residual_skip_steps": 1,
}
edge_level_kwargs = {
    "out_dim": 8,
    "hidden_dims": [8, 9, 10],
    "activation": "relu",
    "last_activation": "none",
    "normalization": "none",
    "dropout": 0.2,
    "name": "LNN",
    "layer_type": FCLayer,
    "residual_type": "none",
    "residual_skip_steps": 1,
}

nodepair_level_kwargs = {
    "out_dim": 8,
    "hidden_dims": [8, 9, 10],
    "activation": "relu",
    "last_activation": "none",
    "normalization": "none",
    "dropout": 0.2,
    "name": "LNN",
    "layer_type": FCLayer,
    "residual_type": "none",
    "residual_skip_steps": 1,
}

task_1_params = {}
task_1_params.update(task_1_kwargs)
task_1_params.update(kwargs)
task_2_params = {}
task_2_params.update(task_2_kwargs)
task_2_params.update(kwargs)
task_3_params = {}
task_3_params.update(task_3_kwargs)
task_3_params.update(kwargs)
task_4_params = {}
task_4_params.update(task_4_kwargs)
task_4_params.update(kwargs)


def toy_test_data(in_dim=7, in_dim_edges=3):
    edge_idx1 = torch.stack([torch.tensor([0, 1, 2, 3, 2]), torch.tensor([1, 2, 3, 0, 0])])
    edge_idx2 = torch.stack([torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])])
    x1 = torch.randn(edge_idx1.max() + 1, in_dim, dtype=torch.float32)
    e1 = torch.randn(edge_idx1.shape[-1], in_dim_edges, dtype=torch.float32)
    x2 = torch.randn(edge_idx2.max() + 1, in_dim, dtype=torch.float32)
    e2 = torch.randn(edge_idx2.shape[-1], in_dim_edges, dtype=torch.float32)
    # edge_idx1, e1 = add_self_loops(edge_idx1, e1)
    # edge_idx2, e2 = add_self_loops(edge_idx2, e2)
    g1 = Data(feat=x1, edge_index=edge_idx1, edge_feat=e1)
    g2 = Data(feat=x2, edge_index=edge_idx2, edge_feat=e2)
    bg = Batch.from_data_list([g1, g2])
    num_nodes_per_graph = bg.batch.bincount()
    batch_node = bg["feat"].size()[0]
    batch_edge = bg["edge_feat"].size()[0]
    batch_graph = 2
    batch_nodepair = ((num_nodes_per_graph.max().item() ** 2) - num_nodes_per_graph.max().item()) // 2
    return bg, batch_node, batch_edge, batch_graph, batch_nodepair


class test_GraphOutputNN(ut.TestCase):
    def generate_test_data(self):
        g_1 = torch.tensor([[2, 3, 4], [0, 1, 0], [0, 1, 1], [1, 0, 0]])
        g_2 = torch.tensor([[3, 4, 0], [0, 1, 1], [1, 0, 0]])
        g_3 = torch.tensor([[3, 4, 0], [0, 1, 1], [1, 0, 0], [2, 2, 2], [4, 4, 4]])

        batch = torch.tensor([0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 2, 2])
        x = torch.concat([g_1, g_2, g_3]).float()

        nan = -1
        expected_result = [
            # graph 1
            [
                [2.0, 4.0, 4.0, 2.0, 2.0, 4.0],
                [2.0, 4.0, 5.0, 2.0, 2.0, 3.0],
                [3.0, 3.0, 4.0, 1.0, 3.0, 4.0],
                [nan, nan, nan, nan, nan, nan],
                [0.0, 2.0, 1.0, 0.0, 0.0, 1.0],
                [1.0, 1.0, 0.0, 1.0, 1.0, 0.0],
                [nan, nan, nan, nan, nan, nan],
                [1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
                [nan, nan, nan, nan, nan, nan],
                [nan, nan, nan, nan, nan, nan],
            ],
            # graph 2
            [
                [3.0, 5.0, 1.0, 3.0, 3.0, 1.0],
                [4.0, 4.0, 0.0, 2.0, 4.0, 0.0],
                [nan, nan, nan, nan, nan, nan],
                [nan, nan, nan, nan, nan, nan],
                [1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
                [nan, nan, nan, nan, nan, nan],
                [nan, nan, nan, nan, nan, nan],
                [nan, nan, nan, nan, nan, nan],
                [nan, nan, nan, nan, nan, nan],
                [nan, nan, nan, nan, nan, nan],
            ],
            # graph 3
            [
                [3.0, 5.0, 1.0, 3.0, 3.0, 1.0],
                [4.0, 4.0, 0.0, 2.0, 4.0, 0.0],
                [5.0, 6.0, 2.0, 1.0, 2.0, 2.0],
                [7.0, 8.0, 4.0, 1.0, 0.0, 4.0],
                [1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
                [2.0, 3.0, 3.0, 2.0, 1.0, 1.0],
                [4.0, 5.0, 5.0, 4.0, 3.0, 3.0],
                [3.0, 2.0, 2.0, 1.0, 2.0, 2.0],
                [5.0, 4.0, 4.0, 3.0, 4.0, 4.0],
                [6.0, 6.0, 6.0, 2.0, 2.0, 2.0],
            ],
        ]

        return x, batch, expected_result

    def test_nodepair_max_num_nodes_not_set(self):
        in_dim = 3
        in_dim_edges = 8

        graph_output_nn_kwargs = {
            "node": node_level_kwargs,
            "edge": edge_level_kwargs,
            "graph": graph_level_kwargs,
            "nodepair": nodepair_level_kwargs,
        }

        graph_output_nn = GraphOutputNN(
            in_dim=in_dim,
            in_dim_edges=in_dim_edges,
            task_level="nodepair",
            graph_output_nn_kwargs=graph_output_nn_kwargs,
        )

        x, batch, expected_result = self.generate_test_data()

        out = graph_output_nn.compute_nodepairs(node_feats=x, batch=batch)
        out = torch.nan_to_num(out, nan=-1)
        self.assertListEqual(expected_result, out.tolist())

    def test_nodepair_with_max_num_nodes(self):
        in_dim = 3
        in_dim_edges = 8

        graph_output_nn_kwargs = {
            "node": node_level_kwargs,
            "edge": edge_level_kwargs,
            "graph": graph_level_kwargs,
            "nodepair": nodepair_level_kwargs,
        }

        graph_output_nn = GraphOutputNN(
            in_dim=in_dim,
            in_dim_edges=in_dim_edges,
            task_level="nodepair",
            graph_output_nn_kwargs=graph_output_nn_kwargs,
        )

        max_num_nodes = 5  # if we change this value, we also have to change the expected result.
        x, batch, expected_result = self.generate_test_data()

        out = graph_output_nn.compute_nodepairs(node_feats=x, batch=batch, max_num_nodes=max_num_nodes)
        out = torch.nan_to_num(out, nan=-1)
        self.assertListEqual(expected_result, out.tolist())


class test_TaskHeads(ut.TestCase):
    def test_task_heads_forward(self):
        in_dim = 8  # Dimension of the incoming data
        in_dim_edges = 8

        task_heads_params = {
            "task_1": task_1_params,
            "task_2": task_2_params,
            "task_3": task_3_params,
            "task_4": task_4_params,
        }
        graph_output_nn_kwargs = {
            "node": node_level_kwargs,
            "edge": edge_level_kwargs,
            "graph": graph_level_kwargs,
            "nodepair": nodepair_level_kwargs,
        }
        # Create the "multitask" network. Really it's just an input going to various FFNNs since there's nothing shared.
        multi_head_nn = TaskHeads(
            in_dim=in_dim,
            in_dim_edges=in_dim_edges,
            task_heads_kwargs=task_heads_params,
            graph_output_nn_kwargs=graph_output_nn_kwargs,
        )

        # Test the sizes of the MLPs for each head
        # Head for task_1
        task_1_head = multi_head_nn.task_heads["task_1"]

        # Check the dimensions
        self.assertEqual(len(task_1_head.layers), len(task_1_kwargs["hidden_dims"]) + 1)
        self.assertEqual(task_1_head.layers[0].in_dim, in_dim)
        self.assertEqual(task_1_head.layers[1].in_dim, task_1_kwargs["hidden_dims"][0])
        self.assertEqual(task_1_head.layers[2].in_dim, task_1_kwargs["hidden_dims"][1])
        self.assertEqual(task_1_head.layers[3].in_dim, task_1_kwargs["hidden_dims"][2])

        # Head for task_2
        task_2_head = multi_head_nn.task_heads["task_2"]

        # Check the dimensions
        self.assertEqual(len(task_2_head.layers), len(task_2_kwargs["hidden_dims"]) + 1)
        self.assertEqual(task_2_head.layers[0].in_dim, in_dim)
        self.assertEqual(task_2_head.layers[1].in_dim, task_2_kwargs["hidden_dims"][0])
        self.assertEqual(task_2_head.layers[2].in_dim, task_2_kwargs["hidden_dims"][1])
        self.assertEqual(task_2_head.layers[3].in_dim, task_2_kwargs["hidden_dims"][2])

        # Head for task_3
        task_3_head = multi_head_nn.task_heads["task_3"]

        # Check the dimensions
        self.assertEqual(len(task_3_head.layers), len(task_3_kwargs["hidden_dims"]) + 1)
        self.assertEqual(task_3_head.layers[0].in_dim, in_dim)
        self.assertEqual(task_3_head.layers[1].in_dim, task_3_kwargs["hidden_dims"][0])
        self.assertEqual(task_3_head.layers[2].in_dim, task_3_kwargs["hidden_dims"][1])
        self.assertEqual(task_3_head.layers[3].in_dim, task_3_kwargs["hidden_dims"][2])

        # Head for task_4
        task_4_head = multi_head_nn.task_heads["task_3"]

        # Check the dimensions
        self.assertEqual(len(task_4_head.layers), len(task_4_kwargs["hidden_dims"]) + 1)
        self.assertEqual(task_4_head.layers[0].in_dim, in_dim)
        self.assertEqual(task_4_head.layers[1].in_dim, task_4_kwargs["hidden_dims"][0])
        self.assertEqual(task_4_head.layers[2].in_dim, task_4_kwargs["hidden_dims"][1])
        self.assertEqual(task_4_head.layers[3].in_dim, task_4_kwargs["hidden_dims"][2])

        # Check the output: It's a per-task prediction!
        bg, batch_node, batch_edge, batch_graph, batch_nodepair = toy_test_data(
            in_dim=in_dim, in_dim_edges=in_dim_edges
        )
        feat_out = multi_head_nn.forward(bg)

        self.assertListEqual(
            list(feat_out["task_1"].shape), [batch_node, task_1_kwargs["out_dim"]]
        )  # node level task
        self.assertListEqual(
            list(feat_out["task_2"].shape), [batch_edge, task_2_kwargs["out_dim"]]
        )  # edge level task
        self.assertListEqual(
            list(feat_out["task_3"].shape), [batch_graph, task_3_kwargs["out_dim"]]
        )  # graph level task
        self.assertListEqual(
            list(feat_out["task_4"].shape), [batch_graph, batch_nodepair, task_4_kwargs["out_dim"]]
        )  # nodepair level task

    def test_task_heads_non_supported_level(self):
        in_dim = 8  # Dimension of the incoming data
        in_dim_edges = 8

        fake_task = deepcopy(task_1_params)
        fake_task["task_level"] = "certainly_not_supported_level"
        task_heads_params = {
            "task_1": fake_task,
            "task_2": task_2_params,
            "task_3": task_3_params,
            "task_4": task_4_params,
        }
        graph_output_nn_kwargs = {
            "certainly_not_supported_level": node_level_kwargs,
            "edge": edge_level_kwargs,
            "graph": graph_level_kwargs,
            "nodepair": nodepair_level_kwargs,
        }

        with self.assertRaises(ValueError):
            TaskHeads(
                in_dim=in_dim,
                in_dim_edges=in_dim_edges,
                task_heads_kwargs=task_heads_params,
                graph_output_nn_kwargs=graph_output_nn_kwargs,
            )


class test_Multitask_NN(ut.TestCase):
    pyg_kwargs = {
        "activation": "relu",
        "last_activation": "none",
        "normalization": "none",
        "dropout": 0.2,
        "name": "LNN",
    }

    in_dim = 7
    fullgraph_out_dim = 11
    in_dim_edges = 3
    hidden_dims = [6, 6, 6, 6, 6]

    virtual_nodes = ["none", "mean", "sum"]
    norms = ["none", None, "batch_norm", "layer_norm"]
    pna_kwargs = {"aggregators": ["mean", "max", "sum"], "scalers": ["identity", "amplification"]}

    gnn_layers_kwargs = {
        "pyg:gin": {},
        "pyg:gine": {"in_dim_edges": in_dim_edges},
        "pyg:gated-gcn": {"in_dim_edges": in_dim_edges, "hidden_dims_edges": hidden_dims},
        "pyg:pna-msgpass#1": {"layer_kwargs": pna_kwargs, "in_dim_edges": 0},
        "pyg:pna-msgpass#2": {"layer_kwargs": pna_kwargs, "in_dim_edges": in_dim_edges},
    }

    def test_full_graph_multitask_forward(self):
        # Params for the network
        temp_dim_1 = 5
        temp_dim_2 = 7
        temp_dim_edges = 21

        default_graph_output_nn_kwargs = {
            "node": dict(out_dim=10, hidden_dims=[3, 3, 3, 3]),
            "edge": dict(out_dim=11, hidden_dims=[4, 4, 4, 4]),
            "graph": dict(out_dim=12, hidden_dims=[5, 5, 5, 5]),
            "nodepair": dict(out_dim=13, hidden_dims=[6, 6, 6, 6]),
        }

        # TODO: Test all combinations?
        all_task_heads_kwargs = dict(
            task_node=task_1_params,
            task_edge=task_2_params,
            task_graph=task_3_params,
            task_nodepair=task_4_params,
        )
        only_node_and_graph_task_heads_kwargs = dict(task_node=task_1_params, task_graph=task_3_params)

        # TODO: Configure properly in pytest best-practice
        options = product(
            [["sum"], ["mean", "max"]],
            [1, 2],
            self.virtual_nodes,
            self.norms,
            self.gnn_layers_kwargs.items(),
            [None, all_task_heads_kwargs, only_node_and_graph_task_heads_kwargs],
            [None, dict(in_dim=self.in_dim, out_dim=temp_dim_1, hidden_dims=[4, 4, 4, 4, 4])],
            [None, dict(in_dim=self.in_dim_edges, out_dim=temp_dim_edges, hidden_dims=[4, 4])],
        )
        for (
            pooling,
            residual_skip_steps,
            virtual_node,
            normalization,
            (layer_name, this_kwargs),
            task_heads_kwargs,
            pre_nn_kwargs,
            pre_nn_edges_kwargs,
        ) in options:
            err_msg = f"pooling={pooling}, virtual_node={virtual_node}, layer_name={layer_name}, residual_skip_steps={residual_skip_steps}, normalization={normalization}, task_heads={task_heads_kwargs}, pre_nn_kwargs={pre_nn_kwargs}, graph_output_nn_kwargs={default_graph_output_nn_kwargs}, pre_nn_edges_kwargs={pre_nn_edges_kwargs}"
            layer_type = layer_name.split("#")[0]

            # TODO: graph_output_nn is currently non-optional, should it be optional?
            if task_heads_kwargs is not None:
                continue

            # TODO: Allow to pass single graph_output_nn_kwargs to apply to all levels?

            bg, num_nodes, num_edges, num_graphs, num_nodepairs = toy_test_data(
                in_dim=self.in_dim, in_dim_edges=self.in_dim_edges
            )

            expected_dim_edges = self.in_dim_edges
            this_kwargs2 = deepcopy(this_kwargs)
            if pre_nn_edges_kwargs is not None:
                this_kwargs2["in_dim_edges"] = expected_dim_edges = temp_dim_edges

            gnn_kwargs = dict(
                in_dim=self.in_dim if pre_nn_kwargs is None else temp_dim_1,
                out_dim=temp_dim_2,
                hidden_dims=self.hidden_dims,
                residual_type="densenet",
                residual_skip_steps=residual_skip_steps,
                layer_type=layer_type,
                **this_kwargs2,
                **self.pyg_kwargs,
            )

            graph_output_nn_kwargs = deepcopy(default_graph_output_nn_kwargs)
            graph_output_nn_kwargs["graph"]["pooling"] = pooling

            expectFailure = False

            if (not LAYERS_DICT[layer_type].layer_supports_edges) and (pre_nn_edges_kwargs is not None):
                expectFailure = True

            if (not LAYERS_DICT[layer_type].layer_supports_edges) and (
                task_heads_kwargs is not None and "task_edge" in task_heads_kwargs
            ):
                expectFailure = True

            if (
                "in_dim_edges" in this_kwargs2
                and this_kwargs2["in_dim_edges"] < 1
                and (task_heads_kwargs is not None and "task_edge" in task_heads_kwargs)
            ):
                expectFailure = True

            if task_heads_kwargs is not None and "task_graph" in task_heads_kwargs:
                expectFailure = True

            if expectFailure:
                with self.assertRaises(ValueError):
                    multitask_graph_nn = FullGraphMultiTaskNetwork(
                        gnn_kwargs=gnn_kwargs,
                        pre_nn_kwargs=pre_nn_kwargs,
                        pre_nn_edges_kwargs=pre_nn_edges_kwargs,
                        task_heads_kwargs=task_heads_kwargs,
                        graph_output_nn_kwargs=graph_output_nn_kwargs,
                    )
                continue

            try:
                multitask_graph_nn = FullGraphMultiTaskNetwork(
                    gnn_kwargs=gnn_kwargs,
                    pre_nn_kwargs=pre_nn_kwargs,
                    pre_nn_edges_kwargs=pre_nn_edges_kwargs,
                    task_heads_kwargs=task_heads_kwargs,
                    graph_output_nn_kwargs=graph_output_nn_kwargs,
                )

                batch_out = multitask_graph_nn.forward(bg)
            except Exception as e:
                exc_type, exc_value, exc_traceback = sys.exc_info()
                msg = err_msg + "\n" + str(traceback.format_exception(exc_type, exc_value, exc_traceback))
                self.fail(msg)

            if task_heads_kwargs is not None:
                if "task_node" in task_heads_kwargs:
                    self.assertListEqual(
                        list(batch_out["task_node"].shape),
                        [num_nodes, task_1_kwargs["out_dim"]],
                        msg=err_msg,
                    )
                if "task_edge" in task_heads_kwargs:
                    self.assertListEqual(
                        list(batch_out["task_edge"].shape),
                        [num_edges, task_2_kwargs["out_dim"]],
                        msg=err_msg,
                    )
                if "task_graph" in task_heads_kwargs:
                    self.assertListEqual(
                        list(batch_out["task_graph"].shape),
                        [num_graphs, task_3_kwargs["out_dim"]],
                        msg=err_msg,
                    )
                if "task_nodepair" in task_heads_kwargs:
                    self.assertListEqual(
                        list(batch_out["task_nodepair"].shape),
                        [num_nodepairs, task_4_kwargs["out_dim"]],
                        msg=err_msg,
                    )
            else:
                feat_out = batch_out["feat"]
                edge_feat_out = batch_out["edge_feat"]

                out_dim_edges = (
                    multitask_graph_nn.gnn.full_dims_edges[-1]
                    if multitask_graph_nn.gnn.full_dims_edges is not None
                    else expected_dim_edges
                )

                if not LAYERS_DICT[layer_type].layer_supports_edges or this_kwargs2["in_dim_edges"] == 0:
                    self.assertEqual(multitask_graph_nn.out_dim_edges, 0)
                else:
                    self.assertEqual(multitask_graph_nn.out_dim_edges, out_dim_edges)
                self.assertListEqual(list(feat_out.shape), [num_nodes, temp_dim_2], msg=err_msg)
                self.assertListEqual(list(edge_feat_out.shape), [num_edges, out_dim_edges], msg=err_msg)

    def test_full_graph_multi_task_from_config(self):
        cfg = graphium.load_config(name="zinc_default_multitask_pyg")

        # Initialize the network
        in_dims = {"feat": self.in_dim, "edge_feat": self.in_dim_edges}
        model_class, model_kwargs = load_architecture(cfg, in_dims=in_dims)

        multitask_full_graph_nn = model_class(**model_kwargs)

        # Test
        bg, num_nodes, num_edges, num_graphs, num_nodepairs = toy_test_data(
            in_dim=self.in_dim, in_dim_edges=self.in_dim_edges
        )
        batch_out = multitask_full_graph_nn.forward(bg)

        self.assertListEqual(
            list(batch_out["task_1"].shape),
            [num_nodes, task_1_kwargs["out_dim"]],
        )
        self.assertListEqual(
            list(batch_out["task_2"].shape),
            [num_edges, task_2_kwargs["out_dim"]],
        )
        self.assertListEqual(
            list(batch_out["task_3"].shape),
            [num_graphs, task_3_kwargs["out_dim"]],
        )
        self.assertListEqual(
            list(batch_out["task_4"].shape),
            [num_graphs, num_nodepairs, task_4_kwargs["out_dim"]],
        )

    def test_full_graph_multi_task_set_max_num_nodes(self):
        cfg = graphium.load_config(name="zinc_default_multitask_pyg")

        # Initialize the network
        in_dims = {"feat": self.in_dim, "edge_feat": self.in_dim_edges}
        model_class, model_kwargs = load_architecture(cfg, in_dims=in_dims)

        multitask_full_graph_nn: FullGraphMultiTaskNetwork = model_class(**model_kwargs)
        multitask_full_graph_nn.set_max_num_nodes_edges_per_graph(10, 4)

        # Probing
        self.assertEqual(
            multitask_full_graph_nn.task_heads.graph_output_nn["nodepair"].max_num_nodes_per_graph, 10
        )
        self.assertEqual(multitask_full_graph_nn.gnn.layers[2].max_num_edges_per_graph, 4)


if __name__ == "__main__":
    ut.main()


# tests/test_multitask_datamodule.py

import unittest as ut
from omegaconf import OmegaConf
import pandas as pd
import numpy as np
import graphium


class Test_Multitask_DataModule(ut.TestCase):
    def test_multitask_fromsmiles_dm(
        self,
    ):  # TODO: I think we can remove this as it tests tiny_zinc which only contain graph level labels
        """Cover similar testing as for the original data module."""
        df = graphium.data.load_tiny_zinc()  # 100 molecules

        # Here we take the microzinc dataset and split the labels up into 'SA', 'logp' and 'score' in order to simulate having multiple single-task datasets
        df_micro_zinc_SA = df[["SMILES", "SA"]]
        df_micro_zinc_logp = df[["SMILES", "logp"]]
        df_micro_zinc_score = df[["SMILES", "score"]]

        # Setup the featurization. This will be the same across all tasks.
        featurization_args = {}
        featurization_args["atom_property_list_float"] = []  # ["weight", "valence"]
        featurization_args["atom_property_list_onehot"] = ["atomic-number", "degree"]
        featurization_args["edge_property_list"] = ["in-ring", "bond-type-onehot"]
        featurization_args["add_self_loop"] = False
        featurization_args["use_bonds_weights"] = False
        featurization_args["explicit_H"] = False

        # Config for multitask datamodule.

        # Per-task arguments.
        dm_task_args_SA = {}
        dm_task_args_SA["df"] = df_micro_zinc_SA
        dm_task_args_SA["task_level"] = "graph"
        dm_task_args_SA["smiles_col"] = "SMILES"
        dm_task_args_SA["label_cols"] = ["SA"]
        dm_task_args_SA["split_val"] = 0.2
        dm_task_args_SA["split_test"] = 0.2
        dm_task_args_SA["seed"] = 19
        dm_task_args_SA["splits_path"] = None  # This may not always be provided
        dm_task_args_SA["sample_size"] = None  # This may not always be provided
        dm_task_args_SA["idx_col"] = None  # This may not always be provided
        dm_task_args_SA["weights_col"] = None  # This may not always be provided
        dm_task_args_SA["weights_type"] = None  # This may not always be provided

        dm_task_args_logp = {}
        dm_task_args_logp["df"] = df_micro_zinc_logp
        dm_task_args_logp["task_level"] = "graph"
        dm_task_args_logp["smiles_col"] = "SMILES"
        dm_task_args_logp["label_cols"] = ["logp"]
        dm_task_args_logp["split_val"] = 0.2
        dm_task_args_logp["split_test"] = 0.2
        dm_task_args_logp["seed"] = 19
        dm_task_args_logp["splits_path"] = None  # This may not always be provided
        dm_task_args_logp["sample_size"] = None  # This may not always be provided
        dm_task_args_logp["idx_col"] = None  # This may not always be provided
        dm_task_args_logp["weights_col"] = None  # This may not always be provided
        dm_task_args_logp["weights_type"] = None  # This may not always be provided

        dm_task_args_score = {}
        dm_task_args_score["df"] = df_micro_zinc_score
        dm_task_args_score["task_level"] = "graph"
        dm_task_args_score["smiles_col"] = "SMILES"
        dm_task_args_score["label_cols"] = ["score"]
        dm_task_args_score["split_val"] = 0.2
        dm_task_args_score["split_test"] = 0.2
        dm_task_args_score["seed"] = 19
        dm_task_args_score["splits_path"] = None  # This may not always be provided
        dm_task_args_score["sample_size"] = None  # This may not always be provided
        dm_task_args_score["idx_col"] = None  # This may not always be provided
        dm_task_args_score["weights_col"] = None  # This may not always be provided
        dm_task_args_score["weights_type"] = None  # This may not always be provided

        dm_task_kwargs = {}
        dm_task_kwargs["SA"] = dm_task_args_SA
        dm_task_kwargs["logp"] = dm_task_args_logp
        dm_task_kwargs["score"] = dm_task_args_score

        dm_args = {}

        # Task-specific arguments for the datamodule
        dm_args["task_specific_args"] = dm_task_kwargs

        # Task-independent arguments
        dm_args["featurization"] = featurization_args
        dm_args["featurization_n_jobs"] = 16
        dm_args["featurization_progress"] = True
        dm_args["featurization_backend"] = "loky"
        dm_args["num_workers"] = 0
        dm_args["pin_memory"] = True
        dm_args["cache_data_path"] = None
        dm_args["batch_size_training"] = 16
        dm_args["batch_size_inference"] = 16

        # Create the data module
        dm = graphium.data.MultitaskFromSmilesDataModule(**dm_args)

        # self.assertEqual(50, dm.num_node_feats)    # Not implemeneted error
        # self.assertEqual(6, dm.num_edge_feats)

        dm.prepare_data()
        dm.setup()

        # self.assertEqual(len(dm), 100)                      # Should this have a fixed value for when it's initialized? MTL dataset only gets created after.
        self.assertEqual(len(dm.train_ds), 60)  # type: ignore
        self.assertEqual(len(dm.val_ds), 20)  # type: ignore
        self.assertEqual(len(dm.test_ds), 20)  # type: ignore
        # assert dm.num_node_feats == 50
        # assert dm.num_edge_feats == 6

        for dl in [dm.train_dataloader(), dm.val_dataloader(), dm.test_dataloader()]:
            it = iter(dl)
            batch = next(it)

            assert set(batch.keys()) == {"labels", "features"}

            # assert batch["labels"].shape == (16, 1)            # Single-task case
            assert batch["labels"]["graph_SA"].shape == (16, 1)
            assert batch["labels"]["graph_logp"].shape == (16, 1)
            assert batch["labels"]["graph_score"].shape == (16, 1)

    def test_multitask_fromsmiles_from_config(self):
        config = graphium.load_config(name="zinc_default_multitask_pyg")

        df = graphium.data.load_tiny_zinc()  # 100 molecules

        # Here we take the microzinc dataset and split the labels up into 'SA', 'logp' and 'score' in order to simulate having multiple single-task datasets
        df_micro_zinc_SA = df[["SMILES", "SA"]]
        df_micro_zinc_logp = df[["SMILES", "logp"]]
        df_micro_zinc_score = df[["SMILES", "score"]]

        # dm_args = dict(config.datamodule.args)
        dm_args = OmegaConf.to_container(config.datamodule.args, resolve=True)
        # dm_args["task_specific_args"]["SA"]["df"] = df
        dm_args["task_specific_args"]["SA"]["df"] = df_micro_zinc_SA
        dm_args["task_specific_args"]["logp"]["df"] = df_micro_zinc_logp
        dm_args["task_specific_args"]["score"]["df"] = df_micro_zinc_score

        dm_args["task_specific_args"]["SA"]["smiles_col"] = "SMILES"
        dm_args["task_specific_args"]["logp"]["smiles_col"] = "SMILES"
        dm_args["task_specific_args"]["score"]["smiles_col"] = "SMILES"

        dm_args["task_specific_args"]["SA"]["label_cols"] = ["SA"]
        dm_args["task_specific_args"]["logp"]["label_cols"] = ["logp"]
        dm_args["task_specific_args"]["score"]["label_cols"] = ["score"]

        dm_args["task_specific_args"]["SA"]["df_path"] = None
        dm_args["task_specific_args"]["logp"]["df_path"] = None
        dm_args["task_specific_args"]["score"]["df_path"] = None

        dm = graphium.data.MultitaskFromSmilesDataModule(**dm_args)

        # assert dm.num_node_feats == 50
        # assert dm.num_edge_feats == 6

        dm.prepare_data()
        dm.setup()

        # self.assertEqual(len(dm), 100)                      # Should this have a fixed value for when it's initialized? MTL dataset only gets created after.
        self.assertEqual(len(dm.train_ds), 60)  # type: ignore
        self.assertEqual(len(dm.val_ds), 20)  # type: ignore
        self.assertEqual(len(dm.test_ds), 20)  # type: ignore
        # assert dm.num_node_feats == 50
        # assert dm.num_edge_feats == 6

        for dl in [dm.train_dataloader(), dm.val_dataloader(), dm.test_dataloader()]:
            it = iter(dl)
            batch = next(it)

            assert set(batch.keys()) == {"labels", "features"}

            # assert batch["labels"].shape == (16, 1)            # Single-task case
            assert batch["labels"]["graph_SA"].shape == (16, 1)
            assert batch["labels"]["graph_logp"].shape == (16, 1)
            assert batch["labels"]["graph_score"].shape == (16, 1)

    def test_multitask_fromsmiles_from_config_csv(self):
        config = graphium.load_config(name="zinc_default_multitask_pyg")

        dm_args = OmegaConf.to_container(config.datamodule.args, resolve=True)
        dm = graphium.data.MultitaskFromSmilesDataModule(**dm_args)

        dm.prepare_data()
        dm.setup()

        # self.assertEqual(len(dm), 100)                      # Should this have a fixed value for when it's initialized? MTL dataset only gets created after.
        self.assertEqual(len(dm.train_ds), 60)  # type: ignore
        self.assertEqual(len(dm.val_ds), 20)  # type: ignore
        self.assertEqual(len(dm.test_ds), 20)  # type: ignore
        # assert dm.num_node_feats == 50
        # assert dm.num_edge_feats == 6

        for dl in [dm.train_dataloader(), dm.val_dataloader(), dm.test_dataloader()]:
            it = iter(dl)
            batch = next(it)

            assert set(batch.keys()) == {"labels", "features"}

            # assert batch["labels"].shape == (16, 1)            # Single-task case
            assert batch["labels"]["graph_SA"].shape == (16, 1)
            assert batch["labels"]["graph_logp"].shape == (16, 1)
            assert batch["labels"]["graph_score"].shape == (16, 1)

    def test_multitask_fromsmiles_from_config_parquet(self):
        config = graphium.load_config(name="fake_multilevel_multitask_pyg")

        dm_args = OmegaConf.to_container(config.datamodule.args, resolve=True)
        dm = graphium.data.MultitaskFromSmilesDataModule(**dm_args)

        dm.prepare_data()
        dm.setup()

        # self.assertEqual(len(dm), 100)                      # Should this have a fixed value for when it's initialized? MTL dataset only gets created after.
        self.assertEqual(len(dm.train_ds), 602)  # type: ignore
        self.assertEqual(len(dm.val_ds), 201)  # type: ignore
        self.assertEqual(len(dm.test_ds), 201)  # type: ignore
        # assert dm.num_node_feats == 50
        # assert dm.num_edge_feats == 6

        for dl in [dm.train_dataloader(), dm.val_dataloader(), dm.test_dataloader()]:
            dl = dm.train_dataloader()
            it = iter(dl)
            batch = next(it)

            assert set(batch.keys()) == {"labels", "features"}

            # assert batch["labels"].shape == (16, 1)            # Single-task case
            assert batch["labels"]["graph_SA"].shape == (16, 1)
            assert batch["labels"]["node_logp"].shape == (
                batch["features"].feat.size(0),
                2,
            )  # test node level
            assert batch["labels"]["edge_score"].shape == (
                batch["features"].edge_feat.size(0),
                2,
            )  # test edge level

    def test_multitask_with_missing_fromsmiles_from_config_parquet(self):
        config = graphium.load_config(name="fake_and_missing_multilevel_multitask_pyg")

        dm_args = OmegaConf.to_container(config.datamodule.args, resolve=True)
        dm = graphium.data.MultitaskFromSmilesDataModule(**dm_args)

        dm.prepare_data()
        dm.setup()

        # self.assertEqual(len(dm), 100)                      # Should this have a fixed value for when it's initialized? MTL dataset only gets created after.
        self.assertEqual(len(dm.train_ds), 602)  # type: ignore
        self.assertEqual(len(dm.val_ds), 201)  # type: ignore
        self.assertEqual(len(dm.test_ds), 201)  # type: ignore
        # assert dm.num_node_feats == 50
        # assert dm.num_edge_feats == 6

        for dl in [dm.train_dataloader(), dm.val_dataloader(), dm.test_dataloader()]:
            dl = dm.train_dataloader()
            it = iter(dl)
            batch = next(it)

            assert set(batch.keys()) == {"labels", "features"}

            # assert batch["labels"].shape == (16, 1)            # Single-task case
            assert batch["labels"]["graph_SA"].shape == (16, 1)
            assert batch["labels"]["node_logp"].shape == (
                batch["features"].feat.size(0),
                2,
            )  # test node level
            assert batch["labels"]["edge_score"].shape == (
                batch["features"].edge_feat.size(0),
                2,
            )  # test edge level

    def test_extract_graph_level_singletask(self):
        df = pd.read_parquet(f"tests/converted_fake_multilevel_data.parquet")
        num_graphs = len(df)
        label_cols = ["graph_label"]
        output = graphium.data.datamodule.extract_labels(df, "graph", label_cols)

        assert isinstance(output, np.ndarray)
        assert len(output.shape) == 2
        assert output.shape[0] == num_graphs
        assert output.shape[1] == 1

    def test_extract_graph_level_multitask(self):
        df = pd.read_parquet(f"tests/converted_fake_multilevel_data.parquet")
        num_graphs = len(df)
        label_cols = ["graph_label", "graph_label"]
        output = graphium.data.datamodule.extract_labels(df, "graph", label_cols)

        assert isinstance(output, np.ndarray)
        assert len(output.shape) == 2
        assert output.shape[0] == num_graphs
        assert output.shape[1] == len(label_cols)

    def test_extract_graph_level_multitask_missing_cols(self):
        df = pd.read_parquet(f"tests/converted_fake_multilevel_data.parquet")
        num_graphs = len(df)
        label_cols = ["graph_label", "graph_label"]

        drop_index = [2, 5, 21, 237, 192, 23, 127, 11]
        for replace in [1, 2]:
            for missing_col in label_cols[:replace]:
                df[missing_col].iloc[drop_index] = None

            output = graphium.data.datamodule.extract_labels(df, "graph", label_cols)

            assert isinstance(output, np.ndarray)
            assert len(output.shape) == 2
            assert output.shape[0] == num_graphs
            assert output.shape[1] == len(label_cols)

    def test_non_graph_level_extract_labels(self):
        df = pd.read_parquet(f"tests/converted_fake_multilevel_data.parquet")

        for level in ["node", "edge", "nodepair"]:
            label_cols = [f"{level}_label_{suffix}" for suffix in ["list", "np"]]
            output = graphium.data.datamodule.extract_labels(df, level, label_cols)

            assert isinstance(output, list)
            assert len(output[0].shape) == 2
            assert output[0].shape[1] == len(label_cols)

    def test_non_graph_level_extract_labels_missing_cols(self):
        df = pd.read_parquet(f"tests/converted_fake_multilevel_data.parquet")

        for level in ["node", "edge", "nodepair"]:
            label_cols = [f"{level}_label_{suffix}" for suffix in ["list", "np"]]
            drop_index = [2, 5, 21, 237, 192, 23, 127, 11]
            for replace in [1, 2]:
                for missing_col in label_cols[:replace]:
                    df.loc[drop_index, missing_col] = None

                output = graphium.data.datamodule.extract_labels(df, level, label_cols)

                for idx in drop_index:
                    assert len(output[idx].shape) == 2
                    assert output[idx].shape[1] == len(label_cols)

                    # Check that number of labels is adjusted correctly
                    if replace == 1:
                        non_missing_col = label_cols[1]
                        assert output[idx].shape[0] == len(df[non_missing_col][idx])


if __name__ == "__main__":
    ut.main()


# tests/test_mup.py

"""
Unit tests for the implementation of mup
"""

import unittest as ut
from copy import deepcopy
import torch.nn as nn
import torch
import yaml

from torch_geometric.data import Batch, Data

from graphium.nn.architectures import FeedForwardNN, FeedForwardPyg, FullGraphMultiTaskNetwork


def get_pyg_graphs(in_dim, in_dim_edges):
    edge_idx1 = torch.stack([torch.tensor([0, 1, 2, 3, 2]), torch.tensor([1, 2, 3, 0, 0])])
    edge_idx2 = torch.stack([torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])])
    x1 = torch.randn(edge_idx1.max() + 1, in_dim, dtype=torch.float32)
    e1 = torch.randn(edge_idx1.shape[-1], in_dim_edges, dtype=torch.float32)
    x2 = torch.randn(edge_idx2.max() + 1, in_dim, dtype=torch.float32)
    e2 = torch.randn(edge_idx2.shape[-1], in_dim_edges, dtype=torch.float32)
    g1 = Data(feat=x1, edge_index=edge_idx1, edge_feat=e1)
    g2 = Data(feat=x2, edge_index=edge_idx2, edge_feat=e2)
    bg = Batch.from_data_list([g1, g2])

    return bg


class test_mup(ut.TestCase):
    kwargs = dict(
        in_dim=12,
        out_dim=60,
        hidden_dims=8 * [84],
        depth=None,
        activation="LeakyReLU",
        last_activation="LeakyReLU",
        dropout=0.1,
        last_dropout=0.2,
        normalization="batch_norm",
        first_normalization="batch_norm",
        last_normalization="batch_norm",
        residual_type="simple",
        residual_skip_steps=2,
        name="testing",
        layer_type="fc",
        layer_kwargs=None,
    )

    def test_feedforwardnn_mup(self):
        kwargs = deepcopy(self.kwargs)
        model = FeedForwardNN(**kwargs, last_layer_is_readout=False)
        model_lastreadout = FeedForwardNN(**kwargs, last_layer_is_readout=True)
        base_1 = model.make_mup_base_kwargs(divide_factor=1)
        base_1_lastreadout = model_lastreadout.make_mup_base_kwargs(divide_factor=1)

        base_2 = model.make_mup_base_kwargs(divide_factor=2)
        base_2_lastreadout = model_lastreadout.make_mup_base_kwargs(divide_factor=2)
        kwargs_2 = deepcopy(base_1)
        kwargs_2.update(dict(out_dim=30, hidden_dims=8 * [42]))
        kwargs_2_lastreadout = deepcopy(base_1_lastreadout)
        kwargs_2_lastreadout.update(dict(hidden_dims=8 * [42]))

        # Check the kwargs matching
        for key in kwargs_2.keys():
            if isinstance(kwargs_2[key], nn.Module):
                # Can't match the random weights
                self.assertEqual(str(kwargs_2[key]), str(base_2[key]), msg=key)
            else:
                self.assertEqual(kwargs_2[key], base_2[key], msg=key)

        # Check the kwargs matching
        for key in kwargs_2_lastreadout.keys():
            if isinstance(kwargs_2_lastreadout[key], nn.Module):
                # Can't match the random weights
                self.assertEqual(str(kwargs_2_lastreadout[key]), str(base_2_lastreadout[key]), msg=key)
            else:
                self.assertEqual(kwargs_2_lastreadout[key], base_2_lastreadout[key], msg=key)

        # Test that the models with divide_factor=1 can be built run a forward pass
        in_features = torch.randn((10, kwargs["in_dim"]))
        model_1 = FeedForwardNN(**base_1)
        model_1.forward(deepcopy(in_features))
        model_1_lastreadout = FeedForwardNN(**base_1_lastreadout)
        model_1_lastreadout.forward(deepcopy(in_features))

        # Test that the models with divide_factor=2 can be built run a forward pass
        model_2 = FeedForwardNN(**base_2)
        model_2.forward(deepcopy(in_features))
        model_2_lastreadout = FeedForwardNN(**base_2_lastreadout)
        model_2_lastreadout.forward(deepcopy(in_features))

    def test_feedforwardgraph_mup(self):
        kwargs = deepcopy(self.kwargs)
        in_dim_edges = kwargs["in_dim"]
        kwargs.update(dict(layer_type="pyg:gine", in_dim_edges=in_dim_edges))
        model = FeedForwardPyg(**kwargs, last_layer_is_readout=False)
        model_lastreadout = FeedForwardPyg(**kwargs, last_layer_is_readout=True)
        base_1 = model.make_mup_base_kwargs(divide_factor=1)
        base_1_lastreadout = model_lastreadout.make_mup_base_kwargs(divide_factor=1)

        base_2 = model.make_mup_base_kwargs(divide_factor=2)
        base_2_lastreadout = model_lastreadout.make_mup_base_kwargs(divide_factor=2)
        kwargs_2 = deepcopy(base_1)
        kwargs_2.update(dict(out_dim=30, hidden_dims=8 * [42]))
        kwargs_2_lastreadout = deepcopy(base_1_lastreadout)
        kwargs_2_lastreadout.update(dict(hidden_dims=8 * [42]))

        # Check the kwargs matching
        for key in kwargs_2.keys():
            if isinstance(kwargs_2[key], nn.Module):
                # Can't match the random weights
                self.assertEqual(str(kwargs_2[key]), str(base_2[key]), msg=key)
            else:
                self.assertEqual(kwargs_2[key], base_2[key], msg=key)

        # Check the kwargs matching
        for key in kwargs_2_lastreadout.keys():
            if isinstance(kwargs_2_lastreadout[key], nn.Module):
                # Can't match the random weights
                self.assertEqual(str(kwargs_2_lastreadout[key]), str(base_2_lastreadout[key]), msg=key)
            else:
                self.assertEqual(kwargs_2_lastreadout[key], base_2_lastreadout[key], msg=key)

        # Test that the models with divide_factor=1 can be built run a forward pass
        in_features = get_pyg_graphs(in_dim=kwargs["in_dim"], in_dim_edges=in_dim_edges)
        model_1 = FeedForwardPyg(**base_1)
        model_1.forward(deepcopy(in_features))
        model_1_lastreadout = FeedForwardPyg(**base_1_lastreadout)
        model_1_lastreadout.forward(deepcopy(in_features))

        # Test that the models with divide_factor=2 can be built run a forward pass
        model_2 = FeedForwardPyg(**base_2)
        model_2.forward(deepcopy(in_features))
        model_2_lastreadout = FeedForwardPyg(**base_2_lastreadout)
        model_2_lastreadout.forward(deepcopy(in_features))

    def test_fullgraphmultitasknetwork(self):
        # Load the configuration file for the model
        CONFIG_FILE = "tests/config_test_ipu_dataloader.yaml"
        with open(CONFIG_FILE, "r") as f:
            cfg = yaml.safe_load(f)

        # Make fake graphs
        in_dim = 12
        in_dim_edges = 12
        pe_indims = {
            "laplacian_eigvec": 5,
            "laplacian_eigval": 5,
            "rw_return_probs": 2,
            "edge_rw_transition_probs": 2,
            "nodepair_rw_return_probs": 1,
            "electrostatic": 6,
            "edge_commute": 1,
            "nodepair_graphormer": 1,
            "positions_3d": 3,
        }
        in_features = get_pyg_graphs(in_dim=in_dim, in_dim_edges=in_dim_edges)
        for key, dim in pe_indims.items():
            if key.startswith("edge_"):
                in_features[key] = torch.randn(in_features.num_edges, dim)
            elif key.startswith("nodepair_"):
                in_features[key] = torch.randn(in_features.num_nodes, in_features.num_nodes, dim)
            else:
                in_features[key] = torch.randn(in_features.num_nodes, dim)

        # Load the model
        kwargs = {}
        for key, val in cfg["architecture"].items():
            if key in ["model_type", "mup_base_path"]:
                continue
            kwargs[key + "_kwargs"] = val
        kwargs["pre_nn_kwargs"]["in_dim"] = in_dim + kwargs["pe_encoders_kwargs"]["out_dim"]
        kwargs["pre_nn_edges_kwargs"]["in_dim"] = in_dim_edges + kwargs["pe_encoders_kwargs"]["edge_out_dim"]
        kwargs["pe_encoders_kwargs"]["in_dims"] = pe_indims

        model = FullGraphMultiTaskNetwork(**kwargs, last_layer_is_readout=True)

        kw_1 = model.make_mup_base_kwargs(divide_factor=1)
        kw_2 = model.make_mup_base_kwargs(divide_factor=2)

        # Check the parameter sizes
        for key, elem in kw_1.items():
            if not isinstance(elem, dict):
                continue
            for subkey, subelem in elem.items():
                if "dim" in subkey:
                    match = f"{key}:{subkey}"
                    if match in ["pre_nn_kwargs:in_dim"]:
                        self.assertNotEqual(subelem, kw_2[key][subkey], msg=match)
                    elif match in [
                        "pre_nn_kwargs:out_dim",
                        "pre_nn_edges_kwargs:out_dim",
                        "gnn_kwargs:in_dim",
                        "graph_output_nn_kwargs:in_dim",
                        "gnn_kwargs:out_dim",
                        "gnn_kwargs:in_dim_edges",
                        "pe_encoders_kwargs:out_dim",
                        "graph_output_nn_kwargs:out_dim",
                    ]:
                        # Divide by 2
                        self.assertEqual(round(subelem / 2), kw_2[key][subkey], msg=match)
                    elif match in [
                        "pre_nn_kwargs:hidden_dims",
                        "pre_nn_edges_kwargs:hidden_dims",
                        "gnn_kwargs:hidden_dims",
                        "graph_output_nn_kwargs:hidden_dims",
                        "gnn_kwargs:hidden_dims_edges",
                    ]:
                        # Arrays divide by 2
                        new_list = [round(e / 2) for e in subelem]
                        self.assertListEqual(new_list, kw_2[key][subkey], msg=match)
                elif subkey in ["homo", "alpha", "cv"]:
                    for subsubkey, subsubelem in subelem.items():
                        match = f"{key}:{subsubkey}"
                        if match in [
                            "task_heads_kwargs:out_dim",
                            "task_heads_kwargs:out_dim",
                            "task_heads_kwargs:out_dim",
                        ]:
                            # No divide
                            self.assertEqual(subsubelem, kw_2[key][subkey][subsubkey], msg=match)
                        elif match in [
                            "task_heads_kwargs:in_dim",
                            "task_heads_kwargs:in_dim",
                            "task_heads_kwargs:in_dim",
                        ]:
                            # Divide by 2
                            self.assertEqual(round(subsubelem / 2), kw_2[key][subkey][subsubkey], msg=match)
                        elif match in [
                            "task_heads_kwargs:hidden_dims",
                            "task_heads_kwargs:hidden_dims",
                            "task_heads_kwargs:hidden_dims",
                        ]:
                            # Divide by 2 a list
                            new_list = [round(e / 2) for e in subsubelem]
                            self.assertListEqual(new_list, kw_2[key][subkey][subsubkey], msg=match)

        # Test that the models with divide_factor=1 can be built run a forward pass
        kw_1["last_layer_is_readout"] = False
        model_1 = FullGraphMultiTaskNetwork(**kw_1)
        model_1.forward(deepcopy(in_features))
        kw_1["last_layer_is_readout"] = True
        model_1 = FullGraphMultiTaskNetwork(**kw_1)
        model_1.forward(deepcopy(in_features))

        # Test that the models with divide_factor=2 can be built run a forward pass
        kw_2["last_layer_is_readout"] = False
        model_2 = FullGraphMultiTaskNetwork(**kw_2)
        model_2.forward(deepcopy(in_features))
        kw_2["last_layer_is_readout"] = True
        model_2 = FullGraphMultiTaskNetwork(**kw_2)
        model_2.forward(deepcopy(in_features))


if __name__ == "__main__":
    ut.main()


# tests/test_packing.py

# General imports
import unittest as ut
import numpy as np

import torch
from torch_geometric.data import Data, Batch

# Current library imports
from graphium.utils.packing import (
    smart_packing,
    get_pack_sizes,
    fast_packing,
    hybrid_packing,
    node_to_pack_indices_mask,
)


def random_packing(num_nodes, batch_size):
    ipu_batch_size = int(len(num_nodes) / batch_size)
    indices = np.arange(len(num_nodes))
    np.random.shuffle(indices)
    indices = np.reshape(indices, (ipu_batch_size, batch_size)).tolist()
    return indices


class test_Packing(ut.TestCase):
    def test_smart_packing(self):
        np.random.seed(42)

        batch_sizes = [2, 4, 8, 16, 32, 64]
        ipu_batch_sizes = [2, 3, 4, 8, 16, 32, 64]

        for batch_size in batch_sizes:
            for ipu_batch_size in ipu_batch_sizes:
                err_msg = f"bz={batch_size}, ipu_bz={ipu_batch_size}"

                # Generate random batch size
                global_batch = batch_size * ipu_batch_size
                num_nodes = np.abs(np.random.gamma(2, 20, size=global_batch)).astype(int)

                # Use the smart packing
                packed_indices = smart_packing(num_nodes=num_nodes, batch_size=batch_size)
                pack_num_nodes = get_pack_sizes(packed_indices, num_nodes)

                # Use the random packing
                rand_packed_indices = random_packing(num_nodes=num_nodes, batch_size=batch_size)
                rand_pack_num_nodes = get_pack_sizes(rand_packed_indices, num_nodes)

                # Assert that the smart packing is better than the random packing
                self.assertLessEqual(max(pack_num_nodes), max(rand_pack_num_nodes), msg=err_msg)
                self.assertGreaterEqual(min(pack_num_nodes), min(rand_pack_num_nodes), msg=err_msg)

                # Assert that the total number of atoms is right
                self.assertEqual(sum(pack_num_nodes), sum(num_nodes), msg=err_msg)
                self.assertEqual(sum(rand_pack_num_nodes), sum(num_nodes), msg=err_msg)

                # Assert that all index are there
                self.assertListEqual(
                    np.sort(np.asarray(packed_indices).flatten()).tolist(), np.arange(len(num_nodes)).tolist()
                )
                self.assertListEqual(
                    np.sort(np.asarray(rand_packed_indices).flatten()).tolist(),
                    np.arange(len(num_nodes)).tolist(),
                )

    def test_fast_packing(self):
        np.random.seed(42)

        # Start at 4 for fast_packing for better statistical significance
        batch_sizes = [4, 8, 16, 32, 64]
        ipu_batch_sizes = [4, 8, 16, 32, 64]

        for batch_size in batch_sizes:
            for ipu_batch_size in ipu_batch_sizes:
                err_msg = f"bz={batch_size}, ipu_bz={ipu_batch_size}"

                # Generate random batch size
                global_batch = batch_size * ipu_batch_size
                num_nodes = np.abs(np.random.gamma(2, 20, size=global_batch)).astype(int)

                # Use the smart packing
                packed_indices = fast_packing(num_nodes=num_nodes, batch_size=batch_size)
                pack_num_nodes = get_pack_sizes(packed_indices, num_nodes)

                # Use the random packing
                rand_packed_indices = random_packing(num_nodes=num_nodes, batch_size=batch_size)
                rand_pack_num_nodes = get_pack_sizes(rand_packed_indices, num_nodes)

                # Assert that the smart packing is better than the random packing
                self.assertLessEqual(max(pack_num_nodes), max(rand_pack_num_nodes), msg=err_msg)
                self.assertGreaterEqual(min(pack_num_nodes), min(rand_pack_num_nodes), msg=err_msg)

                # Assert that the total number of atoms is right
                self.assertEqual(sum(pack_num_nodes), sum(num_nodes), msg=err_msg)
                self.assertEqual(sum(rand_pack_num_nodes), sum(num_nodes), msg=err_msg)

                # Assert that all index are there
                self.assertListEqual(
                    np.sort(np.asarray(packed_indices).flatten()).tolist(), np.arange(len(num_nodes)).tolist()
                )
                self.assertListEqual(
                    np.sort(np.asarray(rand_packed_indices).flatten()).tolist(),
                    np.arange(len(num_nodes)).tolist(),
                )

    def test_hybrid_packing(self):
        np.random.seed(42)

        batch_sizes = [2, 4, 8, 16, 32, 64]
        ipu_batch_sizes = [2, 3, 4, 8, 16, 32, 64]

        for batch_size in batch_sizes:
            for ipu_batch_size in ipu_batch_sizes:
                err_msg = f"bz={batch_size}, ipu_bz={ipu_batch_size}"

                # Generate random batch size
                global_batch = batch_size * ipu_batch_size
                num_nodes = np.abs(np.random.gamma(2, 20, size=global_batch)).astype(int)

                # Use the smart packing
                packed_indices = hybrid_packing(num_nodes=num_nodes, batch_size=batch_size)
                pack_num_nodes = get_pack_sizes(packed_indices, num_nodes)

                # Use the random packing
                rand_packed_indices = random_packing(num_nodes=num_nodes, batch_size=batch_size)
                rand_pack_num_nodes = get_pack_sizes(rand_packed_indices, num_nodes)

                # Assert that the smart packing is better than the random packing
                self.assertLessEqual(max(pack_num_nodes), max(rand_pack_num_nodes), msg=err_msg)
                self.assertGreaterEqual(min(pack_num_nodes), min(rand_pack_num_nodes), msg=err_msg)

                # Assert that the total number of atoms is right
                self.assertEqual(sum(pack_num_nodes), sum(num_nodes), msg=err_msg)
                self.assertEqual(sum(rand_pack_num_nodes), sum(num_nodes), msg=err_msg)

                # Assert that all index are there
                self.assertListEqual(
                    np.sort(np.asarray(packed_indices).flatten()).tolist(), np.arange(len(num_nodes)).tolist()
                )
                self.assertListEqual(
                    np.sort(np.asarray(rand_packed_indices).flatten()).tolist(),
                    np.arange(len(num_nodes)).tolist(),
                )

    def test_node_to_pack_indices_mask(self):
        # Create a dummy batch
        in_dim = 7
        in_dim_edges = 11
        max_num_nodes_per_graph = 20
        batch_size_per_pack = 5

        torch.manual_seed(42)

        # Create a dummy batch of graphs
        batch, all_num_nodes = [], []
        for ii in range(100):
            num_nodes = torch.randint(1, max_num_nodes_per_graph, (1,)).item()
            all_num_nodes.append(num_nodes)
            num_edges = abs(round(2.2 * num_nodes) + torch.randint(-2, 2, (1,)).item()) + 1
            x = torch.randn(num_nodes, in_dim, dtype=torch.float32)
            edge_idx = torch.randint(0, num_nodes, (2, num_edges))
            e = torch.randn(edge_idx.shape[-1], in_dim_edges, dtype=torch.float32)
            g = Data(h=x, edge_index=edge_idx, edge_attr=e)
            batch.append(g)
        batch = Batch.from_data_list(batch)

        # Get the packing
        packed_graph_idx = fast_packing(all_num_nodes, batch_size_per_pack)
        pack_sizes = get_pack_sizes(packed_graph_idx, all_num_nodes)
        max_pack_size = max(pack_sizes)
        num_packs = len(pack_sizes)

        # Get the node to pack indices and the mask
        pack_from_node_idx, pack_attn_mask = node_to_pack_indices_mask(packed_graph_idx, all_num_nodes)

        # Assert that the nodes to pack indices are correct
        h = torch.arange(batch.num_nodes, dtype=torch.float32)
        packed_shape = [num_packs, max_pack_size]
        h_packed = torch.zeros(packed_shape)
        h_packed[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]] = h
        h_packed_unique = torch.sort(torch.unique(h_packed))[0]
        np.testing.assert_array_equal(h_packed_unique, torch.arange(batch.num_nodes))
        self.assertEqual(h_packed.sum(), h.sum())

        # Test again with additional h dimension
        h = batch.h
        packed_shape = [num_packs, max_pack_size] + list(h.shape[1:])
        h_packed = torch.zeros(packed_shape)
        h_packed[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]] = h
        h_packed_unique = torch.sort(torch.unique(h_packed))[0]
        h_packed_unique = h_packed_unique[h_packed_unique != 0]
        np.testing.assert_array_almost_equal(h_packed_unique, torch.unique(h))
        self.assertAlmostEqual(h_packed.sum().item(), h.sum().item(), places=3)

        # Assert that the mask is correct by counting the number of False values (the sum of squared number of nodes per pack)
        num_false = (~pack_attn_mask).sum([1, 2])
        num_expected = torch.as_tensor(
            [sum([all_num_nodes[graph_idx] ** 2 for graph_idx in pack]) for pack in packed_graph_idx]
        )
        np.testing.assert_array_equal(num_false, num_expected)

        # Assert that the mask is correct by counting the number of elements in each row and column
        num_expected = []
        for pack in packed_graph_idx:
            pack_num_expected = []
            for graph_idx in pack:
                num_nodes = all_num_nodes[graph_idx]
                for ii in range(num_nodes):
                    pack_num_expected.append(num_nodes)
            pack_num_expected.extend([0] * (max_pack_size - len(pack_num_expected)))
            num_expected.append(pack_num_expected)
        num_expected = torch.as_tensor(num_expected)
        num_false_row = (~pack_attn_mask).sum([2])
        num_false_col = (~pack_attn_mask).sum([1])
        np.testing.assert_array_equal(num_false_row, num_expected)
        np.testing.assert_array_equal(num_false_col, num_expected)


if __name__ == "__main__":
    ut.main()


# tests/test_pe_nodepair.py

"""
Unit tests for the positional encodings in graphium/features/*
"""

import numpy as np
import networkx as nx
import unittest as ut

from graphium.features.electrostatic import compute_electrostatic_interactions
from graphium.features.commute import compute_commute_distances
from graphium.features.graphormer import compute_graphormer_distances


class test_positional_encodings(ut.TestCase):
    # Test graphs
    adj_dict = {}
    max_dict = {}

    # 6-ring
    adj = np.asarray(
        [
            [0, 1, 0, 0, 0, 1],
            [1, 0, 1, 0, 0, 0],
            [0, 1, 0, 1, 0, 0],
            [0, 0, 1, 0, 1, 0],
            [0, 0, 0, 1, 0, 1],
            [1, 0, 0, 0, 1, 0],
        ]
    )
    adj_dict["6-ring"] = adj
    max_dict["6-ring"] = 3

    # 5-path
    G = nx.path_graph(5)
    adj = nx.to_numpy_array(G)
    adj_dict["5-path"] = adj
    max_dict["5-path"] = 4

    # 4-clique
    adj = 1 - np.eye(4)
    adj_dict["4-clique"] = adj
    max_dict["4-clique"] = 1

    # 4-barbell
    H = nx.barbell_graph(4, 0)
    adj = nx.to_numpy_array(H)
    adj_dict["4-barbell"] = adj
    max_dict["4-barbell"] = 3

    def test_dimensions(self):
        for _, adj in self.adj_dict.items():
            pe, _, _ = compute_electrostatic_interactions(adj, cache={})
            self.assertEqual(pe.shape, adj.shape)

            pe, _, _ = compute_graphormer_distances(adj, adj.shape[0], cache={})
            self.assertEqual(pe.shape, adj.shape)

            pe, _, _ = compute_commute_distances(adj, adj.shape[0], cache={})
            self.assertEqual(pe.shape, adj.shape)

    def test_symmetry(self):
        for _, adj in self.adj_dict.items():
            pe, _, _ = compute_graphormer_distances(adj, adj.shape[0], cache={})
            np.testing.assert_array_almost_equal(pe, pe.T)

            pe, _, _ = compute_commute_distances(adj, adj.shape[0], cache={})
            np.testing.assert_array_almost_equal(pe, pe.T)

    def test_max_dist(self):
        for key, adj in self.adj_dict.items():
            pe, _, _ = compute_graphormer_distances(adj, adj.shape[0], cache={})
            np.testing.assert_array_almost_equal(pe.max(), self.max_dict[key])


if __name__ == "__main__":
    ut.main()


# tests/test_pe_rw.py

"""
Unit tests for the positional encodings in graphium/features/*
"""

import numpy as np
import networkx as nx
import unittest as ut

from graphium.features.rw import compute_rwse


class test_pe_spectral(ut.TestCase):
    def test_caching_and_outputs(self):
        # 4-barbell
        G = nx.barbell_graph(4, 0)
        adj = nx.to_numpy_array(G)
        num_nodes = adj.shape[0]
        cache = {}

        ksteps1 = [4, 6]
        ksteps2 = [2]
        ksteps3 = [6, 7]

        pe1, _, cache = compute_rwse(
            adj.astype(np.float32), ksteps1, num_nodes, cache, pos_type="rw_transition_probs"
        )

        pe2, _, cache = compute_rwse(
            adj.astype(np.float32), ksteps2, num_nodes, cache, pos_type="rw_return_probs"
        )

        pe3, _, cache = compute_rwse(
            adj.astype(np.float32), ksteps3, num_nodes, cache, pos_type="rw_return_probs"
        )

        self.assertTrue(all([k in cache["ksteps"] for k in ksteps1 + ksteps2 + ksteps3]))
        self.assertTrue(pe1.shape, np.zeros((num_nodes, num_nodes, len(ksteps1))))
        self.assertTrue(pe2.shape, np.zeros((num_nodes, len(ksteps2))))
        self.assertTrue(pe3.shape, np.zeros((num_nodes, len(ksteps3))))

        for i in range(len(ksteps1)):
            np.testing.assert_array_almost_equal(pe1[..., i].sum(1), np.ones(num_nodes))

        self.assertGreaterEqual(pe1.min(), 0.0)
        self.assertLessEqual(pe1.max(), 1.0)

        self.assertGreaterEqual(pe2.min(), 0.0)
        self.assertLessEqual(pe2.max(), 1.0)

        self.assertGreaterEqual(pe3.min(), 0.0)
        self.assertLessEqual(pe3.max(), 1.0)


if __name__ == "__main__":
    ut.main()


# tests/test_pe_spectral.py

"""
Unit tests for the positional encodings in graphium/features/*
"""

import numpy as np
import networkx as nx
import unittest as ut

from graphium.features.spectral import compute_laplacian_pe


class test_pe_spectral(ut.TestCase):
    # 2 disconnected 3 cliques
    adj1 = np.zeros((6, 6))
    adj_3clq = 1 - np.eye(3)
    adj1[:3, :3] = adj_3clq
    adj1[3:, 3:] = adj_3clq

    # 3-clique
    adj2 = 1 - np.eye(6)

    def test_for_connected_vs_disconnected_graph(self):
        num_pos = 3

        # test if pe works identically on connected vs disconnected graphs
        eigvals_pe1, _, _, cache = compute_laplacian_pe(self.adj1, num_pos, cache={})
        eigvals_pe1 = np.real(eigvals_pe1).astype(np.float32)
        _, eigvecs_pe1, _, _ = compute_laplacian_pe(self.adj1, num_pos, cache=cache)

        # We expect to cache 4 objects in when running the functon for the first time
        self.assertEqual(len(cache.keys()), 4)

        eigvals_pe2, _, _, _ = compute_laplacian_pe(self.adj2, num_pos, cache={})
        eigvals_pe2 = np.real(eigvals_pe2).astype(np.float32)
        _, eigvecs_pe2, _, _ = compute_laplacian_pe(self.adj2, num_pos, cache={})

        np.testing.assert_array_almost_equal(2 * eigvals_pe1, eigvals_pe2)
        self.assertListEqual(list(eigvals_pe2.shape), [self.adj2.shape[0], num_pos])
        self.assertListEqual(list(eigvecs_pe2.shape), [self.adj2.shape[0], num_pos])


if __name__ == "__main__":
    ut.main()


# tests/test_pos_transfer_funcs.py

"""
Unit tests for the positional encodings in graphium/features/*
"""

import numpy as np
import networkx as nx
import unittest as ut

from graphium.features.spectral import compute_laplacian_pe
from graphium.features.transfer_pos_level import (
    node_to_edge,
    node_to_nodepair,
    edge_to_nodepair,
    nodepair_to_node,
    nodepair_to_edge,
    graph_to_node,
)


class test_pos_transfer_funcs(ut.TestCase):
    # 4-barbell
    G = nx.barbell_graph(4, 0)
    adj = nx.to_numpy_array(G)
    num_nodes, num_feat = 8, 5
    node_pe = np.random.rand(num_nodes, num_feat)

    def test_different_pathways_from_node_to_edge(self):
        edge_pe1, _ = node_to_edge(self.node_pe, self.adj, {})
        nodepair_pe1 = node_to_nodepair(self.node_pe, self.num_nodes)
        edge_pe2, _ = nodepair_to_edge(nodepair_pe1, self.adj, {})
        nodepair_pe2, _ = edge_to_nodepair(edge_pe1, self.adj, self.num_nodes, {})
        edge_pe3, _ = nodepair_to_edge(nodepair_pe2, self.adj, {})
        np.testing.assert_array_almost_equal(edge_pe1, edge_pe2)
        np.testing.assert_array_almost_equal(edge_pe1, edge_pe3)


if __name__ == "__main__":
    ut.main()


# tests/test_positional_encoders.py

"""
Unit tests for the different datasets of graphium/features/featurizer.py
"""

import numpy as np
import unittest as ut
from copy import deepcopy
from rdkit import Chem
import datamol as dm
import torch
from scipy.sparse import coo_matrix

from graphium.features.featurizer import GraphDict
from graphium.features.positional_encoding import graph_positional_encoder
from graphium.nn.encoders import laplace_pos_encoder, mlp_encoder, signnet_pos_encoder

# TODO: Test the MLP_encoder and signnet_pos_encoder


class test_positional_encoder(ut.TestCase):
    smiles = [
        "C",
        "CC",
        "CC.CCCC",
        "CC(C)CC1=CC=C(C=C1)C(C)C(=O)O",
        "OCCc1c(C)[n+](cs1)Cc2cnc(C)nc2N",
        "O1C=C[C@H]([C@H]1O2)c3c2cc(OC)c4c3OC(=O)C5=C4CCC(=O)5",
    ]
    mols = [dm.to_mol(s) for s in smiles]
    adjs = [Chem.rdmolops.GetAdjacencyMatrix(mol) for mol in mols]

    def test_laplacian_eigvec_eigval(self):
        for ii, adj in enumerate(deepcopy(self.adjs)):
            for num_pos in [1, 2, 4]:  # Can't test too much eigs because of multiplicities
                for disconnected_comp in [True, False]:
                    err_msg = f"adj_id={ii}, num_pos={num_pos}, disconnected_comp={disconnected_comp}"

                    # returns a dictionary of computed pe
                    pos_kwargs = {
                        "pos_type": "laplacian_eigvec",
                        "num_pos": num_pos,
                        "disconnected_comp": disconnected_comp,
                        "pos_level": "node",
                    }
                    num_nodes = adj.shape[0]
                    eigvecs, cache = graph_positional_encoder(adj, num_nodes, pos_kwargs=pos_kwargs)
                    pos_kwargs["pos_type"] = "laplacian_eigval"
                    eigvals, cache = graph_positional_encoder(adj, num_nodes, pos_kwargs=pos_kwargs)

                    self.assertEqual(list(eigvecs.shape), [adj.shape[0], num_pos], msg=err_msg)
                    self.assertEqual(list(eigvals.shape), [adj.shape[0], num_pos], msg=err_msg)

                    # Compute eigvals and eigvecs
                    lap = np.diag(np.sum(adj, axis=1)) - adj
                    true_eigvals, true_eigvecs = np.linalg.eig(lap)
                    sort_idx = np.argsort(true_eigvals)
                    true_eigvals, true_eigvecs = true_eigvals[sort_idx], true_eigvecs[:, sort_idx]
                    true_eigvecs = true_eigvecs / (np.sum(true_eigvecs**2, axis=0, keepdims=True) + 1e-8)

                    true_num_pos = min(num_pos, len(true_eigvals))
                    true_eigvals, true_eigvecs = true_eigvals[:true_num_pos], true_eigvecs[:, :true_num_pos]

                    if not ("." in self.smiles[ii]):
                        np.testing.assert_array_almost_equal(
                            np.abs(true_eigvecs),
                            np.abs(eigvecs[:, :true_num_pos]),
                            decimal=6,
                            err_msg=err_msg,
                        )
                        self.assertAlmostEqual(np.sum(true_eigvecs[:, 1:]), 0, places=6, msg=err_msg)
                        np.testing.assert_array_almost_equal(
                            true_eigvals, eigvals[0, :true_num_pos], decimal=6, err_msg=err_msg
                        )

    # didn't actually check the exact computation result because the code was adapted
    def test_rwse(self):
        for ii, adj in enumerate(deepcopy(self.adjs)):
            for ksteps in [1, 2, 4]:
                err_msg = f"adj_id={ii}, ksteps={ksteps}"

                num_nodes = adj.shape[0]
                pos_kwargs = {"pos_type": "rw_return_probs", "ksteps": ksteps, "pos_level": "node"}
                rwse_embed, cache = graph_positional_encoder(adj, num_nodes, pos_kwargs=pos_kwargs)
                self.assertEqual(list(rwse_embed.shape), [num_nodes, ksteps], msg=err_msg)

    # TODO: work in progress

    """
    continue debugging here, see how to adapt the laplace_pos_encoder
    code running now, question is where to add the laplace_pos_encoder
    """

    def test_laplacian_eigvec_with_encoder(self):
        for ii, adj in enumerate(deepcopy(self.adjs)):
            for num_pos in [2, 4, 8]:  # Can't test too much eigs because of multiplicities
                for disconnected_comp in [True, False]:
                    for model_type in ["Transformer", "DeepSet", "MLP"]:
                        err_msg = f"adj_id={ii}, num_pos={num_pos}, disconnected_comp={disconnected_comp}"

                        # returns a dictionary of computed pe
                        pos_kwargs = {
                            "pos_type": "laplacian_eigvec",
                            "num_pos": num_pos,
                            "disconnected_comp": disconnected_comp,
                            "pos_level": "node",
                        }
                        num_nodes = adj.shape[0]
                        eigvecs, cache = graph_positional_encoder(adj, num_nodes, pos_kwargs=pos_kwargs)
                        pos_kwargs["pos_type"] = "laplacian_eigval"
                        eigvals, cache = graph_positional_encoder(adj, num_nodes, pos_kwargs=pos_kwargs)

                        input_keys = ["laplacian_eigvec", "laplacian_eigval"]
                        in_dim = num_pos
                        hidden_dim = 64
                        out_dim = 64
                        num_layers = 1

                        eigvecs = torch.from_numpy(eigvecs)
                        eigvals = torch.from_numpy(eigvals)

                        g = GraphDict(
                            {
                                "adj": coo_matrix(adj),
                                "data": {"laplacian_eigval": eigvals, "laplacian_eigvec": eigvecs},
                            }
                        )
                        batch = g.make_pyg_graph()

                        encoder = laplace_pos_encoder.LapPENodeEncoder(
                            input_keys=input_keys,
                            output_keys=["node"],
                            in_dim=in_dim,  # Size of Laplace PE embedding
                            hidden_dim=hidden_dim,
                            out_dim=out_dim,
                            model_type=model_type,  # 'Transformer' or 'DeepSet'
                            num_layers=num_layers,
                            num_layers_post=2,  # Num. layers to apply after pooling
                            dropout=0.1,
                            first_normalization=None,
                        )

                        hidden_embed = encoder(batch, key_prefix=None)
                        assert "node" in hidden_embed.keys()
                        self.assertEqual(list(hidden_embed["node"].shape), [num_nodes, out_dim], msg=err_msg)


if __name__ == "__main__":
    ut.main()


# tests/test_positional_encodings.py

"""
Unit tests for the positional encodings in graphium/features/*
"""

import numpy as np
import networkx as nx
import unittest as ut

# from graphium.features.spectral import compute_laplacian_positional_eigvecs # TODO: add tests
# from graphium.features.rw import compute_rwse # TODO: add tests
from graphium.features.electrostatic import compute_electrostatic_interactions
from graphium.features.commute import compute_commute_distances
from graphium.features.graphormer import compute_graphormer_distances


class test_positional_encodings(ut.TestCase):
    # Test graphs
    adj_dict = {}
    max_dict = {}

    # 6-ring
    adj = np.asarray(
        [
            [0, 1, 0, 0, 0, 1],
            [1, 0, 1, 0, 0, 0],
            [0, 1, 0, 1, 0, 0],
            [0, 0, 1, 0, 1, 0],
            [0, 0, 0, 1, 0, 1],
            [1, 0, 0, 0, 1, 0],
        ]
    )
    adj_dict["6-ring"] = adj
    max_dict["6-ring"] = 3

    # 5-path
    G = nx.path_graph(5)
    adj = nx.to_numpy_array(G)
    adj_dict["5-path"] = adj
    max_dict["5-path"] = 4

    # 4-clique
    adj = 1 - np.eye(4)
    adj_dict["4-clique"] = adj
    max_dict["4-clique"] = 1

    # 4-barbell
    H = nx.barbell_graph(4, 0)
    adj = nx.to_numpy_array(H)
    adj_dict["4-barbell"] = adj
    max_dict["4-barbell"] = 3

    def test_dimensions(self):
        for _, adj in self.adj_dict.items():
            pe, _, _ = compute_electrostatic_interactions(adj, cache={})
            self.assertEqual(pe.shape, adj.shape)

            pe, _, _ = compute_graphormer_distances(adj, adj.shape[0], cache={})
            self.assertEqual(pe.shape, adj.shape)

            pe, _, _ = compute_commute_distances(adj, adj.shape[0], cache={})
            self.assertEqual(pe.shape, adj.shape)

    def test_symmetry(self):
        for _, adj in self.adj_dict.items():
            pe, _, _ = compute_graphormer_distances(adj, adj.shape[0], cache={})
            np.testing.assert_array_almost_equal(pe, pe.T)

            pe, _, _ = compute_commute_distances(adj, adj.shape[0], cache={})
            np.testing.assert_array_almost_equal(pe, pe.T)

    def test_max_dist(self):
        for key, adj in self.adj_dict.items():
            pe, _, _ = compute_graphormer_distances(adj, adj.shape[0], cache={})
            np.testing.assert_array_almost_equal(pe.max(), self.max_dict[key])


if __name__ == "__main__":
    ut.main()


# tests/test_predictor.py

"""
Unit tests for the file graphium/trainer/predictor.py
"""

import torch
from torch.nn import BCELoss, MSELoss
import unittest as ut

from graphium.trainer.predictor_options import EvalOptions


class test_Predictor(ut.TestCase):
    def test_parse_loss_fun(self):
        losses = ["bce", "mse", "mae", "l1", BCELoss(), MSELoss()]
        preds = torch.rand(10, 5)
        target = (torch.rand(10, 5) > 0.5).to(preds.dtype)
        for this_loss in losses:
            loss_fun = EvalOptions.parse_loss_fun(this_loss)
            loss = loss_fun(preds, target)


if __name__ == "__main__":
    ut.main()


# tests/test_pyg_layers.py

"""
Unit tests for the different layers of graphium/nn/pyg_layers/...

The layers are not thoroughly tested due to the difficulty of testing them
"""

import numpy as np
import torch
import unittest as ut
from torch_geometric.data import Data, Batch
from copy import deepcopy

from graphium.nn.pyg_layers import (
    GINConvPyg,
    GINEConvPyg,
    MPNNPlusPyg,
    GatedGCNPyg,
    PNAMessagePassingPyg,
    GPSLayerPyg,
    VirtualNodePyg,
    DimeNetPyg,
)

from graphium.nn.pyg_layers.utils import (
    PreprocessPositions,
    GaussianLayer,
)


class test_Pyg_Layers(ut.TestCase):
    in_dim = 21
    out_dim = 11
    in_dim_edges = 13
    out_dim_edges = 17

    edge_idx1 = torch.stack([torch.tensor([0, 1, 2, 3, 2]), torch.tensor([1, 2, 3, 0, 0])])
    edge_idx2 = torch.stack([torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])])
    x1 = torch.randn(edge_idx1.max() + 1, in_dim, dtype=torch.float32)
    e1 = torch.randn(edge_idx1.shape[-1], in_dim_edges, dtype=torch.float32)
    x2 = torch.randn(edge_idx2.max() + 1, in_dim, dtype=torch.float32)
    e2 = torch.randn(edge_idx2.shape[-1], in_dim_edges, dtype=torch.float32)
    # edge_idx1, e1 = add_self_loops(edge_idx1, e1)
    # edge_idx2, e2 = add_self_loops(edge_idx2, e2)
    g1 = Data(feat=x1, edge_index=edge_idx1, edge_feat=e1)
    g2 = Data(feat=x2, edge_index=edge_idx2, edge_feat=e2)
    bg = Batch.from_data_list([g1, g2])

    kwargs = {
        "activation": "relu",
        "dropout": 0.1,
        "normalization": "batch_norm",
        "droppath_rate": 0.1,
        "layer_idx": 1,
        "layer_depth": 10,
    }

    def test_gpslayer(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        kwargs = deepcopy(self.kwargs)
        kwargs.pop("droppath_rate")
        kwargs["droppath_rate_attn"] = 0.2
        kwargs["droppath_rate_ffn"] = 0.3
        kwargs["mpnn_kwargs"] = {"droppath_rate_ffn": 0.4}

        layer = GPSLayerPyg(in_dim=self.in_dim, out_dim=self.out_dim, **kwargs)

        # Check the re-implementation of abstract methods
        self.assertTrue(layer.layer_supports_edges)
        self.assertTrue(layer.layer_inputs_edges)
        self.assertFalse(layer.layer_outputs_edges)
        self.assertEqual(layer.out_dim_factor, 1)

        # Apply layer. Should crash due to different dim for nodes vs edges
        with self.assertRaises(ValueError):
            bg = layer.forward(bg)

        # Create new edge attributes with same dim and check that it works
        bg.edge_feat = torch.zeros((bg.edge_feat.shape[0], self.in_dim), dtype=torch.float32)
        bg = layer.forward(bg)
        self.assertEqual(bg.feat.shape[0], feat_in.shape[0])
        self.assertEqual(bg.feat.shape[1], self.out_dim * layer.out_dim_factor)

    def test_ginlayer(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        layer = GINConvPyg(in_dim=self.in_dim, out_dim=self.out_dim, **self.kwargs)

        # Check the re-implementation of abstract methods
        self.assertFalse(layer.layer_supports_edges)
        self.assertFalse(layer.layer_inputs_edges)
        self.assertFalse(layer.layer_outputs_edges)
        self.assertEqual(layer.out_dim_factor, 1)

        # Apply layer
        bg = layer.forward(bg)
        self.assertEqual(bg.feat.shape[0], feat_in.shape[0])
        self.assertEqual(bg.feat.shape[1], self.out_dim * layer.out_dim_factor)

    def test_ginelayer(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        layer = GINEConvPyg(in_dim=self.in_dim, out_dim=self.out_dim, **self.kwargs)

        # Check the re-implementation of abstract methods
        self.assertTrue(layer.layer_supports_edges)
        self.assertTrue(layer.layer_inputs_edges)
        self.assertFalse(layer.layer_outputs_edges)
        self.assertEqual(layer.out_dim_factor, 1)

        # Apply layer. Should crash due to different dim for nodes vs edges
        with self.assertRaises(ValueError):
            bg = layer.forward(bg)

        # Create new edge attributes with same dim and check that it works
        bg.edge_feat = torch.zeros((bg.edge_feat.shape[0], self.in_dim), dtype=torch.float32)
        bg = layer.forward(bg)
        self.assertEqual(bg.feat.shape[0], feat_in.shape[0])
        self.assertEqual(bg.feat.shape[1], self.out_dim * layer.out_dim_factor)

    def test_mpnnlayer(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        # need in_dim = out_dim for skip connection
        # mpnn layer accept different dimension for node and edge features
        layer = MPNNPlusPyg(
            in_dim=self.in_dim,
            out_dim=self.in_dim,
            use_edges=True,
            in_dim_edges=self.in_dim_edges,
            out_dim_edges=self.in_dim_edges,
            **self.kwargs,
        )

        # Check the re-implementation of abstract methods
        self.assertTrue(layer.layer_supports_edges)
        self.assertTrue(layer.layer_inputs_edges)
        self.assertTrue(layer.layer_outputs_edges)
        self.assertEqual(layer.out_dim_factor, 1)

        # Create new edge attributes with same dim and check that it works
        bg.edge_feat = torch.zeros((bg.edge_feat.shape[0], self.in_dim_edges), dtype=torch.float32)
        bg = layer.forward(bg)
        self.assertEqual(bg.feat.shape[0], feat_in.shape[0])
        self.assertEqual(bg.feat.shape[1], self.in_dim * layer.out_dim_factor)

    def test_gatedgcnlayer(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        e_in = bg.edge_feat
        layer = GatedGCNPyg(
            in_dim=self.in_dim,
            out_dim=self.out_dim,
            in_dim_edges=self.in_dim_edges,
            out_dim_edges=self.out_dim_edges,
            **self.kwargs,
        )

        # Check the re-implementation of abstract methods
        self.assertTrue(layer.layer_supports_edges)
        self.assertTrue(layer.layer_inputs_edges)
        self.assertTrue(layer.layer_outputs_edges)
        self.assertIsInstance(layer.layer_supports_edges, bool)
        self.assertIsInstance(layer.layer_inputs_edges, bool)
        self.assertIsInstance(layer.layer_outputs_edges, bool)
        self.assertEqual(layer.out_dim_factor, 1)

        # Apply layer with edges
        bg2 = layer.forward(bg)
        self.assertEqual(bg2.feat.shape[0], feat_in.shape[0])
        self.assertEqual(bg2.feat.shape[1], self.out_dim * layer.out_dim_factor)

    def test_pnamessagepassinglayer(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        aggregators = ["mean", "max", "min", "std", "sum"]
        scalers = ["identity", "amplification", "attenuation"]

        layer = PNAMessagePassingPyg(
            in_dim=self.in_dim, out_dim=self.out_dim, aggregators=aggregators, scalers=scalers, **self.kwargs
        )

        # Check the re-implementation of abstract methods
        self.assertTrue(layer.layer_supports_edges)
        self.assertIsInstance(layer.layer_supports_edges, bool)
        # self.assertFalse(layer.layer_inputs_edges)
        self.assertFalse(layer.layer_outputs_edges)
        self.assertEqual(layer.out_dim_factor, 1)

        # Apply layer
        bg2 = layer.forward(bg)
        self.assertEqual(bg2.feat.shape[0], feat_in.shape[0])
        self.assertEqual(bg2.feat.shape[1], self.out_dim * layer.out_dim_factor)

        # Now try with edges
        bg = deepcopy(self.bg)
        layer = PNAMessagePassingPyg(
            in_dim=self.in_dim,
            out_dim=self.out_dim,
            aggregators=aggregators,
            scalers=scalers,
            in_dim_edges=self.in_dim_edges,
            **self.kwargs,
        )

        # Check the re-implementation of abstract methods
        self.assertTrue(layer.layer_supports_edges)
        # self.assertTrue(layer.layer_inputs_edges)
        self.assertIsInstance(layer.layer_supports_edges, bool)
        # self.assertIsInstance(layer.layer_inputs_edges, bool)
        self.assertFalse(layer.layer_outputs_edges)
        self.assertEqual(layer.out_dim_factor, 1)

        bg2 = layer.forward(bg)
        self.assertEqual(bg2.feat.shape[0], feat_in.shape[0])
        self.assertEqual(bg2.feat.shape[1], self.out_dim * layer.out_dim_factor)
        self.assertTrue((bg2.edge_feat == self.bg.edge_feat).all)

    def test_dimenetlayer(self):
        from graphium.nn.encoders.bessel_pos_encoder import BesselSphericalPosEncoder

        bg = deepcopy(self.bg)
        # dummy input pos
        bg.pos = torch.randn((bg.feat.shape[0], 3), dtype=torch.float32)

        # position 3d encoder
        pos_enc = BesselSphericalPosEncoder(
            input_keys=["pos"],
            output_keys=[
                "node_feat",
                "edge_feat",
                "edge_rbf",
                "triplet_sbf",
                "radius_edge_index",
            ],  # The keys to return
            in_dim=3,
            out_dim=self.in_dim,
            out_dim_edges=self.in_dim_edges,
            num_output_layers=2,
            num_layers=2,
            num_spherical=4,
            num_radial=32,
        )

        enc_output = pos_enc(bg, None)  # forward requires: pos, edge_index, batch

        bg.feat = bg.feat + enc_output["node_feat"]  # [num_nodes, in_dim]
        bg.edge_feat = enc_output["edge_feat"]  # [num_edges, out_dim_edges]
        bg.radius_edge_index = enc_output["radius_edge_index"]  # [2, num_edges]
        bg.edge_rbf = enc_output["edge_rbf"]  # [num_edges, num_radial]
        bg.triplet_sbf = enc_output["triplet_sbf"]  # [num_triplets, num_spherical * num_radial]

        kwargs = deepcopy(self.kwargs)
        kwargs["num_bilinear"] = 32
        kwargs["num_spherical"] = 4
        kwargs["num_radial"] = 32

        layer = DimeNetPyg(
            in_dim=self.in_dim,
            out_dim=self.out_dim,
            in_dim_edges=self.in_dim_edges,
            out_dim_edges=self.out_dim_edges,
            **kwargs,
        )

        # Check the re-implementation of abstract methods
        self.assertTrue(layer.layer_supports_edges)
        self.assertTrue(layer.layer_inputs_edges)
        self.assertTrue(layer.layer_outputs_edges)
        self.assertEqual(layer.out_dim_factor, 1)

        # Apply layer with encoded feats as input
        bg2 = layer.forward(bg)
        # output sanity check
        self.assertEqual(bg2.edge_feat.shape, bg.edge_feat.shape)
        self.assertEqual(bg2.edge_feat.shape[1], self.out_dim_edges)
        self.assertEqual(bg2.feat.shape[1], self.out_dim)
        # not change rbf/sbf embedding
        self.assertTrue((bg2.edge_rbf == bg.edge_rbf).all)
        self.assertTrue((bg2.triplet_sbf == bg.triplet_sbf).all)

    def test_preprocess3Dfeaturelayer(self):
        bg = deepcopy(self.bg)
        num_heads = 2
        num_kernel = 2
        in_dim = 3
        bg.positions_3d = torch.zeros(bg.feat.size()[0], in_dim)
        layer = PreprocessPositions(
            num_heads=num_heads,
            embed_dim=self.out_dim,
            num_kernel=num_kernel,
            in_dim=in_dim,
            first_normalization="layer_norm",
        )
        # bias: [batch, num_heads, nodes, nodes]
        # node_feature: [total_nodes, embed_dim]
        bias, node_feature = layer.forward(
            bg, max_num_nodes_per_graph=4, on_ipu=False, positions_3d_key="positions_3d"
        )
        self.assertEqual(bias.size(), torch.Size([2, num_heads, 4, 4]))
        self.assertFalse(np.isnan(bias.detach().numpy()).any())
        self.assertEqual(node_feature.size(), torch.Size([7, self.out_dim]))

    def test_gaussianlayer(self):
        num_kernels = 3
        input = torch.zeros(2, 4, 4)
        layer = GaussianLayer(num_kernels=num_kernels)
        # tensor_with_kernel: [batch, nodes, nodes, num_kernel]
        tensor_with_kernel = layer.forward(input)
        self.assertEqual(tensor_with_kernel.size(), torch.Size([2, 4, 4, 3]))

    def test_pooling_virtual_node(self):
        bg = deepcopy(self.bg)
        feat_in = bg.feat
        e_in = bg.edge_feat
        vn_h = 0

        vn_types = ["sum", "mean"]
        for vn_type, use_edges, expected_v_node in zip(vn_types, [False, True], [(2, 21), (2, 34)]):
            with self.subTest(
                vn_type=vn_type,
                use_edges=use_edges,
                expected_v_node=expected_v_node,
            ):
                vn_h = 0.0
                print(vn_h, self.out_dim, feat_in.size())
                layer = VirtualNodePyg(
                    in_dim=self.in_dim,
                    out_dim=self.in_dim,
                    in_dim_edges=self.in_dim_edges,
                    out_dim_edges=self.in_dim_edges,
                    vn_type=vn_type,
                    use_edges=use_edges,
                    **self.kwargs,
                )

                feat_out, vn_out, e_out = layer.forward(bg, feat_in, vn_h, edge_feat=e_in)
                assert vn_out.shape == expected_v_node
                assert feat_out.shape == (7, 21)
                assert e_out.shape == (9, 13)
                if use_edges is False:
                    # i.e. that the node features have been updated
                    assert torch.equal(feat_out, feat_in) == False
                    # And that the edge features have not
                    assert torch.equal(e_out, e_in) == True
                else:
                    assert torch.equal(feat_out, feat_in) == False
                    assert torch.equal(e_out, e_in) == False


if __name__ == "__main__":
    ut.main()


# tests/test_residual_connections.py

"""
Unit tests for the file residual_connections.py
"""

import numpy as np
import torch
import unittest as ut

from graphium.nn.residual_connections import (
    ResidualConnectionConcat,
    ResidualConnectionDenseNet,
    ResidualConnectionNone,
    ResidualConnectionSimple,
    ResidualConnectionWeighted,
    ResidualConnectionRandom,
)


class test_ResidualConnectionNone(ut.TestCase):
    def test_get_true_out_dims_none(self):
        full_dims = [4, 6, 8, 10, 12]
        in_dims, out_dims = full_dims[:-1], full_dims[1:]
        rc = ResidualConnectionNone(skip_steps=1)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = out_dims[:-1]

        self.assertListEqual(expected_out_dims, true_out_dims)

    def test_forward_none(self):
        rc = ResidualConnectionNone(skip_steps=1)
        num_loops = 10
        shape = (3, 11)
        h_original = [torch.rand(shape) for _ in range(num_loops)]

        h_prev = None
        for ii in range(num_loops):
            feat, h_prev = rc.forward(h_original[ii], h_prev, step_idx=ii)
            np.testing.assert_array_equal(feat.numpy(), h_original[ii].numpy(), err_msg=f"ii={ii}")
            self.assertIsNone(h_prev)


class test_ResidualConnectionSimple(ut.TestCase):
    def test_get_true_out_dims_simple(self):
        full_dims = [4, 6, 8, 10, 12]
        in_dims, out_dims = full_dims[:-1], full_dims[1:]
        rc = ResidualConnectionSimple(skip_steps=1)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = out_dims[:-1]

        self.assertListEqual(expected_out_dims, true_out_dims)

    def test_forward_simple(self):
        for skip_steps in [1, 2, 3]:
            rc = ResidualConnectionSimple(skip_steps=skip_steps)
            num_loops = 10
            shape = (3, 11)
            h_original = [torch.ones(shape) * (ii + 1) for ii in range(num_loops)]

            h_prev = None
            for ii in range(num_loops):
                feat, h_prev = rc.forward(h_original[ii], h_prev, step_idx=ii)

                if ((ii % skip_steps) == 0) and (ii > 0):
                    h_expected = (
                        torch.sum(torch.stack(h_original[0 : ii + 1 : skip_steps], dim=0), dim=0)
                    ).numpy()
                    h_expected_prev = h_expected
                else:
                    h_expected = h_original[ii].numpy()
                if ii == 0:
                    h_expected_prev = h_expected

                np.testing.assert_array_equal(
                    feat.numpy(), h_expected, err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}"
                )
                np.testing.assert_array_equal(
                    h_prev.numpy(), h_expected_prev, err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}"
                )


class test_ResidualConnectionRandom(ut.TestCase):
    def test_get_true_out_dims_random(self):
        full_dims = [4, 6, 8, 10, 12]
        in_dims, out_dims = full_dims[:-1], full_dims[1:]
        rc = ResidualConnectionRandom(skip_steps=1, out_dims=out_dims)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = out_dims[:-1]

        self.assertListEqual(expected_out_dims, true_out_dims)

        rc = ResidualConnectionRandom(skip_steps=1, num_layers=len(out_dims))
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = out_dims[:-1]

        self.assertListEqual(expected_out_dims, true_out_dims)

        with self.assertRaises(ValueError):
            rc = ResidualConnectionRandom(skip_steps=1, out_dims=None, num_layers=None)

    def test_forward_random(self):
        for skip_steps in [1, 2, 3]:
            num_loops = 10

            if skip_steps > 1:
                with self.assertRaises(ValueError):
                    rc = ResidualConnectionRandom(skip_steps=skip_steps, num_layers=num_loops)
                continue

            rc = ResidualConnectionRandom(skip_steps=skip_steps, num_layers=num_loops + 1)
            shape = (3, 11)
            h_original = [torch.ones(shape) * (ii + 1) for ii in range(num_loops)]

            h_prev = None
            # Not really testing the expected values due to randomness, just testing if it runs
            for ii in range(num_loops):
                print(ii)
                feat, h_prev = rc.forward(h_original[ii], h_prev, step_idx=ii)


class test_ResidualConnectionWeighted(ut.TestCase):
    def test_get_true_out_dims_weighted(self):
        full_dims = [4, 6, 8, 10, 12]
        in_dims, out_dims = full_dims[:-1], full_dims[1:]
        rc = ResidualConnectionWeighted(skip_steps=1, out_dims=full_dims[1:])
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = out_dims[:-1]

        self.assertListEqual(expected_out_dims, true_out_dims)

    def test_forward_weighted(self):
        for skip_steps in [1, 2, 3]:
            num_loops = 10
            shape = (3, 11)
            full_dims = [shape[1]] * (num_loops + 1)
            rc = ResidualConnectionWeighted(
                skip_steps=skip_steps,
                out_dims=full_dims[1:],
                activation="none",
                normalization="none",
                bias=False,
            )

            h_original = [torch.ones(shape) * (ii + 1) for ii in range(num_loops)]
            h_forward = []

            h_prev = None
            step_counter = 0
            for ii in range(num_loops):
                h_prev_backup = h_prev
                feat, h_prev = rc.forward(h_original[ii], h_prev, step_idx=ii)

                if ((ii % skip_steps) == 0) and (ii > 0):
                    h_forward.append(rc.residual_list[step_counter].forward(h_prev_backup))
                    h_expected = (h_forward[-1] + h_original[ii]).detach().numpy()
                    h_expected_prev = h_expected
                    step_counter += 1
                else:
                    h_expected = h_original[ii].detach().numpy()
                if ii == 0:
                    h_expected_prev = h_expected

                np.testing.assert_array_equal(
                    feat.detach().numpy(), h_expected, err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}"
                )
                np.testing.assert_array_equal(
                    h_prev.detach().numpy(),
                    h_expected_prev,
                    err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}",
                )


class test_ResidualConnectionConcat(ut.TestCase):
    def test_get_true_out_dims_concat(self):
        full_dims = [4, 6, 8, 10, 12, 14, 16, 18, 20]
        in_dims, out_dims = full_dims[:-1], full_dims[1:]

        # skip_steps=1
        rc = ResidualConnectionConcat(skip_steps=1)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = [6, 14, 18, 22, 26, 30, 34]
        self.assertListEqual(expected_out_dims, true_out_dims)

        # skip_steps=2
        rc = ResidualConnectionConcat(skip_steps=2)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = [6, 8, 16, 12, 24, 16, 32]
        self.assertListEqual(expected_out_dims, true_out_dims)

        # skip_steps=3
        rc = ResidualConnectionConcat(skip_steps=3)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = [6, 8, 10, 18, 14, 16, 30]
        self.assertListEqual(expected_out_dims, true_out_dims)

    def test_forward_concat(self):
        for skip_steps in [1, 2, 3]:
            rc = ResidualConnectionConcat(skip_steps=skip_steps)
            num_loops = 10
            shape = (3, 11)
            h_original = [torch.ones(shape) * (ii + 1) for ii in range(num_loops)]

            h_prev = None
            for ii in range(num_loops):
                feat, h_prev = rc.forward(h_original[ii], h_prev, step_idx=ii)

                if ((ii % skip_steps) == 0) and (ii > 0):
                    h_expected = (
                        torch.cat(h_original[ii - skip_steps : ii + 1 : skip_steps][::-1], dim=-1)
                    ).numpy()
                    h_expected_prev = h_original[ii].numpy()
                else:
                    h_expected = h_original[ii].numpy()
                if ii == 0:
                    h_expected_prev = h_expected

                np.testing.assert_array_equal(
                    feat.numpy(), h_expected, err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}"
                )
                np.testing.assert_array_equal(
                    h_prev.numpy(), h_expected_prev, err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}"
                )


class test_ResidualConnectionDenseNet(ut.TestCase):
    def test_get_true_out_dims_densenet(self):
        full_dims = [4, 6, 8, 10, 12, 14, 16, 18, 20]
        in_dims, out_dims = full_dims[:-1], full_dims[1:]

        # skip_steps=1
        rc = ResidualConnectionDenseNet(skip_steps=1)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = np.cumsum(out_dims).tolist()[:-1]
        self.assertListEqual(expected_out_dims, true_out_dims)

        # skip_steps=2
        rc = ResidualConnectionDenseNet(skip_steps=2)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = [6, 8, 10 + 6, 12, 14 + 10 + 6, 16, 18 + 14 + 10 + 6]
        self.assertListEqual(expected_out_dims, true_out_dims)

        # skip_steps=3
        rc = ResidualConnectionDenseNet(skip_steps=3)
        true_out_dims = rc.get_true_out_dims(out_dims)
        expected_out_dims = [6, 8, 10, 12 + 6, 14, 16, 18 + 12 + 6]
        self.assertListEqual(expected_out_dims, true_out_dims)

    def test_forward_densenet(self):
        for skip_steps in [1, 2, 3]:
            rc = ResidualConnectionDenseNet(skip_steps=skip_steps)
            num_loops = 10
            shape = (3, 11)
            h_original = [torch.ones(shape) * (ii + 1) for ii in range(num_loops)]

            h_prev = None
            for ii in range(num_loops):
                feat, h_prev = rc.forward(h_original[ii], h_prev, step_idx=ii)

                if ((ii % skip_steps) == 0) and (ii > 0):
                    h_expected = (torch.cat(h_original[0 : ii + 1 : skip_steps][::-1], dim=-1)).numpy()
                    h_expected_prev = h_expected
                else:
                    h_expected = h_original[ii].numpy()
                if ii == 0:
                    h_expected_prev = h_expected

                np.testing.assert_array_equal(
                    feat.numpy(), h_expected, err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}"
                )
                np.testing.assert_array_equal(
                    h_prev.numpy(), h_expected_prev, err_msg=f"Error at: skip_steps={skip_steps}, ii={ii}"
                )


if __name__ == "__main__":
    ut.main()


# tests/test_utils.py

"""
Unit tests for the metrics and wrappers of graphium/utils/...
"""

from graphium.utils.tensor import nan_mad, nan_mean, nan_std, nan_var, nan_median
from graphium.utils.safe_run import SafeRun
import torch
import numpy as np
import scipy as sp
import unittest as ut


class test_nan_statistics(ut.TestCase):
    torch.manual_seed(42)

    dims = [
        None,
        (0),
        (1),
        (2),
        (-1),
        (-2),
        (-3),
        (0, 1),
        (0, 2),
    ]
    # Create tensor
    sz = (10, 6, 7)
    tensor = torch.randn(sz, dtype=torch.float32) ** 2 + 3
    is_nan = torch.rand(sz) > 0.4
    tensor[is_nan] = float("nan")

    def test_nan_mean(self):
        for keepdim in [False, True]:
            for dim in self.dims:
                err_msg = f"Error for :\n dim = {dim}\n keepdim = {keepdim}"

                tensor = self.tensor.clone()

                # Prepare the arguments for numpy vs torch
                if dim is not None:
                    torch_kwargs = {"dim": dim, "keepdim": keepdim}
                    numpy_kwargs = {"axis": dim, "keepdims": keepdim}
                else:
                    torch_kwargs = {}
                    numpy_kwargs = {}

                # Compare the nan-mean
                torch_mean = nan_mean(tensor, **torch_kwargs)
                numpy_mean = np.nanmean(tensor.numpy(), **numpy_kwargs)
                np.testing.assert_almost_equal(torch_mean.numpy(), numpy_mean, decimal=6, err_msg=err_msg)

    def test_nan_std_var(self):
        for unbiased in [True, False]:
            for keepdim in [False, True]:
                for dim in self.dims:
                    err_msg = f"Error for :\n\tdim = {dim}\n\tkeepdim = {keepdim}\n\tunbiased = {unbiased}"

                    tensor = self.tensor.clone()

                    # Prepare the arguments for numpy vs torch
                    if dim is not None:
                        torch_kwargs = {"dim": dim, "keepdim": keepdim, "unbiased": unbiased}
                        numpy_kwargs = {"axis": dim, "keepdims": keepdim, "ddof": float(unbiased)}
                    else:
                        torch_kwargs = {"unbiased": unbiased}
                        numpy_kwargs = {"ddof": float(unbiased)}

                    # Compare the std
                    torch_std = nan_std(tensor, **torch_kwargs)
                    numpy_std = np.nanstd(tensor.numpy(), **numpy_kwargs)
                    np.testing.assert_almost_equal(torch_std.numpy(), numpy_std, decimal=6, err_msg=err_msg)

                    # Compare the variance
                    torch_var = nan_var(tensor, **torch_kwargs)
                    numpy_var = np.nanvar(tensor.numpy(), **numpy_kwargs)
                    np.testing.assert_almost_equal(torch_var.numpy(), numpy_var, decimal=6, err_msg=err_msg)

    def test_nan_median(self):
        for keepdim in [False, True]:
            # Cannot test
            for dim in self.dims:  # in [d for d in self.dims if not isinstance(d, Tuple)]:
                err_msg = f"Error for :\n dim = {dim}\n keepdim = {keepdim}"

                tensor = torch.randn(
                    (7, 9, 11, 13)
                )  # Need odd number of values to properly compare torch to numpy

                # Prepare the arguments for numpy vs torch
                if dim is not None:
                    torch_kwargs = {"dim": dim, "keepdim": keepdim}
                    numpy_kwargs = {"axis": dim, "keepdims": keepdim}
                else:
                    torch_kwargs = {}
                    numpy_kwargs = {}

                # Compare the nan-median
                torch_med = nan_median(tensor, **torch_kwargs)
                numpy_med = np.nanmedian(tensor.numpy(), **numpy_kwargs)
                torch_sum = torch.nansum(tensor, **torch_kwargs)
                np.testing.assert_almost_equal(torch_med.numpy(), numpy_med, decimal=4, err_msg=err_msg)
                self.assertListEqual(list(torch_med.shape), list(torch_sum.shape))

    def test_nan_mad(self):
        for normal in [False, True]:
            # Cannot test
            for dim in self.dims:  # in [d for d in self.dims if not isinstance(d, Tuple)]:
                err_msg = f"Error for :\n dim = {dim}\n normal = {normal}"

                tensor = torch.randn(
                    (7, 9, 11, 13)
                )  # Need odd number of values to properly compare torch to numpy

                # Prepare the arguments for numpy vs torch
                if dim is not None:
                    torch_kwargs = {"dim": dim, "keepdim": False, "normal": normal}
                    numpy_kwargs = {"axis": dim, "nan_policy": "omit", "scale": 1 / 1.4826 if normal else 1.0}
                else:
                    torch_kwargs = {"normal": normal}
                    numpy_kwargs = {"axis": dim, "nan_policy": "omit", "scale": 1 / 1.4826 if normal else 1.0}

                # Compare the nan-median
                torch_mad = nan_mad(tensor, **torch_kwargs)
                numpy_mad = sp.stats.median_abs_deviation(tensor.numpy(), **numpy_kwargs)
                np.testing.assert_almost_equal(torch_mad.numpy(), numpy_mad, decimal=4, err_msg=err_msg)


class test_SafeRun(ut.TestCase):
    def test_safe_run(self):
        # Error is caught
        with SafeRun(name="bob", raise_error=False, verbose=0):
            raise ValueError("This is an error")

        # Error is caught
        with SafeRun(name="bob", raise_error=False, verbose=0):
            2 + None

        # Error is not caught
        with self.assertRaises(ValueError):
            with SafeRun(name="bob", raise_error=True, verbose=0):
                raise ValueError("This is an error")

        # Error is not caught
        with self.assertRaises(TypeError):
            with SafeRun(name="bob", raise_error=True, verbose=0):
                2 + None

        # No error. Runs correctly
        with SafeRun(name="bob", raise_error=True, verbose=0):
            print("This is not an error")

        # No error. Runs correctly
        with SafeRun(name="bob", raise_error=False, verbose=0):
            print("This is not an error")


if __name__ == "__main__":
    ut.main()


# graphium/utils/fs.py
from typing import Union
from typing import Optional

import os
import io
import platformdirs
import pathlib

from tqdm.auto import tqdm
import fsspec


def get_cache_dir(suffix: str = None, create: bool = True) -> pathlib.Path:
    """Get a local cache directory. You can append a suffix folder to it and optionnaly create
    the folder if it doesn't exist.
    """

    cache_dir = pathlib.Path(platformdirs.user_cache_dir(appname="graphium"))

    if suffix is not None:
        cache_dir /= suffix

    if create:
        cache_dir.mkdir(exist_ok=True, parents=True)

    return cache_dir


def get_mapper(path: Union[str, os.PathLike]):
    """Get the fsspec mapper.
    Args:
        path: a path supported by `fsspec` such as local, s3, gcs, etc.
    """
    return fsspec.get_mapper(str(path))


def get_basename(path: Union[str, os.PathLike]):
    """Get the basename of a file or a folder.
    Args:
        path: a path supported by `fsspec` such as local, s3, gcs, etc.
    """
    path = str(path)
    mapper = get_mapper(path)
    clean_path = path.rstrip(mapper.fs.sep)
    return str(clean_path).split(mapper.fs.sep)[-1]


def get_extension(path: Union[str, os.PathLike]):
    """Get the extension of a file.
    Args:
        path: a path supported by `fsspec` such as local, s3, gcs, etc.
    """
    basename = get_basename(path)
    return basename.split(".")[-1]


def exists(path: Union[str, os.PathLike, fsspec.core.OpenFile, io.IOBase]):
    """Check whether a file exists.
    Args:
        path: a path supported by `fsspec` such as local, s3, gcs, etc.
    """

    if isinstance(path, fsspec.core.OpenFile):
        return path.fs.exists(path.path)

    elif isinstance(path, (str, pathlib.Path)):
        mapper = get_mapper(str(path))
        return mapper.fs.exists(path)

    else:
        # NOTE(hadim): file-like objects always exist right?
        return True


def exists_and_not_empty(path: Union[str, os.PathLike]):
    """Check whether a directory exists and is not empty."""

    if not exists(path):
        return False

    fs = get_mapper(path).fs

    return len(fs.ls(path)) > 0


def mkdir(path: Union[str, os.PathLike], exist_ok: bool = True):
    """Create directory including potential parents."""
    fs = get_mapper(path).fs
    fs.mkdirs(path, exist_ok=exist_ok)


def rm(path: Union[str, os.PathLike], recursive=False, maxdepth=None):
    """Delete a file or a directory with all nested files."""
    fs = get_mapper(path).fs
    fs.rm(path, recursive=recursive, maxdepth=maxdepth)


def join(*paths):
    """Join paths together. The first element determine the
    filesystem to use (and so the separator.
    Args:
        paths: a list of paths supported by `fsspec` such as local, s3, gcs, etc.
    """
    paths = [str(path) for path in paths]
    source_path = paths[0]
    fs = get_mapper(source_path).fs
    full_path = fs.sep.join(paths)
    return full_path


def get_size(file: Union[str, os.PathLike, io.IOBase, fsspec.core.OpenFile]) -> Optional[int]:
    """Get the size of a file given its path. Return None if the
    size can't be retrieved.
    """

    if isinstance(file, io.IOBase) and hasattr(file, "name"):
        fs_local = fsspec.filesystem("file")
        file_size = fs_local.size(getattr(file, "name"))

    elif isinstance(file, (str, pathlib.Path)):
        fs = get_mapper(str(file)).fs
        file_size = fs.size(str(file))

    elif isinstance(file, fsspec.core.OpenFile):
        file_size = file.fs.size(file.path)

    else:
        file_size = None

    return file_size


def copy(
    source: Union[str, os.PathLike, io.IOBase, fsspec.core.OpenFile],
    destination: Union[str, os.PathLike, io.IOBase, fsspec.core.OpenFile],
    chunk_size: int = None,
    force: bool = False,
    progress: bool = False,
    leave_progress: bool = True,
):
    """Copy one file to another location across different filesystem (local, S3, GCS, etc).

    Args:
        source: path or file-like object to copy from.
        destination: path or file-like object to copy to.
        chunk_size: the chunk size to use. If progress is enabled the chunk
            size is `None`, it is set to 2048.
        force: whether to overwrite the destination file it it exists.
        progress: whether to display a progress bar.
        leave_progress: whether to hide the progress bar once the copy is done.
    """

    if progress and chunk_size is None:
        chunk_size = 2048

    if isinstance(source, (str, os.PathLike)):
        source_file = fsspec.open(str(source), "rb")
    else:
        source_file = source

    if isinstance(destination, (str, os.PathLike)):
        # adapt the file mode of the destination depending on the source file.
        destination_mode = "wb"
        if hasattr(source_file, "mode"):
            destination_mode = "wb" if "b" in getattr(source_file, "mode") else "w"
        elif isinstance(source_file, io.BytesIO):
            destination_mode = "wb"
        elif isinstance(source_file, io.StringIO):
            destination_mode = "w"

        destination_file = fsspec.open(str(destination), destination_mode)
    else:
        destination_file = destination

    if not exists(source_file):
        raise ValueError(f"The file being copied does not exist: {source}")

    if not force and exists(destination_file):
        raise ValueError(f"The destination file to copy already exists: {destination}")

    with source_file as source_stream:
        with destination_file as destination_stream:
            if chunk_size is None:
                # copy without chunks
                destination_stream.write(source_stream.read())

            else:
                # copy with chunks

                # determine the size of the source file
                source_size = None
                if progress:
                    source_size = get_size(source)

                # init progress bar
                pbar = tqdm(
                    total=source_size,
                    leave=leave_progress,
                    disable=not progress,
                    unit="B",
                    unit_divisor=1024,
                    unit_scale=True,
                )

                # start the loop
                while True:
                    data = source_stream.read(chunk_size)
                    if not data:
                        break
                    destination_stream.write(data)
                    pbar.update(chunk_size)

                pbar.close()
