__all__ = ["identify_associations"]

from functools import reduce
from os.path import exists
from pathlib import Path
from typing import Literal, Sized, Union, cast

import hydra
import numpy as np
import pandas as pd
import torch
from omegaconf import OmegaConf
from scipy.stats import ks_2samp, pearsonr  # type: ignore
from torch.utils.data import DataLoader

from move.analysis.metrics import get_2nd_order_polynomial
from move.conf.schema import (
    IdentifyAssociationsBayesConfig,
    IdentifyAssociationsConfig,
    IdentifyAssociationsKSConfig,
    IdentifyAssociationsTTestConfig,
    MOVEConfig,
)
from move.core.logging import get_logger
from move.core.typing import BoolArray, FloatArray, IntArray
from move.data import io
from move.data.dataloaders import MOVEDataset, make_dataloader
from move.data.perturbations import (
    ContinuousPerturbationType,
    perturb_categorical_data,
    perturb_continuous_data_extended,
)
from move.data.preprocessing import one_hot_encode_single
from move.models.vae import VAE
from move.visualization.dataset_distributions import (
    plot_correlations,
    plot_cumulative_distributions,
    plot_feature_association_graph,
    plot_reconstruction_movement,
)

TaskType = Literal["bayes", "ttest", "ks"]
CONTINUOUS_TARGET_VALUE = ["minimum", "maximum", "plus_std", "minus_std"]


def _get_task_type(
    task_config: IdentifyAssociationsConfig,
) -> TaskType:
    task_type = OmegaConf.get_type(task_config)
    if task_type is IdentifyAssociationsBayesConfig:
        return "bayes"
    if task_type is IdentifyAssociationsTTestConfig:
        return "ttest"
    if task_type is IdentifyAssociationsKSConfig:
        return "ks"
    raise ValueError("Unsupported type of task!")


def _validate_task_config(
    task_config: IdentifyAssociationsConfig, task_type: TaskType
) -> None:
    if not (0.0 <= task_config.sig_threshold <= 1.0):
        raise ValueError("Significance threshold must be within [0, 1].")
    if task_type == "ttest":
        task_config = cast(IdentifyAssociationsTTestConfig, task_config)
        if len(task_config.num_latent) != 4:
            raise ValueError("4 latent space dimensions required.")


def prepare_for_categorical_perturbation(
    config: MOVEConfig,
    interim_path: Path,
    baseline_dataloader: DataLoader,
    cat_list: list[FloatArray],
) -> tuple[
    list[DataLoader],
    BoolArray,
    BoolArray,
]:
    """
    This function creates the required dataloaders and masks
    for further categorical association analysis.

    Args:
        config: main configuration file
        interim_path: path where the intermediate outputs are saved
        baseline_dataloader: reference dataloader that will be perturbed
        cat_list: list of arrays with categorical data

    Returns:
        dataloaders: all dataloaders, including baseline appended last.
        nan_mask: mask for Nans
        feature_mask: masks the column for the perturbed feature.
    """

    # Read original data and create perturbed datasets
    task_config = cast(IdentifyAssociationsConfig, config.task)
    logger = get_logger(__name__)

    # Loading mappings:
    mappings = io.load_mappings(interim_path / "mappings.json")
    target_mapping = mappings[task_config.target_dataset]
    target_value = one_hot_encode_single(target_mapping, task_config.target_value)
    logger.debug(
        f"Target value: {task_config.target_value} => {target_value.astype(int)[0]}"
    )

    dataloaders = perturb_categorical_data(
        baseline_dataloader,
        config.data.categorical_names,
        task_config.target_dataset,
        target_value,
    )
    dataloaders.append(baseline_dataloader)

    baseline_dataset = cast(MOVEDataset, baseline_dataloader.dataset)

    assert baseline_dataset.con_all is not None
    orig_con = baseline_dataset.con_all
    nan_mask = (orig_con == 0).numpy()  # NaN values encoded as 0s
    logger.debug(f"# NaN values: {np.sum(nan_mask)}/{orig_con.numel()}")

    target_dataset_idx = config.data.categorical_names.index(task_config.target_dataset)
    target_dataset = cat_list[target_dataset_idx]
    feature_mask = np.all(target_dataset == target_value, axis=2)  # 2D: N x P
    feature_mask |= np.sum(target_dataset, axis=2) == 0

    return (
        dataloaders,
        nan_mask,
        feature_mask,
    )


