#!/usr/bin/env python3
"""Combined figure: (a) Motif Preservation (A1 + A2 stacked) and (b) JS heatmap, in one row."""

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl
import seaborn as sns
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec

# ── Data ──
df_summary = pd.read_csv(
    "/data/n5huang/dna_token/tokenizer_evaluation/motif_eval_outputs_new/summary_by_tokenizer.csv"
)
df_js = pd.read_csv(
    "/data/n5huang/dna_token/enrichment_analysis/js_vs_random_matrix.csv",
    index_col=0,
)

# ── Method config (shared) ──
METHODS = ["DNAbert2", "Grover", "cCRE_region_BPE", "motif_region_BPE", "DNAMotifTokenizer"]
DISPLAY_NAMES = {
    "DNAbert2": "DNABERT-2",
    "Grover": "Grover",
    "cCRE_region_BPE": "cCRE-Region BPE",
    "motif_region_BPE": "Motif-Region BPE",
    "DNAMotifTokenizer": "DNAMotif",
}
METHOD_COLORS = {
    "DNAbert2": "#6C5CE7",          # purple (from C)
    "Grover": "#E74C3C",            # red (from T)
    "cCRE_region_BPE": "#52BE80",   # green (from A)
    "motif_region_BPE": "#F4D03F",  # yellow (Token Characterization)
    "DNAMotifTokenizer": "#E67E22", # orange (Fine-tuning Ability)
}
HEATMAP_REGIONS = ["promoter", "enhancer", "exon"]
JS_INDEX_DISPLAY = [DISPLAY_NAMES.get(n, n) for n in df_js.index]

# ── Plot style ──
mpl.rcParams.update({
    "font.family": "sans-serif",
    "font.size": 58,
    "axes.labelsize": 64,
    "axes.titlesize": 66,
    "axes.titleweight": "bold",
    "xtick.labelsize": 58,
    "ytick.labelsize": 58,
})


def get_value(method, source_df, col_name):
    row = source_df[source_df["Tokenizer"] == method]
    if len(row) == 0:
        return float("nan")
    return row[col_name].values[0]


def draw_bar_panel(ax, title, ylabel, col_name, higher_better, title_pad=30, ylabel_y=0.5):
    vals = [get_value(m, df_summary, col_name) for m in METHODS]
    colors = [METHOD_COLORS[m] for m in METHODS]
    x_labels = [DISPLAY_NAMES[m] for m in METHODS]
    x = np.arange(len(METHODS)) * 0.3
    bars = ax.bar(x, vals, color=colors, edgecolor="white", linewidth=0.5, alpha=0.85, width=0.2)

    for i, v in enumerate(vals):
        if np.isnan(v):
            continue
        y_off = max(vals) * 0.02
        fmt = f"{v:.2f}" if abs(v) >= 1 else f"{v:.4f}"
        ax.text(x[i], v + y_off, fmt, ha="center", va="bottom", fontsize=62, fontweight="bold")

    ax.set_xticks(x)
    ax.set_xticklabels(x_labels, rotation=15, ha="right")
    ax.set_xlim(x[0] - 0.2, x[-1] + 0.2)
    ax.set_ylabel(ylabel)
    # Pin ylabel position in axes coords so A1 and A2 line up regardless of tick-label widths.
    ax.yaxis.set_label_coords(-0.08, ylabel_y)
    ax.set_title(title, pad=title_pad)
    ax.tick_params(axis="x", pad=10)
    ax.tick_params(axis="y", pad=15)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.grid(axis="y", alpha=0.3, linewidth=0.8)

    valid = [v for v in vals if not np.isnan(v)]
    if valid:
        best = (np.nanargmax if higher_better else np.nanargmin)(vals)
        bars[best].set_edgecolor("black")
        bars[best].set_linewidth(2)


# ── Figure layout: 1 row, 2 top-level columns ──
fig = plt.figure(figsize=(50, 26))
gs = GridSpec(1, 2, figure=fig, width_ratios=[1, 0.9], wspace=0.30)

# Left column (panel a): 2 stacked subplots
gs_a = GridSpecFromSubplotSpec(2, 1, subplot_spec=gs[0, 0], hspace=0.8)
ax_a1 = fig.add_subplot(gs_a[0, 0])
ax_a2 = fig.add_subplot(gs_a[1, 0])

# Right column (panel b): heatmap
ax_b = fig.add_subplot(gs[0, 1])

# ── Draw panel a (A1 + A2) ──
draw_bar_panel(ax_a1, "Perfect Match Rate", "Perfect match rate (%)",
               "Perfect_Match_Rate", higher_better=True, title_pad=80, ylabel_y=0.35)
draw_bar_panel(ax_a2, "Motif Fragmentation", "Avg tokens per motif",
               "Avg_Tokens_Per_Motif", higher_better=False, title_pad=80)

# ── Draw panel b (heatmap) ──
df_disp = df_js.copy()
df_disp.index = JS_INDEX_DISPLAY
cmap = LinearSegmentedColormap.from_list("evo_heat", ["white", "#AED6F1", "#5DADE2", "#2E86C1"])
sns.heatmap(
    df_disp,
    ax=ax_b,
    annot=True,
    fmt=".4f",
    annot_kws={"size": 62, "fontweight": "bold"},  # matches bar-value annotation size
    cmap=cmap,
    vmin=0.0,
    vmax=float(df_disp.values.max()) or 0.001,
    linewidths=0.5,
    cbar_kws={"label": "Jensen-Shannon distance"},
)
ax_b.set_title("Region vs random-genome separation",
               fontsize=66, fontweight="bold", pad=80)
ax_b.set_xlabel("Genomic region", fontsize=64, labelpad=15)
ax_b.set_ylabel("")
ax_b.tick_params(axis="x", labelsize=58)
ax_b.tick_params(axis="y", labelsize=58, rotation=0)
cbar = ax_b.collections[0].colorbar
cbar.ax.tick_params(labelsize=58)
cbar.set_label("Jensen-Shannon distance", size=64)

fig.subplots_adjust(left=0.06, right=0.98, top=0.88, bottom=0.10)

# ── Panel labels (a), (b) on same line as titles ──
# Use axes-relative y so they line up with titles (which sit just above each axes top).
ax_a1.text(-0.08, 1.234, "(a)", transform=ax_a1.transAxes,
           fontsize=66, fontweight="bold", va="center", ha="center")
ax_b.text(-0.30, 1.083, "(b)", transform=ax_b.transAxes,
          fontsize=66, fontweight="bold", va="center", ha="center")

out_dir = "/data/n5huang/dna_token/paper_figures"
os.makedirs(out_dir, exist_ok=True)
out = os.path.join(out_dir, "combined_panel_a_b.pdf")
fig.savefig(out, dpi=300)
fig.savefig(out.replace(".pdf", ".png"), dpi=80)
print(f"Saved: {out}")
print(f"Saved: {out.replace('.pdf', '.png')}")
plt.close(fig)