"""
=============================
Main EasyLink Pipeline Schema
=============================
"""

from easylink.graph_components import (
    EdgeParams,
    InputSlot,
    InputSlotMapping,
    OutputSlot,
    OutputSlotMapping,
)
from easylink.step import (
    HierarchicalStep,
    InputStep,
    LoopStep,
    OutputStep,
    ParallelStep,
    Step,
)
from easylink.utilities.validation_utils import (
    dont_validate,
    validate_blocks,
    validate_clusters,
    validate_dataset_dir,
    validate_ids_to_remove,
    validate_input_dataset_or_known_clusters,
    validate_links,
    validate_records,
)

NODES = [
    # NOTE: In our pipeline schema as documented, there are two inputs: input datasets and known clusters
    # However, due to limitations currently in EasyLink, we can't have multiple output slots on the InputStep.
    # Instead we have a single undifferentiated slot and make it the *implementation's* problem to differentiate
    # based on filename.
    InputStep(),
    LoopStep(
        template_step=HierarchicalStep(
            step_name="entity_resolution",
            input_slots=[
                InputSlot(
                    name="input_datasets",
                    env_var="INPUT_DATASETS_AND_INPUT_KNOWN_CLUSTERS_FILE_PATHS",
                    # NOTE: Since this originates from the InputStep, it will be a *list*
                    # of files, and this validator will be called on *each*
                    # TODO: Change this when https://jira.ihme.washington.edu/browse/MIC-6070 is implemented
                    validator=validate_input_dataset_or_known_clusters,
                ),
                InputSlot(
                    name="known_clusters",
                    env_var="KNOWN_CLUSTERS_AND_MAYBE_INPUT_DATASETS_FILE_PATHS",
                    validator=validate_input_dataset_or_known_clusters,
                ),
            ],
            output_slots=[OutputSlot("clusters")],
            nodes=[
                ParallelStep(
                    # NOTE: Splitters/aggregators on the ParallelStep are implicit!
                    template_step=HierarchicalStep(
                        step_name="determining_exclusions_and_removing_records",
                        directly_implemented=False,
                        input_slots=[
                            InputSlot(
                                name="input_datasets",
                                env_var="INPUT_DATASETS_AND_INPUT_KNOWN_CLUSTERS_FILE_PATHS",
                                validator=validate_input_dataset_or_known_clusters,
                            ),
                            InputSlot(
                                name="known_clusters",
                                env_var="KNOWN_CLUSTERS_AND_MAYBE_INPUT_DATASETS_FILE_PATHS",
                                validator=validate_input_dataset_or_known_clusters,
                            ),
                        ],
                        output_slots=[OutputSlot("datasets")],
                        nodes=[
                            Step(
                                step_name="determining_exclusions",
                                input_slots=[
                                    InputSlot(
                                        name="input_datasets",
                                        env_var="INPUT_DATASETS_AND_INPUT_KNOWN_CLUSTERS_FILE_PATHS",
                                        validator=validate_input_dataset_or_known_clusters,
                                    ),
                                    InputSlot(
                                        name="known_clusters",
                                        env_var="KNOWN_CLUSTERS_AND_MAYBE_INPUT_DATASETS_FILE_PATHS",
                                        validator=validate_input_dataset_or_known_clusters,
                                    ),
                                ],
                                output_slots=[OutputSlot("ids_to_remove")],
                            ),
                            Step(
                                step_name="removing_records",
                                input_slots=[
                                    InputSlot(
                                        name="input_datasets",
                                        env_var="INPUT_DATASETS_AND_INPUT_KNOWN_CLUSTERS_FILE_PATHS",
                                        validator=validate_input_dataset_or_known_clusters,
                                    ),
                                    InputSlot(
                                        name="ids_to_remove",
                                        env_var="IDS_TO_REMOVE_FILE_PATH",
                                        validator=validate_ids_to_remove,
                                    ),
                                ],
                                output_slots=[OutputSlot("dataset")],
                            ),
                        ],
                        edges=[
                            EdgeParams(
                                source_node="determining_exclusions",
                                target_node="removing_records",
                                output_slot="ids_to_remove",
                                input_slot="ids_to_remove",
                            )
                        ],
                        input_slot_mappings=[
                            # NOTE: This is the edge that would normally be split,
                            # but it won't be here, because we don't want it to split
                            # the known clusters to be a separate thing!
                            InputSlotMapping(
                                parent_slot="input_datasets",
                                child_node="determining_exclusions",
                                child_slot="input_datasets",
                            ),
                            InputSlotMapping(
                                parent_slot="known_clusters",
                                child_node="determining_exclusions",
                                child_slot="known_clusters",
                            ),
                            InputSlotMapping(
                                parent_slot="input_datasets",
                                child_node="removing_records",
                                child_slot="input_datasets",
                            ),
                        ],
                        output_slot_mappings=[
                            OutputSlotMapping(
                                # Becomes multiple, after implicit cloneable aggregator
                                parent_slot="datasets",
                                child_node="removing_records",
                                child_slot="dataset",
                            )
                        ],
                    )
                ),
                HierarchicalStep(
                    step_name="clustering",
                    input_slots=[
                        InputSlot(
                            name="datasets",
                            env_var="DATASETS_FILE_PATHS",
                            validator=validate_dataset_dir,
                        ),
                        InputSlot(
                            name="known_clusters",
                            env_var="KNOWN_CLUSTERS_AND_MAYBE_INPUT_DATASETS_FILE_PATHS",
                            validator=validate_input_dataset_or_known_clusters,
                        ),
                    ],
                    output_slots=[OutputSlot("new_clusters")],
                    nodes=[
                        Step(
                            step_name="clusters_to_links",
                            input_slots=[
                                InputSlot(
                                    name="known_clusters",
                                    env_var="KNOWN_CLUSTERS_AND_MAYBE_INPUT_DATASETS_FILE_PATHS",
                                    validator=validate_input_dataset_or_known_clusters,
                                ),
                            ],
                            output_slots=[OutputSlot("known_links")],
                        ),
                        LoopStep(
                            template_step=HierarchicalStep(
                                step_name="linking",
                                input_slots=[
                                    InputSlot(
                                        name="datasets",
                                        env_var="DATASETS_FILE_PATHS",
                                        validator=validate_dataset_dir,
                                    ),
                                    InputSlot(
                                        name="known_links",
                                        env_var="KNOWN_LINKS_FILE_PATH",
                                        validator=validate_links,
                                    ),
                                ],
                                output_slots=[OutputSlot("links")],
                                nodes=[
                                    ParallelStep(
                                        template_step=LoopStep(
                                            template_step=Step(
                                                step_name="pre-processing",
                                                input_slots=[
                                                    InputSlot(
                                                        # NOTE: No splitter here, because
                                                        # not supported by EasyLink;
                                                        # the implementation must do the splitting itself.
                                                        name="dataset",
                                                        env_var="DATASET_DIR_PATHS",
                                                        validator=validate_dataset_dir,
                                                    ),
                                                ],
                                                output_slots=[OutputSlot("dataset")],
                                            ),
                                            self_edges=[
                                                EdgeParams(
                                                    source_node="pre-processing",
                                                    target_node="pre-processing",
                                                    output_slot="dataset",
                                                    input_slot="dataset",
                                                ),
                                            ],
                                        )
                                    ),
                                    Step(
                                        step_name="schema_alignment",
                                        input_slots=[
                                            InputSlot(
                                                name="datasets",
                                                env_var="DATASETS_DIR_PATHS",
                                                validator=validate_dataset_dir,
                                            ),
                                        ],
                                        output_slots=[OutputSlot("records")],
                                    ),
                                    Step(
                                        step_name="blocking_and_filtering",
                                        input_slots=[
                                            InputSlot(
                                                name="records",
                                                env_var="RECORDS_FILE_PATH",
                                                validator=validate_records,
                                            ),
                                            InputSlot(
                                                name="known_links",
                                                env_var="KNOWN_LINKS_FILE_PATH",
                                                validator=validate_links,
                                            ),
                                        ],
                                        output_slots=[OutputSlot("blocks")],
                                    ),
                                    Step(
                                        step_name="evaluating_pairs",
                                        input_slots=[
                                            InputSlot(
                                                name="blocks",
                                                env_var="BLOCKS_DIR_PATH",
                                                validator=validate_blocks,
                                            ),
                                            InputSlot(
                                                name="known_links",
                                                env_var="KNOWN_LINKS_FILE_PATH",
                                                validator=validate_links,
                                            ),
                                        ],
                                        output_slots=[OutputSlot("links")],
                                    ),
                                ],
                                edges=[
                                    EdgeParams(
                                        source_node="pre-processing",
                                        target_node="schema_alignment",
                                        output_slot="dataset",
                                        # NOTE: The implicit ParallelStep aggregator has
                                        # made this multiple (a list)
                                        input_slot="datasets",
                                    ),
                                    EdgeParams(
                                        source_node="schema_alignment",
                                        target_node="blocking_and_filtering",
                                        output_slot="records",
                                        input_slot="records",
                                    ),
                                    EdgeParams(
                                        source_node="blocking_and_filtering",
                                        target_node="evaluating_pairs",
                                        output_slot="blocks",
                                        input_slot="blocks",
                                    ),
                                ],
                                input_slot_mappings=[
                                    InputSlotMapping(
                                        parent_slot="datasets",
                                        child_node="pre-processing",
                                        child_slot="dataset",
                                    ),
                                    InputSlotMapping(
                                        parent_slot="known_links",
                                        child_node="blocking_and_filtering",
                                        child_slot="known_links",
                                    ),
                                    InputSlotMapping(
                                        parent_slot="known_links",
                                        child_node="evaluating_pairs",
                                        child_slot="known_links",
                                    ),
                                ],
                                output_slot_mappings=[
                                    OutputSlotMapping(
                                        parent_slot="links",
                                        child_node="evaluating_pairs",
                                        child_slot="links",
                                    )
                                ],
                            ),
                            self_edges=[
                                EdgeParams(
                                    source_node="linking",
                                    target_node="linking",
                                    output_slot="links",
                                    input_slot="known_links",
                                )
                            ],
                        ),
                        Step(
                            step_name="links_to_clusters",
                            input_slots=[
                                InputSlot(
                                    name="links",
                                    env_var="LINKS_FILE_PATH",
                                    validator=validate_links,
                                ),
                            ],
                            output_slots=[OutputSlot("clusters")],
                        ),
                    ],
                    edges=[
                        EdgeParams(
                            source_node="clusters_to_links",
                            target_node="linking",
                            output_slot="known_links",
                            input_slot="known_links",
                        ),
                        EdgeParams(
                            source_node="linking",
                            target_node="links_to_clusters",
                            output_slot="links",
                            input_slot="links",
                        ),
                    ],
                    input_slot_mappings=[
                        InputSlotMapping(
                            parent_slot="datasets",
                            child_node="linking",
                            child_slot="datasets",
                        ),
                        InputSlotMapping(
                            parent_slot="known_clusters",
                            child_node="clusters_to_links",
                            child_slot="known_clusters",
                        ),
                    ],
                    output_slot_mappings=[
                        OutputSlotMapping(
                            parent_slot="new_clusters",
                            child_node="links_to_clusters",
                            child_slot="clusters",
                        ),
                    ],
                ),
                Step(
                    step_name="updating_clusters",
                    input_slots=[
                        InputSlot(
                            name="new_clusters",
                            env_var="NEW_CLUSTERS_FILE_PATH",
                            validator=validate_clusters,
                        ),
                        InputSlot(
                            name="known_clusters",
                            env_var="KNOWN_CLUSTERS_AND_MAYBE_INPUT_DATASETS_FILE_PATHS",
                            validator=validate_input_dataset_or_known_clusters,
                        ),
                    ],
                    output_slots=[OutputSlot("clusters")],
                ),
            ],
            edges=[
                EdgeParams(
                    source_node="determining_exclusions_and_removing_records",
                    target_node="clustering",
                    output_slot="datasets",
                    input_slot="datasets",
                ),
                EdgeParams(
                    source_node="clustering",
                    target_node="updating_clusters",
                    output_slot="new_clusters",
                    input_slot="new_clusters",
                ),
            ],
            input_slot_mappings=[
                InputSlotMapping(
                    parent_slot="input_datasets",
                    child_node="determining_exclusions_and_removing_records",
                    child_slot="input_datasets",
                ),
                InputSlotMapping(
                    parent_slot="known_clusters",
                    child_node="determining_exclusions_and_removing_records",
                    child_slot="known_clusters",
                ),
                InputSlotMapping(
                    parent_slot="known_clusters",
                    child_node="clustering",
                    child_slot="known_clusters",
                ),
                InputSlotMapping(
                    parent_slot="known_clusters",
                    child_node="updating_clusters",
                    child_slot="known_clusters",
                ),
            ],
            output_slot_mappings=[
                OutputSlotMapping(
                    child_node="updating_clusters",
                    child_slot="clusters",
                    parent_slot="clusters",
                ),
            ],
        ),
        self_edges=[
            EdgeParams(
                source_node="entity_resolution",
                target_node="entity_resolution",
                output_slot="clusters",
                input_slot="known_clusters",
            )
        ],
    ),
    Step(
        step_name="canonicalizing_and_downstream_analysis",
        input_slots=[
            InputSlot(
                name="input_datasets",
                env_var="INPUT_DATASETS_AND_INPUT_KNOWN_CLUSTERS_FILE_PATHS",
                validator=validate_input_dataset_or_known_clusters,
            ),
            InputSlot(
                name="clusters",
                env_var="CLUSTERS_FILE_PATH",
                validator=validate_clusters,
            ),
        ],
        output_slots=[OutputSlot("analysis_output")],
    ),
    OutputStep(
        input_slots=[
            InputSlot(name="analysis_output", env_var=None, validator=dont_validate)
        ],
    ),
]
EDGES = [
    EdgeParams(
        source_node="input_data",
        target_node="entity_resolution",
        output_slot="all",
        input_slot="input_datasets",
    ),
    EdgeParams(
        source_node="input_data",
        target_node="entity_resolution",
        output_slot="all",
        input_slot="known_clusters",
    ),
    EdgeParams(
        source_node="input_data",
        target_node="canonicalizing_and_downstream_analysis",
        output_slot="all",
        input_slot="input_datasets",
    ),
    EdgeParams(
        source_node="entity_resolution",
        target_node="canonicalizing_and_downstream_analysis",
        output_slot="clusters",
        input_slot="clusters",
    ),
    EdgeParams(
        source_node="canonicalizing_and_downstream_analysis",
        target_node="results",
        output_slot="analysis_output",
        input_slot="analysis_output",
    ),
]
SCHEMA_PARAMS = (NODES, EDGES)