def prepare_for_continuous_perturbation(
    config: MOVEConfig,
    output_subpath: Path,
    baseline_dataloader: DataLoader,
) -> tuple[
    list[DataLoader],
    BoolArray,
    BoolArray,
]:
    """
    This function creates the required dataloaders and masks
    for further continuous association analysis.

    Args:
        config:
            main configuration file.
        output_subpath:
            path where the output plots for continuous analysis are saved.
        baseline_dataloader:
            reference dataloader that will be perturbed.

    Returns:
        dataloaders:
            list with all dataloaders, including baseline appended last.
        nan_mask:
            mask for NaNs
        feature_mask:
            same as `nan_mask`, in this case.
    """

    # Read original data and create perturbed datasets
    logger = get_logger(__name__)
    task_config = cast(IdentifyAssociationsConfig, config.task)

    dataloaders = perturb_continuous_data_extended(
        baseline_dataloader,
        config.data.continuous_names,
        task_config.target_dataset,
        cast(ContinuousPerturbationType, task_config.target_value),
        output_subpath,
    )
    dataloaders.append(baseline_dataloader)

    baseline_dataset = cast(MOVEDataset, baseline_dataloader.dataset)

    assert baseline_dataset.con_all is not None
    orig_con = baseline_dataset.con_all
    nan_mask = (orig_con == 0).numpy()  # NaN values encoded as 0s
    logger.debug(f"# NaN values: {np.sum(nan_mask)}/{orig_con.numel()}")
    feature_mask = nan_mask

    return (dataloaders, nan_mask, feature_mask)


