"""Render the MuleGuard system architecture + data-flow diagrams to PNG."""
from __future__ import annotations

import matplotlib
matplotlib.use("Agg")
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, FancyBboxPatch

from src import config

NAVY, BLUE, TEAL, RED, GREY = "#1f4e79", "#2e86c1", "#16a085", "#c0392b", "#7f8c8d"


def _box(ax, x, y, w, h, text, color, tcolor="white", fs=10):
    ax.add_patch(FancyBboxPatch((x, y), w, h, boxstyle="round,pad=0.02,rounding_size=0.08",
                                fc=color, ec="white", lw=1.5))
    ax.text(x + w / 2, y + h / 2, text, ha="center", va="center",
            color=tcolor, fontsize=fs, weight="bold", wrap=True)


def _arrow(ax, x1, y1, x2, y2, color=GREY):
    ax.add_patch(FancyArrowPatch((x1, y1), (x2, y2), arrowstyle="-|>",
                                 mutation_scale=16, color=color, lw=1.8))


def architecture() -> None:
    fig, ax = plt.subplots(figsize=(12, 6.8))
    ax.set_xlim(0, 12); ax.set_ylim(0, 7); ax.axis("off")
    ax.text(6, 6.7, "MuleGuard — System Architecture", ha="center", fontsize=15, weight="bold", color=NAVY)

    # Ingestion
    _box(ax, 0.3, 4.6, 2.4, 1.6,
         "INGESTION\n\n• Financial txns\n• FMS / TMS alerts\n• Govt cyber-fraud\n  tickets (I4C/NCRP)\n• Cross-channel data", NAVY, fs=8.5)
    # Feature pipeline
    _box(ax, 3.2, 4.7, 2.3, 1.4,
         "FEATURE PIPELINE\n\nclean · impute ·\none-hot · select\n+ anomaly score", BLUE, fs=9)
    # Model
    _box(ax, 6.0, 4.7, 2.3, 1.4,
         "ML MODEL\n\nLightGBM +\nIsolationForest\ncalibrated · SHAP", TEAL, fs=9)
    # Artifacts
    _box(ax, 6.0, 2.7, 2.3, 1.2, "ARTIFACTS\nmodel · threshold\nfeatures · explainer", GREY, fs=8.5)
    # Scoring API
    _box(ax, 8.9, 4.7, 2.7, 1.4,
         "SCORING API (FastAPI)\n\n/score → risk 0-100,\ntier, decision,\nreason codes, alert", RED, fs=9)
    # Simulator
    _box(ax, 3.2, 2.6, 2.3, 1.3, "FEED SIMULATOR\nstreams accounts\n& tickets → API", BLUE, fs=9)
    # Dashboard
    _box(ax, 8.9, 2.5, 2.7, 1.5,
         "ANALYST CONSOLE (Streamlit)\nalert queue · drill-down\nexplanations · model perf", NAVY, fs=8.5)
    # Outcome
    _box(ax, 4.4, 0.5, 3.2, 1.1,
         "OUTCOME: block circulation of\nfraudulent proceeds via mule accounts", RED, fs=9.5)

    _arrow(ax, 2.7, 5.4, 3.2, 5.4)
    _arrow(ax, 5.5, 5.4, 6.0, 5.4)
    _arrow(ax, 8.3, 5.4, 8.9, 5.4)
    _arrow(ax, 7.15, 4.7, 7.15, 3.9)          # model -> artifacts
    _arrow(ax, 5.5, 3.25, 6.0, 3.25, BLUE)    # simulator -> artifacts/api path
    _arrow(ax, 4.35, 4.7, 4.35, 3.9, GREY)    # pipeline <-> simulator
    _arrow(ax, 10.25, 4.7, 10.25, 4.0, RED)   # api -> dashboard
    _arrow(ax, 10.25, 2.5, 7.6, 1.6, NAVY)    # dashboard -> outcome
    _arrow(ax, 8.9, 3.25, 8.3, 5.0, GREY)     # artifacts -> api (load)

    fig.tight_layout()
    out = config.ROOT / "docs" / "architecture.png"
    fig.savefig(out, dpi=140, bbox_inches="tight"); plt.close(fig)
    print("Wrote", out)


def pipeline_flow() -> None:
    fig, ax = plt.subplots(figsize=(12, 2.6))
    ax.set_xlim(0, 12); ax.set_ylim(0, 2); ax.axis("off")
    steps = [("Raw\n3,924 feats", NAVY), ("Drop leakage\n(F3912, F2230)", RED),
             ("Clean +\nencode", BLUE), ("Impute +\nanomaly score", TEAL),
             ("CV feature\nselection → 103", BLUE), ("LightGBM +\ncalibration", TEAL),
             ("Threshold\ntuning (F2)", GREY), ("Risk 0-100\n+ SHAP", NAVY)]
    w = 1.32; gap = (12 - len(steps) * w) / (len(steps) + 1)
    x = gap
    for i, (t, c) in enumerate(steps):
        _box(ax, x, 0.6, w, 0.9, t, c, fs=7.5)
        if i < len(steps) - 1:
            _arrow(ax, x + w, 1.05, x + w + gap, 1.05)
        x += w + gap
    ax.text(6, 1.8, "Modeling Pipeline", ha="center", fontsize=13, weight="bold", color=NAVY)
    fig.tight_layout()
    out = config.ROOT / "docs" / "pipeline_flow.png"
    fig.savefig(out, dpi=140, bbox_inches="tight"); plt.close(fig)
    print("Wrote", out)


if __name__ == "__main__":
    architecture()
    pipeline_flow()