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, multiplier=5, col="median_value",  
    res_df=None, res_name_col="measurement", res_p_col="p_mwu"
):
    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")

    # 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()}
        
    parts = []
    for m in measurements:
        sub = d[d["measurement"] == m].copy()
        sub = apply_plausible_range(sub, m)
        ctrl = clean_measurement(sub[sub["group"] == "Controls"], col=col, multiplier=multiplier)
        case = sub[sub["group"] == "Cases"].copy()  # don't IQR-clean cases
        parts.append(pd.concat([ctrl, case], ignore_index=True))
    p = pd.concat(parts, ignore_index=True)

    fig, axes = plt.subplots(2, 3, figsize=(4, 3))
    axes = axes.flatten()

    for i, m in enumerate(measurements):
        ax = axes[i]
        sub = p[p["measurement"] == m].copy()
        sub[col] = pd.to_numeric(sub[col], errors="coerce")
        sub = sub.dropna(subset=[col])
        sns.violinplot(data=sub, x="group", y=col, order=["Controls", "Cases"],
                       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

            y_line = ymax + 0.07 * span
            y_star = y_line + 0.02 * span
            y_p = y_star - 0.03 * span  # right below stars

            if not pd.isna(pval):
                ax.text(0.5, y_p, fmt_p(pval), ha="center", va="top")

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

    plt.tight_layout()
    return fig, axes, p
