# Work from this notebook's folder (content/), also when the notebook is executed from the project root
# (myst build --execute, NeuroLibre), so the local modules (helpers, stats_test, ...) and ../data/ resolve.
import os
import sys

if os.path.isdir("content") and not os.path.isfile("helpers.py"):
    os.chdir("content")
if os.getcwd() not in sys.path:
    sys.path.insert(0, os.getcwd())

import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

import numpy as np

from statsmodels.stats.multitest import multipletests

import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots

# Setting text svg exports as exports
plt.rcParams['svg.fonttype'] = 'none'

# rpy2 configuration
import rpy2.robjects as ro
from rpy2.robjects.packages import importr
from rpy2.robjects import pandas2ri
from rpy2.robjects.conversion import localconverter

try:
    effectsize = importr('effectsize')
    marginaleffects = importr('marginaleffects')
    car = importr('car')
    print("✅ All R dependencies are loaded and ready!")
except Exception as e:
    print(f"❌ Missing R package: {e}")

# Import auxiliary functions bundled in this notebook
import stats_test as st
import helpers

# Static (matplotlib/seaborn) figure versions for the Word/PDF/LaTeX exports and the journal figures.
# Written to static/figures/ only when STATIC_FIGURE_FORMATS is set (pixi tasks static-figures-png/-pdf/-svg,
# export-docx, export-tex). A plain `pixi run run-notebooks` does not write them.
create_Static_figures = bool(helpers.STATIC_FIGURE_FORMATS)

# This color pallet matches plotly figures with seaborn defaults.
color_palette = ["#4C72B0", "#DD8452", "#55A868", "#C44E52", "#8172B2", "#937860", "#DA8BC3", "#8C8C8C", "#CCB974", "#64B5CD"]

# read MRI metrics data: HippUnfold surface metrics averaged over the vertices of each participant, hemisphere and
# subfield (n_vertices = number of vertices averaged)
df = pd.read_csv(helpers.data_path("hyades_hippunfoldAD_surfmetrics.csv"))

# read volumetric data
df_vol = pd.read_csv(helpers.data_path("hyades_hippunfoldAD_volumes_normalized.tsv"),sep='\t')

# Volumes are fractions of total intracranial volume (TIV). Express them as percentages of TIV: the smallest subfields (e.g. CA2)
# otherwise fall below the numerical tolerance of car::Anova (residual sum of squares treated as 0).
# F statistics, effect sizes and p-values are unaffected by this change of units.
volume_cols = ['Sub', 'CA1', 'CA2', 'CA3', 'CA4', 'DG', 'SRLM', 'Cyst']
df_vol[volume_cols] = df_vol[volume_cols] * 100

# Exclusion list obtained from QC the segemntation and registration of the anatomical scans. Participants with a QC level of 0 were excluded from the analyses.
exclusion_list = pd.read_csv(helpers.data_path("exclusion_list.csv")) 
print(exclusion_list)
exclusion_list = exclusion_list["participant_id"].to_list()

# rename column from dataframe (just for stylec onsitency with paper)
df = df.rename(columns={'D_dki': 'MD'})

# Drop subjects with other diagnosis than dementia
nonDemGroups = ['PDD','PSP','VCI','PD','Conc','aMCI','TBD']
df = df.query("group != @nonDemGroups")

# Drop subjects in the exclusion list
df = df.query("participant_id != @exclusion_list")
df_vol = df_vol.query("participant_id != @exclusion_list")

# Define the desired order of categories
group_order = ['HC', 'LBD', 'FTD', 'AD']

# Convert to ordered categorical type (the data files code healthy controls as 'Control'; the paper uses HC)
df['group'] = df['group'].replace({'Control': 'HC'})
df['group'] = pd.Categorical(df['group'], categories=group_order, ordered=True)

# harmonize color palette
group_colors = dict(zip(group_order, color_palette))

# dropping srlm and cyst subfields
df = df.query("subfield != 'SRLM' and subfield != 'Cyst'")

