"""Load the connectomes and transform to be consistent."""


from pathlib import Path
import pandas as pd
from ptpython.ipython import embed
from scipy.spatial.distance import squareform
from nilearn.connectome import vec_to_sym_matrix, sym_matrix_to_vec


def _prepare_xcp_fc(connectome):

    square = squareform(connectome)
    vec = sym_matrix_to_vec(square, discard_diagonal=True)
    return vec


def _prepare_junifer_fc(connectome):

    square = vec_to_sym_matrix(connectome)
    # remove subcortical
    square = square[:-32, :-32].copy()
    # back to vec
    vec = sym_matrix_to_vec(square, discard_diagonal=True)
    return vec


def load_data(dataset, session, atlas_xcp, atlas_junifer, smoothed_or_empirical):
    """Load the connectomes for junifer and xcp output."""
    junifer_path = (
        Path(__file__).parent
        / ".."
        / ".."
        / ".."
        / "junifer_storage"
        / "JUNIFER_AOMIC_TSV_CONNECTOMES"
        / f"{dataset}"
        / (
            f"{dataset}_BOLD_parccortical-{atlas_junifer}_"
            "parcsubcortical-TianxS2x3TxMNInonlinear2009cAsym_"
            f"marker-{smoothed_or_empirical}FC_{session}.tsv"
        )
    )
    xcp_path = (
        Path(__file__).parent
        / ".."
        / ".."
        / ".."
        / ".."
        / "aomic-fc_xcp-engine"
        / "xcpengine_storage"
        / "group_level"
        / f"{dataset}_{session}_{atlas_xcp}.tsv"
    )
    junifer_connectomes = pd.read_csv(junifer_path, index_col=0, sep="\t")
    xcpengine_connectomes = pd.read_csv(xcp_path, index_col=0, sep="\t")
    junifer_connectomes = junifer_connectomes.apply(
        _prepare_junifer_fc, axis=1, result_type="expand"
    ).dropna()
    new_ind = [f"{dataset}_sub-{x:04d}" for x in junifer_connectomes.index]
    junifer_connectomes.index = new_ind

    xcpengine_connectomes = xcpengine_connectomes.apply(
        _prepare_xcp_fc, axis=1, result_type="expand"
    )
    xcpengine_connectomes = xcpengine_connectomes.loc[
        junifer_connectomes.index
    ]
    return junifer_connectomes, xcpengine_connectomes


if __name__ == "__main__":
    junifer_conn, xcp_conn = load_data(
        "PIOP2", "restingstate", "schaefer400x17", "Schaefer400x17FSLMNI"
    )
    embed()
