High-Density Ridge / Joyplot

#!/usr/bin/env python3
"""FigureCraft Chart Engine: High-Density Ridge / Joyplot (c03_ridge_joyplot)."""

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 = "c03_ridge_joyplot"


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=["cohort", "value"])

    cohorts = list(df["cohort"].unique())
    k = len(cohorts)

    # Compute global x range
    x_min = df["value"].min() - (df["value"].max() - df["value"].min()) * 0.1
    x_max = df["value"].max() + (df["value"].max() - df["value"].min()) * 0.1
    x_eval = np.linspace(x_min, x_max, 250)

    fig, axes = plt.subplots(
        nrows=k,
        ncols=1,
        figsize=(6.5, 4.8),
        sharex=True,
        gridspec_kw={"hspace": -0.45},
    )

    if k == 1:
        axes = [axes]

    for i, (ax, cohort) in enumerate(zip(axes, cohorts)):
        cohort_vals = df[df["cohort"] == cohort]["value"].values
        color = palette[i % len(palette)]

        if len(cohort_vals) >= 2:
            kde = stats.gaussian_kde(cohort_vals, bw_method="silverman")
            y_density = kde(x_eval)
        else:
            y_density = np.zeros_like(x_eval)

        # Baseline white mask to prevent bleed-through
        ax.fill_between(x_eval, 0, y_density, color="#FFFFFF", alpha=1.0, zorder=2)
        # Colored KDE fill
        ax.fill_between(x_eval, 0, y_density, color=color, alpha=0.75, zorder=3)
        # Top boundary line
        ax.plot(x_eval, y_density, color=color, linewidth=1.2, zorder=4)

        # Baseline reference line
        ax.axhline(0, color="#D1D5DB", linewidth=0.6, zorder=1)

        # Left label for cohort
        ax.text(
            x_min,
            np.max(y_density) * 0.25 if np.max(y_density) > 0 else 0.05,
            cohort.replace("_", " "),
            fontsize=7.5,
            fontweight="bold",
            color="#1F2937",
            ha="right",
            va="bottom",
        )

        # Clean background & hide unnecessary spines
        ax.patch.set_alpha(0.0)
        ax.set_ylim(bottom=0)
        ax.set_yticks([])
        for spine in ["top", "right", "left", "bottom"]:
            ax.spines[spine].set_visible(False)

    # Enable bottom spine and x-ticks on bottom axis
    axes[-1].spines["bottom"].set_visible(True)
    axes[-1].spines["bottom"].set_color("#111827")
    axes[-1].spines["bottom"].set_linewidth(0.5)
    axes[-1].tick_params(bottom=True, labelbottom=True, labelsize=7.5)
    axes[-1].set_xlabel("Quantitative Response Metric", fontsize=8.0)

    fig.suptitle("High-Density Ridge Probability Distribution Across Cohorts", fontsize=9.5, fontweight="bold", y=0.98)

    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 High-Density Ridge / Joyplot")
    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}")