# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.19.3
#   kernelspec:
#     display_name: acore-dev
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Metabolomics drift correction
#

# %% [markdown]
# Within-batch correction of metabolomics data. This notebook shows how to correct for
# instrumental drift based on pooled QC data samples.

# %% [markdown]
# ## 1. LOESS smoothing-based drift correction

# %% [markdown]
# First, let's load our packages, including acore.

# %% tags=["hide-output"]
# %pip install acore

# %%
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler

from acore import drift_correction as dc


# %% tags=["hide-input"]
def plot_loess_example_curve(
    df: pd.DataFrame,
    feature_idx: int,
    samples: list,
    qcs: list,
    sample_order: pd.DataFrame,
    show_corrected: bool = True,
    alpha: float = None,  # fixed smoothing span; if None, selected by LOOCV
):
    """
    Plot the raw intensities, LOESS drift curve, and optionally the corrected
    intensities for a single feature according to the loess drift correction function.
    Useful for inspecting drift behaviour before running full drift correction.

    The drift curve is estimated with the same method used by
    ldc.run_drift_correction: LOESS is fit to the QC points and then
    interpolated across all injection positions via a cubic spline
    (qc_rlsc_loess).

    Parameters
    ----------
    df : pd.DataFrame
        Feature matrix with samples as rows and features as columns.
        Metadata columns should have been removed already.
    feature_idx : int
        Column index of the feature to plot.
    samples : list of str
        Row index labels of the biological samples.
    qcs : list of str
        Row index labels of the pooled QC samples.
    sample_order : pd.DataFrame
        Injection-order table with columns `File Name` and `Sample ID`
        (integer run order).
    show_corrected : bool, optional
        If True (default), overlays drift-corrected sample intensities as
        diamond markers.
    alpha : float, optional
        LOESS smoothing span (0 < α ≤ 1). If None (default), the optimal span
        is selected automatically by leave-one-out cross-validation over
        α ∈ [0.40, 1.00]. The selected value is shown in the legend.
    """

    all_rows = samples + qcs
    feature_col = df.iloc[:, feature_idx]

    order_dict = sample_order.set_index("File Name")["Sample ID"].to_dict()
    x_all = np.array([order_dict.get(r, np.nan) for r in all_rows])
    y_all = feature_col.loc[all_rows].astype(float).values

    n_s = len(samples)
    x_sample, y_sample = x_all[:n_s], y_all[:n_s]
    x_qc_arr, y_qc_arr = x_all[n_s:], y_all[n_s:]

    valid_sample = ~np.isnan(x_sample) & ~np.isnan(y_sample)
    valid_qc = ~np.isnan(x_qc_arr) & ~np.isnan(y_qc_arr)
    x_s_v, y_s_v = x_sample[valid_sample], y_sample[valid_sample]
    x_qc_v, y_qc_v = x_qc_arr[valid_qc], y_qc_arr[valid_qc]

    feature_name = df.columns[feature_idx]

    fig, ax = plt.subplots(figsize=(11, 5))
    ax.scatter(x_s_v, y_s_v, label="Samples", color="steelblue", alpha=0.7, zorder=3)
    ax.scatter(
        x_qc_v, y_qc_v, label="QC", color="firebrick", edgecolor="k", s=60, zorder=4
    )

    if len(x_qc_v) >= 4:
        if alpha is not None:
            # Pass a single-element candidate list to skip LOOCV and use the given alpha
            # directly
            drift_curve, best_alpha = dc.qc_rlsc_loess(
                x_qc_v, y_qc_v, x_all, always_use_default=True, default=alpha
            )
        else:
            drift_curve, best_alpha = dc.qc_rlsc_loess(x_qc_v, y_qc_v, x_all)

        valid_curve = ~np.isnan(x_all) & ~np.isnan(drift_curve)
        sort_idx = np.argsort(x_all[valid_curve])
        ax.plot(
            x_all[valid_curve][sort_idx],
            drift_curve[valid_curve][sort_idx],
            label=f"LOESS drift curve (α={best_alpha:.2f})",
            color="black",
            lw=2,
            zorder=5,
        )

        if show_corrected:
            median_qc = np.median(y_qc_v)
            drift_at_samples = drift_curve[:n_s][valid_sample]
            corrected = (y_s_v / drift_at_samples) * median_qc
            ax.scatter(
                x_s_v,
                corrected,
                label="Corrected samples",
                color="lightsteelblue",
                marker="D",
                s=40,
                alpha=0.9,
                zorder=3,
            )
    else:
        print(f"Not enough valid QC points for LOESS ({len(x_qc_v)} found, need ≥4).")

    ax.set_xlabel("Injection Order")
    ax.set_ylabel("Intensity")
    ax.set_title(f"Drift correction example ({feature_name})")
    ax.legend()
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()


