统计中介与调节效应路径图 (Mediation & Moderation Path)

#!/usr/bin/env python3
"""FigureCraft Chart Engine: c18_mediation_path (统计中介效应与调节效应路径模型图)."""

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
import matplotlib.patches as patches

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


def render(
    data_path: Optional[str] = None,
    style: str = "thesis_cn",
    output_dir: str = "output",
    formats: Optional[List[str]] = None,
    dpi: int = 600,
) -> Dict[str, str]:
    if formats is None:
        formats = ["svg", "pdf", "png", "tiff"]

    apply_style(style)
    palette = get_palette(style)

    if data_path is None:
        data_path = str(Path(__file__).resolve().parents[1] / "data" / f"{CHART_ID}.csv")

    df = pd.read_csv(data_path)
    os.makedirs(output_dir, exist_ok=True)

    fig, ax = plt.subplots(figsize=(7.0, 4.0), dpi=dpi)
    ax.set_xlim(0, 10)
    ax.set_ylim(0, 6)
    ax.axis('off')

    nodes = {
        'X': {'pos': (1.5, 1.8), 'label': '自变量 (X)\n组织支持感', 'color': palette[0]},
        'M': {'pos': (5.0, 4.2), 'label': '中介变量 (M)\n心理安全感', 'color': palette[1]},
        'Y': {'pos': (8.5, 1.8), 'label': '因变量 (Y)\n员工创新绩效', 'color': palette[2]},
    }

    box_w, box_h = 2.4, 1.2
    for k, n in nodes.items():
        x, y = n['pos']
        rect = patches.FancyBboxPatch(
            (x - box_w / 2, y - box_h / 2), box_w, box_h,
            boxstyle="round,pad=0.1,rounding_size=0.2",
            facecolor='white', edgecolor=n['color'], linewidth=1.8, zorder=3
        )
        ax.add_patch(rect)
        ax.text(x, y, n['label'], ha='center', va='center', fontsize=8.5, fontweight='bold', zorder=4, color='#111111')

    path_a = df[df['path_name'] == 'Path_a'].iloc[0] if len(df[df['path_name'] == 'Path_a']) else None
    path_b = df[df['path_name'] == 'Path_b'].iloc[0] if len(df[df['path_name'] == 'Path_b']) else None
    path_c = df[df['path_name'] == 'Path_c_direct'].iloc[0] if len(df[df['path_name'] == 'Path_c_direct']) else None

    # Path a (X -> M)
    ax.annotate(
        "", xy=(3.8, 3.8), xytext=(2.4, 2.5),
        arrowprops=dict(arrowstyle="->,head_width=0.35,head_length=0.5", color='#333333', lw=1.2)
    )
    if path_a is not None:
        ax.text(2.7, 3.4, f"a = {path_a['beta']:.3f}***\n(SE={path_a['std_err']:.3f})",
                ha='center', va='center', fontsize=8.0, fontweight='semibold', color=palette[0])

    # Path b (M -> Y)
    ax.annotate(
        "", xy=(7.6, 2.5), xytext=(6.2, 3.8),
        arrowprops=dict(arrowstyle="->,head_width=0.35,head_length=0.5", color='#333333', lw=1.2)
    )
    if path_b is not None:
        ax.text(7.3, 3.4, f"b = {path_b['beta']:.3f}***\n(SE={path_b['std_err']:.3f})",
                ha='center', va='center', fontsize=8.0, fontweight='semibold', color=palette[1])

    # Path c direct (X -> Y)
    ax.annotate(
        "", xy=(7.2, 1.8), xytext=(2.8, 1.8),
        arrowprops=dict(arrowstyle="->,head_width=0.35,head_length=0.5", color='#333333', lw=1.2, linestyle='--')
    )
    if path_c is not None:
        ax.text(5.0, 1.35, f"直接效应 c' = {path_c['beta']:.3f}* (p = {path_c['p_val']:.3f})\n总效应 c = 0.357***, 间接效应 ab = 0.165***",
                ha='center', va='top', fontsize=8.0, fontweight='semibold', color='#222222',
                bbox=dict(boxstyle='square,pad=0.3', facecolor='#F5F5F5', edgecolor='#CCCCCC', lw=0.5))

    ax.set_title('结构方程与统计中介效应路径模型分析', y=1.03, fontsize=10.5, fontweight='bold')
    plt.tight_layout()

    out_paths = {}
    base_name = f"{CHART_ID}_{style}"
    for fmt in formats:
        p = os.path.join(output_dir, f"{base_name}.{fmt}")
        fig.savefig(p, format=fmt, dpi=dpi, bbox_inches='tight')
        out_paths[fmt] = p

    plt.close(fig)
    return out_paths


def main():
    parser = argparse.ArgumentParser(description="Render Mediation & Moderation Path Diagram")
    parser.add_argument("--data", type=str, default=None)
    parser.add_argument("--style", type=str, default="thesis_cn")
    parser.add_argument("--output-dir", type=str, default="output")
    parser.add_argument("--formats", type=str, default="svg,pdf,png,tiff")
    parser.add_argument("--dpi", type=int, default=600)
    args = parser.parse_args()

    fmts = [f.strip() for f in args.formats.split(",")]
    render(args.data, args.style, args.output_dir, fmts, args.dpi)


if __name__ == "__main__":
    main()