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, 7, 9, 32), 'os.path.dirname', 'os.path.dirname', ({(9, 23, 9, 31): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((13, 19, 13, 52), 'nilearn.datasets.fetch_adhd', 'datasets.fetch_adhd', (), '', False, 'from nilearn import datasets\n'), ((22, 35, 22, 52), 'numpy.random.rand', 'np.random.rand', ({(22, 50, 22, 51): '2'}, {}), '(2)', True, 'import numpy as np\n'), ((23, 21, 23, 57), 'pandas.DataFrame', 'pd.DataFrame', (), '', True, 'import pandas as pd\n'), ((42, 19, 42, 52), 'nilearn.datasets.fetch_adhd', 'datasets.fetch_adhd', (), '', False, 'from nilearn import datasets\n'), ((51, 35, 51, 52), 'numpy.random.rand', 'np.random.rand', ({(51, 50, 51, 51): '2'}, {}), '(2)', True, 'import numpy as np\n'), ((52, 21, 52, 57), 'pandas.DataFrame', 'pd.DataFrame', (), '', True, 'import pandas as pd\n'), ((72, 30, 79, 5), 'sbfc.parser.seed_base_connectivity', 'seed_base_connectivity', (), '', False, 'from sbfc.parser import seed_base_connectivity\n'), ((89, 30, 96, 5), 'sbfc.parser.seed_base_connectivity', 'seed_base_connectivity', (), '', False, 'from sbfc.parser import seed_base_connectivity\n'), ((108, 30, 115, 5), 'sbfc.parser.seed_base_connectivity', 'seed_base_connectivity', (), '', False, 'from sbfc.parser import seed_base_connectivity\n'), ((14, 22, 14, 59), 'pandas.DataFrame', 'pd.DataFrame', ({(14, 35, 14, 58): 'adhd_dataset.phenotypic'}, {}), '(adhd_dataset.phenotypic)', True, 'import pandas as pd\n'), ((18, 26, 18, 63), 'pandas.DataFrame', 'pd.DataFrame', ({(18, 39, 18, 62): 'adhd_dataset.phenotypic'}, {}), '(adhd_dataset.phenotypic)', True, 'import pandas as pd\n'), ((43, 22, 43, 59), 'pandas.DataFrame', 'pd.DataFrame', ({(43, 35, 43, 58): 'adhd_dataset.phenotypic'}, {}), '(adhd_dataset.phenotypic)', True, 'import pandas as pd\n'), ((47, 26, 47, 63), 'pandas.DataFrame', 'pd.DataFrame', ({(47, 39, 47, 62): 'adhd_dataset.phenotypic'}, {}), '(adhd_dataset.phenotypic)', True, 'import pandas as pd\n')]