def pca_for_cpca_drift(
    df: pd.DataFrame,
    samples,  # list of row index labels, OR dict {group_name: [row index labels]}
    qcs: list,
    log_transform: bool = True,
    title: str = "PCA",
):
    """
    PCA of samples and QC samples.

    Parameters
    ----------
    df      : feature matrix with samples as rows and features as columns.
              Metadata columns should have been removed already.
    samples : list of row index labels of the biological samples (all one group),
              or dict {group_name: [row index labels]} for multiple groups
    qcs     : list of row index labels of the pooled QC samples
    log_transform : apply log1p before scaling
    title   : plot title
    """
    # Normalise samples to a dict
    if isinstance(samples, list):
        sample_groups = {"Samples": samples}
    else:
        sample_groups = samples

    # Build ordered row + label arrays
    all_rows, labels = [], []
    for group, rows in sample_groups.items():
        for r in rows:
            if r in df.index:
                all_rows.append(r)
                labels.append(group)
    for r in qcs:
        if r in df.index:
            all_rows.append(r)
            labels.append("QC")

    # Coerce to numeric, drop features (columns) with any NaN
    X = df.loc[all_rows].apply(pd.to_numeric, errors="coerce").dropna(axis=1)

    if log_transform:
        X = np.log1p(X.clip(lower=0))

    X_scaled = StandardScaler().fit_transform(X.values.astype(float))

    pca = PCA(n_components=2)
    coords = pca.fit_transform(X_scaled)
    pc1_var = pca.explained_variance_ratio_[0] * 100
    pc2_var = pca.explained_variance_ratio_[1] * 100

    palette = plt.cm.tab10.colors
    color_map = {g: palette[i % len(palette)] for i, g in enumerate(sample_groups)}
    color_map["QC"] = "black"

    fig, ax = plt.subplots(figsize=(8, 6))
    for group in list(sample_groups) + ["QC"]:
        mask = [_label == group for _label in labels]
        pts = coords[mask]
        is_qc = group == "QC"
        ax.scatter(
            pts[:, 0],
            pts[:, 1],
            label=group,
            color=color_map[group],
            marker="D" if is_qc else "o",
            s=70 if is_qc else 50,
            edgecolors="k" if is_qc else "none",
            alpha=0.9 if is_qc else 0.75,
            zorder=5 if is_qc else 3,
        )

    ax.set_xlabel(f"PC1 ({pc1_var:.1f}%)")
    ax.set_ylabel(f"PC2 ({pc2_var:.1f}%)")
    ax.set_title(title)
    ax.axhline(0, color="grey", lw=0.5, ls="--")
    ax.axvline(0, color="grey", lw=0.5, ls="--")
    ax.legend(framealpha=0.8)
    ax.grid(True, alpha=0.2)
    plt.tight_layout()
    plt.show()


# %% [markdown]
# ### Load in data

# %% [markdown]
# We will use some example data that can be found in this repo.

# %%
data_path = (
    "https://raw.githubusercontent.com/Multiomics-Analytics-Group/acore/"
    "refs/heads/main/"
)

df_original = pd.read_csv(
    "../../example_data/DidacMauricio_hilic/DM_FIS2018_Hilic_pos_results2023_filled_imputed.csv",
    index_col=0,
)

# %% [markdown]
# This is what the data frame, an intensity table, looks like.

# %% tags=["hide-inputs"]
df_original

# %% [markdown]
# In order to run our further analysis, including the filtering functions,
# we have to transform the data and remove metadata such as mass and retention time.

# %%
df = df_original.T
df = df.drop(
    ["Qidx", "SOIidx", "rtmed", "start", "end", "mass", "MaxInt", "formula", "anot"]
)

# %% tags=["hide-input"]
df

# %% [markdown]
# We also have data that contains the order in which our samples were run. This
# information is crucial for the drift correction algorithm; it cannot be performed
# without it.
#
# The data needs to contain the columns `File Name` and `Sample ID`, referring to the
# name of the sample and the index of the sample, meaning the order in which the samples
# were run.
#
# Note: If you do not have metadata that contains the order in which samples and QCs
# were run, you can use the second method explained in this notebook,
# [CPCA](#common-principal-components-analysis-based-drift-correction).

