Bar Plot with Error Bars & Jittered Raw Data Points

#!/usr/bin/env python3
"""FigureCraft Chart Engine: Bar Plot with Error Bars & Jittered Raw Data (c04_bar_jitter)."""

import os
import sys
import argparse
from pathlib import Path
from typing import List, Dict, Optional
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy import stats

try:
    from engines.styles import apply_style, get_palette
except ImportError:
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from engines.styles import apply_style, get_palette

CHART_ID = "c04_bar_jitter"


def get_pvalue_asterisks(p_val: float) -> str:
    if p_val < 0.0001:
        return "****"
    elif p_val < 0.001:
        return "***"
    elif p_val < 0.01:
        return "**"
    elif p_val < 0.05:
        return "*"
    else:
        return "ns"


def draw_significance_bracket(ax, x1: float, x2: float, y: float, h: float, text: str, color: str = "#111827"):
    ax.plot([x1, x1, x2, x2], [y, y + h, y + h, y], lw=0.8, c=color)
    ax.text((x1 + x2) * 0.5, y + h * 1.1, text, ha='center', va='bottom', color=color, fontsize=7.5, fontweight='bold')


def render(
    data_path: Optional[str] = None,
    style: str = "nature",
    output_dir: str = "output",
    formats: Optional[List[str]] = None,
    dpi: int = 600,
) -> Dict[str, str]:
    if formats is None:
        formats = ["svg", "pdf", "png", "tiff"]
    if data_path is None:
        data_path = str(Path(__file__).resolve().parents[1] / "data" / f"{CHART_ID}.csv")

    apply_style(style)
    palette = get_palette(style)
    df = pd.read_csv(data_path)
    df = df.dropna(subset=["group", "value"])

    groups = list(df["group"].unique())
    k = len(groups)

    fig, ax = plt.subplots(figsize=(6.0, 4.5))

    means = []
    sems = []
    max_val = df["value"].max()

    for i, grp in enumerate(groups):
        grp_vals = df[df["group"] == grp]["value"].values
        mean = np.mean(grp_vals)
        sem = stats.sem(grp_vals) if len(grp_vals) > 1 else 0.0
        means.append(mean)
        sems.append(sem)
        color = palette[i % len(palette)]

        # Draw Bar
        ax.bar(
            i,
            mean,
            yerr=sem,
            capsize=4.0,
            color=color,
            alpha=0.65,
            edgecolor=color,
            linewidth=1.0,
            width=0.55,
            zorder=2,
            error_kw=dict(lw=1.0, capthick=1.0, ecolor="#111827"),
        )

        # Draw jittered raw points
        np.random.seed(101 + i)
        jitter = np.random.uniform(-0.16, 0.16, size=len(grp_vals))
        ax.scatter(
            i + jitter,
            grp_vals,
            s=28,
            color=color,
            alpha=0.85,
            edgecolors="#FFFFFF",
            linewidth=0.6,
            zorder=4,
        )

    # Statistical significance brackets across comparisons
    y_max_error = max([m + s for m, s in zip(means, sems)] + [max_val])
    curr_bracket_y = y_max_error * 1.08
    bracket_h = y_max_error * 0.03

    # Test baseline (group 0) against all others
    for i in range(1, k):
        g0 = df[df["group"] == groups[0]]["value"].values
        gi = df[df["group"] == groups[i]]["value"].values
        if len(g0) >= 3 and len(gi) >= 3:
            _, p_val = stats.ttest_ind(g0, gi, equal_var=False)
            stars = get_pvalue_asterisks(p_val)
            draw_significance_bracket(ax, 0, i, curr_bracket_y, bracket_h, stars)
            curr_bracket_y += y_max_error * 0.10

    ax.set_xticks(range(k))
    formatted_groups = [g.replace("_", " ") for g in groups]
    ax.set_xticklabels(formatted_groups)
    ax.set_ylabel("Quantified Signal Intensity (Mean ± SEM)")
    ax.set_xlabel("Experimental Group")
    ax.set_title("Group Comparison: Mean ± SEM with Individual Replicates")
    ax.set_ylim(bottom=0, top=curr_bracket_y + y_max_error * 0.05)

    os.makedirs(output_dir, exist_ok=True)
    generated_files = {}
    for fmt in formats:
        out_path = os.path.join(output_dir, f"{CHART_ID}_{style}.{fmt}")
        fig.savefig(out_path, format=fmt, dpi=dpi, bbox_inches="tight")
        generated_files[fmt] = str(Path(out_path).resolve())

    plt.close(fig)
    return generated_files


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Render Bar Plot with Error Bars & Jitter")
    parser.add_argument("--data", type=str, default=None, help="Path to input CSV dataset")
    parser.add_argument("--style", type=str, default="nature", choices=["nature", "cell", "lancet", "ieee", "economist"], help="Journal style pack")
    parser.add_argument("--output-dir", type=str, default="output", help="Output directory")
    parser.add_argument("--formats", type=str, default="svg,pdf,png,tiff", help="Comma-separated export formats")
    parser.add_argument("--dpi", type=int, default=600, help="Raster DPI")
    args = parser.parse_args()

    fmt_list = [f.strip() for f in args.formats.split(",") if f.strip()]
    results = render(args.data, style=args.style, output_dir=args.output_dir, formats=fmt_list, dpi=args.dpi)
    for fmt, path in results.items():
        print(f"Generated {fmt.upper()}: {path}")