from importer import *

sN = slice(None)
now = datetime.now()
# results_uid = now.strftime("%Y%m%d-%H%M%S")
results_uid = '20240807-135247-postRebuttal2-Night2'
results_dir = get_root(f'results', results_uid)
figdir = results_dir.append('figures')
results_file = results_dir.append('result_dict.pkl')
restats_file = results_dir.append('result_and_stats_dict.pkl')
dl_aggregate_file = results_dir.append('aggregate_dict.pkl')  # Stores models
os.makedirs(figdir, exist_ok=True)
print('Results stored at:', results_dir)

path_roi = get_root('data', 'ENIGMA-MDD_patients-included_data-subcortical.csv')
path_dbc = get_root('data', 'ENIGMA-MDD_patients-included_data-clinical_wk2.csv')
data = load_and_join_dfs(path_dbc, path_roi, index_col='SubjID', decimal=',')

path_dbc_wk1 = get_root('data', 'ENIGMA-MDD_patients-included_data-clinical_wk1.csv')
wk1_subjs = read_any_csv(path_dbc_wk1, index_col='SubjID').index
print(f'Number of patients excluded if scan threshold is week 1 instead of week 2: {len([i for i in data.index if i not in wk1_subjs])}')

# Remove edges and extra channel, and make it into a nice DataFrame with x,y pixel coordinates as column labels
dim = 32
X_thc, X_thc_arr = load_proj_df(data, 'Thickness', dim=dim)
X_are, X_are_arr = load_proj_df(data, 'Area', dim=dim)

# Find zero variance (corpus callosum (CC))
bool_var = [var < 0.005 for _, var in X_thc.var().iteritems()]
zero_var = [col for col, var in X_thc.var().iteritems() if var < 0.002]

# Drop CC
X_thc = X_thc.drop(columns=zero_var)
X_are = X_are.drop(columns=zero_var)

# Visualization of loaded data
fig, axes = plt.subplots(1, 3)
img_arrays = X_thc_arr.mean(axis=0)[0], X_are_arr.mean(axis=0)[0], np.array(bool_var).reshape(dim - 2, dim - 2)
titles = 'Mean Thickness', 'Mean Area', 'Zero Variance (white)'
cmaps = 'jet', 'jet', 'gray'
for ax, arr, title, cmap in zip(axes, img_arrays, titles, cmaps):
    ax.imshow(arr, cmap=cmap)
    ax.set(title=title, xticks=[], yticks=[])

scoring_methods = 'MADRS', 'HDRS', 'BDI'
data, norm_fig = clean_data_df(data, drop_old_symptom_scores=True)
norm_fig.savefig(figdir.append('NormalizedSeverityScores.png'))

score_df = pd.DataFrame(columns=scoring_methods)
for cohort in cohorts:
    cohort_data = pd.DataFrame(data=data[data.cohort == cohort.key], columns=data.columns)
    if len(cohort_data):
        contains = lambda scoring_method: cohort_data[f'{scoring_method}_pre'].notna().sum()
        score_df.loc[f'{cohort.name} (N={len(cohort_data)})'] = [
            contains(score) if contains(score) else '' for score in scoring_methods]
print(score_df)

# Split up the treatment information into separate boolean columns
drug_types = set(flatten(data['Treatment_type_category'].values))
for drug_type in drug_types:
    data[f'uses_{drug_type}'] = [drug_type in i for i in data['Treatment_type_category']]

q1, q3 = np.percentile(data['Response_percentage'], [25, 75])
data.insert(loc=0, column='is_extreme',
            value=[(data['Response_percentage'] < q1) | (data['Response_percentage'] > q3)][0])
same_responders = 'AFFDIS', 'Hiroshima', 'Melb'
data.insert(loc=0, column='is_same_responder', value=[c in same_responders for c in data.cohort])

long_treated = 'MOODS', 'Melb', 'Minnesota'
q1, q3 = np.percentile(data['Response_percentage'], [25, 75])
populations = {
    'All': data.index,
    'Hiroshima': data[data.cohort == 'Hiroshima'].index,
    'SameResponders': pd.Index([i for i in data.index if i.split('_')[0] in same_responders]),
    'Extremes': data[data.is_extreme].index,
    'LongTreated': pd.Index([i for i in data.index if i.split('_')[0] in long_treated]),
}

# Plot histograms of response percentage for each subpopulation
resp_fig = subpop_hists(data, populations)
resp_fig.savefig(figdir.append('ResponseHistogram.png'))
plt.show()

duration_based_cohort_printer = lambda in_only: '\n  '.join([''.join([
    c.name.ljust(20),
    f'n={data.cohort.value_counts()[c.key]}'.ljust(7),
    f'{data[data.cohort == c.key].is_responder.mean():.1%}  ',
    f'{str(c.treatment_duration_mu).rjust(4)}+{c.treatment_duration_sd}'
]) for c in cohorts if c.key in set(data.cohort) and (c.key in long_treated if in_only else c.key not in long_treated)])

print('Long treated:\n ', duration_based_cohort_printer(True), '\nShort treated:\n ',
      duration_based_cohort_printer(False))

expl = {
    'All': '# samples that this cohort contributed to the entire population and the fraction of the entire population that comes from this sample.',
    'All_Responders': '# Responders that this cohort contributed to the entire population.',
    'F_of_All_Responders': "% all responders that comes from this cohort",
    'Extremes': '# samples that this cohort contributed to the Extremes subpopulation and the fraction of the extremes population that comes from this sample.',
    'Extreme_Responders': '# Responders that this cohort contributed to the Extremes subpopulation.',
    'F_of_Extreme_Responders': "% Extreme responders that comes from this cohort",
}
extremes_contribution_df = pd.DataFrame(columns=list(expl.keys()))
extremes_contribution_df.loc['Total'] = [
    f'{len(data)} (100%)',
    f'{data.is_responder.sum()} ({data.is_responder.mean():.1%})',
    '100%',
    f'{len(data[data.is_extreme])} ({len(data[data.is_extreme]) / len(data):.1%})',
    f'{data[data.is_extreme].is_responder.sum()} ({data[data.is_extreme].is_responder.mean():.1%})',
    '100%',
]