# %%
sample_order = pd.read_csv(
    "../../example_data/DidacMauricio_hilic/DidacMauricio_hilic_pos_injectionorder.csv"
)
sample_order["File Name"] = sample_order["SampleName"]
sample_order["Sample ID"] = sample_order["injectionOrder"]

sample_order

# %% [markdown]
# ### Run drift correction with LOESS smoothing

# %% [markdown]
# We can now correct our data for experimental drift.
#
# With the acore LOESS drift-correction function, a LOESS (locally estimated regression)
# smoother is applied separately to the features in the data to model slow temporal
# trends, and the resulting smooth trend is used to correct the data.
#
# Before the function estimation and correction, the data can filtered, to remove
# features that have too many missing values in the QC samples.

# %% [markdown]
# First, we can create a dictionary for our sample names, ordering them
# into groups, to make the upcoming function call easier.

# %%
column_list = list(df_original.columns.values)
sample_cols = []
qc_cols = []
for col in column_list:
    if col.startswith("AAA"):
        sample_cols.append(col)
    elif col.startswith("QC"):
        qc_cols.append(col)

# %% [markdown]
# Now, are sample column names are summarised in the sample_cols variable:

# %% tags=["hide-output"]
sample_cols

# %% [markdown]
# Our columns corresponding to QC data are saved in the qc_cols variable:

# %% tags=["hide-output"]
qc_cols

# %% [markdown]
# Now we can run the drift correction, using the acore `run_loess_drift_correction`
# function.
#
# Note: For this demonstration, we will use the parameter `always_use_default`, so that the leave-one-out cross validation, during which the best smoothing parameter for the curve is determined, is skipped. This is just for efficiency purposes; so that this does not take too long to load. However, it is recommended to not use this option, so that the best parameter can be determined for the curve of every feature.

# %% tags=["hide-input"]
help(dc.run_loess_drift_correction)

# %%
corrected_df, correction_info = dc.run_loess_drift_correction(
    df,
    qc_cols,
    sample_cols,
    sample_order=sample_order,
    filter_percent=0.5,
    always_use_default=True,
)

# %% [markdown]
# Explanation of the parameters chosen:
# - `filter_percent` is  the minimum percentage of values that must be present for this
#   feature to be retained. If the percentage of non-missing is below this, the feature
#   will be filtered out. If this parameter is set to "None", no filtering will be done. In this case, the filter_percent parameter does not do anything, as the data is already imputed and there are no missing values in the QCs.
# - `always_use_default` disables the leave-one-out cross validation which calculates
#   the ideal smoothing parameter for the LOESS curve for every feature. It will speed up the computation, but the curve which decides the value correction is less optimal.
# - `default`(not used here) changes the default smoothing value. If the parameter is
#   not used, the default is 0.75. The default smoothing value is either used if `always_use_default`=True, or if no better smoothing parameter can be found during leave-one-out cross validation.
#

# %% [markdown]
# Now we can inspect our results. First, the corrected output dataframe.

# %% tags=["hide-input"]
corrected_df

# %% [markdown]
# We can also look further into the correction_info object, to see the parameters that
# were chosen for each feature.
#
# For example, let's check the parameters used for the 200th feature.

# %% tags=["hide-output"]
print(correction_info[200])

# %% [markdown]
# ### Plot an example curve for one feature

# %% [markdown]
# We can also plot an example feature, to see how the values have changed
# and what the LOESS curve would look like for the data of this feature.

# %% tags=["hide-input"]
plot_loess_example_curve(
    df=df,
    feature_idx=300,
    samples=sample_cols,
    qcs=qc_cols,
    sample_order=sample_order,
)

# %% [markdown]
# In this plot, we can see all of the data points of this feature, ordered by the
# injection time. The red points are our QC samples, so we can see whether there was any
# instrumental drift over time. The LOESS curve is also calculated, with the smoothing
# value alpha chosen with leave-one-out cross validation, just like in the
# run_drift_correction function.
#
# Alternative values for the smoothing parameter alpha can be tested by adding the
# argument "alpha" and choosing a value, just like in the example below.

# %% tags=["hide-input"]
plot_loess_example_curve(
    df=df,
    feature_idx=300,
    samples=sample_cols,
    qcs=qc_cols,
    sample_order=sample_order,
    alpha=0.6,
)

# %% [markdown]
# ## 2. Common Principal Components Analysis-based drift correction

# %% [markdown]
# We can also use another method for correcting drift which is based on Common Principal
# Components Analysis (CPCA). This method is based on common principal components in
# defined groups of the data. It assumes that when calculating common principal
# components of QC samples, the drift contribution can be identified as the direction
# capturing maximum variance that simultaneously diagonalizes the covariance matrices of
# a set of classes.
#
# Therefore, the variability in the identified direction can be explained as caused by
# experimental drift and subtracted from all samples.
#
# Let's use different example data for demonstrating this method.