def _bayes_approach(
    config: MOVEConfig,
    task_config: IdentifyAssociationsBayesConfig,
    train_dataloader: DataLoader,
    baseline_dataloader: DataLoader,
    dataloaders: list[DataLoader],
    models_path: Path,
    num_perturbed: int,
    num_samples: int,
    num_continuous: int,
    nan_mask: BoolArray,
    feature_mask: BoolArray,
) -> tuple[Union[IntArray, FloatArray], ...]:

    assert task_config.model is not None
    device = torch.device("cuda" if task_config.model.cuda else "cpu")

    # Train models
    logger = get_logger(__name__)
    logger.info("Training models")
    mean_diff = np.zeros((num_perturbed, num_samples, num_continuous))
    normalizer = 1 / task_config.num_refits

    # Last appended dataloader is the baseline
    baseline_dataset = cast(MOVEDataset, baseline_dataloader.dataset)

    for j in range(task_config.num_refits):
        # Initialize model
        model: VAE = hydra.utils.instantiate(
            task_config.model,
            continuous_shapes=baseline_dataset.con_shapes,
            categorical_shapes=baseline_dataset.cat_shapes,
        )
        if j == 0:
            logger.debug(f"Model: {model}")

        # Train/reload model
        model_path = models_path / f"model_{task_config.model.num_latent}_{j}.pt"
        if model_path.exists():
            logger.debug(f"Re-loading refit {j + 1}/{task_config.num_refits}")
            model.load_state_dict(torch.load(model_path))
            model.to(device)
        else:
            logger.debug(f"Training refit {j + 1}/{task_config.num_refits}")
            model.to(device)
            hydra.utils.call(
                task_config.training_loop,
                model=model,
                train_dataloader=train_dataloader,
            )
            if task_config.save_refits:
                torch.save(model.state_dict(), model_path)
        model.eval()

        # Calculate baseline reconstruction
        _, baseline_recon = model.reconstruct(baseline_dataloader)
        min_feat, max_feat = np.zeros((num_perturbed, num_continuous)), np.zeros(
            (num_perturbed, num_continuous)
        )
        min_baseline, max_baseline = np.min(baseline_recon, axis=0), np.max(
            baseline_recon, axis=0
        )

        # Calculate perturb reconstruction => keep track of mean difference
        for i in range(num_perturbed):
            _, perturb_recon = model.reconstruct(dataloaders[i])
            diff = perturb_recon - baseline_recon  # 2D: N x C
            mean_diff[i, :, :] += diff * normalizer

            min_perturb, max_perturb = np.min(perturb_recon, axis=0), np.max(
                perturb_recon, axis=0
            )
            min_feat[i, :], max_feat[i, :] = np.min(
                [min_baseline, min_perturb], axis=0
            ), np.max([max_baseline, max_perturb], axis=0)

    # Calculate Bayes factors
    logger.info("Identifying significant features")
    bayes_k = np.empty((num_perturbed, num_continuous))
    bayes_mask = np.zeros(np.shape(bayes_k))
    for i in range(num_perturbed):
        mask = feature_mask[:, [i]] | nan_mask  # 2D: N x C
        diff = np.ma.masked_array(mean_diff[i, :, :], mask=mask)  # 2D: N x C
        prob = np.ma.compressed(np.mean(diff > 1e-8, axis=0))  # 1D: C
        bayes_k[i, :] = np.log(prob + 1e-8) - np.log(1 - prob + 1e-8)
        if task_config.target_value in CONTINUOUS_TARGET_VALUE:
            bayes_mask[i, :] = (
                baseline_dataloader.dataset.con_all[0, :]
                - dataloaders[i].dataset.con_all[0, :]
            )

    bayes_mask[bayes_mask != 0] = 1
    bayes_mask = np.array(bayes_mask, dtype=bool)

    # Calculate Bayes probabilities
    bayes_abs = np.abs(bayes_k)
    bayes_p = np.exp(bayes_abs) / (1 + np.exp(bayes_abs))  # 2D: N x C
    bayes_abs[bayes_mask] = np.min(
        bayes_abs
    )  # Bring feature_i feature_i associations to minimum
    sort_ids = np.argsort(bayes_abs, axis=None)[::-1]  # 1D: N x C
    prob = np.take(bayes_p, sort_ids)  # 1D: N x C
    logger.debug(f"Bayes proba range: [{prob[-1]:.3f} {prob[0]:.3f}]")

    # Sort Bayes
    bayes_k = np.take(bayes_k, sort_ids)  # 1D: N x C

    # Calculate FDR
    fdr = np.cumsum(1 - prob) / np.arange(1, prob.size + 1)  # 1D
    idx = np.argmin(np.abs(fdr - task_config.sig_threshold))
    logger.debug(f"FDR range: [{fdr[0]:.3f} {fdr[-1]:.3f}]")

    return sort_ids[:idx], prob[:idx], fdr[:idx], bayes_k[:idx]