for cohort in [c for c in cohorts if c.key in set(data.cohort)]:
    cohort_data = data[data.cohort == cohort.key]
    extreme_cohort_data = cohort_data[cohort_data.is_extreme]
    values = [
        f'{len(cohort_data)}' + f'({len(cohort_data) / len(data):.0%})'.rjust(7),
        bool_fmt(cohort_data.is_responder),
        f'{cohort_data.is_responder.sum() / data.is_responder.sum():.1%}',
        f'{len(extreme_cohort_data)}' + f'({len(extreme_cohort_data) / data.is_extreme.sum():.0%})'.rjust(6),
        bool_fmt(extreme_cohort_data.is_responder),
        f'{cohort_data.is_responder.sum() / data.is_responder.sum():.1%}',
    ]
    extremes_contribution_df.loc[cohort.name] = values
print(extremes_contribution_df)
print('\n'.join(['', 'Explained:'] + [f'{k}:'.ljust(np.max([len(i) for i in expl]) + 4) + v for k, v in expl.items()]))

# Load subjects in Stage 1
path_stg1 = get_root('data', 'ENIGMA-MDD_patients-all_data-clinical.csv')
df_stg1 = read_any_csv(path_stg1, index_col='SubjID').iloc[:, 1:]
# Load subjects in Stage 2
df_stg2 = read_any_csv(path_dbc, index_col='SubjID')
# We already have Stage 3
df_stg3 = data
# For Stage 4 we need to edit our projections data frame a bit
df_stg4 = X_thc.loc[[i for i in X_thc.index if 'LH' in i]]
df_stg4.index = df_stg4.index.map(lambda tag: tag.split('|')[0])

# Get the sites, and the data that are in each of the three steps
dfs = df_stg1, df_stg2, data, df_stg4,
sts = [sorted(set([i.split('_')[0] for i in df.index])) for df in dfs]

# Print a verbose overview of the subject selection process
print(f'Prior to subject selection (stage 1) we have {len(dfs[0])} subjects from {len(sts[0])} cohorts.')
print(f'After subject selection (stage 2) we have {len(dfs[1])} subjects from {len(sts[1])} cohorts.')
print(
    f'After checking for completeness of data for ROI analysis (stage 3) we still have {len(dfs[2])} subjects from {len(sts[2])} cohorts.')
print(f'For Vec and 2D analysis, we have {len(dfs[3])} subjects from {len(sts[3])} cohorts')

# Print a table header
n_dfs = [
    {'Total': len(df), '--': 0, **{f'{c.name} ({c.key})': len([i for i in df.index if c.key in i]) for c in cohorts}}
    for df in dfs]
n_all, n_cov, n_dat, _ = n_dfs

# We can split this up further
n_series = [pd.Series(data=n_df, name=f'Stage {i}') for i, n_df in enumerate(n_dfs, 1)]
series_list = flatten([[ser, None] for ser in n_series])[:-1]

# Set the even numbered columns
for i, (df_before, df_after) in enumerate(zip(dfs[:-1], dfs[1:])):
    series_list[1 + i * 2] = pd.Series(data={
        'Total': len(df_before) - len(df_after),
        '--': 0,
        **{f'{cohorts[k].name} ({k})': len(
            [i for i in df_before.index if i not in df_after.index and i.split('_')[0] == k]) for k in cohorts.keys()}
    }, dtype=float, name=f'--excl-->')
print(pd.concat(series_list, axis=1).fillna(0).astype(int).replace(0, '').head(12))

# Uncomment this line if you would like to see which subjects do NOT have Cortical Data
# print("Subjects without Raw Cortigal data are:\n  -", "\n  - ".join([i for i in df_stg3.index if i not in df_stg4.index]))

# Load structured data
dkt_atlas_features = [i for i in data.columns if any([j in i for j in dkt_atlas_lut])]
subcortical_rois = [i for i in pd.read_csv(path_roi, sep=';', index_col='SubjID').columns if
                    i[1] != '_' and not 'Site' in i][:-2]
subcortical_rois = [i for i in subcortical_rois if 'Thick' not in i and 'Surf' not in i]

print(
    f'The number of cortical features we found is {len(dkt_atlas_features)} ({32 * 2 * 3} = 32 ROIs × 2 maps (Thickness & Sufrace) × 3 hemispheres (Left, Right, Mean). That means there are {len(data.columns) - len(dkt_atlas_features)} non-surface columns.')
subc_rois = subcortical_rois
subc_no_icv = subc_rois[:-1]
print(f'\nWe also have ICV and {len(subc_rois) - 1} subcortical_rois:', ', '.join(subc_rois[:-1]))
unique_subc_rois = sorted(set([i[1:].capitalize() if i[0] in 'LRM' else i for i in subc_rois]))
print(f'Of these {len(unique_subc_rois)} are unique:', ', '.join(unique_subc_rois),
      'Vent, LatVent and ICV are not included, so 7 remain.')
# data[dkt_atlas_features].head()  # You can take a look at the DKT Atlas features by uncommenting this line
data[subc_no_icv] = data[subc_no_icv].astype(float).div(data['ICV'], axis=0)

site_stats, site_means = collect_stats_per_site(data)
print(site_means.T)

responder_stats, responder_means = collect_stats_per_site(data.loc[data.is_responder])
non_responder_stats, non_responder_means = collect_stats_per_site(data.loc[~data.is_responder])
non_vs_responders = pd.concat([site_means.T.Total, responder_means.T.Total, non_responder_means.T.Total], axis=1)
non_vs_responders.columns = ['Total', 'Responders', 'Non-responders']
print(non_vs_responders)

covariates = 'Age', 'Treatment_duration', 'is_responder'
variables = ('is_responder',), ('is_responder',), ('uses_SSRI', 'uses_SNRI', 'uses_ATYP')
response_correlates = pd.DataFrame(columns=['Population', 'Property', '+', '-', 'n+', 'n-'])
for covariate, variable in zip(covariates, variables):
    for var in variable:
        group_a = data[data[var]]
        group_b = data[~data[var]]
        num_a = group_a[covariate].mean()
        num_b = group_b[covariate].mean()
        fmt = '{:.1%}' if num_a < 1 else '{:.1f}'
        response_correlates.loc[len(response_correlates)] = var, covariate, fmt.format(num_a), fmt.format(num_b), len(
            group_a), len(group_b)
print(response_correlates)

