Scatter Plot with Marginal Histograms / KDE

#!/usr/bin/env python3
"""FigureCraft Chart Engine: Scatter Plot with Marginal Histograms / KDE (c09_scatter_marginal)."""

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


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=["x_val", "y_val"])

    has_group = "group" in df.columns
    groups = list(df["group"].unique()) if has_group else ["All"]

    fig = plt.figure(figsize=(6.5, 6.0))
    gs = fig.add_gridspec(
        nrows=2,
        ncols=2,
        width_ratios=[4.5, 1.2],
        height_ratios=[1.2, 4.5],
        hspace=0.08,
        wspace=0.08,
    )

    ax_marg_x = fig.add_subplot(gs[0, 0])
    ax_main = fig.add_subplot(gs[1, 0], sharex=ax_marg_x)
    ax_marg_y = fig.add_subplot(gs[1, 1], sharey=ax_main)

    # Main scatter & group regressions
    for i, grp in enumerate(groups):
        color = palette[i % len(palette)]
        sub_df = df[df["group"] == grp] if has_group else df
        x = sub_df["x_val"].values
        y = sub_df["y_val"].values

        # Scatter
        ax_main.scatter(
            x,
            y,
            c=color,
            label=grp.replace("_", " "),
            alpha=0.75,
            s=32,
            edgecolors="#FFFFFF",
            linewidth=0.5,
            zorder=3,
        )

        # Marginal X (Top KDE)
        if len(x) >= 3:
            kde_x = stats.gaussian_kde(x)
            x_eval = np.linspace(df["x_val"].min(), df["x_val"].max(), 100)
            ax_marg_x.plot(x_eval, kde_x(x_eval), color=color, lw=1.2)
            ax_marg_x.fill_between(x_eval, 0, kde_x(x_eval), color=color, alpha=0.3)

        # Marginal Y (Right KDE)
        if len(y) >= 3:
            kde_y = stats.gaussian_kde(y)
            y_eval = np.linspace(df["y_val"].min(), df["y_val"].max(), 100)
            ax_marg_y.plot(kde_y(y_eval), y_eval, color=color, lw=1.2)
            ax_marg_y.fill_betweenx(y_eval, 0, kde_y(y_eval), color=color, alpha=0.3)

    # Overall linear regression fit across all data
    x_all = df["x_val"].values
    y_all = df["y_val"].values
    if len(x_all) >= 3:
        slope, intercept, r_val, p_val, std_err = stats.linregress(x_all, y_all)
        x_line = np.linspace(x_all.min(), x_all.max(), 100)
        y_line = slope * x_line + intercept
        ax_main.plot(x_line, y_line, color="#1F2937", linestyle="--", linewidth=1.4, zorder=2)

        # 95% Confidence band for regression
        n = len(x_all)
        t_crit = stats.t.ppf(0.975, df=n - 2)
        s_err = np.sqrt(np.sum((y_all - (slope * x_all + intercept)) ** 2) / (n - 2))
        x_mean = np.mean(x_all)
        ci = t_crit * s_err * np.sqrt(1 / n + (x_line - x_mean) ** 2 / np.sum((x_all - x_mean) ** 2))
        ax_main.fill_between(x_line, y_line - ci, y_line + ci, color="#9CA3AF", alpha=0.25, zorder=1)

        stat_text = f"Linear Fit: $r = {r_val:.2f}$\n$p = {p_val:.3e}$" if p_val < 0.001 else f"Linear Fit: $r = {r_val:.2f}$\n$p = {p_val:.3f}$"
        ax_main.text(
            0.05,
            0.92,
            stat_text,
            transform=ax_main.transAxes,
            fontsize=7.5,
            va="top",
            bbox=dict(boxstyle="round,pad=0.3", facecolor="#F3F4F6", edgecolor="#D1D5DB", alpha=0.85),
        )

    ax_main.set_xlabel("Independent Predictor ($X$)")
    ax_main.set_ylabel("Dependent Response ($Y$)")
    ax_main.legend(loc="lower right", frameon=False, fontsize=7.5)

    # Clean marginal axes
    ax_marg_x.axis("off")
    ax_marg_y.axis("off")

    fig.suptitle("Bivariate Association with Marginal Kernel Densities", 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 Scatter with Marginals")
    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}")