"""dahlin wt vs kit-w41 lineage composition and log2fc."""
from __future__ import annotations
import warnings, sys
warnings.filterwarnings("ignore")
from pathlib import Path
import numpy as np, pandas as pd
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt

sys.path.insert(0, str(Path(__file__).parent))
from palette import apply_style, color_for, GENOTYPE_COLORS
apply_style()

# figure font sizes (style spec): small titles, larger body text
SUPTITLE_FS = 15   # figure suptitle
TITLE_FS = 14      # panel/axes titles
LABEL_FS = 12      # axis labels
TICK_FS = 11       # tick labels
LEGEND_FS = 10     # legend text
ANNOT_FS = 10      # in-plot annotations
ANNOT_EM_FS = 11   # emphasized in-plot annotations
CBAR_FS = 12       # colorbar labels
PANEL_FS = 12      # bold panel letters
plt.rcParams.update({
    "figure.titlesize": SUPTITLE_FS,
    "axes.titlesize": TITLE_FS,
    "axes.labelsize": LABEL_FS,
    "xtick.labelsize": TICK_FS,
    "ytick.labelsize": TICK_FS,
    "legend.fontsize": LEGEND_FS,
    "legend.title_fontsize": LEGEND_FS + 1,
})

import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
ROOT = Path(str(PANDA_ROOT))
CSV = ROOT / "discovery/hematopoiesis/marker/dahlin_predictions.csv"
OUT = ROOT / "figures/biology/biology_05_dahlin_composition.pdf"

GT_MAP = {"SIGAB1":"WT","SIGAC1":"WT","SIGAD1":"WT","SIGAF1":"WT",
          "SIGAG1":"WT","SIGAH1":"WT","SIGAG8":"Kit_W41","SIGAH8":"Kit_W41"}

LINEAGE_ORDER = ["LT-HSC", "MPP", "lymphoid", "erythroid", "megakaryocyte",
                 "myeloid", "monocyte", "macrophage", "basophil-mast"]


def main():
    df = pd.read_csv(CSV)
    df["sample"] = df["cell_id"].str.split("_").str[0]
    df["genotype"] = df["sample"].map(GT_MAP).fillna("other")

    df = df[df["genotype"].isin(["WT", "Kit_W41"])]
    tab = pd.crosstab(df["pred_label"], df["genotype"]).reindex(LINEAGE_ORDER).fillna(0)
    print("[fig5] counts:")
    print(tab)

    frac = tab.div(tab.sum(axis=0), axis=1)  # columns are genotypes
    log2fc = np.log2((frac["Kit_W41"] + 1e-4) / (frac["WT"] + 1e-4))

    fig, axes = plt.subplots(1, 2, figsize=(17, 7.2), gridspec_kw={"width_ratios": [1.1, 1]})

    ax = axes[0]
    genos = ["WT", "Kit_W41"]
    bottoms = np.zeros(len(genos))
    for lin in LINEAGE_ORDER:
        vals = np.array([frac[g].loc[lin] for g in genos])
        ax.bar(genos, vals, bottom=bottoms,
               color=color_for(lin), edgecolor="white",
               linewidth=0.7, label=lin)
        for j, v in enumerate(vals):
            if v > 0.035:
                ax.text(j, bottoms[j] + v / 2, f"{v*100:.1f}%",
                        ha="center", va="center", fontsize=ANNOT_FS, color="white",
                        fontweight="bold")
        bottoms += vals
    ax.set_ylabel("Within-genotype fraction")
    ax.set_ylim(0, 1)
    ax.set_title(f"Lineage composition — WT (n={int(tab['WT'].sum())}) vs "
                 f"Kit-W41 (n={int(tab['Kit_W41'].sum())})\n"
                 f"9 largest lineages shown; 91 cells in 4 minor classes omitted "
                 f"(pro-B 66, T-cell 14, endothelial 6, naive-B 5)", fontsize=TITLE_FS)
    ax.legend(loc="center left", bbox_to_anchor=(1.02, 0.5),
              frameon=False, ncol=1, title="Lineage")
    ax.text(-0.12, 1.05, "(a)", transform=ax.transAxes,
            fontsize=PANEL_FS, fontweight="bold", va="bottom", ha="right")

    ax = axes[1]
    lins = LINEAGE_ORDER
    vals = [log2fc.loc[l] for l in lins]
    colors = [GENOTYPE_COLORS["Kit_W41"] if v > 0 else GENOTYPE_COLORS["WT"] for v in vals]
    y = np.arange(len(lins))
    ax.barh(y, vals, color=colors, edgecolor="black", linewidth=0.6)
    ax.axvline(0, color="black", lw=0.8)
    # value labels always sit just to the RIGHT of the bar tip regardless of
    # sign, so negative-bar labels never crash into the y-tick labels on the
    # left. Positive bars: label extends outward to the right. Negative bars:
    # label sits inside the axis (toward zero) on the right of the tip.
    for i, (l, v) in enumerate(zip(lins, vals)):
        n_kit = int(tab["Kit_W41"].loc[l]); n_wt = int(tab["WT"].loc[l])
        ax.annotate(
            f"{v:+.2f}  (Kit={n_kit}, WT={n_wt})",
            xy=(v, i), xycoords="data",
            xytext=(6, 0), textcoords="offset points",
            ha="left", va="center", fontsize=ANNOT_EM_FS, clip_on=False,
        )
    ax.set_yticks(y); ax.set_yticklabels(lins)
    ax.tick_params(axis="y", pad=6)
    ax.invert_yaxis()
    ax.set_xlabel("log2(Kit-W41 fraction / WT fraction)")
    ax.set_title("Compositional shift — Kit-W41 expands erythroid, contracts lymphoid")
    max_abs = max(abs(min(vals)), abs(max(vals)))
    # Right-side padding accommodates the value+n text at annotation size
    ax.set_xlim(-max_abs * 1.15, max_abs * 2.8)
    ax.text(-0.14, 1.05, "(b)", transform=ax.transAxes,
            fontsize=PANEL_FS, fontweight="bold", va="bottom", ha="right")

    plt.suptitle("Dahlin — Kit-W41 alters lineage output balance",
                 fontsize=SUPTITLE_FS, y=1.02, fontweight="bold")
    plt.tight_layout()
    plt.savefig(OUT, bbox_inches="tight", dpi=180)
    plt.close()
    print(f"[fig5] wrote {OUT}")


if __name__ == "__main__":
    main()