def _ttest_approach(
    task_config: IdentifyAssociationsTTestConfig,
    train_dataloader: DataLoader,
    baseline_dataloader: DataLoader,
    dataloaders: list[DataLoader],
    models_path: Path,
    interim_path: Path,
    num_perturbed: int,
    num_samples: int,
    num_continuous: int,
    nan_mask: BoolArray,
    feature_mask: BoolArray,
) -> tuple[Union[IntArray, FloatArray], ...]:

    from scipy.stats import ttest_rel

    assert task_config.model is not None
    device = torch.device("cuda" if task_config.model.cuda else "cpu")

    # Train models
    logger = get_logger(__name__)
    logger.info("Training models")
    pvalues = np.empty(
        (
            len(task_config.num_latent),
            task_config.num_refits,
            num_perturbed,
            num_continuous,
        )
    )

    # Last appended dataloader is the baseline
    baseline_dataset = cast(MOVEDataset, baseline_dataloader.dataset)

    for k, num_latent in enumerate(task_config.num_latent):
        for j in range(task_config.num_refits):

            # Initialize model
            model: VAE = hydra.utils.instantiate(
                task_config.model,
                continuous_shapes=baseline_dataset.con_shapes,
                categorical_shapes=baseline_dataset.cat_shapes,
                num_latent=num_latent,
            )
            if j == 0:
                logger.debug(f"Model: {model}")

            # Train model
            model_path = models_path / f"model_{num_latent}_{j}.pt"
            if model_path.exists():
                logger.debug(f"Re-loading refit {j + 1}/{task_config.num_refits}")
                model.load_state_dict(torch.load(model_path))
                model.to(device)
            else:
                logger.debug(f"Training refit {j + 1}/{task_config.num_refits}")
                model.to(device)
                hydra.utils.call(
                    task_config.training_loop,
                    model=model,
                    train_dataloader=train_dataloader,
                )
                if task_config.save_refits:
                    torch.save(model.state_dict(), model_path)
            model.eval()

            # Get baseline reconstruction and baseline difference
            _, baseline_recon = model.reconstruct(baseline_dataloader)
            baseline_diff = np.empty((10, num_samples, num_continuous))
            for i in range(10):
                _, recon = model.reconstruct(baseline_dataloader)
                baseline_diff[i, :, :] = recon - baseline_recon
            baseline_diff = np.mean(baseline_diff, axis=0)  # 2D: N x C
            baseline_diff = np.where(nan_mask, np.nan, baseline_diff)

            # T-test between baseline and perturb difference
            for i in range(num_perturbed):
                _, perturb_recon = model.reconstruct(dataloaders[i])
                perturb_diff = perturb_recon - baseline_recon
                mask = feature_mask[:, [i]] | nan_mask  # 2D: N x C
                _, pvalues[k, j, i, :] = ttest_rel(
                    a=np.where(mask, np.nan, perturb_diff),
                    b=np.where(mask, np.nan, baseline_diff),
                    axis=0,
                    nan_policy="omit",
                )

    # Correct p-values (Bonferroni)
    pvalues = np.minimum(pvalues * num_continuous, 1.0)
    np.save(interim_path / "pvals.npy", pvalues)

    # Find significant hits
    overlap_thres = task_config.num_refits // 2
    reject = pvalues <= task_config.sig_threshold  # 4D: L x R x P x C
    overlap = reject.sum(axis=1) >= overlap_thres  # 3D: L x P x C
    sig_ids = overlap.sum(axis=0) >= 3  # 2D: P x C
    sig_ids = np.flatnonzero(sig_ids)  # 1D

    # Report median p-value
    masked_pvalues = np.ma.masked_array(pvalues, mask=~reject)  # 4D
    masked_pvalues = np.ma.median(masked_pvalues, axis=1)  # 3D
    masked_pvalues = np.ma.median(masked_pvalues, axis=0)  # 2D
    sig_pvalues = np.ma.compressed(np.take(masked_pvalues, sig_ids))  # 1D

    return sig_ids, sig_pvalues