# creating a subfield type to give a right order to values
subfield_category_order = ['Sub','CA1','CA2','CA3','CA4','DG']
# keep only the six subfields (drops unlabeled vertices, label '0') before ordering them
df = df[df['subfield'].isin(subfield_category_order)].copy()
df['subfield'] = pd.Categorical(df['subfield'], categories=subfield_category_order, ordered=True)

metadata = ["participant_id", "group", "age",'sex',"hemisphere","subfield",'subfield_ID']
#metrics = ["FA","D_uFA", "uFA", "Ki", "Vol_mm3_norm", "MTsat"]
metrics = ["FA","MD", "uFA", "Ki", "MTsat", "thickness","curvature","gyrification"]
df = df[metadata + metrics + ['n_vertices']]

# group by subfields
df_Hemi_grouped = df.drop(columns=['n_vertices']).groupby(['participant_id','group','sex','age','subfield','subfield_ID','hemisphere'],observed=True).mean()

# group data by subject: whole hippocampus = mean over all vertices (weighted by the number of vertices)
def vertex_weighted_mean(x):
    w = x['n_vertices']
    return pd.Series({m: np.average(x[m], weights=w) if x[m].notna().all() else np.nan for m in metrics})

df_Subject_grouped = df.groupby(['participant_id','group','sex','age'], observed=True).apply(vertex_weighted_mean)

# re-arrage data into same format as our dataframe
df_vol_melt = df_vol.drop(columns=['SRLM', 'Cyst'])
df_vol_melt = df_vol_melt.melt(id_vars=['participant_id','hemisphere'], value_vars=subfield_category_order, value_name='volume')
df_vol_melt = df_vol_melt.set_axis(['participant_id', 'hemisphere', 'subfield','volume'], axis=1)
df_vol_melt = df_vol_melt.set_index(['participant_id', 'hemisphere', 'subfield'])

# arrange index of original dataframe so we can merge them both
df_Hemi_grouped = df_Hemi_grouped.reset_index()
df_Hemi_grouped = df_Hemi_grouped.set_index(['participant_id', 'hemisphere', 'subfield'])
df_Hemi_grouped= pd.concat([df_Hemi_grouped,df_vol_melt], axis=1, ignore_index=False)
df_Hemi_grouped = df_Hemi_grouped.reset_index()
#df_Hemi_grouped

df_vol_melID = df_vol_melt.groupby('participant_id').sum()

df_Subject_grouped = df_Subject_grouped.reset_index()
df_Subject_grouped = df_Subject_grouped.set_index(['participant_id'])
df_Subject_grouped= pd.concat([df_Subject_grouped,df_vol_melID], axis=1, ignore_index=False)
df_Subject_grouped = df_Subject_grouped.reset_index()
#df_Subject_grouped

df_Hemi_grouped_melted = df_Hemi_grouped.melt(id_vars=['participant_id','group','sex','age', 'hemisphere','subfield'], value_vars=metrics)
#df_Hemi_grouped_melted

df_Subject_grouped_melted = df_Subject_grouped.melt(id_vars=['participant_id','group','sex','age'], value_vars=metrics)
#df_Subject_grouped_melted

# whole hippocampus: one row per participant (metric means over subfields and hemispheres; volume = summed subfield volumes, % of TIV)
metrics_grid = [
    ["volume", "FA", "MD"],
    ["uFA", "Ki", "MTsat"],
]

fig = make_subplots(
    rows=2, cols=3,
    subplot_titles=[st.metric_axis_label(m) for row in metrics_grid for m in row],
    horizontal_spacing=0.05,
    vertical_spacing=0.08,
)

for r, row_metrics in enumerate(metrics_grid, start=1):
    for c, metric in enumerate(row_metrics, start=1):
        for group in group_order:
            data_group = df_Subject_grouped[df_Subject_grouped["group"] == group]
            fig.add_trace(
                go.Box(
                    x=data_group["group"],
                    y=data_group[metric],
                    name=group,
                    legendgroup=group,
                    marker_color=group_colors[group],
                    showlegend=(r == 1 and c == 1),
                    width=0.6,  # wider boxes since each has its own x slot
                ),
                row=r, col=c,
            )

