
from pathlib import Path
import csv
import sys
import numpy as np
import matplotlib.pyplot as plt

PKG = Path(__file__).resolve().parent.parent
IMG_DIR = PKG / "图片"
COMPOSITE_DIR = IMG_DIR / "组合图"
SUBPLOT_DIR = IMG_DIR / "独立子图"
TABLE_DIR = PKG / "表格"
IMG_DIR.mkdir(parents=True, exist_ok=True)
COMPOSITE_DIR.mkdir(parents=True, exist_ok=True)
SUBPLOT_DIR.mkdir(parents=True, exist_ok=True)
TABLE_DIR.mkdir(parents=True, exist_ok=True)

sys.path.insert(0, r"C:\Users\User\.codex\skills\scientific-paper-figures\assets")
from nature_figure_helpers import setup_nature_style
setup_nature_style()

plt.rcParams.update({
    "font.family": "sans-serif",
    "font.sans-serif": ["Microsoft YaHei", "Arial", "DejaVu Sans"],
    "axes.unicode_minus": False,
    "text.color": "#111111",
    "axes.labelcolor": "#111111",
    "xtick.color": "#111111",
    "ytick.color": "#111111",
    "axes.edgecolor": "#333333",
    "legend.frameon": False,
    "axes.titlesize": 8.5,
    "axes.labelsize": 8,
    "xtick.labelsize": 7,
    "ytick.labelsize": 7,
    "legend.fontsize": 7,
})

def half_frame(ax, grid_axis="y"):
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.tick_params(which="both", top=False, right=False, direction="in")
    ax.grid(axis=grid_axis, color="#E5E7EB", linewidth=0.55, zorder=0)
    ax.set_axisbelow(True)

def save_all(fig, stem):
    fig.savefig(stem.with_suffix(".png"), dpi=301, bbox_inches="tight", pad_inches=0.04, facecolor="white")
    fig.savefig(stem.with_suffix(".pdf"), bbox_inches="tight", pad_inches=0.04, facecolor="white")
    fig.savefig(stem.with_suffix(".tif"), dpi=600, bbox_inches="tight", pad_inches=0.04,
                facecolor="white", pil_kwargs={"compression": "tiff_lzw"})
    plt.close(fig)

# ---------------- data ----------------
agent_names = ["FinRobot", "Dexter", "TradingAgents"]
agent_metrics = ["Macro F1", "第一层 F1", "第二层 F1", "第三层 F1"]
agent_values = np.array([
    [4.99, 10.93, 4.36, 1.07],
    [8.46, 12.16, 4.28, 8.42],
    [5.01, 10.54, 3.30, 2.08],
])
agent_colors = ["#008DB8", "#00A878", "#FF7A00"]
agent_hatches = ["///", "\\\\\\", "xx"]

methods = ["Baseline", "Prompt", "QLoRA", "LoRA"]
class_metrics = ["等级准确率", "平衡准确率", "Macro F1"]
class_values = np.array([
    [50.00, 45.83, 41.87],
    [50.00, 45.83, 43.06],
    [83.33, 79.17, 80.32],
    [83.33, 79.17, 80.32],
])
mae = np.array([11.06, 7.26, 4.71, 4.39])
rmse = np.array([12.11, 8.79, 7.08, 6.92])
method_colors = ["#6B7280", "#F2B701", "#F2385A", "#7B2CBF"]
method_hatches = ["", "///", "\\\\\\", "xx"]

years = np.array([2023, 2024, 2025, 2026])
companies = [
    ("迪瑞医疗", "300396", [97.99, 70.00, 44.00, 16.67], (12.29, 21.06), "#F2385A", "o"),
    ("华铭智能", "300462", [74.00, 84.00, 70.00, 72.00], (67.61, 76.39), "#008DB8", "s"),
    ("汇金科技", "300561", [94.04, 70.00, 70.00, 53.97], (49.59, 58.36), "#00A878", "^"),
]

# CSV files
with (TABLE_DIR / "表1_金融智能体结果.csv").open("w", newline="", encoding="utf-8-sig") as f:
    w = csv.writer(f)
    w.writerow(["方法", "Macro F1", "第一层F1", "第二层F1", "第三层F1", "格式合规率", "标签覆盖率", "平均调用次数", "平均耗时（秒/条）", "总成本（美元）", "失败率"])
    w.writerows([
        ["FinRobot", 0.0499, 0.1093, 0.0436, 0.0107, 1.00, 1.00, 3.0, 103.3, 9.05, 0.00],
        ["Dexter", 0.0846, 0.1216, 0.0428, 0.0842, 1.00, 0.9992, 3.0, 109.0, 9.71, 0.00],
        ["TradingAgents", 0.0501, 0.1054, 0.0330, 0.0208, 0.96, 1.00, 4.0, 117.6, 10.72, 0.00],
    ])