def _ks_approach(
    config: MOVEConfig,
    task_config: IdentifyAssociationsKSConfig,
    train_dataloader: DataLoader,
    baseline_dataloader: DataLoader,
    dataloaders: list[DataLoader],
    models_path: Path,
    num_perturbed: int,
    num_samples: int,
    num_continuous: int,
    con_names: list[list[str]],
    output_path: Path,
) -> tuple[Union[IntArray, FloatArray], ...]:
    """
    Find associations between continuous features using Kolmogorov-Smirnov distances.
    When perturbing feature A, this function measures the shift of the reconstructed
    distribution for feature B (over samples) from 1) the baseline reconstruction to 2)
    the reconstruction when perturbing A.

    If A and B are related the perturbation of A in the input will lead to a change in
    feature B's reconstruction, that will be measured by KS distance.

    Associations are then ranked according to KS distance (absolute value).


    Args:
        config: MOVE main configuration.
        task_config: IdentifyAssociationsKSConfig configuration.
        train_dataloader: training DataLoader.
        baseline_dataloader: unperturbed DataLoader.
        dataloaders: list of DataLoaders where DataLoader[i] is obtained by perturbing
                     feature i in the target dataset.
        models_path: path to the models.
        num_perturbed: number of perturbed features.
        num_samples: total number of samples
        num_continuous: number of continuous features
                        (all continuous datasets concatenated).
        con_names: list of lists where eah inner list
                   contains the feature names of a specific continuous dataset
        output_path: path where QC summary metrics will be saved.

    Returns:
        sort_ids: list with flattened IDs of the associations
                  above the significance threshold.
        ks_distance: Ordered list with signed KS scores. KS scores quantify the
                    direction and magnitude of the shift in feature B's reconstruction
                    when perturbing feature A.


    !!! Note !!!:

    The sign of the KS score can be misleading: negative sign means positive shift.
    since the cumulative distribution starts growing later and is found below
    the reference (baseline). Hence:
    a) with plus_std, negative sign means a positive correlation.
    b) with minus_std, negative sign means a negative correlation.
    """

    assert task_config.model is not None
    device = torch.device("cuda" if task_config.model.cuda else "cpu")
    figure_path = output_path / "figures"
    figure_path.mkdir(exist_ok=True, parents=True)

    # Data containers
    stats = np.empty((task_config.num_refits, num_perturbed, num_continuous))
    stat_signs = np.empty_like(stats)
    rec_corr, slope = np.empty((task_config.num_refits, num_continuous)), np.empty(
        (task_config.num_refits, num_continuous)
    )
    ks_mask = np.zeros((num_perturbed, num_continuous))
    latent_matrix = np.empty(
        (num_samples, task_config.model.num_latent, len(dataloaders))
    )

    # Last appended dataloader is the baseline
    baseline_dataset = cast(MOVEDataset, baseline_dataloader.dataset)

    # Train models
    logger = get_logger(__name__)
    logger.info("Training models")

    target_dataset_idx = config.data.continuous_names.index(task_config.target_dataset)
    perturbed_names = con_names[target_dataset_idx]

    for j in range(task_config.num_refits):  # Train num_refits models

        # Initialize model
        model: VAE = hydra.utils.instantiate(
            task_config.model,
            continuous_shapes=baseline_dataset.con_shapes,
            categorical_shapes=baseline_dataset.cat_shapes,
        )
        if j == 0:
            logger.debug(f"Model: {model}")

        # Train/reload model
        model_path = models_path / f"model_{task_config.model.num_latent}_{j}.pt"
        if model_path.exists():
            logger.debug(f"Re-loading refit {j + 1}/{task_config.num_refits}")
            model.load_state_dict(torch.load(model_path))
            model.to(device)
        else:
            logger.debug(f"Training refit {j + 1}/{task_config.num_refits}")
            model.to(device)
            hydra.utils.call(
                task_config.training_loop,
                model=model,
                train_dataloader=train_dataloader,
            )
            if task_config.save_refits:
                torch.save(model.state_dict(), model_path)
        model.eval()

        # Calculate baseline reconstruction
        _, baseline_recon = model.reconstruct(baseline_dataloader)
        min_feat = np.zeros((num_perturbed, num_continuous))
        max_feat = np.zeros((num_perturbed, num_continuous))
        min_baseline = np.min(baseline_recon, axis=0)
        max_baseline = np.max(baseline_recon, axis=0)

        # QC of feature's reconstruction ##############################
        logger.debug("Calculating quality control of the feature reconstructions")
        # Correlation and slope for each feature's reconstruction
        feature_names = reduce(list.__add__, con_names)

        for k in range(num_continuous):
            x = baseline_dataloader.dataset.con_all.numpy()[:, k]  # baseline_recon[:,i]
            y = baseline_recon[:, k]
            x_pol, y_pol, (a2, a1, a) = get_2nd_order_polynomial(x, y)
            slope[j, k] = a1
            rec_corr[j, k] = pearsonr(x, y).statistic

            if (
                feature_names[k] in task_config.perturbed_feature_names
                or feature_names[k] in task_config.target_feature_names
            ):

                # Plot correlations
                fig = plot_correlations(x, y, x_pol, y_pol, a2, a1, a, k)
                fig.savefig(
                    figure_path
                    / f"Input_vs_reconstruction_correlation_feature_{k}_refit_{j}.png",
                    dpi=50,
                )

        # Calculate perturbed reconstruction and shifts #############################
        logger.debug("Computing KS scores")

        # Save original latent space for first refit:
        if j == 0:
            latent = model.project(baseline_dataloader)
            latent_matrix[:, :, -1] = latent

        for i, pert_feat in enumerate(perturbed_names):
            _, perturb_recon = model.reconstruct(dataloaders[i])
            min_perturb = np.min(perturb_recon, axis=0)
            max_perturb = np.max(perturb_recon, axis=0)
            min_feat[i, :] = np.min([min_baseline, min_perturb], axis=0)
            max_feat[i, :] = np.max([max_baseline, max_perturb], axis=0)

            # Save latent representation for perturbed samples
            if j == 0:
                latent_pert = model.project(dataloaders[i])
                latent_matrix[:, :, i] = latent_pert

            for k, targ_feat in enumerate(feature_names):
                # Calculate ks factors: measure distance between baseline and perturbed
                # reconstruction distributions per feature (k)
                res = ks_2samp(perturb_recon[:, k], baseline_recon[:, k])
                stats[j, i, k] = res.statistic
                stat_signs[j, i, k] = res.statistic_sign

                if (
                    pert_feat in task_config.perturbed_feature_names
                    and targ_feat in task_config.target_feature_names
                ):

                    # Plotting preliminary results:
                    n_bins = 50
                    hist_base, edges = np.histogram(
                        baseline_recon[:, k],
                        bins=np.linspace(min_feat[i, k], max_feat[i, k], n_bins),
                        density=True,
                    )
                    hist_pert, edges = np.histogram(
                        perturb_recon[:, k],
                        bins=np.linspace(min_feat[i, k], max_feat[i, k], n_bins),
                        density=True,
                    )

                    # Cumulative distribution:
                    fig = plot_cumulative_distributions(
                        edges,
                        hist_base,
                        hist_pert,
                        title=f"Cumulative_perturbed_{i}_measuring_"
                        f"{k}_stats_{stats[j, i, k]}",
                    )
                    fig.savefig(
                        figure_path
                        / (
                            f"Cumulative_refit_{j}_perturbed_{i}_"
                            f"measuring_{k}_stats_{stats[j, i, k]}.png"
                        )
                    )

                    # Feature changes:
                    fig = plot_reconstruction_movement(baseline_recon, perturb_recon, k)
                    fig.savefig(
                        figure_path / f"Changes_pert_{i}_on_feat_{k}_refit_{j}.png"
                    )

    # Save latent space matrix:
    np.save(output_path / "latent_location.npy", latent_matrix)
    np.save(output_path / "perturbed_features_list.npy", np.array(perturbed_names))

    # Creating a mask for self associations
    logger.debug("Creating self-association mask")
    for i in range(num_perturbed):
        if task_config.target_value in CONTINUOUS_TARGET_VALUE:
            ks_mask[i, :] = (
                baseline_dataloader.dataset.con_all[0, :]
                - dataloaders[i].dataset.con_all[0, :]
            )
    ks_mask[ks_mask != 0] = 1
    ks_mask = np.array(ks_mask, dtype=bool)

    # Take the median of KS values (with sign) over refits.
    final_stats = np.nanmedian(stats * stat_signs, axis=0)
    final_stats[ks_mask] = (
        0.0  # Zero all masked values, placing them at end of the ranking
    )

    # KS-threshold:
    ks_thr = np.sqrt(-np.log(task_config.sig_threshold / 2) * 1 / (num_samples))
    logger.info(f"Suggested absolute KS threshold is: {ks_thr}")

    # Sort associations by absolute KS value
    sort_ids = np.argsort(abs(final_stats), axis=None)[::-1]  # 1D: N x C
    ks_distance = np.take(final_stats, sort_ids)  # 1D: N x C

    # Writing Quality control csv file.
    # Mean slope and correlation over refits as qc metrics.
    logger.info("Writing QC file")
    qc_df = pd.DataFrame({"Feature names": feature_names})
    qc_df["slope"] = np.nanmean(slope, axis=0)
    qc_df["reconstruction_correlation"] = np.nanmean(rec_corr, axis=0)
    qc_df.to_csv(output_path / "QC_summary_KS.tsv", sep="\t", index=False)

    # Return first idx associations: redefined for reasonable threshold

    return sort_ids[abs(ks_distance) >= ks_thr], ks_distance[abs(ks_distance) >= ks_thr]


