# -*- coding: utf-8 -*-
from __future__ import annotations

import logging

from dataiku import Dataset, Folder
from dataiku.customrecipe import get_input_names_for_role, get_output_names_for_role, get_recipe_config

from graph.exploration.algorithms.backend_service import (
    GraphAnalyticsExecutionLocation,
    build_recipe_execution_note,
    execute_graph_algorithm,
)
from graph.exploration.algorithms.graph_algorithm_registry import (
    GRAPH_CLUSTERING_ALGORITHM_DEFINITIONS,
)
from graph.exploration.algorithms.graph_clustering import GraphClusteringParams, ProjectedGraph
from recipes.dataiku_recipe_helpers import (
    GraphAnalyticsRecipeBackendConfig,
    parse_graph_algorithm_selection,
    parse_recipe_bool,
    validate_graph_analytics_recipe_backend,
    write_batches_to_dataset,
)

logger = logging.getLogger(__name__)


db_folder_names = get_input_names_for_role("graph_db_folder")
db_folder = Folder(db_folder_names[0]) if db_folder_names else None
output_ds = Dataset(get_output_names_for_role("output_ds")[0])

config = get_recipe_config()
backend_config = GraphAnalyticsRecipeBackendConfig.from_recipe_config(
    config,
    default_location=GraphAnalyticsExecutionLocation.IN_MEMORY,
)

is_directed = parse_recipe_bool(config.get("directed_graph"), default=False)
selected_algorithms = set(
    parse_graph_algorithm_selection(
        config,
        definitions=GRAPH_CLUSTERING_ALGORITHM_DEFINITIONS,
        directed_graph=is_directed,
    )
)

params = GraphClusteringParams(
    directed_graph=is_directed,
    output_type=config.get("output_type", "output_edges"),
    weight_property=config.get("weight_property") or None,
    fastgreedy="fastgreedy" in selected_algorithms,
    multilevel="multilevel" in selected_algorithms,
    infomap="infomap" in selected_algorithms,
    walktrap="walktrap" in selected_algorithms,
)

projected_graph = ProjectedGraph(set(config["node_groups"]), set(config["edge_groups"]))

effective_database_type, conn_info, db_handler = backend_config.resolve_backend(db_folder=db_folder)
validate_graph_analytics_recipe_backend(
    "graph_clustering",
    effective_database_type=effective_database_type,
    backend_config=backend_config,
)

logger.info(build_recipe_execution_note(backend_config.location, backend_config.sampling_mode))
with execute_graph_algorithm(
    "graph_clustering",
    conn_info,
    db_handler,
    projected_graph,
    params,
    batch_size=int(config.get("batch_size", 10000)),
    location=backend_config.location,
    sampling_mode=backend_config.sampling_mode,
    sampling_limit=backend_config.sampling_limit,
    readonly=True,
) as execution:
    logger.info("Computing graph clustering on DB '%s'.", execution.connection_info.uid)
    write_batches_to_dataset(
        output_ds,
        execution.batches,
        empty_error_factory=lambda: ValueError("The selected graph is empty; cannot compute graph clustering."),
    )

logger.info("Done computing graph clustering.")