adolescent_df = pd.DataFrame(columns=('N', '%'))
for cohort in [c for c in cohorts if c.name in site_stats.index]:
    adolescents = np.sum(site_stats.loc[cohort.name].Age < 20)
    total = site_stats.loc[cohort.name].Numel
    table_head = f'{cohort.name} (N={total})'
    adolescent_df.loc[table_head] = (adolescents, f'{adolescents / total * 100:.0f}')
print(adolescent_df)

outcomes = 'is_female', 'is_responder', 'is_remitter', 'is_extreme'
outcomes_to_plot = outcomes[1], outcomes[3]
n_otp = len(outcomes_to_plot)
pie_fig, axes = plt.subplots(n_otp, len(set(data.cohort)), figsize=(12, 2 * n_otp))
#axes = axes.reshape(-1)
for u, outcome in enumerate(outcomes_to_plot):  # 'is_remitter']):
    v = 0
    for cohort in cohorts:
        df_cohort = data[data.cohort == cohort.key]
        if not any(df_cohort.index):
            continue
        axes[u, v].pie(
            x=[df_cohort[outcome].sum(), len(df_cohort) - df_cohort[outcome].sum()],
            colors=(cohort.color, colors.dark),
            autopct='%.0f%%',
        )
        if not u:
            axes[u, v].set_title(f'{cohort.name}\n(n={len(df_cohort)})')
        if not v:
            axes[u, v].set_ylabel(outcome)
        v += 1
plt.tight_layout()
pie_fig.savefig(figdir.append('ResponseRateBySite.png'))
plt.show()

trend_fig, axes = plt.subplots(1, 3, figsize=(10, 4))
for ax, score in zip(axes.reshape(-1), scoring_methods):
    for cohort in cohorts:
        # Get the values for a specific cohort and a specific symptom scoring method
        cohort_data = pd.DataFrame(data=data[data.cohort == cohort.key], columns=data.columns)
        pre = cohort_data[score + '_pre'].notna()
        post = cohort_data[score + '_post'].notna()
        subjs = cohort_data[(pre & post)]
        if not len(subjs):
            continue
        vals = subjs[[score + '_pre', score + '_post']].to_numpy().astype(int).T
        ax.plot(vals, color=cohort.color, lw=1)  # Plot the found individual lines
        ax.plot(vals.mean(axis=1), color=colors.dark, lw=5.0, zorder=3)  # And plot the mean in bold
    ax.set(xticks=(), xlim=(0, 1), title=f'{score} (n={data[score + "_pre"].notna().sum()})')
trend_fig.suptitle('Severity Scores from Time Point 1 to 2 for Each Metric')
# Format the legend
line_per_cohort = [Line2D([], [], color=c.color, label=c.name, lw=3) for c in cohorts if c.name in site_stats.index]
axes[1].legend(handles=line_per_cohort, ncol=len(site_stats.index), loc='center', bbox_to_anchor=(0.5, -.1))
trend_fig.savefig(figdir.append('ResponseTrends.png'))

hist_fig, axes = plt.subplots(2, 2, figsize=(16, 8))
stacked_hist(data, site_stats, ax=axes[0, 0], make_legend=False, name='Age')
stacked_hist(data, site_stats, ax=axes[0, 1], make_legend=False, name='Age_of_Onset')
stacked_hist(data, site_stats, ax=axes[1, 0], make_legend=False, name='Normalized_Pretreatment_Severity')
stacked_hist(data, site_stats, ax=axes[1, 1], make_legend=True, name='Response_percentage')
axes[1, 1].set_xticklabels([f'{t:.0%}' for t in axes[1, 1].get_xticks()])
hist_fig.savefig(figdir.append('ResponseHistogram.png'))
plt.show()

_, ss_means = collect_stats_per_site(data.loc[populations['Hiroshima']])
print(ss_means.T)

_, rrs_means = collect_stats_per_site(data.loc[populations['SameResponders']])
print(rrs_means.T)

extr_stats, extr_means = collect_stats_per_site(data.loc[populations['Extremes']])
print(extr_means.T)

extremes = data.loc[populations['Extremes']]
extr_lo_stats, extr_lo_means = collect_stats_per_site(extremes[~extremes.is_responder])
extr_hi_stats, extr_hi_means = collect_stats_per_site(extremes[extremes.is_responder])
extr_non_vs_responders = pd.concat([extr_means.T.Total, extr_lo_means.T.Total, extr_hi_means.T.Total], axis=1)
extr_non_vs_responders.columns = ['Extremes Total', 'Extreme Non-responders', 'Extreme Responders']
print(extr_non_vs_responders)

print('Mean ± 1 SD treatment duration for the following populations:')
relevant_cohorts = [c for c in cohorts if c.name in site_stats.index]
for stats_desc, stats in zip(
        ['All', 'Responders', 'Non-responders', 'Extremes', 'Responding Extremes', 'Non-responding Extremes'],
        [site_stats, responder_stats, non_responder_stats, extr_stats, extr_hi_stats, extr_lo_stats]):
    print(stats_desc.ljust(24),
          f'{np.sum([stats.loc[c.name].Numel * c.treatment_duration_mu for c in relevant_cohorts if c.name in stats.index]) / stats.Numel.sum():.2f} ± '
          f'{np.sum([stats.loc[c.name].Numel * c.treatment_duration_sd for c in relevant_cohorts if c.name in stats.index]) / stats.Numel.sum():.2f}')

print('\n'.join([f'Extremes in {c.name} with ADcur == 0:'.ljust(50) + str(
    (data[(data.cohort == c.key) & data.is_extreme].ADcur == 0).sum()) for c in relevant_cohorts]))

# Preallocate dictionary for storage
if os.path.isfile(results_file):
    results_dict = pickle_in(results_file)
    print(f'Loaded results file: {results_file}.\n The file contained {len(nested_dict_to_df(results_dict))} results')
    dl_aggregates = pickle_in(dl_aggregate_file)
    print(
        f'Loaded aggregates file: {results_file}.\n The file contained {len(nested_dict_to_df(dl_aggregates))} trained Deep Learning models')
else:
    results_dict, dl_aggregates = {}, {}
    print(f'Created new results file: {results_file}')

# Some settings to our machine learning method:
target_labels = 'is_responder',  # 'is_remitter', 'is_female', 'is_extreme',
n_splits = 10
n_repeats = 1
n_permutations = 100
n_iter = 25
n_epochs = 20

