# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.17.2
#   kernelspec:
#     display_name: .venv
#     language: python
#     name: python3
# ---

# %% [markdown]
# # `ABaCo` tutorial: Anaerobic Digestion benchmark dataset

# %% [markdown]
# In this notebook we will implement ABaCo for batch correction on an Anaerobic Digestion (AD) dataset under different phenol concentration conditions. 
#
# -----
# **Data Description:**
# - The AD dataset generated by ([Chapleur et al., 2016](https://doi.org/10.1007/s10532-015-9751-4)) and shared in R package [PLSDAbatch](https://github.com/EvaYiwenWang/PLSDAbatch) by [Wang and Cao (2023)](https://doi.org/10.1093/bib/bbac622) is a dataset composed of 75 samples with 567 identified taxonomic groups. 
# - The samples were treated with two different phenol concentrations, accounting for the biological source of variation. 
# - Samples were processed on 5 different dates over the lapse of 2 years, accounting for the technical source of variation (i.e., batch effect).
#
# **Goal:**
#
# With ABaCo the aim is to remove the technical variation (batch effect) while retaining the biological variation (phenol group effect). 
#
# -----

# %% [markdown]
# ## Dataset Requirements for ABaCo
#
# The dataset should contain the following to be compatible with the ABaCo framework:
#
# | sample | batch      | trt  | feat1 | feat2 | ... |
# |--------|------------|------|-------|-------|-----|
# | A      | 24/07/2025 | low  | #     | #     | ... |
# | B      | 15/06/2024 | low  | #     | #     | ... |
# | C      | 24/07/2025 | high | #     | #     | ... |
#
# - The data must have 3 categorical columns: 
#     1. unique ids to identify the observations/samples e.g. sample ids `sample`
#     2. ids for the batch/factor groupings to be corrected by abaco. e.g. dates of sample analysis `batch`
#     3. biological/experimental factor variation for abaco to retain when correcting batch effect e.g., phenol concentration condition `trt`
#
# - And the features (numeric type) to be trained on. 
#
# `abaco` provides a `BatchEffectDataLoader.dataPreprocess()` function to help convert a plain text file (e.g, csv, tsv) into a compatible pd.DataFrame format, and a `onehotencoding()` function to encode categorical columns. 

# %%
from abaco.dataloader import DataPreprocess, one_hot_encoding

# Load AD count
path_to_dataset = "data/dataset_ad.csv"
batch_col = "batch"
id_col = "sample"
bio_col = "trt"

# Convert data path into compatible pd.DataFrame with ABaCo framework
df_ad = DataPreprocess(
    path_to_dataset,
    factors=[
        batch_col, 
        id_col, 
        bio_col
    ]
)

# see that there are 3 categorical and n numeric columns
print(df_ad.info())

# Isolating the features (numeric) to be inputed into the model
df_features = df_ad.drop(columns=[batch_col, id_col, bio_col]).values
print(f"\n The feature inputs have shape: {df_features.shape}")

# %%
# One-hot encode the categorical columns
# the batch factor
ohe_batches, ohe_batch_classes = one_hot_encoding(df_ad[batch_col])
# and the biological factor
ohe_bio, ohe_bio_classes = one_hot_encoding(df_ad[bio_col])

print(f"Batch classes: {ohe_batch_classes}\n")
print(f"Biological classes: {ohe_bio_classes}")

# %% [markdown]
# ### The AD dataset in brief
# `abaco` also provides plotting functions to easily explore the data. Here we use `BatchEffectPlots.plotPCoA()` to visualize the batch and biological effects on the data.

# %%
from abaco.plots import plotPCoA
import plotly.io as pio
pio.renderers.default = "notebook_connected"
# pcoa
plotPCoA(
    data=df_ad, 
    sample_label=id_col, 
    batch_label=batch_col, 
    experiment_label=bio_col,
)

# %% [markdown]
# - Batch effect (colours): 
#     - Samples were processed on 5 different dates over the lapse of 2 years, accounting for the technical source of variation 
#     - The points clustering by colour demonstrate that there is a batch effect.
# - Biological effect (shapes): 
#     - The samples were treated with 2 different phenol concentrations, accounting for the biological source of variation. 
#     - Within the batches the circle points (lower phenol conc.) are to the left of the square points (higher conc).
#
#
# ### The goal 
#
# The aim of **ABaCo** is to: 
# 1) correct the batch effect (e.g., the points should no longer cluster by colour in the PCoA) while
# 2) maintaining biological variance (e.g., keep circles and squares separated).
#
# Ideally, after using AbaCo to transform the data, the resulting PCoA will look like a colourful mixture of points, where the circle points remain to the left of the square ones.
#
#
# ---

# %% [markdown]
# # Using `ABaCo`

# %% [markdown]
# ## Training the ABaCo model
#
# To train ABaCo on the prepared AD dataset, we create the `abaco.metaABaCo()` class and pass the required parameters shown in the cell below. Thus, running abaco with the default optional parameters.
#
# Usually, setup of the parameters is required, which are explained in brief in the documentation e.g. `help(metaABaCo)`

# %%
from abaco.ABaCo import metaABaCo
import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Create ABaCo model
model = metaABaCo(
    data=df_ad, # Pre-processed dataframe
    n_bios=2, # Number of biological groups in the data
    bio_label=bio_col, # Column where biological groups are labeled in the dataframe
    n_batches=5, # Number of batch groups in the data
    batch_label=batch_col, # Column where batch groups are labeled in the dataframe
    n_features=567, # Number of features (taxonomic groups)
    prior="MoG", # Prior distribution 
    device=device, # Device
)

# %% tags=["hide-output"]
help(metaABaCo)

# %% [markdown]
# For training the model we can call the method `metaABaCo.fit()` with the default parameters.

# %% tags=["hide-output"]
model.fit(seed=42)

# %% [markdown]
# After training the abaco model we can use it to reconstruct the AD dataset without the batch variance while keeping the biological variance.
#
# -----

# %% [markdown]
# ## ABaCo data reconstruction

# %% [markdown]
# To reconstruct the dataset we use the method `metaABaco.correct()`.

# %%
# Reconstruct the dataset using the trained ABaCo model
corrected_dataset = model.correct(seed=42)

# Plot the PCoA of the reconstructed dataset
plotPCoA(
    data = corrected_dataset, 
    sample_label=id_col, 
    batch_label=batch_col, 
    experiment_label=bio_col
)

# %% [markdown]
# ### Conclusion
#
# The goal was to: 
#
# &#x2705; correct the batch effect (reduce clustering of points by colour)
#
# &#x2705; maintain biological variance (maintain separation of circles and sequares)
#
# A brief visual inspection of the PCoA of the reconstructed AD data suggests that ABaCo reduced the batch effect associated with processing the samples on different days, while still retaining the variance due to the experimental condition of lower vs higher phenol concentration.
#
#
# ---
