import os

import numpy as np
import pandas as pd
from nilearn import datasets

from sbfc.parser import seed_base_connectivity

seed = os.path.dirname(__file__) + "/data/difumo64_pcc.nii.gz"


def _make_data_single_run(confound=True):
    adhd_dataset = datasets.fetch_adhd(n_subjects=2)
    group_confounds = pd.DataFrame(adhd_dataset.phenotypic)[
        ["Subject", "MeanFD", "age", "sex"]
    ]
    group_confounds = group_confounds.rename(columns={"Subject": "subject_label"})
    group_design_matrix = pd.DataFrame(adhd_dataset.phenotypic)[["Subject"]]
    group_design_matrix = group_design_matrix.rename(
        columns={"Subject": "subject_label"}
    )
    group_design_matrix["pheno"] = np.random.rand(2)
    group_contrast = pd.DataFrame([1], columns=["pheno"])
    if confound:
        func_img = {
            f"{sub_id}": {"func": [func], "confound": [confound]}
            for func, confound, sub_id in zip(
                adhd_dataset.func, adhd_dataset.confounds, group_confounds.index
            )
        }
    else:
        func_img = {
            f"{sub_id}": {"func": [func], "confound": [None]}
            for func, confound, sub_id in zip(
                adhd_dataset.func, adhd_dataset.confounds, group_confounds.index
            )
        }
    return func_img, group_design_matrix, group_confounds, group_contrast


def _make_data_multi_run():
    adhd_dataset = datasets.fetch_adhd(n_subjects=2)
    group_confounds = pd.DataFrame(adhd_dataset.phenotypic)[
        ["Subject", "MeanFD", "age", "sex"]
    ]
    group_confounds = group_confounds.rename(columns={"Subject": "subject_label"})
    group_design_matrix = pd.DataFrame(adhd_dataset.phenotypic)[["Subject"]]
    group_design_matrix = group_design_matrix.rename(
        columns={"Subject": "subject_label"}
    )
    group_design_matrix["pheno"] = np.random.rand(2)
    group_contrast = pd.DataFrame([1], columns=["pheno"])

    func_img = {
        f"{sub_id}": {"func": [func, func], "confound": [confound, confound]}
        for func, confound, sub_id in zip(
            adhd_dataset.func, adhd_dataset.confounds, group_confounds.index
        )
    }
    return func_img, group_design_matrix, group_confounds, group_contrast


def test_sbfc_single_run(tmpdir):
    (
        func_img,
        group_design_matrix,
        group_confounds,
        group_contrast,
    ) = _make_data_single_run()
    # Prepare seed
    pcc_coords = (0, -53, 26)
    first_m, first_con, s_m = seed_base_connectivity(
        func_img,
        pcc_coords,
        group_confounds,
        group_design_matrix,
        group_contrast,
        write_dir=tmpdir,
    )
    assert len(first_m) == 2

    (
        func_img,
        group_design_matrix,
        group_confounds,
        group_contrast,
    ) = _make_data_single_run(confound=False)
    # mask seed
    first_m, first_con, s_m = seed_base_connectivity(
        func_img,
        seed,
        group_confounds,
        group_design_matrix,
        group_contrast,
        write_dir=tmpdir,
    )
    assert len(first_m) == 2


def test_sbfc_mutli_run(tmpdir):
    (
        func_img,
        group_design_matrix,
        group_confounds,
        group_contrast,
    ) = _make_data_multi_run()
    # mask seed
    first_m, first_con, s_m = seed_base_connectivity(
        func_img,
        seed,
        group_confounds,
        group_design_matrix,
        group_contrast,
        write_dir=tmpdir,
    )
    assert len(first_m) == 2

[(9, 'os.path.dirname', 'os.path.dirname', False, 'import os\n'), (13, 'nilearn.datasets.fetch_adhd', 'datasets.fetch_adhd', False, 'from nilearn import datasets\n'), (22, 'numpy.random.rand', 'np.random.rand', True, 'import numpy as np\n'), (23, 'pandas.DataFrame', 'pd.DataFrame', True, 'import pandas as pd\n'), (42, 'nilearn.datasets.fetch_adhd', 'datasets.fetch_adhd', False, 'from nilearn import datasets\n'), (51, 'numpy.random.rand', 'np.random.rand', True, 'import numpy as np\n'), (52, 'pandas.DataFrame', 'pd.DataFrame', True, 'import pandas as pd\n'), (72, 'sbfc.parser.seed_base_connectivity', 'seed_base_connectivity', False, 'from sbfc.parser import seed_base_connectivity\n'), (89, 'sbfc.parser.seed_base_connectivity', 'seed_base_connectivity', False, 'from sbfc.parser import seed_base_connectivity\n'), (108, 'sbfc.parser.seed_base_connectivity', 'seed_base_connectivity', False, 'from sbfc.parser import seed_base_connectivity\n'), (14, 'pandas.DataFrame', 'pd.DataFrame', True, 'import pandas as pd\n'), (18, 'pandas.DataFrame', 'pd.DataFrame', True, 'import pandas as pd\n'), (43, 'pandas.DataFrame', 'pd.DataFrame', True, 'import pandas as pd\n'), (47, 'pandas.DataFrame', 'pd.DataFrame', True, 'import pandas as pd\n')]