def save_results(
    config: MOVEConfig,
    con_shapes: list[int],
    cat_names: list[list[str]],
    con_names: list[list[str]],
    output_path: Path,
    sig_ids,
    extra_cols,
    extra_colnames,
) -> None:
    """
    This function saves the obtained associations in a TSV file containing
    the following columns:
        feature_a_id
        feature_b_id
        feature_a_name
        feature_b_name
        feature_b_dataset
        proba/p_value: number quantifying the significance of the association

    Args:
        config: main config
        con_shapes: tuple with the number of features per continuous dataset
        cat_names: list of lists of names for the categorical features.
                   Each inner list corresponds to a separate dataset.
        con_names: list of lists of names for the continuous features.
                   Each inner list corresponds to a separate dataset.
        output_path: path where the results will be saved
        sig_ids: ids for the significat features
        extra_cols: extra data when calling the approach function
        extra_colnames: names for the extra data columns
    """
    logger = get_logger(__name__)
    logger.info(f"Significant hits found: {sig_ids.size}")
    task_config = cast(IdentifyAssociationsConfig, config.task)
    task_type = _get_task_type(task_config)

    num_continuous = sum(con_shapes)  # C

    if sig_ids.size > 0:
        sig_ids = np.vstack((sig_ids // num_continuous, sig_ids % num_continuous)).T
        logger.info("Writing results")
        results = pd.DataFrame(sig_ids, columns=["feature_a_id", "feature_b_id"])

        # Check if the task is for continuous or categorical data
        if task_config.target_value in CONTINUOUS_TARGET_VALUE:
            target_dataset_idx = config.data.continuous_names.index(
                task_config.target_dataset
            )
            a_df = pd.DataFrame(dict(feature_a_name=con_names[target_dataset_idx]))
        else:
            target_dataset_idx = config.data.categorical_names.index(
                task_config.target_dataset
            )
            a_df = pd.DataFrame(dict(feature_a_name=cat_names[target_dataset_idx]))
        a_df.index.name = "feature_a_id"
        a_df.reset_index(inplace=True)
        feature_names = reduce(list.__add__, con_names)
        b_df = pd.DataFrame(dict(feature_b_name=feature_names))
        b_df.index.name = "feature_b_id"
        b_df.reset_index(inplace=True)
        results = results.merge(a_df, on="feature_a_id", how="left").merge(
            b_df, on="feature_b_id", how="left"
        )
        results["feature_b_dataset"] = pd.cut(
            results["feature_b_id"],
            bins=cast(list[int], np.cumsum([0] + con_shapes)),
            right=False,
            labels=config.data.continuous_names,
        )
        for col, colname in zip(extra_cols, extra_colnames):
            results[colname] = col
        results.to_csv(
            output_path / f"results_sig_assoc_{task_type}.tsv", sep="\t", index=False
        )


def identify_associations(config: MOVEConfig) -> None:
    """
    Leads to the execution of the appropriate association
    identification tasks. The function is organized in three
    blocks:
        1) Prepare the data and create the dataloaders with their masks.
        2) Evaluate associations using bayes or ttest approach.
        3) Save results.
    """
    # DATA PREPARATION ######################
    # Read original data and create perturbed datasets####

    logger = get_logger(__name__)
    task_config = cast(IdentifyAssociationsConfig, config.task)
    task_type = _get_task_type(task_config)
    _validate_task_config(task_config, task_type)

    interim_path = Path(config.data.interim_data_path)

    models_path = interim_path / "models"
    if task_config.save_refits:
        models_path.mkdir(exist_ok=True)

    output_path = Path(config.data.results_path) / "identify_associations"
    output_path.mkdir(exist_ok=True, parents=True)

    # Load datasets:
    cat_list, cat_names, con_list, con_names = io.load_preprocessed_data(
        interim_path,
        config.data.categorical_names,
        config.data.continuous_names,
    )

    train_dataloader = make_dataloader(
        cat_list,
        con_list,
        shuffle=True,
        batch_size=task_config.batch_size,
        drop_last=True,
    )

    con_shapes = [con.shape[1] for con in con_list]

    num_samples = len(cast(Sized, train_dataloader.sampler))  # N
    num_continuous = sum(con_shapes)  # C
    logger.debug(f"# continuous features: {num_continuous}")

    # Creating the baseline dataloader:
    baseline_dataloader = make_dataloader(
        cat_list, con_list, shuffle=False, batch_size=task_config.batch_size
    )

    # Indentify associations between continuous features:
    logger.info(f"Perturbing dataset: '{task_config.target_dataset}'")
    if task_config.target_value in CONTINUOUS_TARGET_VALUE:
        logger.info(f"Beginning task: identify associations continuous ({task_type})")
        logger.info(f"Perturbation type: {task_config.target_value}")
        output_subpath = Path(output_path) / "perturbation_visualization"
        output_subpath.mkdir(exist_ok=True, parents=True)
        (
            dataloaders,
            nan_mask,
            feature_mask,
        ) = prepare_for_continuous_perturbation(
            config, output_subpath, baseline_dataloader
        )

    # Identify associations between categorical and continuous features:
    else:
        logger.info("Beginning task: identify associations categorical")
        (
            dataloaders,
            nan_mask,
            feature_mask,
        ) = prepare_for_categorical_perturbation(
            config, interim_path, baseline_dataloader, cat_list
        )

    num_perturbed = len(dataloaders) - 1  # P
    logger.debug(f"# perturbed features: {num_perturbed}")

    # APPROACH EVALUATION ##########################

    if task_type == "bayes":
        task_config = cast(IdentifyAssociationsBayesConfig, task_config)
        sig_ids, *extra_cols = _bayes_approach(
            config,
            task_config,
            train_dataloader,
            baseline_dataloader,
            dataloaders,
            models_path,
            num_perturbed,
            num_samples,
            num_continuous,
            nan_mask,
            feature_mask,
        )

        extra_colnames = ["proba", "fdr", "bayes_k"]

    elif task_type == "ttest":
        task_config = cast(IdentifyAssociationsTTestConfig, task_config)
        sig_ids, *extra_cols = _ttest_approach(
            task_config,
            train_dataloader,
            baseline_dataloader,
            dataloaders,
            models_path,
            interim_path,
            num_perturbed,
            num_samples,
            num_continuous,
            nan_mask,
            feature_mask,
        )

        extra_colnames = ["p_value"]

    elif task_type == "ks":
        task_config = cast(IdentifyAssociationsKSConfig, task_config)
        sig_ids, *extra_cols = _ks_approach(
            config,
            task_config,
            train_dataloader,
            baseline_dataloader,
            dataloaders,
            models_path,
            num_perturbed,
            num_samples,
            num_continuous,
            con_names,
            output_path,
        )

        extra_colnames = ["ks_distance"]

    else:
        raise ValueError()

    # RESULTS ################################
    save_results(
        config,
        con_shapes,
        cat_names,
        con_names,
        output_path,
        sig_ids,
        extra_cols,
        extra_colnames,
    )

    if exists(output_path / f"results_sig_assoc_{task_type}.tsv"):
        association_df = pd.read_csv(
            output_path / f"results_sig_assoc_{task_type}.tsv", sep="\t"
        )
        _ = plot_feature_association_graph(association_df, output_path)
        _ = plot_feature_association_graph(association_df, output_path, layout="spring")