# 1. Define data type for full population we need :X = train data, y = label data, c = stratification data
get_y_for_target = lambda tl: pd.Series(data=[data[tl].loc[sub.split('|')[0]] for sub in X_thc.index],
                                        index=X_thc.index, name=tl)
dtypes = {'vec': (X_thc,
                  pd.DataFrame.from_dict(
                      data={target_label: get_y_for_target(target_label) for target_label in target_labels}),
                  pd.Series(data={sub: data.loc[sub.split('|')[0]].cohort_idx for sub in X_thc.index},
                            index=X_thc.index, name='Stratification Label'),),
          'roi': (data[dkt_atlas_features],
                  pd.DataFrame.from_dict(data={target_label: data[target_label] for target_label in target_labels}),
                  data.cohort_idx,),
          }

# 2.1 Presence of Demographic, behavioural and clinical (DBC) data:
dbc_cols = ['is_ad', 'is_recurrent', 'is_female', 'Age', 'Normalized_Pretreatment_Severity', 'Age_of_Onset']
dbc_prj = data[dbc_cols].loc[[s.split('|')[0] for s in X_thc.index]]
sub_prj = data[dbc_cols].loc[[s.split('|')[0] for s in X_thc.index]]
dbc_prj.insert(loc=0, value=X_thc.index, column='SubjID|Hem')
sub_prj.insert(loc=0, value=X_thc.index, column='SubjID|Hem')
dbc_prj = dbc_prj.set_index('SubjID|Hem')
dbc_options = {
    'dbc_no': {'roi': pd.DataFrame(), 'vec': pd.DataFrame()},  # add nothing to cortical data
    'dbc_yes': {'roi': data[dbc_cols], 'vec': dbc_prj},  # add DBC data to cortical data
    'sub_add': {'roi': data[subc_rois], 'vec': None},  # add subcortical data to cortical data
    'sub_nly': {'roi': pd.DataFrame(), 'vec': None},  # add nothing to subcortical data
    'sub_dbc': {'roi': data[dbc_cols], 'vec': None},  # add DBC data to subcortical data
    'sub_all': {'roi': data[subc_rois + dbc_cols], 'vec': None},  # add subcortical and DBC data to cortical data
    'dbc_nly': {'roi': pd.DataFrame(), 'vec': None},  # only use DBC
}

# 2.2 Get indices each population, both for our roi data frame (which we already defined at 2.3) but also for projection data
get_indices = lambda key, indices: {'roi': indices, 'vec': [i for i in X_thc.index if i.split('|')[0] in indices]}
population_indices = {k: get_indices(k, i) for k, i in populations.items()}

# 3. Define the data for each of our previously defined populations
data_dict = make_empty_nest(proj_type=dtypes, dbc_presence=dbc_options, population=populations)
for dtype_name, dtype_values in dtypes.items():
    X, y, c = dtype_values
    for dbc_option, dbc_values in dbc_options.items():
        # These conditionals are a bit hacky because we replace the cortical
        # ROIs with something else (subc or dbc data), but it save a ton of code
        if dbc_option in ('sub_nly', 'sub_dbc',):
            X = data[subc_rois]
        elif dbc_option in ('dbc_nly',):
            X = data[dbc_cols]

        if dbc_values[dtype_name] is None:
            del data_dict[dtype_name][dbc_option]
            continue
        X_cort_and_dbc = pd.concat([dbc_values[dtype_name], X], axis=1)
        for pop_key, pop_idx in population_indices.items():
            get_subpop_subjs = lambda sub_df: sub_df.loc[pop_idx[dtype_name]]
            data_dict[dtype_name][dbc_option][pop_key] = [get_subpop_subjs(sub_df) for sub_df in (X_cort_and_dbc, y, c)]

# 4. Define our cross-validation methods:
cv_schemes = {'Fold': (RepeatedStratifiedKFold(n_splits=n_splits, n_repeats=n_repeats, random_state=0),
                       CombatWrapper(data=data, discrete_covariates=['is_female'],
                                     continuous_covariates=['Age', 'Age2']),),
              'Site': (LeaveOneGroupOut(), None,)}

# 5. Define classifiers and hyperparameter search space to use:
from hyper_space import svc_space, gbc_space, log_space, rfc_space
scorer = Scorer({'accuracy': accuracy_score, 'balanced_accuracy': balanced_accuracy_score})
classifiers = \
    (SVC(), svc_space), \
    (GradientBoostingClassifier(), gbc_space), \
    (LogisticRegression(max_iter=500), log_space), \
    (RandomForestClassifier(), rfc_space)

# Define other pipeline components
imputer = PipeWrapper(KNNImputer)
regressor = RegressorWrapper(data=data, continuous_covariates=['Age', 'Age2'])
selector = SelectFromModel(LinearSVC(penalty="l1", dual=False, max_iter=20000), threshold=0)
scaler = PipeWrapper(StandardScaler)

# Define Deep Learning model
torch_scorer = {'accuracy': torch_accuracy, 'balanced_accuracy': torch_balanced_accuracy, }
resnet_model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
resnet_model.fc = nn.Sequential(nn.Linear(in_features=512, out_features=2, bias=True), Squeeze())
model = TorchTrainer(
    subestimator=resnet_model,
    batch_size=32,
    epochs=n_epochs,
    verbose=False,
    random_state=0,
    metrics=torch_scorer
)

# Sensitivity testing
subj_classes = {'wk_1': wk1_subjs, 'wk_2': data.index}

# Print progress
n_tasks = count_values(data_dict) * len(cv_schemes) * len(classifiers) * len(target_labels)
# n_tasks = 1512
n_tasks = 412
pb = ProgressBar(n_tasks, desc='models trained')

