script:
    "../scripts/R/ancombc.R"

script:
    "../scripts/R/aldex2.R"

shell:
    """
    songbird multinomial \
        --input-biom {input.table} \
        --metadata-file {input.metadata} \
        --formula "{params.formula}" \
        --epochs {params.epochs} \
        --differential-prior {params.diff_prior} \
        --summary-interval 1 \
        --min-feature-count 0 \
        --min-sample-count 0 \
        --random-seed 1 \
        --summary-dir {params.outdir} > {log} 2>&1
    """

script:
    "../scripts/R/maaslin2.R"

script:
    "../scripts/R/corncob.R"

script:
    "../scripts/process_differentials.py"

script:
    "../scripts/concatenate_differentials.py"

script:
    "../scripts/process_pvalues.py"

script:
    "../scripts/concatenate_pvalues.py"

script:
    "../scripts/concatenate_all_results.py"

script:
    "../scripts/run_pca.py"

script:
    "../scripts/get_pctile_feats.py"

script:
    "../scripts/get_pctile_feats.py"

script:
    "../scripts/get_log_ratios.py"

script:
    "../scripts/logistic_regression.py"

script:
    "../scripts/plot_differentials.py"

script:
    "../scripts/plot_diff_kendall_heatmap.py"

script:
    "../scripts/plot_pvalue_kendall_heatmap.py"

script:
    "../scripts/plot_upset.py"

script:
    "../scripts/plot_roc.py"

script:
    "../scripts/plot_prc.py"

script:
    "../scripts/plot_pca.py"

script:
    "../scripts/plot_rank_comparison.py"

script:
    "../scripts/plot_pvalue_pw_comparisons.py"

script:
    "../scripts/plot_pvalue_volcanoes.py"

shell:
    """
    qurro \
        --ranks {input.ranks} \
        --table {input.table} \
        --sample-metadata {input.metadata} \
        --output-dir {output} > {log} 2>&1
    """

shell:
    """
    empress tree-plot \
        --tree {input.tree} \
        --feature-metadata {input.ranks} \
        --output-dir {output} > {log} 2>&1
    """

script:
    "../scripts/create_diff_table.py"

script:
    "../scripts/create_pval_table.py"

import logging
import pandas as pd

logger = logging.getLogger("qadabra")
logger.setLevel(logging.INFO)
fh = logging.FileHandler(snakemake.log[0], mode="w")
formatter = logging.Formatter(
    f"[%(asctime)s - {snakemake.rule}] :: %(message)s"
)
fh.setFormatter(formatter)
logger.addHandler(fh)

logging.captureWarnings(True)
logging.getLogger("py.warnings").addHandler(fh)

logger.info("Loading differentials...")
diffs_files = [
    pd.read_table(x, sep="\t", index_col=0) for x in snakemake.input
]
concat_diffs = pd.concat(diffs_files, axis=1)

concat_diffs.to_csv(snakemake.output[0], sep="\t", index=True)
logger.info(f"Saved to {snakemake.output[0]}")

import logging
import pandas as pd

logger = logging.getLogger("qadabra")
logger.setLevel(logging.INFO)
fh = logging.FileHandler(snakemake.log[0], mode="w")
formatter = logging.Formatter(
    f"[%(asctime)s - {snakemake.rule}] :: %(message)s"
)
fh.setFormatter(formatter)
logger.addHandler(fh)

logging.captureWarnings(True)
logging.getLogger("py.warnings").addHandler(fh)

logger.info("Loading p-values...")
pvalue_files = [
    pd.read_table(x, sep="\t", index_col=0) for x in snakemake.input
]
concat_pvalues = pd.concat(pvalue_files, axis=1)

concat_pvalues.to_csv(snakemake.output[0], sep="\t", index=True)
logger.info(f"Saved to {snakemake.output[0]}")

import logging
import pandas as pd
from bokeh.plotting import output_file, save
from bokeh.models import DataTable, TableColumn, ColumnDataSource


logger = logging.getLogger("qadabra")
logger.setLevel(logging.INFO)
fh = logging.FileHandler(snakemake.log[0], mode="w")
formatter = logging.Formatter(
    f"[%(asctime)s - {snakemake.rule}] :: %(message)s"
)
fh.setFormatter(formatter)
logger.addHandler(fh)

output_file(filename=snakemake.output[0], title="Table")
diff_df = pd.read_table(snakemake.input[0], sep="\t", index_col=0)

diff_df = diff_df.round(2)
diff_df = diff_df.reset_index()