with (TABLE_DIR / "表2_四模型比较.csv").open("w", newline="", encoding="utf-8-sig") as f:
    w = csv.writer(f)
    w.writerow(["方法", "解析率", "等级准确率", "平衡准确率", "Macro F1", "MAE", "RMSE"])
    w.writerows([
        ["Baseline", 1.00, 0.5000, 0.4583, 0.4187, 11.06, 12.11],
        ["Prompt", 1.00, 0.5000, 0.4583, 0.4306, 7.26, 8.79],
        ["QLoRA", 1.00, 0.8333, 0.7917, 0.8032, 4.71, 7.08],
        ["LoRA", 1.00, 0.8333, 0.7917, 0.8032, 4.39, 6.92],
    ])
with (TABLE_DIR / "表3_三家公司预测.csv").open("w", newline="", encoding="utf-8-sig") as f:
    w = csv.writer(f)
    w.writerow(["企业", "股票代码", "2023", "2024", "2025", "2026预测", "2026下限", "2026上限", "趋势判断"])
    w.writerows([
        ["迪瑞医疗", "300396", 97.99, 70.00, 44.00, 16.67, 12.29, 21.06, "持续下行，需优先复核"],
        ["华铭智能", "300462", 74.00, 84.00, 70.00, 72.00, 67.61, 76.39, "总体稳定，维持关注"],
        ["汇金科技", "300561", 94.04, 70.00, 70.00, 53.97, 49.59, 58.36, "趋势转弱，建议跟踪"],
    ])

# ---------------- Figure 1: agents ----------------
fig, ax = plt.subplots(figsize=(7.1, 3.05), dpi=301)
x = np.arange(len(agent_metrics)); width = 0.23
for i, (name, vals, color, hatch) in enumerate(zip(agent_names, agent_values, agent_colors, agent_hatches)):
    ax.bar(x + (i - 1) * width, vals, width, label=name, color=color,
           edgecolor="#222222", linewidth=0.45, hatch=hatch, zorder=3)
ax.set_xticks(x, agent_metrics); ax.set_ylabel("F1（%）")
ax.set_ylim(0, 14); ax.set_yticks(np.arange(0, 15, 2))
ax.legend(ncol=3, loc="upper center", bbox_to_anchor=(0.5, 1.12),
          handlelength=1.5, columnspacing=1.6)
half_frame(ax)
save_all(fig, COMPOSITE_DIR / "图1_金融智能体F1比较")

# 图1本身为单面板，同时作为第1张独立子图另存，便于单独引用。
fig, ax = plt.subplots(figsize=(7.1, 3.05), dpi=301)
x = np.arange(len(agent_metrics)); width = 0.23
for i, (name, vals, color, hatch) in enumerate(zip(agent_names, agent_values, agent_colors, agent_hatches)):
    ax.bar(x + (i - 1) * width, vals, width, label=name, color=color,
           edgecolor="#222222", linewidth=0.45, hatch=hatch, zorder=3)
ax.set_xticks(x, agent_metrics); ax.set_ylabel("F1（%）")
ax.set_ylim(0, 14); ax.set_yticks(np.arange(0, 15, 2))
ax.legend(ncol=3, loc="upper center", bbox_to_anchor=(0.5, 1.12),
          handlelength=1.5, columnspacing=1.6)
half_frame(ax)
save_all(fig, SUBPLOT_DIR / "独立子图1_金融智能体F1比较")

# ---------------- Figure 2: model combination ----------------
def draw_classification(ax):
    x = np.arange(len(class_metrics)); width = 0.19
    for i, (name, vals, color, hatch) in enumerate(zip(methods, class_values, method_colors, method_hatches)):
        ax.bar(x + (i - 1.5) * width, vals, width, label=name, color=color,
               edgecolor="#222222", linewidth=0.4, hatch=hatch, zorder=3)
    ax.set_xticks(x, class_metrics); ax.set_ylabel("指标值（%）"); ax.set_ylim(0, 100)
    ax.legend(ncol=2, loc="upper left", handlelength=1.35, columnspacing=1.2)
    half_frame(ax)

