"""Plot correlations between subjects obtained in two separate pipelines."""

from itertools import product
from plot_utils import load_data
from pathlib import Path
import matplotlib.pyplot as plt
import seaborn as sns


def plot_correlation_distribution():
    """Plot the correlations between same subjects."""
    datasets = {
        #"ID1000": ["moviewatching"],
        #"PIOP2": ["restingstate", "emomatching", "workingmemory"],
        "PIOP1": ["restingstate", "faces", "emomatching", "workingmemory"],
    }
    atlases = [
        (f"Schaefer{n_roi}x17FSLMNI", n_roi) for n_roi in [100, 200, 400]
    ] + [
        (f"Schaefer{n_roi}x17MNI152NLin2009cAsym", n_roi)
        for n_roi in [400]
    ]

    for suffix in ["empirical", "smoothed"]:
        for dataset, sessions in datasets.items():
            for session, atlas in product(sessions, atlases):
                atlas_name, n_roi = atlas
                outpath = Path("..") / "images" / dataset
                if not outpath.is_dir():
                    outpath.mkdir(exist_ok=True, parents=True)
    
                junifer_conn, xcp_conn = load_data(
                    dataset, session, f"schaefer{n_roi}x17", atlas_name, suffix,
                )
    
                print(junifer_conn)
                print(xcp_conn)
                corr_dist = junifer_conn.corrwith(xcp_conn, axis=1)
                print(corr_dist)
    
                # plot corr distributions
                cm = 1 / 2.54
                figsize = (18 * cm, 10 * cm)
                grid_rows, grid_cols = 1, 1
    
                with plt.style.context("./style.mplstyle"):
                    fig = plt.figure(figsize=figsize)
    
                    grid = fig.add_gridspec(grid_rows, grid_cols)
    
                    ax1 = fig.add_subplot(grid[0, 0])
                    sns.displot(x=corr_dist, ax=ax1)
    
                    outfile = (
                        outpath / f"{atlas_name}_self_correlations_{session}_{suffix}.pdf"
                    )
    
                    plt.savefig(outfile)


if __name__ == "__main__":
    plot_correlation_distribution()