fig.update_xaxes(showticklabels=False)

fig.update_layout(
    width=1300,
    height=800,
    boxmode='group',
    boxgap=0.15,
    margin=dict(l=60, r=180, t=60, b=40),
    legend=dict(
        orientation='v',
        yanchor='top', y=1,
        xanchor='left', x=1.02,
    ),
)

fig

# whole hippocampus: one row per participant (metric means over subfields and hemispheres; volume = summed subfield volumes, % of TIV)
if create_Static_figures:
    helpers.static_theme()
    metrics_grid = ["volume", "FA", "MD", "uFA", "Ki", "MTsat"]
    fig, axes = plt.subplots(2, 3, figsize=(helpers.TEXT_WIDTH_IN, 3.8))
    for ax, metric in zip(axes.flatten(), metrics_grid):
        sns.boxplot(ax=ax, data=df_Subject_grouped, y=metric, hue='group', palette=group_colors)
        ax.set_ylabel(st.metric_axis_label(metric))
        if ax is not axes[0, 2]:
            ax.legend_.remove()
    sns.move_legend(axes[0, 2], "upper left", bbox_to_anchor=(1, 1), title="Group", frameon=False)
    sns.despine(fig)
    fig.tight_layout()
    helpers.save_static_figure(fig, "fig_hippocampus_metrics")
    plt.close(fig)
    sns.set_theme()

fig = px.scatter(
    df_Subject_grouped_melted,
    x='age',
    y='value',
    color='group',
    facet_col='variable',
    facet_col_wrap=2,
    category_orders={
        'group': group_order,
    },
    color_discrete_map=group_colors,
    trendline='ols',
    labels={
        'age': 'Age',
        'value': 'Value',
        'group': 'Group',
    },
)

# sharey=False equivalent: let each facet have its own y-axis range
fig.update_yaxes(matches=None, showticklabels=True)

# px.scatter + trendline creates a duplicate legend entry per group
# (one for points, one for the "...,trendline" line) — collapse to one
seen = set()
for trace in fig.data:
    base_name = trace.name.split(',')[0]
    trace.legendgroup = base_name
    if 'trendline' in trace.name:
        trace.showlegend = False
    elif base_name in seen:
        trace.showlegend = False
    else:
        trace.name = base_name
        seen.add(base_name)

# clean up facet titles ("structure=Hippocampus" -> "Hippocampus")
fig.for_each_annotation(lambda a: a.update(text=st.metric_axis_label(a.text.split('=')[-1])))

fig.update_layout(
    width=1100,
    height=1600,
    legend=dict(
        orientation='v',
        yanchor='middle', y=0.5,
        xanchor='left', x=1.02,
    ),
)

fig

if create_Static_figures:
    helpers.static_theme()
    # FacetGrid + regplot (same as lmplot) so the bootstrapped confidence bands can use a fixed seed
    p = sns.FacetGrid(df_Subject_grouped_melted, hue='group', palette=group_colors, col="variable", col_wrap=4,
                      sharey=False, height=1.55, aspect=1.0)
    p.map_dataframe(sns.regplot, x='age', y='value', seed=0, scatter_kws={"s": 6})
    p.add_legend(title="Group")
    for name, ax in p.axes_dict.items():
        ax.set_title(st.metric_axis_label(name))
    p.set_axis_labels("Age (years)", "")
    helpers.save_static_figure(p.figure, "fig_hippocampus_ageMetrics")
    plt.close(p.figure)
    sns.set_theme()

qMRI_metrics = ["volume", "FA", "MD", "uFA", "Ki", "MTsat"]
subfield_order = subfield_category_order

default_metric = "uFA"
default_idx = qMRI_metrics.index(default_metric)

fig = go.Figure()

traces_per_group = 2  # band + line
traces_per_metric = len(group_order) * traces_per_group  # 4 groups * 2 = 8

