import pandas as pd
# plotting imports 
import seaborn as sns
import matplotlib
import matplotlib.pyplot as plt
matplotlib.rcParams['pdf.fonttype'] = 42
matplotlib.rcParams['ps.fonttype'] = 42
# matplotlib.rcParams['font.sans-serif'] = "Arial" # missing fonts:: https://alexanderlabwhoi.github.io/post/2021-03-missingfont/
# Then, "ALWAYS use sans-serif fonts"
matplotlib.rcParams['font.family'] = "sans-serif"
matplotlib.rcParams.update({'font.size': 6, 'axes.linewidth': 1, 'xtick.major.width': 1, 'xtick.major.size': 5, 'ytick.major.width': 1, 'ytick.major.size': 5})


def percent_bar_with_n(df, group_col, cat_col, title=None, top_n=None,
                       control_label="Controls", case_label="Cases",
                       figsize=(8, 6),
                       ax=None, show=True,
                       palette=None):
    
    d = df[[group_col, cat_col]].copy()
    d[group_col] = (d[group_col]
                    .fillna("Unknown")
                    .astype(str)
                    .replace({"": "Unknown", "nan": "Unknown", "None": "Unknown"}))

    d[cat_col] = (d[cat_col]
                  .fillna("Unknown")
                  .astype(str)
                  .replace({"": "Unknown", "nan": "Unknown", "None": "Unknown"}))
    # Optional: collapse rare categories (global frequency)
    if top_n is not None:
        top = d[cat_col].value_counts().head(top_n).index
        d[cat_col] = d[cat_col].where(d[cat_col].isin(top), "Other")
    # counts + within-group percent
    ct = d.groupby([group_col, cat_col]).size().reset_index(name="n")
    ct["pct"] = ct["n"] / ct.groupby(group_col)["n"].transform("sum") * 100
    # enforce hue order
    present = ct[group_col].unique().tolist()
    hue_order = [control_label, case_label]
    if control_label not in present or case_label not in present:
        hue_order = sorted(present)
    # category order: sort by CONTROL pct desc
    ctrl = ct[ct[group_col] == hue_order[0]]
    if ctrl.empty:
        cat_order = d[cat_col].value_counts().index.tolist()
    else:
        cat_order = ctrl.sort_values("pct", ascending=False)[cat_col].tolist()
    # append categories missing from controls (keep at end)
    cat_order += [c for c in d[cat_col].unique().tolist() if c not in cat_order]
    # FORCE categorical order BEFORE plotting
    ct[cat_col] = pd.Categorical(ct[cat_col], categories=cat_order, ordered=True)
    ct[group_col] = pd.Categorical(ct[group_col], categories=hue_order, ordered=True)
    ct = ct.sort_values([cat_col, group_col])
    # plotting
    if ax is None:
        fig = plt.figure(figsize=figsize)
        ax = plt.gca()
    else:
        fig = ax.figure

    sns.barplot(
        data=ct,
        x=cat_col, y="pct",
        hue=group_col,
        gap=0.1,
        order=cat_order,
        hue_order=hue_order,
        estimator=sum,
        errorbar=None,
        ax=ax,
        palette=palette
    )

    ax.set_title(title or f"{cat_col} (% within group)")
    ax.set_xlabel("")
    ax.set_ylabel("Percent (%)")
    ax.tick_params(axis="x", rotation=25)
    for lbl in ax.get_xticklabels():
        lbl.set_ha("right")

    # annotate n
    for container, grp in zip(ax.containers, hue_order):
        ns = (ct[ct[group_col] == grp]
              .set_index(cat_col)
              .reindex(cat_order)["n"]
              .fillna(0)
              .astype(int)
              .tolist())

        labels = [f"{n}" if rect.get_height() > 0 else "" for rect, n in zip(container, ns)]
        if hasattr(ax, "bar_label"):
            ax.bar_label(container, labels=labels, padding=3, rotation=90)
        else:
            for rect, lab in zip(container, labels):
                if not lab:
                    continue
                h = rect.get_height()
                ax.text(rect.get_x() + rect.get_width()/2, h, lab,
                        ha="center", va="bottom", rotation=90)

    ax.spines[['top', 'right']].set_visible(False)

    if ax.get_legend() is not None:
        ax.get_legend().set_title("")

    if show and ax is not None and ax.figure is not None and ax.figure.get_axes()[0] == ax:
        plt.tight_layout()
        plt.show()

    return fig, ax


def demographics_plot(person_df,
                        group_col="group",
                        control_label="Controls",
                        case_label="Cases",
                        sex_col="sex_at_birth",
                        ethnicity_col="ethnicity",
                        ancestry_col="ancestry_pred_other",
                        age_col="age_at_cdr",
                        top_n_sex=2,
                        top_n_ethnicity=2,
                        top_n_ancestry=None,
                        figsize=(4, 5),
                        palette=None,
                        share_legend=True,
                        savepath=None,
                        dpi=300):

    if palette is None:
        palette = {
            control_label: "#2F6690",  # BLUE
            case_label:    "#D1495B",  # RED
        }

    fig = plt.figure(figsize=figsize)
    gs = fig.add_gridspec(nrows=2, ncols=4, wspace=1, hspace=0.5)

    ax_age = fig.add_subplot(gs[0, 0:1])
    ax_ancestry      = fig.add_subplot(gs[0, 1:4])
    ax_sex      = fig.add_subplot(gs[1, 0:2])
    ax_eth      = fig.add_subplot(gs[1, 2:4])

    # Top row
    percent_bar_with_n(
        person_df, group_col=group_col, cat_col=ancestry_col,
        title="Ancestry", top_n=top_n_ancestry,
        control_label=control_label, case_label=case_label,
        ax=ax_ancestry, show=False, palette=palette
    )

    sns.violinplot(
        data=person_df, y=age_col, hue=group_col,
        ax=ax_age, palette=palette,# cut=0, inner="quartile"
    )
    ax_age.set_title("Age")
    ax_age.set_xlabel("")
    ax_age.set_ylabel("Age")
    ax_age.spines[['top', 'right']].set_visible(False)
    if ax_age.get_legend() is not None:
        ax_age.get_legend().set_title("")

    # Bottom row
    percent_bar_with_n(
        person_df, group_col=group_col, cat_col=sex_col,
        title="Sex at Birth", top_n=top_n_sex,
        control_label=control_label, case_label=case_label,
        ax=ax_sex, show=False, palette=palette
    )

    percent_bar_with_n(
        person_df, group_col=group_col, cat_col=ethnicity_col,
        title="Ethnicity", top_n=top_n_ethnicity,
        control_label=control_label, case_label=case_label,
        ax=ax_eth, show=False, palette=palette
    )

    axes = [ax_ancestry, ax_age, ax_sex, ax_eth]

    # One shared legend
    if share_legend:
        handles_labels = None
        for ax in axes:
            leg = ax.get_legend()
            if leg is not None and handles_labels is None:
                handles_labels = ax.get_legend_handles_labels()
            if leg is not None:
                leg.remove()

        if handles_labels is not None:
            handles, labels = handles_labels
            fig.legend(handles, labels, loc="upper center",
                       ncol=len(labels), frameon=False, bbox_to_anchor=(0.5, 1.0))

    fig.tight_layout()

    if savepath is not None:
        fig.savefig(savepath, dpi=dpi, bbox_inches="tight")

    return fig, {
        "ancestry": ax_ancestry,
        "age": ax_age,
        "sex_at_birth": ax_sex,
        "ethnicity": ax_eth
    }