# For explanation on these six nested loops please see the list at 3.
with Timer():
    for subj_class, subj_idxs in subj_classes.items():
        for target_label in target_labels:
            for dtype, data_a in data_dict.items():
                for dbc_presence, data_b in data_a.items():
                    for population_name, (X, ys, c) in data_b.items():
                        subj_incl = [i for i in X.index if i.split('|')[0] in subj_idxs]
                        # Define True label
                        X_safe = X.loc[subj_incl].drop(columns=[target_label], errors='ignore')
                        y = ys.loc[subj_incl][target_label]
                        c = c.loc[subj_incl]
                        for cv_name, (cv, harmonizer) in cv_schemes.items():
                            n_groups = (1, len(set(c))) if cv_name == 'Site' else (n_repeats, n_splits)
                            for classifier, param_space in classifiers:
                                if population_name == 'Extremes' and target_label == 'is_extreme' or \
                                        len(c.unique()) == 1 and cv_name == 'Site':
                                    # Can't predict on own population or do LSO-CV on one site
                                    continue

                                if subj_class == 'wk_1' and (
                                        dbc_presence != 'dbc_no' or
                                        population_name != 'All' or
                                        cv_name != 'Fold' or
                                        target_label != 'is_responder'):
                                    # We only want to run the sensitivity analysis of the weeks of cutoff
                                    # for runs that are relevant to the primary analysis
                                    continue

                                # 1) Classical Machine Learning: Specify which Classical ML experiment we are going to run
                                experiment_specifier = dtype, dbc_presence, population_name, cv_name, classifier.__class__.__name__, target_label, subj_class, 'done'
                                previous_result = safe_dict_get(results_dict, *experiment_specifier)

                                if previous_result['done'] is None:  # Skip if results exist
                                    pipeline = make_pipeline(imputer, regressor, harmonizer, scaler, imputer, selector, classifier, )
                                    opt = BayesSearchCV(estimator=pipeline, search_spaces=param_space, n_iter=n_iter, cv=5, n_jobs=-1)
                                    try:
                                        _, _, pvalue = permutation_test_score(opt, X_safe, y, groups=c, cv=cv, n_permutations=n_permutations, scoring=scorer)
                                        previous_result['pvalue'] = pvalue
                                        for score_name, score_value in scorer.items():
                                            score_tensor = np.reshape(score_value, [1 + n_permutations, *n_groups])
                                            previous_result[score_name] = score_tensor[0]
                                            previous_result[f'null_{score_name}'] = score_tensor[1:]
                                    finally:
                                        scorer.reset()  # Clear the scorer when we are done
                                    # Save results
                                    previous_result['done'] = True
                                    pickle_out(results_dict, results_file)
                                pb()

                                
                                if dbc_presence == 'dbc_no' and dtype == 'vec':
                                    for X_arr, letter in zip((X_thc_arr, X_are_arr), 'TA'):
                                        # 2) Deep Learning: Specify which Deep Learning experiment we are going to run
                                        dl_specifier = f'2D{letter}', dbc_presence, population_name, cv_name, 'ResNet', target_label, subj_class, 'done'
                                        previous_dl_result = safe_dict_get(results_dict, *dl_specifier)
                                        if previous_dl_result['done'] is None:  # Skip if results exist
                                            pop_idx = [idx for idx, subj_id in enumerate(X_thc.index) if
                                                       subj_id in y.index]

                                            # Get CV-scores, score and a priori chance
                                            previous_dl_result.update({''.join(null_metric): [] for null_metric in
                                                                       product(('null_', ''), torch_scorer)})
                                            for do_permute in range(1 + n_permutations):
                                                tag = 'null_' if do_permute else ''
                                                torch_results = torch_val_score(model, X_arr[pop_idx], y, cv,
                                                                                groups=c, verbose=False,
                                                                                return_pipeline='aggregate',
                                                                                do_permute=bool(do_permute))
                                                for metric in torch_scorer:
                                                    dl_score_matrix = np.reshape(
                                                        [r[metric] for r in torch_results['test']], n_groups)
                                                    previous_dl_result[tag + metric].append(dl_score_matrix)
                                            # Convert the list of dl_permutations into a properly formatted array
                                            for metric, tag in product(torch_scorer, ('', 'null_')):
                                                permutes = n_permutations if tag else 1
                                                sc = previous_dl_result[tag + metric]
                                                sc = np.reshape(sc, (permutes, *sc[0].shape))
                                                previous_dl_result[tag + metric] = sc

                                            # Computation of permutation p-value, copied from scipy (see permutation_test)
                                            dl_pvalue = ((previous_dl_result[f'null_accuracy'].mean(2).mean(1) >=
                                                          previous_dl_result[f'accuracy'].mean(2).mean(
                                                              1)).sum() + 1.0) / (n_permutations + 1)
                                            previous_dl_result['pvalue'] = dl_pvalue

                                            dl_aggregate = safe_dict_get(dl_aggregates, *dl_specifier[:-1])
                                            dl_aggregate[dl_specifier[-2]] = torch_results['pipeline']
                                            # Save results
                                            previous_dl_result['done'] = True
                                            pickle_out(results_dict, results_file)
                                            pickle_out(dl_aggregates, dl_aggregate_file)
                                        pb()

results_table = nested_dict_to_df(results_dict)
# The table is ordered by the order of the nesting of the calculation.
# Reorder the index to be consistent with the order described in the manuscript.
results_table = results_table.reorder_levels([5, 0, 1, 4, 3, 2, 6])

# Specify which DL result to look at as an deep learning configuration:
a_dl_res = dl_aggregates['2DA']['dbc_no']['Extremes']['Fold']['ResNet']['is_responder'][
    'wk_2']  # A Deep Learning Result
# Show it
lw = 4
dl_fig, axes = plt.subplots(3, 2, figsize=(9, 6))
line_colors = colors.blue, colors.orange,
for u, (partition_name, partition) in enumerate(a_dl_res.items()):
    for v, (metric_name, metric_values) in enumerate(partition.items()):
        # Plot all learning curves
        axes[v, u].plot(metric_values.T, color='#aaa')
        # Plot the mean learning curve
        axes[v, u].plot(metric_values.mean(0), color=line_colors[u], lw=lw)
        axes[v, u].set(
            title=f'{partition_name.capitalize()}' if u else '',
            ylabel=metric_name if not u else None,
            xlabel='Batch' if v else None,
            yscale='log' if 'Log' in metric_name else 'linear',
        )
        # Only draw a legend once
        if not u and v == 2:
            axes[v, u].legend(frameon=False, handles=[
                Line2D([], [], color=line_colors[0], label='Mean loss', linewidth=lw, ),
                Line2D([], [], color=line_colors[1], label='Mean accuracy', linewidth=lw, ),
                Line2D([], [], color='#aaaaaa', label='Individual folds')],
                              )