source = ColumnDataSource(diff_df)
columns = [
    TableColumn(field=x, title=x)
    for x in diff_df.columns
]
data_table = DataTable(source=source, columns=columns,
                       frozen_columns=1,
                       sizing_mode="stretch_height")
save(data_table)
logger.info(f"Saved table to {snakemake.output[0]}")

import logging
import pandas as pd
from bokeh.plotting import output_file, save
from bokeh.models import DataTable, TableColumn, ColumnDataSource


logger = logging.getLogger("qadabra")
logger.setLevel(logging.INFO)
fh = logging.FileHandler(snakemake.log[0], mode="w")
formatter = logging.Formatter(
    f"[%(asctime)s - {snakemake.rule}] :: %(message)s"
)
fh.setFormatter(formatter)
logger.addHandler(fh)

output_file(filename=snakemake.output[0], title="P-value table")
pval_df = pd.read_table(snakemake.input[0], sep="\t", index_col=0)

pval_df = pval_df.round(2)
pval_df = pval_df.reset_index()

source = ColumnDataSource(pval_df)
columns = [
    TableColumn(field=x, title=x)
    for x in pval_df.columns
]
data_table = DataTable(source=source, columns=columns,
                       frozen_columns=1,
                       sizing_mode="stretch_height")
save(data_table)
logger.info(f"Saved table to {snakemake.output[0]}")

import logging

import pandas as pd


logger = logging.getLogger("qadabra")
logger.setLevel(logging.INFO)
fh = logging.FileHandler(snakemake.log[0], mode="w")
formatter = logging.Formatter(
    f"[%(asctime)s - {snakemake.rule}] :: %(message)s"
)
fh.setFormatter(formatter)
logger.addHandler(fh)

logger.info("Loading differentials...")
diffs = pd.read_table(snakemake.input[0], sep="\t", index_col=0)
if snakemake.wildcards.get("tool") is not None:
    tool = snakemake.wildcards["tool"]
else:
    tool = "PC1"
pctile = int(snakemake.wildcards["pctile"])

sorted_diffs = diffs.sort_values(by=tool, ascending=False)

n = int(diffs.shape[0]*pctile/100)
top_feats = sorted_diffs.head(n).index.tolist()
bot_feats = sorted_diffs.tail(n).index.tolist()
top_feats = pd.DataFrame(
    dict(feature_id=top_feats, location="numerator")
)
bot_feats = pd.DataFrame(
    dict(feature_id=bot_feats, location="denominator")
)
df = pd.concat([top_feats, bot_feats])
df["pctile"] = pctile
df["num_feats"] = n

df.to_csv(snakemake.output[0], sep="\t", index=False)
logger.info(f"Saved to {snakemake.output[0]}")

import logging

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


logger = logging.getLogger("qadabra")
logger.setLevel(logging.INFO)
fh = logging.FileHandler(snakemake.log[0], mode="w")
formatter = logging.Formatter(
    f"[%(asctime)s - {snakemake.rule}] :: %(message)s"
)
fh.setFormatter(formatter)
logger.addHandler(fh)

logging.captureWarnings(True)
logging.getLogger("py.warnings").addHandler(fh)

plt.style.use(snakemake.config["stylesheet"])

logger.info("Loading differentials...")
diffs_df = pd.read_table(snakemake.input[0], sep="\t", index_col=0)

corr = diffs_df.corr("kendall")
mask = np.zeros_like(corr)
mask[np.triu_indices_from(mask)] = True

fig, ax = plt.subplots(1, 1)
sns.heatmap(
    corr,
    annot=True,
    mask=mask,
    square=True,
    ax=ax
)
ax.set_title("Kendall Rank Correlations")
plt.savefig(snakemake.output[0])
logger.info(f"Saved to {snakemake.output[0]}")

import logging

import matplotlib.pyplot as plt
from matplotlib.patches import Patch
import numpy as np
import pandas as pd
import seaborn as sns


logger = logging.getLogger("qadabra")
logger.setLevel(logging.INFO)
fh = logging.FileHandler(snakemake.log[0], mode="w")
formatter = logging.Formatter(
    f"[%(asctime)s - {snakemake.rule}] :: %(message)s"
)
fh.setFormatter(formatter)
logger.addHandler(fh)

logging.captureWarnings(True)
logging.getLogger("py.warnings").addHandler(fh)

plt.style.use(snakemake.config["stylesheet"])

