import pandas as pd
import os
from tqdm import tqdm
import numpy as np
import ast
# plotting imports 
import textwrap
import seaborn as sns
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator
import matplotlib.patches as mpatches
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})
from matplotlib.colors import TwoSlopeNorm
from matplotlib.backends.backend_pdf import PdfPages


def percent_bar_with_n(
    df, group_col, cat_col, cat_order=None, 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 given categories if present else by control pct desc
    if cat_order is not None:
        present_cats = d[cat_col].unique().tolist()
        cat_order = [c for c in cat_order if c in present_cats]          # keep only those that exist
        cat_order += [c for c in present_cats if c not in cat_order]     # append anything else (Unknown/Other/etc.)
    else:
        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()
        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)", pad=20)
    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

# Create a single survey figure
def plot_survey_questions(
    df,
    surveys=None,
    group_col="cases",
    cat_col="answer_category",      # or "answer"
    survey_col="survey",
    question_col="question",
    control_label="Controls",
    case_label="Cases",
    top_n=10,
    figsize=(8, 4),
    palette=None):

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

    d = df.copy()

    # optional: restrict to chosen surveys
    if surveys is not None:
        d = d[d[survey_col].isin(surveys)].copy()

    # map True/False -> nice labels (so your function's hue_order works)
    d[group_col] = np.where(d[group_col].astype(bool), case_label, control_label)

    groups = list(d.groupby([survey_col, question_col], dropna=False))
    n = len(groups)
    ncols = 2
    nrows = int(np.ceil(n / ncols))
    
    fig, axes = plt.subplots(nrows, ncols, figsize=figsize)
    axes = axes.ravel()
    handles = labels = None

    for i, ((svy, q), sub) in enumerate(groups):
        percent_bar_with_n(
            sub,
            group_col=group_col,
            cat_col=cat_col,
            cat_order=["Poor","Fair","Good","Very Good","Excellent"],
            title=f"{q}",
            top_n=top_n,
            control_label=control_label,
            case_label=case_label,
            ax=axes[i],
            show=False,          # << key
            palette=palette
        )

        # remove per-axis legends; we'll add one global legend
        if handles is None:
            handles, labels = axes[i].get_legend_handles_labels()
        leg = axes[i].get_legend()
        if leg is not None:
            leg.remove()

    # hide unused axes
    for j in range(n, len(axes)):
        axes[j].axis("off")

    # global legend from first axis handles
    if handles:
        fig.legend(handles, labels, loc="upper right", frameon=False)

    fig.tight_layout()
    return fig, axes