dl_fig.suptitle('Deep Learning Training Progression')
dl_fig.tight_layout()
dl_fig.savefig(figdir.append(f'dl_training.png'))
dl_fig.show()

results_table_with_alternate_classifiers = deepcopy(results_table)
results_table = results_table[~(results_table.index.get_level_values(3) == 'LogisticRegression')]
results_table = results_table[~(results_table.index.get_level_values(3) == 'RandomForestClassifier')]

statistics_calculator = calc_stats_wrapper(population_indices)
stats_table = statistics_calculator(results_table)

# Horizontal concat the statistics table to the results table
results_and_stats = pd.concat((results_table, stats_table), axis=1)
pickle_out(results_and_stats, restats_file)

# Split our results for primary and exploratory analyses
cortical_results = results_and_stats.loc[(sN, sN, ['dbc_no', 'dbc_yes']), :]  # Limit primary analysis to cortical data

distiller = TableDistiller(cortical_results, 'is_responder', None, 'dbc_no', None, 'Fold', 'All', 'wk_2',
                           is_strict=True, verbose=True, n_splits=n_splits)
distiller(1)

distiller.is_strict = False
distiller(2, 3, 4)

distiller(5)

added_data = results_and_stats.index.get_level_values(2)
subc_compare = results_and_stats[np.logical_or(added_data == 'dbc_no', added_data == 'sub_add')]
distiller = TableDistiller(subc_compare, 'is_responder', None, 'dbc_no', None, 'Fold', None, 'wk_2', n_splits=n_splits)
distiller(6)

distiller = TableDistiller(subc_compare, 'is_responder', None, None, None, 'Fold', 'All', 'wk_2', n_splits=n_splits)
distiller(3)

distiller = TableDistiller(cortical_results, 'is_responder', None, 'dbc_no', None, 'Fold', 'Extremes', 'wk_2',
                           is_strict=True, spacing=15, n_splits=n_splits)
distiller(1)
distiller.is_strict = False
distiller(2, 3, 4, 5)

distiller = TableDistiller(results_and_stats, 'is_responder', None, None, None, 'Fold', 'All', 'wk_2')
distiller(3)

distiller = TableDistiller(results_and_stats, 'is_responder', None, None, None, 'Fold', 'Extremes', 'wk_2')
distiller(3)

distiller = TableDistiller(cortical_results, 'is_responder', None, 'dbc_no', None, 'Fold', 'LongTreated', 'wk_2',
                           is_strict=True, verbose=False, spacing=15, n_splits=n_splits)
distiller(1)
distiller.is_strict = False
distiller(2, 3, 4, 5, 6)

all_clf_distiller = TableDistiller(pd.concat(
    (results_table_with_alternate_classifiers, statistics_calculator(results_table_with_alternate_classifiers)),
    axis=1), 'is_responder', None, 'dbc_no', 'RandomForestClassifier', 'Fold', 'All', 'wk_2', is_strict=True,
    n_splits=n_splits)
all_clf_distiller(1)
all_clf_distiller.is_strict = False
all_clf_distiller(4)

distiller = TableDistiller(cortical_results, 'is_responder', None, 'dbc_no', None, 'Fold', 'All', 'wk_1', is_strict=False,
                           n_splits=n_splits)
distiller(7)

# Human-readable Title Dictionaries
hr_data = OrderedDict({
    'dbc_no': 'cortical',
    'dbc_nly': 'clinical',
    'sub_nly': 'sub-cortical',
    'dbc_yes': 'cortical and clinical',
    'sub_dbc': 'sub-cortical and clinical',
    'sub_add': '(sub-)cortical',
    'sub_all': '(sub-)cortical and clinical',
})
hr_pop = {
    'All': 'Full\nPopulation',
    'Extremes': 'Extremes\nSubpopulation',
}

bins = np.linspace(-80, 80, 17)

queries = flatten(
    [[('is_responder', 'roi', d, 'GradientBoostingClassifier', 'Fold', p, 'wk_2',) for d in hr_data] for p in hr_pop])

n_rows = 2
n_cols = len(queries) // n_rows
xlim = 80
ylim = 0.05  # max(ax.get_ylim()) # example.population.mean() #

cs = [colors.blue] * n_cols + [colors.orange] * n_cols
bacc_fig, axes = plt.subplots(n_rows, n_cols, figsize=(14, 6))

for ax_n, query, ax, c in zip(range(n_rows * n_cols), queries, axes.reshape(-1), cs):
    # Get data
    example = results_and_stats.loc[query]
    score, null = example['balanced_accuracy'], example['null_balanced_accuracy']
    score, null = score.ravel(), null.ravel()

    # Plot
    ax.hist((score.mean() - null) * 100, bins, density=True, color=c, label=hr_pop[query[-2]])
    ax.plot([0, 0], [0, ylim], c='k')
    acc_explainer = '' if ax_n else '(bAcc)'
    if not ax_n % n_cols:
        ax.legend(frameon=False)
    ax.text(-xlim, ylim, f'{np.mean(score):.1%}\n{acc_explainer}', va='top', size=14)
    # Format the plot
    ax.set(
        xlim=(-xlim, xlim),
        ylim=(0, ylim),
        title=hr_data[query[2]].capitalize(),
    )
    if not ax_n % n_cols:
        ax.set_ylabel('Normalized frequency')
    ax.set_xlabel('bAcc Δ to null')
    if ax_n < n_cols:
        ax.get_xaxis().set_visible(False)
    ax.xaxis.set_major_formatter(mtick.PercentFormatter())
    plt.suptitle('Predicting Response Using Various Sets of ROI Data in two populations', fontsize=18)
bacc_fig.tight_layout()
plt.show()
bacc_fig.savefig(figdir.append(f'bAcc_histograms.png'))

# Retrieve the accuracy score from our results
queries = ('Fold', 'Extremes'), ('Site', 'All')
measured_sites_accs = [np.mean(flatten(results_and_stats.loc[('is_responder', sN, sN, sN, *q)]['null_accuracy'])) for q in
                       queries]

