"""This module provides functions to handle connectome networks / graphs generated by CMP3."""

import os
import warnings
from glob import glob

import networkx as nx
import numpy as np

from cmtklib.bids.io import __cmp_directory__

warnings.simplefilter("ignore")


def load_graphs(output_dir, subjects, parcellation_scheme, weight):
    """Return a dictionary of connectivity matrices (graph adjacency matrices).

    Still in development

    Parameters
    ----------
    output_dir : string
        Output/derivatives directory

    subjects : list
        List of subject

    parcellation_scheme : ['NativeFreesurfer', 'Lausanne2018', 'Custom']
        Parcellation scheme

    weight : ['number_of_fibers','fiber_density',...]
        Edge metric to extract from the graph

    Returns
    -------
    connmats: dict
        Dictionary of connectivity matrices

    """
    if parcellation_scheme == "Lausanne2018":
        bids_atlas_label = "L2018"
    elif parcellation_scheme == "NativeFreesurfer":
        bids_atlas_label = "Desikan"

    if parcellation_scheme == "NativeFreesurfer":
        for subj in subjects:
            subj_dir = os.path.join(output_dir, subj)
            subj_session_dirs = glob(os.path.join(subj_dir, "ses-*"))
            subj_sessions = [
                "ses-{}".format(subj_session_dir.split("-")[-1])
                for subj_session_dir in subj_session_dirs
            ]

            if len(subj_sessions) > 0:  # Session structure
                for subj_session in subj_sessions:
                    conn_derivatives_dir = os.path.join(
                        output_dir, __cmp_directory__, subj, subj_session, "connectivity"
                    )

                    # Extract the connectivity matrix
                    # self.subject+'_label-'+bids_atlas_label+'_desc-scale5_conndata-snetwork_connectivity'
                    connmat_fname = os.path.join(
                        conn_derivatives_dir,
                        "{}_{}_atlas-{}_conndata-snetwork_connectivity.gpickle".format(
                            subj, subj_session, bids_atlas_label
                        ),
                    )
                    connmat_gp = nx.read_gpickle(connmat_fname)
                    connmat = nx.to_numpy_matrix(
                        connmat_gp, weight=weight, dtype=np.float32
                    )
    else:
        # For each parcellation scale
        for scale in np.arange(1, 6):
            for subj in subjects:
                subj_dir = os.path.join(output_dir, subj)
                subj_session_dirs = glob(os.path.join(subj_dir, "ses-*"))
                subj_sessions = [
                    "ses-{}".format(subj_session_dir.split("-")[-1])
                    for subj_session_dir in subj_session_dirs
                ]

                if len(subj_sessions) > 0:  # Session structure
                    for subj_session in subj_sessions:
                        conn_derivatives_dir = os.path.join(
                            output_dir, __cmp_directory__, subj, subj_session, "connectivity"
                        )

                        # Extract the connectivity matrix
                        # self.subject+'_label-'+bids_atlas_label+'_desc-scale5_conndata-snetwork_connectivity'
                        connmat_fname = os.path.join(
                            conn_derivatives_dir,
                            "{}_{}_atlas-{}_res-scale{}_conndata-snetwork_connectivity.gpickle".format(
                                subj, subj_session, bids_atlas_label, scale
                            ),
                        )
                        connmat_gp = nx.read_gpickle(connmat_fname)
                        connmat = nx.to_numpy_matrix(
                            connmat_gp, weight=weight, dtype=np.float32
                        )
                # TODO: finalize condition and append all conmat to a list
    return connmat

