import time
import warnings

import numpy as np
import pandas as pd

from graph_based_intrusion_detection.utils import constants


def run_analysis_step(nodes: list[str], analysis_step: dict, network_state: dict) -> dict:
    """
    Run an analysis step on a dataset
    :param nodes: nodes to analyze
    :param analysis_step: dictionary with analysis step configuration
    :param network_state: network state dictionary
    :return: analysis results
    """

    # extract network graph from network state
    network_state_data = network_state[analysis_step["network_state_key"]]

    # apply dataset function to extract dataset for node
    datasets = [analysis_step["dataset_function"](network_state_data, node) for node in nodes]

    # remove identifier columns from datasets
    datasets = [dataset.drop(columns=analysis_step["dataset_identifiers"]) for dataset in datasets]

    # scale datasets
    scaled_datasets = [analysis_step["scaler"].transform(dataset) for dataset in datasets]

    # apply preprocessing function, if given
    if "preprocessing_function" in analysis_step:
        scaled_datasets = [analysis_step["preprocessing_function"](scaled_dataset) for scaled_dataset in
                           scaled_datasets]

    # combine datasets, get indices to split them later
    combined_dataset = np.concatenate(scaled_datasets)
    split_indices = list()
    for dataset in scaled_datasets:
        split_indices.append(len(dataset) + split_indices[-1] if split_indices else len(dataset))

    # predict and interpret results
    prediction_function = analysis_step["predict_function"]
    if constants.SURPRESS_WARNINGS:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            predictions = prediction_function(combined_dataset)
    else:
        predictions = prediction_function(combined_dataset)

    # split predictions
    predictions = np.split(predictions, split_indices[:-1])

    results = [analysis_step["interpretation_function"](prediction) for prediction in predictions]

    # compose results with node names to dict
    result_dict = {node: result for node, result in zip(nodes, results)}

    return result_dict
