"""Fig 2: two-stage stacking framework (schematic). Pure matplotlib patches."""
from pathlib import Path
import sys

import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch

sys.path.insert(0, str(Path(__file__).resolve().parent))
from plot_style import apply, save, PALETTE_DEEP as C  # noqa: E402

apply()
ROOT = Path(__file__).resolve().parents[2]
FIG = ROOT / "reports" / "figures"

fig, ax = plt.subplots(figsize=(13, 5))
ax.set_xlim(0, 13)
ax.set_ylim(0, 5)
ax.axis("off")


def box(x, y, w, h, text, color, fc_alpha=0.18, fs=10):
    p = FancyBboxPatch((x, y), w, h, boxstyle="round,pad=0.08,rounding_size=0.15",
                       fc=color, ec=color, alpha=fc_alpha, lw=2)
    ax.add_patch(p)
    ax.text(x + w / 2, y + h / 2, text, ha="center", va="center", fontsize=fs, fontweight="bold")


def arrow(x1, y1, x2, y2):
    ax.add_patch(FancyArrowPatch((x1, y1), (x2, y2), arrowstyle="-|>", mutation_scale=18, lw=2, color="black"))


# Stage 1: score producers
ax.text(1.25, 4.7, "Stage 1: Score producers", ha="center", fontsize=11, fontweight="bold")
for i, (t, col) in enumerate([("LightGCN", C[0]), ("BPR-MF", C[1]), ("DeepWalk", C[2]), ("Node2Vec", C[3])]):
    box(0.2, 3.4 - i * 0.8, 2.1, 0.6, t, col, fs=9.5)

# Feature engineering
ax.text(5.9, 4.7, "Feature engineering (259-d)", ha="center", fontsize=11, fontweight="bold")
for i, (t, col) in enumerate([("higher-order prop. (69)", C[3]), ("random-walk blocks (88)", C[2]),
                              ("explicit graph / meta-path (18)", C[0]), ("content-rich + BPR + ranks (84)", C[1])]):
    box(4.6, 3.4 - i * 0.8, 2.6, 0.6, t, col, fs=8.6)

# Stage 2
ax.text(9.9, 4.7, "Stage 2 / Decision", ha="center", fontsize=11, fontweight="bold")
box(8.3, 3.0, 3.2, 1.0, "LightGBM\nmeta-learner", C[4], fs=10)
box(8.3, 1.4, 3.2, 1.0, "rank-cutoff decision\n(top-50% + known pos.)", C[5], fs=9)

# arrows
arrow(2.3, 3.1, 4.6, 3.1)
arrow(7.2, 3.5, 8.3, 3.5)
arrow(9.9, 3.0, 9.9, 2.4)

ax.set_title("Two-stage stacking framework", fontsize=13)
save(fig, "fig2_framework", FIG)
print("saved fig2_framework")