"""Unit test class for Node-label prediction pipeline."""
import os
from unittest import TestCase

import numpy as np
import pandas as pd
import pytest
from ensmallen.datasets.kgobo import CIO, HP
from tqdm.auto import tqdm

from embiggen import (get_available_models_for_edge_embedding,
                      get_available_models_for_edge_label_prediction,
                      get_available_models_for_node_embedding)
from embiggen.edge_label_prediction import edge_label_prediction_evaluation
from embiggen.edge_label_prediction.edge_label_prediction_model import \
    AbstractEdgeLabelPredictionModel
from embiggen.embedders.ensmallen_embedders.degree_spine import DegreeSPINE
from embiggen.feature_preprocessors import GraphConvolution
from embiggen.utils import AbstractEmbeddingModel
from glob import glob

from .cached_tests import cache_or_store


class TestEvaluateEdgeLabelPrediction(TestCase):
    """Unit test class for edge-label prediction pipeline."""

    def setUp(self):
        """Setup objects for running tests on edge-label prediction pipeline class."""
        self._number_of_holdouts = 2
        self.graph = HP().remove_singleton_nodes().remove_parallel_edges()

    def test_evaluate_embedding_for_edge_label_prediction(self):
        """Test graph visualization."""
        if cache_or_store([
            "./tests/test_edge_label_prediction_pipeline.py",
            "./embiggen/embedders/ensmallen_embedders/degree_spine.py",
            *glob("./embiggen/edge_label_prediction/**/*.py", recursive=True),
            *glob("./embiggen/edge_label_prediction/*.py",),
            "./embiggen/utils/abstract_models/abstract_classifier_model.py",
        ]):
            return
        df = get_available_models_for_edge_label_prediction()
        feature = DegreeSPINE(embedding_size=5)
        red = self.graph.set_all_edge_types("red")
        blue = CIO().remove_singleton_nodes().set_all_edge_types("blue")
        binary_graph = red | blue
        for evaluation_schema in AbstractEdgeLabelPredictionModel.get_available_evaluation_schemas():
            if "Stratified" not in evaluation_schema:
                # TODO! properly test this!
                continue
            holdouts = edge_label_prediction_evaluation(
                holdouts_kwargs={
                    "train_size": 0.8
                },
                node_features=feature,
                models=df.model_name,
                library_names=df.library_name,
                graphs=[self.graph, binary_graph],
                number_of_holdouts=self._number_of_holdouts,
                evaluation_schema=evaluation_schema,
                verbose=True,
                smoke_test=True
            )
        self.assertEqual(holdouts.shape[0],
                         self._number_of_holdouts*2*2*df.shape[0])

    def test_edge_label_prediction_models_apis(self):
        if cache_or_store([
            "./tests/test_edge_label_prediction_pipeline.py",
            "./embiggen/embedders/ensmallen_embedders/degree_spine.py",
            *glob("./embiggen/edge_label_prediction/**/*.py", recursive=True),
            *glob("./embiggen/edge_label_prediction/*.py",),
            "./embiggen/utils/abstract_models/abstract_classifier_model.py",
        ]):
            return
        df = get_available_models_for_edge_label_prediction()
        graph = self.graph
        red = graph.set_all_edge_types("red")
        blue = CIO().remove_singleton_nodes().set_all_edge_types("blue")
        binary_graph = red | blue
        bar = tqdm(
            df.model_name,
            total=df.shape[0],
            leave=False,
        )
        for g in (graph, binary_graph):
            node_features = DegreeSPINE(embedding_size=10).fit_transform(g)
            edge_features=np.random.uniform(
                size=(g.get_number_of_edges(), 5)
            )
            for ef in (edge_features, None):
                for model_name in bar:
                    bar.set_description(
                        f"Testing API of {model_name}")
                    model_class = AbstractEdgeLabelPredictionModel.get_model_from_library(model_name)
                    model = model_class()
                    use_edge_metrics = "use_edge_metrics" in model.parameters()
                    model = model_class(
                        **{
                            **model_class.smoke_test_parameters(),
                            **dict(use_edge_metrics=use_edge_metrics)
                        }
                    )

                    model.fit(
                        g,
                        node_features=node_features,
                        edge_features=ef
                    )

                    if model.library_name() in ("TensorFlow", "scikit-learn", "LightGBM", "CatBoost", "XGBoost"):
                        path = f"model_{model.library_name()}_{model.model_name()}.pkl.gz"
                    elif model.library_name() == "Ensmallen":
                        path = f"model_{model.library_name()}_{model.model_name()}.json"
                    else:
                        raise NotImplementedError(
                            f"The model {model.model_name()} from library {model.library_name()} "
                            "is not currently covered by the test suite!"
                        )
                    if os.path.exists(path):
                        os.remove(path)
                    
                    model.dump(path)
                    restored_model = model.load(path)

                    # We get the restored model parameters and the
                    # original ones and check that they are the same.
                    # We must be careful to remove the np.Nan values
                    # from both parameters, as they are not equal to
                    # themselves.

                    restored_model_parameters = {
                        k: v for k, v in restored_model.parameters().items()
                        if isinstance(v, (list, tuple, np.ndarray)) or pd.notna(v)
                    }

                    model_parameters = {
                        k: v for k, v in model.parameters().items()
                        if isinstance(v, (list, tuple, np.ndarray)) or pd.notna(v)
                    }

                    assert restored_model_parameters == model_parameters

                    model.predict(
                        g,
                        node_features=node_features,
                        edge_features=ef
                    )
                    model.predict_proba(
                        g,
                        node_features=node_features,
                        edge_features=ef
                    )

    def test_evaluate_edge_label_prediction_with_node_types_features(self):
        if cache_or_store([
            "./tests/test_edge_label_prediction_pipeline.py",
            "./embiggen/embedders/ensmallen_embedders/degree_spine.py",
            *glob("./embiggen/edge_label_prediction/**/*.py", recursive=True),
            *glob("./embiggen/edge_label_prediction/*.py",),
            "./embiggen/utils/abstract_models/abstract_classifier_model.py",
        ]):
            return
        df = get_available_models_for_edge_label_prediction()
        graph = self.graph
        holdouts = edge_label_prediction_evaluation(
            holdouts_kwargs=dict(train_size=0.8),
            models=df.model_name,
            library_names=df.library_name,
            node_features=DegreeSPINE(embedding_size=5),
            node_type_features=np.random.uniform(
                size=(graph.get_number_of_node_types(), 5)
            ),
            evaluation_schema="Stratified Monte Carlo",
            graphs=graph,
            number_of_holdouts=self._number_of_holdouts,
            verbose=True,
            smoke_test=True
        )
        self.assertEqual(
            holdouts.shape[0],
            self._number_of_holdouts*2*df.shape[0]
        )

    def test_evaluate_edge_label_prediction_with_node_types_features(self):
        # We expect this to fail because of the wrong number of edge
        # features. Specifically, this graph is undirected, but the
        # edge features are not, so they are double the expected size.
        with pytest.raises(ValueError):
            df = get_available_models_for_edge_label_prediction()
            graph = self.graph
            edge_label_prediction_evaluation(
                holdouts_kwargs=dict(train_size=0.8),
                models=df.model_name,
                library_names=df.library_name,
                node_features=DegreeSPINE(embedding_size=5),
                edge_features=np.random.uniform(
                    size=(graph.get_number_of_directed_edges(), 5)
                ),
                evaluation_schema="Stratified Monte Carlo",
                graphs=graph,
                number_of_holdouts=self._number_of_holdouts,
                verbose=True,
                smoke_test=True
            )
            

    def test_model_recreation(self):
        """Test graph visualization."""
        df = get_available_models_for_edge_label_prediction()

        for _, row in df.iterrows():
            model = AbstractEdgeLabelPredictionModel.get_model_from_library(
                model_name=row.model_name,
                task_name=AbstractEdgeLabelPredictionModel.task_name(),
                library_name=row.library_name
            )()
            try:
                AbstractEdgeLabelPredictionModel.get_model_from_library(
                    model_name=row.model_name,
                    task_name=AbstractEdgeLabelPredictionModel.task_name(),
                    library_name=row.library_name
                )(**model.parameters())
            except Exception as exception:
                raise ValueError(
                    f"Found an error in model {row.model_name} "
                    f"implemented in library {row.library_name}."
                ) from exception

    def test_all_node_embedding_models_as_feature(self):
        """Test graph visualization."""
        if cache_or_store([
            "./tests/test_edge_label_prediction_pipeline.py",
            "./embiggen/embedders/ensmallen_embedders/degree_spine.py",
            *glob("./embiggen/embedders/**/*.py", recursive=True),
            *glob("./embiggen/edge_label_prediction/**/*.py", recursive=True),
            *glob("./embiggen/edge_label_prediction/*.py",),
            "./embiggen/utils/abstract_models/abstract_classifier_model.py",
        ]):
            return
        df = get_available_models_for_node_embedding()
        bar = tqdm(
            df.iterrows(),
            total=df.shape[0],
            leave=False,
            desc="Testing node embedding methods"
        )
        graph = HP().remove_components(top_k_components=1).sort_by_decreasing_outbound_node_degree()

        for _, row in bar:
            if row.requires_edge_weights or row.requires_edge_types or row.requires_node_types:
                continue

            bar.set_description(
                f"Testing {row.model_name} from library {row.library_name}")
            
            edge_label_prediction_evaluation(
                holdouts_kwargs=dict(train_size=0.8),
                models="Decision Tree Classifier",
                node_features=AbstractEmbeddingModel.get_model_from_library(
                    model_name=row.model_name,
                    library_name=row.library_name,
                    task_name="Node Embedding"
                )(),
                node_features_preprocessing_steps=GraphConvolution(),
                evaluation_schema="Stratified Monte Carlo",
                graphs=graph,
                number_of_holdouts=self._number_of_holdouts,
                verbose=False,
                smoke_test=True,
            )

    def test_all_edge_embedding_models_as_feature(self):
        """Test graph visualization."""
        if cache_or_store([
            "./tests/test_edge_label_prediction_pipeline.py",
            "./embiggen/embedders/ensmallen_embedders/degree_spine.py",
            *glob("./embiggen/embedders/**/*.py", recursive=True),
            *glob("./embiggen/edge_label_prediction/**/*.py", recursive=True),
            *glob("./embiggen/edge_label_prediction/*.py",),
            "./embiggen/utils/abstract_models/abstract_classifier_model.py",
        ]):
            return
        df = get_available_models_for_edge_embedding()
        bar = tqdm(
            df.iterrows(),
            total=df.shape[0],
            leave=False,
            desc="Testing edge embedding methods"
        )
        graph = HP().remove_components(top_k_components=1).sort_by_decreasing_outbound_node_degree()

        for _, row in bar:
            if row.requires_edge_weights or row.requires_edge_types or row.requires_node_types:
                continue

            bar.set_description(
                f"Testing {row.model_name} from library {row.library_name}")

            edge_label_prediction_evaluation(
                holdouts_kwargs=dict(train_size=0.8),
                models="Decision Tree Classifier",
                edge_features=AbstractEmbeddingModel.get_model_from_library(
                    model_name=row.model_name,
                    library_name=row.library_name,
                    task_name="Edge Embedding"
                )(),
                evaluation_schema="Stratified Monte Carlo",
                graphs=graph,
                number_of_holdouts=self._number_of_holdouts,
                verbose=False,
                smoke_test=True,
            )