# %% [markdown]
# ### Load in data

# %%
# Load data
df_original = pd.read_csv(
    "../../example_data/DidacMauricio_hilic/DM_FIS2018_Hilic_pos_results2023_filled_imputed.csv",
    index_col=0,
)

df = df_original.T
df = df.drop(
    ["Qidx", "SOIidx", "rtmed", "start", "end", "mass", "MaxInt", "formula", "anot"]
)

# Define sample columns and qc columns
collist = list(df_original.columns.values)
sample_cols = []
qc_cols = []
for col in collist:
    if col.startswith("AAA"):
        sample_cols.append(col)
    elif col.startswith("QC"):
        qc_cols.append(col)
df

# %% [markdown]
# Seeing as this method is based on calculating principal components, as with PCA, there
# must not be any missing data.
#
# We will therefore first calculate missingness (NAs).
#
# We need to check for missing values in both the sample columns and the QC columns.

# %% tags=["hide-input"]
if dc.check_missingness(df, sample_cols + qc_cols):
    print(
        "There are missing values. Consider imputing first, or use LOESS drift correction instead."
    )
else:
    print("\nThere is no missingness. We can proceed with the CPCA drift correction.")

# %% [markdown]
# ### Visualise non-corrected data
#
# Now that we know we can proceed, let's visualise our data before drift correction with
# a PCA.

# %% tags=["hide-input"]
pca_for_cpca_drift(
    df,
    sample_cols,  # list of col names, OR dict {group_name: [col names]}
    qc_cols,
    log_transform=True,
    title="PCA",
)

# %% [markdown]
# We see some clear indication of instrumental drift in the QC samples.

# %% [markdown]
# ### Run drift correction based on CPCA

# %%
df_corrected = dc.run_cpca_drift_correction(df, sample_cols, qc_cols, n_comps=1)

# %% [markdown]
# Let's plot the corrected data.

# %% tags=["hide-input"]
pca_for_cpca_drift(  # This function is defined in the beginning of this notebook.
    df_corrected,
    sample_cols,  # list of col names, OR dict {group_name: [col names]}
    qc_cols,
    log_transform=True,
    title="PCA",
)

# %% [markdown]
# There is some change, but still a significant amount of drift is clearly visible from
# the PCA plot.
#
# We can play around with the n_comps variable which decides the number of components.

# %%
df_corrected_2comps = dc.run_cpca_drift_correction(df, sample_cols, qc_cols, n_comps=2)
df_corrected_3comps = dc.run_cpca_drift_correction(df, sample_cols, qc_cols, n_comps=3)
df_corrected_4comps = dc.run_cpca_drift_correction(df, sample_cols, qc_cols, n_comps=4)

# %% tags=["hide-input"]
pca_for_cpca_drift(
    df_corrected_2comps,
    sample_cols,  # list of col names, OR dict {group_name: [col names]}
    qc_cols,
    log_transform=True,
    title="PCA with 2 components",
)

pca_for_cpca_drift(
    df_corrected_3comps,
    sample_cols,  # list of col names, OR dict {group_name: [col names]}
    qc_cols,
    log_transform=True,
    title="PCA with 3 components",
)

pca_for_cpca_drift(
    df_corrected_4comps,
    sample_cols,  # list of col names, OR dict {group_name: [col names]}
    qc_cols,
    log_transform=True,
    title="PCA with 4 components",
)

# %% [markdown]
# The correction methods with three and four components are already looking better.
# Let's calculate the centroids of the QC principal components and the distance of the
# QC points to them, to objectively decide which number of n_comps is most favourable.

# %% tags=["hide-input"]
print(
    f"1 component: {dc.cpca_centroid(df_corrected, sample_cols, qc_cols, log_transform=True)}"
)
print(
    f"2 components: {dc.cpca_centroid(df_corrected_2comps, sample_cols, qc_cols, log_transform=True)}"
)
print(
    f"3 components: {dc.cpca_centroid(df_corrected_3comps, sample_cols, qc_cols, log_transform=True)}"
)
print(
    f"4 components: {dc.cpca_centroid(df_corrected_4comps, sample_cols, qc_cols, log_transform=True)}"
)

# %% [markdown]
# According to this, the CPCA method with three principal components is most favourable
# in this case.
#
# We can go ahead and continue our metabolomics data analysis with this data set.

# %%
df_corrected_3comps