logger.info("Loading PCA results...")
feat_df = pd.read_table(snakemake.input["features"], sep="\t", index_col=0)
tool_df = pd.read_table(snakemake.input["tools"], sep="\t", index_col=0)
prop_exp = pd.read_table(snakemake.input["prop_exp"], sep="\t", index_col=0)
prop_exp = prop_exp.squeeze()

pc_cols = tool_df.columns.tolist()
tool_list = tool_df.index.tolist()

arrow_palette = dict(zip(
    tool_list,
    sns.color_palette("colorblind", len(tool_list)).as_hex()
))

fig, ax = plt.subplots(1, 1)

logger.info("Plotting scatterplot...")
sns.scatterplot(
    data=feat_df,
    x="PC1",
    y="PC2",
    color="lightgray",
    edgecolor="black",
    linewidth=0.1,
    ax=ax
)

xlim = ax.get_xlim()
ylim = ax.get_ylim()

xrange = xlim[1] - xlim[0]
yrange = ylim[1] - ylim[0]

logger.info("Plotting arrows...")
for i, row in tool_df.iterrows():
    tool_name = row.name
    color = arrow_palette[tool_name]
    ax.arrow(
        x=0, y=0,
        dx=row["PC1"], dy=row["PC2"],
        facecolor=color,
        linewidth=0.5,
        head_width=xrange*0.005*5,
        width=xrange*0.005,
        edgecolor="black",
    )


patches = []
for i, row in tool_df.iterrows():
    tool_name = row.name
    color = arrow_palette[tool_name]
    ax.arrow(
        x=0, y=0,
        dx=row["PC1"], dy=row["PC2"],
        facecolor=color,
        linewidth=0.5,
        head_width=xrange*0.005*5,
        width=xrange*0.005,
        edgecolor="black",
    )
    tool_patch = Patch(
        facecolor=color,
        edgecolor="black",
        label=tool_name,
        linewidth=0.5,
    )
    patches.append(tool_patch)

ax.legend(
    handles=patches,
    bbox_to_anchor=[1, 1],
    loc="upper left",
    frameon=False
)

prop_exp_labels = [
    f"{i} ({x*100:.2f}%)"
    for i, x in prop_exp.items()
]
ax.set_xlabel(prop_exp_labels[0])
ax.set_ylabel(prop_exp_labels[1])

ax.set_title("PCA")
plt.savefig(snakemake.output[0])
logger.info(f"Saved to {snakemake.output[0]}")

from collections import defaultdict
import logging

import joblib
from matplotlib.lines import Line2D
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns


logger = logging.getLogger("qadabra")
logger.setLevel(logging.INFO)
fh = logging.FileHandler(snakemake.log[0], mode="w")
formatter = logging.Formatter(
    f"[%(asctime)s - {snakemake.rule}] :: %(message)s"
)
fh.setFormatter(formatter)
logger.addHandler(fh)

logging.captureWarnings(True)
logging.getLogger("py.warnings").addHandler(fh)

plt.style.use(snakemake.config["stylesheet"])

tools = snakemake.config["tools"] + ["pca_pc1"]
palette = dict(zip(
    tools, sns.color_palette("colorblind", len(tools))
))

pctile = snakemake.wildcards["pctile"]

models = defaultdict(dict)
for tool, model_loc in zip(tools, snakemake.input):
    logger.info(f"Loading {model_loc}...")
    models[tool] = joblib.load(model_loc)

fig, ax = plt.subplots(1, 1)

best_tool = None
best_auc = 0
leg_lines = []
for tool in tools:
    this_model_data = models[tool]
    mean_recalls = this_model_data["recalls"]
    mean_precs = np.mean(this_model_data["precisions"], axis=0)

    pr_aucs = this_model_data["pr_aucs"]
    mean_auc = np.mean(pr_aucs)
    std_auc = np.std(pr_aucs)

    if mean_auc > best_auc:
        best_auc = mean_auc
        best_tool = tool

    color = palette[tool]
    line = Line2D([0], [0], color=color, lw=2,
                  label=f"{tool} ({mean_auc:.2f} $\pm$ {std_auc:.2f})")
    leg_lines.append(line)

    ax.plot(mean_recalls, mean_precs, lw=2, color=color)
    logger.info(f"{tool} Mean AUC = {mean_auc:.2f} +- {std_auc:.2f}")

leg = ax.legend(
    handles=leg_lines,
    loc="upper left",
    frameon=False,
    bbox_to_anchor=[1, 1]
)
for text in leg.get_texts():
    if best_tool in text._text:
        text.set_weight("bold")