print('We can say the following about the null accuracy score:')
for i, ((cv_name, (cv_scheme, _)), example_population, measured_sites_acc) in enumerate(
        zip(cv_schemes.items(), ('Extremes', 'All'), measured_sites_accs), 1):
    # Simulate simple
    Xt, yt, ct, = data_dict['roi']['dbc_no'][example_population]
    yt = yt.is_responder

    sites_acc = []
    dominant_class_mismatches, n_splits = 0, 0
    for train_idx, test_idx in cv_scheme.split(yt, y=yt, groups=ct):
        n_splits += 1
        dominant_class_in_train = yt.iloc[train_idx].mode()[0].item()
        dominant_class_in_test = yt.iloc[test_idx].mode()[0].item()
        if dominant_class_in_train != dominant_class_in_test:
            dominant_class_mismatches += 1
        test_samples = yt.iloc[test_idx]
        sites_acc.append(np.mean(test_samples == dominant_class_in_train))

    print(
        f'{i}.  In Leave-{cv_name}-Out CV of {example_population} the dominant class was mismatched {dominant_class_mismatches}/{n_splits} times:\n'
        f'    - Expected: {np.mean(sites_acc):.1%}\n'
        f'    - Measured: {measured_sites_acc:.1%}\n')

s = results_and_stats.loc['is_responder', '2DT', 'dbc_no', 'ResNet', 'Fold', 'All', 'wk_2']['accuracy']
n = results_and_stats.loc['is_responder', '2DT', 'dbc_no', 'ResNet', 'Fold', 'All', 'wk_2']['null_accuracy']
print(f'3.  The mean-balanced-accuracy (bAcc) is not the same as balanced-accuracy of the means:\n'
      f'    Mean accuracy score: {np.mean(s):.1%} ± {np.std(s):.1%}, Mean null score: {np.mean(n):.1%} ± {np.std(n):.1%}\n'
      f'    - Balanced accuracy of means:  {np.mean(s) / np.mean(n) / 2:.1%}\n'
      f'    - Mean of balanced accuracies: {np.mean([i / j / 2 for i, j in zip(s, n)]):.1%}\n')

# Parameters of our simulation
n_splits = 40
n_repeats = 2

# Define the machine learning pipeline
combat = list(cv_schemes.values())[-1][-1]
pipeline = make_pipeline(imputer, combat, scaler, selector, LogisticRegression, )

# Preallocate two dictionaries: One to store the simulation bAccs, one to store results
coef_bacc = make_empty_nest(['All', 'Extremes'], ['dbc_no', 'sub_all'], bottom=[])
coef_dict = deepcopy(coef_bacc)

# We use a ROS to balanced training data. Test data is balanced through RSKF-CV.
ros = RandomOverSampler()

for population_name, coef_dict_b in coef_dict.items():
    for data_type in coef_dict_b:
        X, y, _ = data_dict['roi'][data_type][population_name]
        y = y['is_responder']
        cv = RepeatedStratifiedKFold(n_splits=n_splits, n_repeats=n_repeats, random_state=0)
        coef_df = pd.DataFrame(columns=X.columns)
        b_accs = []
        for train_index, test_idx in tqdm(cv.split(X, y)):
            X_train, y_train = X.iloc[train_index], y.iloc[train_index]
            X_train, y_train, = ros.fit_resample(X_train, y_train)
            X_test, y_test = X.iloc[test_idx], y.iloc[test_idx]
            X_test, y_test = ros.fit_resample(X_test, y_test)

            pipeline = make_pipeline(imputer, scaler, selector, classifiers[-1], )
            pipeline.fit(X_train, y_train)
            clf = pipeline[-1]
            y_pred = pipeline.predict(X_test)

            # We have already used ROS, so Acc = bAcc, but just for good measure:
            bacc = balanced_accuracy_score(y_test, y_pred)

            coef_df.loc[len(coef_df)] = clf.feature_importances_
            b_accs.append(bacc)

        coef_dict[population_name][data_type] = coef_df.mean()
        coef_bacc[population_name][data_type] = np.mean(b_accs)

# Flip the atlas that matches regions to indices
inv_dkt_atlas_lut = {v: k for k, v in dkt_atlas_lut.items()}

# Assign storage location
lut_dir = get_root('data', 'fsaverage', 'manual_labels')
os.makedirs(lut_dir, exist_ok=True)

dfs = {'All': [], 'Extremes': []}
# pp stands for projection_property: area or thickness
for pop, hemi, map_prop in product(('All', 'Extremes'), ("Left", "Right"), ('surf', 'thick')):
    coef_means = coef_dict[pop]['dbc_no']
    lut_fname = lut_dir.append(f'{hemi[0]}H-LRC_{map_prop}-coefs-{pop}.pkl')
    if not os.path.isfile(lut_fname):
        idx2coef = make_coef_lut(coef_means, hm=hemi[0], property=map_prop, sd=True)
        pickle_out(idx2coef, lut_fname)
        lut_status = 'created'
    else:
        idx2coef = pickle_in(lut_fname)
        lut_status = 'loaded'
    dfs[pop].append(pd.Series({inv_dkt_atlas_lut[k]: v for k, v in idx2coef.items()}, name=f'{hemi}_{map_prop}'))
    print(f'Coefficient Look-Up-Table {lut_status} for {map_prop} in {hemi} hemisphere. ({lut_fname})')
pd.concat(dfs['Extremes'], axis=1)

# inverse the dkt_atlas
population = 'Extremes'
dkt_atlas_lut_ud = {v: k for k, v in dkt_atlas_lut.items()}
colnames = []
feature_importances = []
for hm in 'LR':
    for letter in 'surf', 'thick':
        # load the LUT
        lut_fname = lut_dir.append(f'{hm}H-LRC_{letter}-coefs-{population}.pkl')
        feature_importances.append({dkt_atlas_lut_ud[k]: v for k, v in pickle_in(lut_fname).items()})
        colnames.append(f'{hm}H_{letter}')
print(f'Coefficients as provided for imaging of the {population} population:')
print(pd.DataFrame(feature_importances, index=colnames).T)

#show colormap
n_shades = 200
cbar_lim = -2, 2
colorbar = get_rgb_cbar(n_shades=n_shades)

# #Plot with label brain (annot)
# imgs = 'lh_annot', 'lh_surf', 'lh_thick', 'colorbar', 'rh_annot', 'rh_surf', 'rh_thick',
# imgs_fpath = r"D:\repositories\ENIGMA\documents\graphics\Blender\{}.png"
# img_titles = 'Labels', 'Surface Coefficients', 'Thickness Coefficients', '',\
#              'Labels', 'Surface Coefficients', 'Thickness Coefficients',

