#!/usr/bin/env python3
"""Supplementary Fig. 2a: perfect match rate (left) and motif fragmentation (right)."""

import os
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.ticker import PercentFormatter

df = pd.read_csv(
    "results/summary_by_tokenizer.csv"
).set_index("Tokenizer")

METHODS = ["DNAbert2", "Grover", "cCRE_region_BPE", "motif_region_BPE", "DNAMotifTokenizer"]
LABELS = [
    "BPE(orig. DNABERT-2)",
    "BPE(orig. Grover)",
    "BPE(cCRE-guided)",
    "BPE(motif-guided)",
    "DNAMotifTokenizer",
]
COLORS = ["#54aa8c", "#ce7f51", "#cca47f", "#c590bd", "#e6bed6"]

mpl.rcParams.update({"font.family": "sans-serif", "font.size": 13})

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5.2))


def draw(ax, col, title, ylabel, fmt):
    vals = df.loc[METHODS, col].values
    ax.bar(range(len(vals)), vals, color=COLORS, edgecolor="#555555", linewidth=0.6, width=0.65)
    for i, v in enumerate(vals):
        ax.text(i, v + max(vals) * 0.015, fmt(v), ha="center", va="bottom", fontsize=13)
    ax.set_xticks(range(len(vals)))
    ax.set_xticklabels(LABELS, rotation=45, ha="right")
    ax.set_title(title, fontsize=14, pad=14)
    ax.set_ylabel(ylabel)
    ax.set_ylim(0, max(vals) * 1.12)
    for s in ["top", "right"]:
        ax.spines[s].set_visible(False)
    ax.spines["left"].set_color("#bbbbbb")
    ax.spines["bottom"].set_color("#bbbbbb")
    ax.grid(axis="y", color="#eeeeee", linewidth=0.8)
    ax.set_axisbelow(True)


# Perfect match rate is a percentage
draw(ax1, "Perfect_Match_Rate", "Perfect Match Rate", "",
     lambda v: f"{v:.2f}" if v >= 1 else f"{v:.4f}")
ax1.yaxis.set_major_formatter(PercentFormatter(decimals=0))

# Fragmentation is a token count, not a percentage
draw(ax2, "Avg_Tokens_Per_Motif", "Motif Fragmentation", "Avg. tokens per motif",
     lambda v: f"{v:.2f}")

fig.text(0.01, 0.97, "a", fontsize=15, fontweight="bold", va="top")
fig.tight_layout(rect=(0.02, 0, 1, 0.97), w_pad=4)

out_dir = "figures"
os.makedirs(out_dir, exist_ok=True)
out = os.path.join(out_dir, "supp_fig2a_motif.pdf")
fig.savefig(out, dpi=600)
fig.savefig(out.replace(".pdf", ".png"), dpi=300)
print(f"Saved: {out}")
plt.close(fig)