logger.info(f"Best model: {best_tool} ({best_auc:.2f})")

ax.set_xlabel("Recall")
ax.set_ylabel("Precision")
ax.set_title(f"{pctile}% Features")

plt.savefig(snakemake.output[0])
logger.info(f"Saved to {snakemake.output[0]}")

import logging

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


logger = logging.getLogger("qadabra")
logger.setLevel(logging.INFO)
fh = logging.FileHandler(snakemake.log[0], mode="w")
formatter = logging.Formatter(
    f"[%(asctime)s - {snakemake.rule}] :: %(message)s"
)
fh.setFormatter(formatter)
logger.addHandler(fh)

logging.captureWarnings(True)
logging.getLogger("py.warnings").addHandler(fh)

plt.style.use(snakemake.config["stylesheet"])

logger.info("Loading pvalues...")
diffs_df = pd.read_table(snakemake.input[0], sep="\t", index_col=0)

corr = diffs_df.corr("kendall")
mask = np.zeros_like(corr)
mask[np.triu_indices_from(mask)] = True

fig, ax = plt.subplots(1, 1)

palette = sns.color_palette("rocket", as_cmap=True)

sns.heatmap(
    corr,
    annot=True,
    mask=mask,
    square=True,
    ax=ax,
    cmap=palette,
    vmin=0,
    vmax=1
)
ax.set_title("Kendall Rank Correlations of P-values")
plt.savefig(snakemake.output[0])
logger.info(f"Saved to {snakemake.output[0]}")

import logging

import numpy as np
import pandas as pd
import seaborn as sns

from bokeh.plotting import figure, output_file, save
from bokeh.layouts import column, row
from bokeh.models import ColumnDataSource, Select, HoverTool
from bokeh.models.callbacks import CustomJS


logger = logging.getLogger("qadabra")
logger.setLevel(logging.INFO)
fh = logging.FileHandler(snakemake.log[0], mode="w")
formatter = logging.Formatter(
    f"[%(asctime)s - {snakemake.rule}] :: %(message)s"
)
fh.setFormatter(formatter)
logger.addHandler(fh)

logging.captureWarnings(True)
logging.getLogger("py.warnings").addHandler(fh)

output_file(filename=snakemake.output[0], title="Rank Comparison")

pval_df = pd.read_table(snakemake.input[0], sep="\t", index_col=0)
rank_df = (
    pval_df.rename(columns={f"{x}": f"{x}_pvalue" for x in pval_df.columns})
    .reset_index()
)

tool_list = pval_df.columns.tolist()
tool_rank_list = [x + "_pvalue" for x in tool_list]

# Rank Comparison
# logger.info("Creating rank comparison panel...")
chosen_tool_1 = Select(options=tool_rank_list, value=tool_rank_list[0], title="x-axis")
chosen_tool_2 = Select(options=tool_rank_list, value=tool_rank_list[1], title="y-axis")

rank_df["x"] = rank_df[chosen_tool_1.value]
rank_df["y"] = rank_df[chosen_tool_2.value]

source = ColumnDataSource(rank_df)

hover = HoverTool(mode="mouse", tooltips=[("Name", "points")], attachment="below")

hover.tooltips = (
    [("Feature ID", "@feature_id")] +
    [(f"{x}", f"@{x}") for x in rank_df.columns if x not in ["feature_id", "x", "y"]]
)

tools = [
    hover,
    "box_zoom",
    "reset",
    "pan",
    "save"
]

# plot = figure(tools=tools)
plot = figure(x_range=(0, 1), y_range=(0, 1), tools=tools)

plot.circle(
    source=source,
    x="x",
    y="y",
    color="lightblue",
    line_width=0.5,
    line_color="black",
    size=10,
    name="points"
)
plot.line(
    x=np.arange(0, rank_df.shape[0]),
    y=np.arange(0, rank_df.shape[0]),
    line_width=3,
    line_dash="dashed",
    color="black"
)

tool_callback = CustomJS(args=dict(source=source, plot=plot), code="""
const data = source.data[cb_obj.value];
if (cb_obj.title == 'x-axis') {
    plot.below[0].axis_label = cb_obj.value;
    source.data.x = data;
} else {
    plot.left[0].axis_label = cb_obj.value;
    source.data.y = data;
}

source.change.emit();
""")
chosen_tool_1.js_on_change("value", tool_callback)
chosen_tool_2.js_on_change("value", tool_callback)

plot.title.text = "P-value Comparison"
plot.title.text_font_size = "20pt"

