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

import logging

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

from graph.exploration.projected_graph import compute_projected_graph
from recipes.dataiku_recipe_helpers import parse_recipe_bool

logger = logging.getLogger(__name__)


input_ds = Dataset(get_input_names_for_role("input_ds")[0])
output_ds = Dataset(get_output_names_for_role("output_ds")[0])

config = get_recipe_config()
graph_of_column = config["create_graph_of"]
linked_by_column = config["linked_by"]
weighted = parse_recipe_bool(config.get("weighted"), default=False)

logger.info(
    "Building projected graph from '%s' linked by '%s' (weighted=%s).",
    graph_of_column,
    linked_by_column,
    weighted,
)

input_df = input_ds.get_dataframe(columns=[graph_of_column, linked_by_column])
output_df = compute_projected_graph(
    input_df,
    graph_of_column=graph_of_column,
    linked_by_column=linked_by_column,
    weighted=weighted,
)

if output_df.empty:
    logger.info(
        "Projected graph is empty: no projected edges were found for '%s' linked by '%s'.",
        graph_of_column,
        linked_by_column,
    )

output_ds.write_with_schema(output_df)

logger.info("Done building projected graph.")