def draw_error(ax):
    y = np.arange(len(methods))
    for i, color in enumerate(method_colors):
        ax.plot([mae[i], rmse[i]], [y[i], y[i]], color=color, linewidth=1.4, zorder=2)
        ax.scatter(mae[i], y[i], s=28, color=color, edgecolor="#222222", linewidth=0.45, marker="o", zorder=3)
        ax.scatter(rmse[i], y[i], s=28, facecolor="white", edgecolor=color, linewidth=1.1, marker="s", zorder=3)
    ax.set_yticks(y, methods); ax.invert_yaxis(); ax.set_xlabel("CFHS分数误差（分）"); ax.set_xlim(3.5, 12.8)
    ax.plot([], [], "o", color="#444444", label="MAE")
    ax.plot([], [], "s", markerfacecolor="white", markeredgecolor="#444444", label="RMSE")
    ax.legend(ncol=2, loc="upper center", bbox_to_anchor=(0.5, 1.11), handletextpad=0.4)
    half_frame(ax)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(7.1, 3.25), dpi=301,
                               gridspec_kw={"width_ratios": [1.32, 1], "wspace": 0.34})
draw_classification(ax1); draw_error(ax2)
save_all(fig, COMPOSITE_DIR / "图2_四模型比较_组合")

fig, ax = plt.subplots(figsize=(4.6, 3.25), dpi=301)
draw_classification(ax)
save_all(fig, SUBPLOT_DIR / "独立子图2_四模型分类指标比较")

fig, ax = plt.subplots(figsize=(4.3, 3.25), dpi=301)
draw_error(ax)
save_all(fig, SUBPLOT_DIR / "独立子图3_四模型CFHS分数误差比较")

# ---------------- Figure 3: company combination and subplots ----------------
def draw_company(ax, company, show_ylabel=True, show_legend=False):
    name, code, vals, interval, color, marker = company
    vals = np.asarray(vals, dtype=float)
    ax.plot(years[:3], vals[:3], color=color, linewidth=1.45, marker=marker,
            markersize=4.0, markeredgecolor="#222222", markeredgewidth=0.4,
            label="模型回溯评分", zorder=3)
    ax.plot(years[2:], vals[2:], color=color, linewidth=1.25, linestyle="--",
            marker=marker, markersize=4.6, markerfacecolor="white",
            markeredgecolor=color, markeredgewidth=1.0, label="2026趋势预测", zorder=3)
    ax.fill_between([2025.82, 2026.18], [interval[0], interval[0]],
                    [interval[1], interval[1]], color=color, alpha=0.18, edgecolor="none", zorder=1)
    ax.vlines(2026, interval[0], interval[1], color=color, linewidth=0.9, zorder=2)
    ax.hlines([interval[0], interval[1]], 2025.94, 2026.06, color=color, linewidth=0.9, zorder=2)
    ax.axvline(2025.5, color="#9CA3AF", linewidth=0.7, linestyle=":", zorder=0)
    ax.text(0.5, 1.06, f"{name}\n({code})", transform=ax.transAxes, ha="center", va="bottom",
            fontsize=8.2, fontweight="bold", color="#111111")
    ax.set_xticks(years); ax.set_xlim(2022.75, 2026.25); ax.set_ylim(0, 105)
    if show_ylabel: ax.set_ylabel("CFHS评分（分）")
    if show_legend: ax.legend(loc="lower left", ncol=1, handlelength=1.6)
    half_frame(ax)

fig, axes = plt.subplots(1, 3, figsize=(7.1, 2.95), dpi=301, sharey=True)
for i, (ax, company) in enumerate(zip(axes, companies)):
    draw_company(ax, company, show_ylabel=(i == 0), show_legend=(i == 0))
axes[1].set_xlabel("年度")
save_all(fig, COMPOSITE_DIR / "图3_三家公司预测_组合")

for idx, company in enumerate(companies, start=1):
    fig, ax = plt.subplots(figsize=(4.25, 3.25), dpi=301)
    draw_company(ax, company, show_ylabel=True, show_legend=True)
    ax.set_xlabel("年度")
    save_all(fig, SUBPLOT_DIR / f"独立子图{idx + 3}_{company[0]}")

print(IMG_DIR)