plot.xaxis.axis_label = chosen_tool_1.value
plot.yaxis.axis_label = chosen_tool_2.value

for ax in [plot.xaxis, plot.yaxis]:
    ax.axis_label_text_font_size = "14pt"
    ax.axis_label_text_font_style = "normal"
    ax.major_label_text_font_size = "10pt"

layout = row(column(chosen_tool_1, chosen_tool_2), plot)
save(layout)

logger.info(f"Saved rank comparisons to {snakemake.output[0]}")

import logging
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from bokeh.plotting import figure, save
from bokeh.models import HoverTool, ColumnDataSource
from bokeh.io import output_file

# Add logging 
logger = logging.getLogger("qadabra")
logger.setLevel(logging.INFO)
fh = logging.FileHandler(snakemake.log[0], mode="w")
formatter = logging.Formatter(
    f"[%(asctime)s - {snakemake.rule}] :: %(message)s"
)
fh.setFormatter(formatter)
logger.addHandler(fh)

logging.captureWarnings(True)
logging.getLogger("py.warnings").addHandler(fh)

plt.style.use(snakemake.config["stylesheet"])

# Read in tsv containing p-values and LFCs
logger.info("Loading p-values and LFCs...")
data = pd.read_table(snakemake.input[0], sep="\t", index_col=0).squeeze()
pval_col = snakemake.params["pval_col"]
diff_ab_col = snakemake.params["diff_ab_col"]

# Take negative log of pval_col column and add as new column
data['negative_log_pval'] = -np.log10(data[pval_col])

# Create a ColumnDataSource object
source_blue = ColumnDataSource(data=dict(
    diff_ab_col=data[diff_ab_col],
    pval=data['negative_log_pval'],
    feature=data.index
))

source_red = ColumnDataSource(data=dict(
    diff_ab_col=data[diff_ab_col],
    pval=data['negative_log_pval'],
    feature=data.index
))

# Create a figure
tool_name = snakemake.wildcards["tool"]
p = figure(title=f"{tool_name} Volcano Plot", x_axis_label='coefficient', y_axis_label='-log10(pval_col)')

p.title.text_font_size = '20pt'

# Define the desired y-axis range
y_min = data['negative_log_pval'].min() - 1  # Minimum value
y_max = data['negative_log_pval'].max() + 1  # Maximum value

# Set the y-axis range
p.y_range.start = y_min
p.y_range.end = y_max

# Define the desired x-axis range
x_min = data[diff_ab_col].min() - 1  # Minimum value
x_max = data[diff_ab_col].max() + 1  # Maximum value

# Set the y-axis range
p.x_range.start = x_min
p.x_range.end = x_max

# # Add a vertical line at x = -1
p.line(x=[-1, -1], y=[y_min, y_max], line_color='gray', line_width=1, line_dash='dashed')
p.line(x=[1, 1], y=[y_min, y_max], line_color='gray', line_width=1, line_dash='dashed')
p.line(x=[x_min,x_max], y=[-np.log10(0.05), -np.log10(0.05)], line_color='gray', line_width=1, line_dash='dashed')

# Add hover tool
hover = HoverTool(
    tooltips=[
        ('Feature', '@feature'),
        ('coefficient', '@diff_ab_col'),
        ('-log10(pval_col)', '@pval'),
    ]
)
p.add_tools(hover)

# Select points where y >= 1 and x >= -log10(0.05)
pos_sig_diff = (data[diff_ab_col] >= 1) & (data['negative_log_pval'] >= -np.log10(0.05))
# Select points where y <= -1 and x >= -log10(0.05)
neg_sig_diff = (data[diff_ab_col] <= -1) & (data['negative_log_pval'] >= -np.log10(0.05))

# Apply selection to the data source
source_blue.selected.indices = np.where(neg_sig_diff)[0]
# Apply selection to the data source for negative significant differences (blue)
p.circle('diff_ab_col', 'pval', size=5, fill_alpha=0.6, nonselection_line_color='black', nonselection_fill_color='black', source=source_blue, selection_color='blue')

# Apply selection to the data source
source_red.selected.indices = np.where(pos_sig_diff)[0]
# Apply selection to the data source for positive significant differences (red)
p.circle('diff_ab_col', 'pval', size=5, fill_alpha=0.6, nonselection_line_color='black', nonselection_fill_color='black', source=source_red, selection_color='red')


output_file(snakemake.output[0])
save(p)

logger.info(f"Saved volocano plot to {snakemake.output[0]}")