for i, metric in enumerate(qMRI_metrics):
    for group in group_order:
        data_group = df_Hemi_grouped[df_Hemi_grouped["group"] == group]

        stats = (
            data_group.groupby("subfield")[metric]
            .agg(mean="mean", sem="sem")
            .reindex(subfield_order)
        )
        ci95 = 1.96 * stats["sem"]
        upper = stats["mean"] + ci95
        lower = stats["mean"] - ci95

        color = group_colors[group]
        is_default = (i == default_idx)

        # band (always showlegend=False, only visibility toggles)
        fig.add_trace(
            go.Scatter(
                x=list(subfield_order) + list(subfield_order)[::-1],
                y=list(upper) + list(lower)[::-1],
                fill='toself',
                fillcolor=color,
                opacity=0.2,
                line=dict(width=0),
                legendgroup=group,
                showlegend=False,
                hoverinfo='skip',
                visible=is_default,
            )
        )

        # mean line
        fig.add_trace(
            go.Scatter(
                x=subfield_order,
                y=stats["mean"],
                mode='lines',
                name=group,
                legendgroup=group,
                line=dict(color=color),
                showlegend=is_default,
                visible=is_default,
            )
        )

fig.update_xaxes(categoryorder='array', categoryarray=subfield_order)

# build dropdown: for each metric, only its 8 traces visible;
# legend only shown on the 4 "line" traces (odd positions within each metric's block)
buttons = []
n_total = len(qMRI_metrics) * traces_per_metric
for i, metric in enumerate(qMRI_metrics):
    visible_mask = [False] * n_total
    showlegend_mask = [False] * n_total
    start = i * traces_per_metric
    for j in range(traces_per_metric):
        visible_mask[start + j] = True
        if j % 2 == 1:  # line traces only (band=0, line=1, band=2, line=3, ...)
            showlegend_mask[start + j] = True

    buttons.append(dict(
        label=st.metric_name(metric),
        method='update',
        args=[
            {'visible': visible_mask, 'showlegend': showlegend_mask},
            {'yaxis': {'title': st.metric_axis_label(metric)}, 'title': f"{st.metric_name(metric)} across subfields"},
        ],
    ))

fig.update_layout(
    width=750,
    height=500,
    title=f"{st.metric_name(default_metric)} across subfields",
    yaxis_title=st.metric_axis_label(default_metric),
    margin=dict(l=60, r=200, t=80, b=60),
    legend=dict(
        orientation='v',
        yanchor='middle', y=0.5,
        xanchor='left', x=1.02,
    ),
    updatemenus=[dict(
        type='dropdown',
        buttons=buttons,
        direction='down',
        x=1.02, xanchor='left',
        y=1.15, yanchor='top',
        showactive=True,
        active=default_idx,
    )],
)

fig

if create_Static_figures:
    helpers.static_theme()
    metrics_grid = ["volume", "FA", "MD", "uFA", "Ki", "MTsat"]
    fig, axes = plt.subplots(2, 3, figsize=(helpers.TEXT_WIDTH_IN, 3.8))
    for ax, metric in zip(axes.flatten(), metrics_grid):
        # band: mean ± 1.96 × standard error, as in the interactive figure and the caption
        sns.lineplot(ax=ax, data=df_Hemi_grouped, x="subfield", y=metric, hue='group', palette=group_colors,
                     errorbar=("se", 1.96))
        ax.set_ylabel(st.metric_axis_label(metric))
        ax.set_xlabel("Subfield")
        if ax is not axes[0, 2]:
            ax.legend_.remove()
    sns.move_legend(axes[0, 2], "upper left", bbox_to_anchor=(1, 1), title="Group", frameon=False)
    sns.despine(fig)
    fig.tight_layout()
    helpers.save_static_figure(fig, "fig_micrometricsSubfield")
    plt.close(fig)
    sns.set_theme()

macro_metrics = ["thickness", "curvature", "gyrification"]
subfield_order = subfield_category_order