#Plot without label brain (annot)
imgs = 'lh_surf', 'lh_thick', 'colorbar', 'rh_surf', 'rh_thick',
imgs_fpath = figdir.append('{}.png')
img_titles = 'Surface Area', 'Cortical Thickness', '', \
    'Surface Area', 'Cortical Thickness',

n_rows = 2
n_cols = len(imgs) // n_rows
for img_pop in ('Extremes',):
    brain_fig, axes = plt.subplots(nrows=n_rows,
                                   ncols=len(imgs) // n_rows + 1,
                                   figsize=(n_cols * 3 + 2, 2 * n_rows),
                                   gridspec_kw={'width_ratios': n_cols * [1] + [0.1]})
    axes = axes.reshape(-1)
    for ax, title, img_path in zip(axes, img_titles, imgs):
        # Select and display image
        formatted_img_path = imgs_fpath.format(img_path if 'annot' in img_path else f'{img_path}_{img_pop}')
        img = colorbar if img_path == 'colorbar' else mpimg.imread(formatted_img_path)
        ax.imshow(img)

        # Set axes properties
        ax.set(xticks=[], title=title)
        if img_path == 'colorbar':
            ax.set_yticks(np.linspace(0, n_shades, 5))
            ax.grid(False)
            ax.set_yticklabels(np.linspace(cbar_lim[-1], cbar_lim[0], 5))
            ax.set_ylabel('Norm. Coefficient (σ)')
        else:
            ax.set_yticks([])
    axes[-1].axis('off')
    axes[0].set_ylabel('Left Hemisphere')
    axes[n_cols + 1].set_ylabel('Right Hemisphere')
    perf = np.mean(results_dict['roi']['dbc_no'][img_pop]['Fold']['GradientBoostingClassifier']['is_responder']['wk_2'][
                       'balanced_accuracy'])
    n_part = len(data_dict['roi']['dbc_no'][img_pop][0])
    # plt.suptitle(f'Coefficients for {img_pop} Population (n = {n_part:.0f}, bAcc = {bacc:.1%})')
    plt.suptitle(f'Coefficients for the Extreme (Non-)responders Subpopulation (n = {n_part:.0f}, bAcc = {perf:.1%})')
    brain_fig.savefig(figdir.append(f'{img_pop}-3Dbrains.tiff'))
plt.show()

# We never ended up using this because it was hard to retrace 
# which "pixels" were dropped in preprocessing, so they might 
# not be placed correctly (23 out of 900 were dropped)
X, y, _ = data_dict['vec']['dbc_no']['Extremes']
clf = LogisticRegression()
clf.fit(X, y)
coef = np.squeeze(clf.coef_)

X_thc2, _ = load_proj_df(data, 'Thickness')

imps = np.zeros(30 * 30)
counter = 0
for i, pixel in enumerate(X_thc2.var().to_numpy().ravel()):
    if pixel >= 0.002:
        imps[i] = clf.coef_[0][counter]
        counter += 1

titles = 'Median Imporance', 'Median thickness', 'aparc.a2009s.annot'

map_imp = np.reshape(imps, (30, 30))
map_thk = np.reshape(np.median(X_thc2, axis=0), (30, 30))
map_img = np.load(results_dir.append('aggregated_label_annot_res-32.npy'))

fig, axes = plt.subplots(1, len(titles), figsize=(10, 4))
h_thk = axes[0].imshow(map_imp, cmap='viridis', clim=(-map_imp.std() * 2, map_imp.std() * 2), )
h_imp = axes[1].imshow(map_thk, cmap='jet')
h_img = axes[2].imshow(map_img)

for ax, mappable, title in zip(axes, [h_thk, h_imp, None], titles):
    ax.grid(False)
    ax.set(title=title, xticks=[], yticks=[])
    if mappable:
        fig.colorbar(mappable, ax=ax)
fig.tight_layout()
fig.show()

# A function to strip subcortical ROIs from their hemisphere information
rm_hem = lambda x: x[1:] if x[0] in hems else x

# Get hemisphere options, and ROI options so we can strip them from column names
hems = 'LRM'

# Features of interest, pick one
feature_set = ['sub_all', 'dbc_no'][0]

imp_dfs = {}
for population, population_coef_dict in coef_dict.items():
    bacc = np.mean(coef_bacc[population][feature_set])

    # Subcortical importances
    subc = {k: v for k, v in population_coef_dict[feature_set].items() if k in subcortical_rois}

    # Apply the subcortical strip function to the index to get all unique ROIs
    no_hem_subc_rois = set(population_coef_dict[feature_set].index.map(rm_hem))

    # Preallocate a numpty array to store our
    subc_coefs = np.empty((len(no_hem_subc_rois), len(hems)))
    subc_coefs[:] = np.nan

    # Collect coef by ROI and hemisphere, and appoint directly into a numpy array
    for i, v in subc.items():
        roi_name = rm_hem(i)
        hemi = i.replace(roi_name, '')
        hemi = hemi if hemi else 'M'
        x = list(no_hem_subc_rois).index(roi_name)
        y = hems.index(hemi)
        subc_coefs[x, y] = v
    subc_coefs = np.concatenate([subc_coefs, np.nansum(np.abs(subc_coefs), 1)[:, np.newaxis]], axis=1)

    # Convert the array to a table, and do formatting
    subc_relevance_table = pd.DataFrame(data=subc_coefs, columns=list(hems) + ['SAR'],
                                        index=no_hem_subc_rois)  # We use Sum of Absolute Relevances
    subc_relevance_table = subc_relevance_table.sort_values(by='SAR', ascending=False)  # for sorting
    subc_relevance_table = subc_relevance_table.drop(columns=['SAR'])  # and drop it after.
    subc_relevance_table = subc_relevance_table.replace(np.nan, pd.NA)
    subc_relevance_table = subc_relevance_table.rename(columns={'L': 'Left', 'R': 'Right', 'M': 'Middle / NA'})

    # Store the table
    imp_dfs[population] = subc_relevance_table

print('start time:', results_uid)
print('finish time:', datetime.now().strftime("%Y%m%d-%H%M%S"))