import pandas as pd
import numpy as np

# plotting imports 
import textwrap
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 p_to_stars(p):
    if pd.isna(p): return ""
    if p < 1e-3: return "***"
    if p < 1e-2: return "**"
    if p < 5e-2: return "*"
    return ""

def fmt_p(p):
    if pd.isna(p): return ""
    return f"P={p:.2e}" if p < 0.001 else f"P={p:.3f}"

def plot_measurements(
    df_long,
    measurements,
    col="median_value",
    res_df=None,
    res_name_col="measurement",
    res_p_col="pvalue",
    ncols=3,
    figsize=(6, 4),
    palette=None
    ):
    d = df_long[df_long["measurement"].isin(measurements)].copy()
    d[col] = pd.to_numeric(d[col], errors="coerce")
    d = d.dropna(subset=[col])
    d["group"] = np.where(d["cases"], "Cases", "Controls")

    if palette is None:
        palette = {
            "Controls": "#2F6690",  # BLUE
            "Cases":    "#D1495B",  # RED
        }
    # star map
    p_map, star_map = {}, {}
    if res_df is not None:
        tmp = res_df.copy()
        tmp[res_p_col] = pd.to_numeric(tmp[res_p_col], errors="coerce")
        p_map = dict(zip(tmp[res_name_col], tmp[res_p_col]))
        star_map = {k: p_to_stars(v) for k, v in p_map.items()}
        
    # Layout: dynamic grid
    n = len(measurements)
    nrows = int(np.ceil(n / ncols))
    fig, axes = plt.subplots(nrows, ncols, figsize=figsize)
    axes = np.array(axes).ravel()

    for i, m in enumerate(measurements):
        ax = axes[i]
        sub = d[d["measurement"] == m].copy()

        sns.violinplot(
            data=sub, x="group", y=col, order=["Controls", "Cases"],
            hue="group", hue_order=["Controls", "Cases"], legend=False,
            palette=palette, inner="quartile", cut=0, linewidth=1, ax=ax
            )

        ax.set_title(textwrap.fill(m, width=20,  max_lines=3))
        ax.set_xlabel("")
        ax.set_ylabel("")
        ax.spines[['top', 'right']].set_visible(False)
        
        stars = star_map.get(m, "")
        pval = p_map.get(m, np.nan)

        if stars or (not pd.isna(pval)):
            ymax = np.nanmax(sub[col].values)
            ymin = np.nanmin(sub[col].values)
            span = (ymax - ymin) if ymax > ymin else 1.0

            # put annotation above violins
            y_txt = ymax + 0.08 * span
            
            # x=0.5 is between the two categories (0 and 1)
            if stars:
                ax.text(0.5, y_txt + 0.03 * span, stars, ha="center", va="bottom")
            if not pd.isna(pval):
                ax.text(0.5, y_txt, fmt_p(pval), ha="center", va="bottom")

            # give headroom so text doesn't clip
            ax.set_ylim(ymin, y_txt + 0.12 * span)

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

    plt.tight_layout()
    return fig, axes