fig = make_subplots(
    rows=1, cols=3,
    subplot_titles=[st.metric_axis_label(m) for m in macro_metrics],
    horizontal_spacing=0.06,
)

for c, metric in enumerate(macro_metrics, start=1):
    for group in group_order:
        data_group = df_Hemi_grouped[df_Hemi_grouped["group"] == group]

        stats = (
            data_group.groupby("subfield")[metric]
            .agg(mean="mean", sem="sem")
            .reindex(subfield_order)
        )
        ci95 = 1.96 * stats["sem"]
        upper = stats["mean"] + ci95
        lower = stats["mean"] - ci95

        color = group_colors[group]
        show_legend = (c == 1)  # matches seaborn keeping only one legend, though placed on axes[2] there

        fig.add_trace(
            go.Scatter(
                x=list(subfield_order) + list(subfield_order)[::-1],
                y=list(upper) + list(lower)[::-1],
                fill='toself',
                fillcolor=color,
                opacity=0.2,
                line=dict(width=0),
                legendgroup=group,
                showlegend=False,
                hoverinfo='skip',
            ),
            row=1, col=c,
        )

        fig.add_trace(
            go.Scatter(
                x=subfield_order,
                y=stats["mean"],
                mode='lines',
                name=group,
                legendgroup=group,
                line=dict(color=color),
                showlegend=show_legend,
            ),
            row=1, col=c,
        )

fig.update_xaxes(categoryorder='array', categoryarray=subfield_order)

fig.update_layout(
    width=1750,
    height=500,
    margin=dict(l=60, r=180, t=60, b=40),
    legend=dict(
        orientation='v',
        yanchor='top', y=1,
        xanchor='left', x=1.02,
    ),
)

fig

if create_Static_figures:
    helpers.static_theme()
    macro_metrics = ["thickness", "curvature", "gyrification"]
    fig, axes = plt.subplots(1, 3, figsize=(helpers.TEXT_WIDTH_IN, 2.0))
    for ax, metric in zip(axes, macro_metrics):
        # band: mean ± 1.96 × standard error, as in the interactive figure and the caption
        sns.lineplot(ax=ax, data=df_Hemi_grouped, x="subfield", y=metric, hue='group', palette=group_colors,
                     errorbar=("se", 1.96))
        ax.set_ylabel(st.metric_axis_label(metric))
        ax.set_xlabel("Subfield")
        if ax is not axes[2]:
            ax.legend_.remove()
    sns.move_legend(axes[2], "upper left", bbox_to_anchor=(1, 1), title="Group", frameon=False)
    sns.despine(fig)
    fig.tight_layout()
    helpers.save_static_figure(fig, "fig_macrometricsSubfield")
    plt.close(fig)
    sns.set_theme()

# Define general variables
subfields =  df_Hemi_grouped["subfield"].unique()
mc_alpha = 0.05
metrics_stats = metrics + ['volume']

df_Hemi_grouped_ancova_subfield = st.ancova_tests_subfield(df_Hemi_grouped, metrics_stats, subfields, covar='age + hemisphere + sex')

df_Hemi_grouped_ancova_subfield_Pcorrected = st.post_ancova_correctMultipleComparison(df_Hemi_grouped_ancova_subfield, mc_method='fdr_by', mc_alpha=mc_alpha, structure_col = 'subfield')

cols = ["subfield", "metric", "F value", "Df_residual",
        "Omega2_partial", "CI_low.y", "pvals_adj", "level"]

df_sig = (
    df_Hemi_grouped_ancova_subfield_Pcorrected
    .query("pvals_adj < @mc_alpha")
    [cols]
    .rename(columns={
        "F value":        "F(3, df_err)",
        "Df_residual":    "df_err",
        "Omega2_partial": "ω²ₚ",
        "CI_low.y":       "95% CI lower",
        "pvals_adj":      "p_adj"
    })
    .sort_values(["subfield", "metric"])
).reset_index(drop=True)

st.display_table(st.format_table(df_sig, two_decimals=["F(3, df_err)", "ω²ₚ", "95% CI lower"], p_values=["p_adj"]))