# -*- coding: utf-8 -*-
from __future__ import annotations
import csv, json, math
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Arc, Circle, FancyArrowPatch, Polygon, Rectangle, Wedge

ROOT = Path(__file__).resolve().parents[1]
FIG = ROOT / "docs" / "figures"
FIG.mkdir(parents=True, exist_ok=True)
INK = "#234C70"
GRAY = "#676767"
LIGHT = "#9DB7CA"
PALE = "#D9E4EC"
GOLD = "#B58B3A"
GREEN = "#4E7A55"
plt.rcParams.update({"font.sans-serif": ["DejaVu Sans"], "axes.unicode_minus": False,
    "figure.dpi": 120, "savefig.dpi": 300, "axes.edgecolor": GRAY,
    "axes.labelcolor": "#222222", "text.color": "#222222"})

def save(fig, name):
    fig.savefig(FIG / name, bbox_inches="tight", facecolor="white")
    plt.close(fig)

def load_formal():
    rows = []
    path = ROOT / "docs" / "paper" / "data" / "formal-results.csv"
    with path.open(encoding="utf-8-sig") as f:
        for row in csv.DictReader(f):
            for key in ("problem_no", "formal_index", "cleared_jammer_count",
                        "measure_accepted_count", "channel_switch_count",
                        "clear_failure_count"):
                row[key] = int(row[key])
            for key in ("virtual_time_s", "program_run_duration_s",
                        "inferred_movement_time_s"):
                row[key] = float(row[key])
            rows.append(row)
    return rows

def fig_formal_total():
    rows = load_formal()
    labels = ["Q%d-%d" % (r["problem_no"], r["formal_index"]) for r in rows]
    total = [r["virtual_time_s"] for r in rows]
    runtime = [r["program_run_duration_s"] for r in rows]
    colors = [INK, LIGHT, PALE, GOLD, "#C98963", GREEN]
    fig, axes = plt.subplots(1, 2, figsize=(10.5, 3.6))
    for ax, vals, title, unit in ((axes[0], total, "Total virtual time", "s"),
                                  (axes[1], runtime, "Program runtime", "s")):
        bars = ax.barh(labels[::-1], vals[::-1], color=colors[::-1])
        ax.set_xlabel(title + " / " + unit)
        ax.grid(axis="x", color="#E5E5E5", linewidth=0.7)
        ax.set_axisbelow(True)
        for bar, val in zip(bars, vals[::-1]):
            ax.text(bar.get_width() + max(vals) * 0.015,
                    bar.get_y() + bar.get_height() / 2,
                    "%.3f" % val, va="center", fontsize=8)
        ax.margins(x=0.18)
    fig.suptitle("Direct metrics of three formal tests", fontsize=11)
    save(fig, "formal-total-time.png")


def fig_formal_time_bars():
    rows = load_formal()
    labels = ["Q%d-%d" % (r["problem_no"], r["formal_index"]) for r in rows]
    move = np.array([r["inferred_movement_time_s"] for r in rows])
    measure = np.array([5.0 * r["measure_accepted_count"] for r in rows])
    switch = np.array([1.0 * r["channel_switch_count"] for r in rows])
    clear = np.array([5.0 * r["cleared_jammer_count"] + 3.0 * r["clear_failure_count"] for r in rows])
    fig, ax = plt.subplots(figsize=(9.0, 3.8))
    ax.barh(labels, move, color=INK, label="Move")
    ax.barh(labels, measure, left=move, color=LIGHT, label="Measure")
    ax.barh(labels, switch, left=move + measure, color=GRAY, label="Switch")
    ax.barh(labels, clear, left=move + measure + switch, color=GOLD, label="Clear")
    ax.set_xlabel("Cumulative virtual time / s")
    ax.grid(axis="x", color="#E5E5E5", linewidth=0.7)
    ax.set_axisbelow(True)
    ax.legend(loc="lower right", frameon=False, ncol=4)
    save(fig, "formal-time-bars.png")

def box(ax, x, y, w, h, text, fc=PALE, ec=INK):
    ax.add_patch(Rectangle((x, y), w, h, facecolor=fc, edgecolor=ec, linewidth=1.0))
    ax.text(x + w/2, y + h/2, text, ha="center", va="center", fontsize=8, wrap=True)


def arrow(ax, p, q, color=GRAY):
    ax.add_patch(FancyArrowPatch(p, q, arrowstyle="-|>", mutation_scale=10,
                                 linewidth=1.0, color=color))


def fig_flow_overall():
    fig, ax = plt.subplots(figsize=(8.0, 4.8))
    ax.axis("off")
    box(ax, 0.5, 4.0, 7.0, 0.8, "L0  Problem abstraction: area, motion, timing, channels", fc="#EFEFEF")
    box(ax, 0.5, 3.0, 2.1, 0.7, "L1  Evidence\ncandidate sets")
    box(ax, 2.9, 3.0, 2.1, 0.7, "L1  Decision\ntask tables")
    box(ax, 5.3, 3.0, 2.1, 0.7, "L1  Execution\nactions")
    box(ax, 0.5, 2.0, 1.6, 0.7, "L2  Q1\ngeometry", fc="#E8EEF4")
    box(ax, 2.3, 2.0, 1.6, 0.7, "L2  Q2\nactive measure", fc="#E8EEF4")
    box(ax, 4.1, 2.0, 1.6, 0.7, "L2  Q3\nomni loop", fc="#E8EEF4")
    box(ax, 5.9, 2.0, 1.6, 0.7, "L2  Q4\nmixed loop", fc="#E8EEF4")
    box(ax, 1.4, 0.8, 5.2, 0.7, "L3  Discovery  ->  Localisation  ->  Clearance", fc="#E6F0E7", ec=GREEN)
    for px in (1.55, 3.95, 6.35):
        arrow(ax, (px, 4.0), (px, 3.7))
    for px in (1.3, 3.1, 4.9, 6.7):
        arrow(ax, (px, 2.7), (3.9, 1.5))
    ax.set_xlim(0, 8); ax.set_ylim(0, 5.2)
    save(fig, "flow-overall-architecture.png")


def fig_framework():
    fig, ax = plt.subplots(figsize=(5.8, 4.2))
    ax.axis("off")
    box(ax, 0.3, 2.8, 2.0, 0.7, "Observation\nbearing / near / no-signal")
    box(ax, 3.1, 2.8, 2.0, 0.7, "Candidate update\ngeometry / grid / heading")
    box(ax, 3.1, 1.4, 2.0, 0.7, "Task selection\ncover / locate / clear")
    box(ax, 0.3, 1.4, 2.0, 0.7, "Action execution\nmove and timing")
    arrow(ax, (2.3, 3.15), (3.1, 3.15)); arrow(ax, (4.1, 2.8), (4.1, 2.1))
    arrow(ax, (3.1, 1.75), (2.3, 1.75)); arrow(ax, (1.3, 2.1), (1.3, 2.8))
    box(ax, 3.1, 0.3, 2.0, 0.6, "Exit / report", fc="#EFEFEF")
    arrow(ax, (4.1, 1.4), (4.1, 0.9))
    ax.set_xlim(0, 5.6); ax.set_ylim(0, 4.0)
    save(fig, "framework-decision-loop.png")


def fig_state_machine():
    fig, ax = plt.subplots(figsize=(7.0, 2.2))
    ax.axis("off")
    xs = [(0.1, "IDLE"), (1.8, "ENTER\n5 s countdown"), (3.8, "OPEN\ninterface"),
          (5.8, "CLOSED\nexit / timeout")]
    for x, t in xs:
        box(ax, x, 0.6, 1.5, 0.8, t)
    for x in (1.6, 3.6, 5.6):
        arrow(ax, (x, 1.0), (x + 0.2, 1.0))
    ax.set_xlim(0, 7.4); ax.set_ylim(0, 2.0)
    save(fig, "steps-state-machine.png")


def fig_timing_attachment():
    steps = ["step 1", "step 2", "step 3", "step 4"]
    move = np.array([100.0, 0.0, 80.0, 0.0])
    action = np.array([5.0, 5.0, 3.0, 5.0])
    switch = np.array([0.0, 1.0, 0.0, 0.0])
    fig, ax = plt.subplots(figsize=(6.0, 3.2))
    ax.bar(steps, move, color=INK, label="Move")
    ax.bar(steps, action, bottom=move, color=LIGHT, label="Measure / clear")
    ax.bar(steps, switch, bottom=move + action, color=GRAY, label="Switch")
    for i, total in enumerate(move + action + switch):
        ax.text(i, total + 2, "%.0f" % total, ha="center", fontsize=8)
    ax.set_ylabel("Virtual time / s")
    ax.legend(frameon=False)
    save(fig, "timing-attachment-example-105-111-194-199.png")

def clip_halfplane(poly, a, b, c):
    if len(poly) == 0:
        return poly
    out = []
    for i in range(len(poly)):
        p = poly[i]; q = poly[(i + 1) % len(poly)]
        fp = a * p[0] + b * p[1] - c
        fq = a * q[0] + b * q[1] - c
        if fp <= 1e-9:
            out.append(p)
        if (fp < -1e-9 and fq > 1e-9) or (fp > 1e-9 and fq < -1e-9):
            t = fp / (fp - fq)
            out.append(p + t * (q - p))
    return np.array(out, dtype=float)


def wedge_halfplanes(sx, sy, theta_deg, delta_deg):
    theta = math.radians(theta_deg); delta = math.radians(delta_deg)
    a1 = math.sin(theta - delta); b1 = -math.cos(theta - delta)
    c1 = a1 * sx + b1 * sy
    a2 = -math.sin(theta + delta); b2 = math.cos(theta + delta)
    c2 = a2 * sx + b2 * sy
    a3 = -math.cos(theta); b3 = -math.sin(theta)
    c3 = a3 * sx + b3 * sy
    return [(a1, b1, c1), (a2, b2, c2), (a3, b3, c3)]


def intersection_polygon(meas, bound=2500.0):
    poly = np.array([[-bound, -bound], [bound, -bound], [bound, bound], [-bound, bound]], dtype=float)
    for m in meas:
        for a, b, c in wedge_halfplanes(m["x"], m["y"], m["bearing_deg"], m["half_angle_deg"]):
            poly = clip_halfplane(poly, a, b, c)
            if len(poly) == 0:
                return poly
    return poly


def pair_circle(p, q):
    c = (p + q) / 2.0
    r = np.linalg.norm(p - q) / 2.0
    return c, r


def contains_all(c, r, pts, tol=1e-6):
    return all(np.linalg.norm(p - c) <= r + tol for p in pts)


def circumcircle(p, q, r):
    ax, ay = p; bx, by = q; cx, cy = r
    d = 2.0 * (ax * (by - cy) + bx * (cy - ay) + cx * (ay - by))
    if abs(d) < 1e-12:
        return None
    ux = ((ax*ax + ay*ay) * (by - cy) + (bx*bx + by*by) * (cy - ay) + (cx*cx + cy*cy) * (ay - by)) / d
    uy = ((ax*ax + ay*ay) * (cx - bx) + (bx*bx + by*by) * (ax - cx) + (cx*cx + cy*cy) * (bx - ax)) / d
    c = np.array([ux, uy])
    return c, np.linalg.norm(c - p)


def min_enclosing_circle(pts):
    pts = np.asarray(pts, dtype=float)
    best_c = pts[0]; best_r = 0.0
    for i in range(len(pts)):
        for j in range(i + 1, len(pts)):
            c, r = pair_circle(pts[i], pts[j])
            if contains_all(c, r, pts) and r < best_r or (best_r == 0.0 and contains_all(c, r, pts)):
                best_c, best_r = c, r
    for i in range(len(pts)):
        for j in range(i + 1, len(pts)):
            for k in range(j + 1, len(pts)):
                cc = circumcircle(pts[i], pts[j], pts[k])
                if cc is None:
                    continue
                c, r = cc
                if contains_all(c, r, pts) and r < best_r:
                    best_c, best_r = c, r
    return best_c, best_r


def fig_q1_counterexample():
    A = np.array([0.0, 0.0]); B = np.array([40.0, 0.0]); C = np.array([20.0, 20.0 * math.sqrt(3)])
    fig, ax = plt.subplots(figsize=(4.6, 4.2))
    ax.add_patch(Polygon([A, B, C], closed=True, fill=False, edgecolor=INK, linewidth=1.4))
    dc, dr = pair_circle(A, B); ax.add_patch(Circle(dc, dr, fill=False, edgecolor=GOLD, linestyle="--", linewidth=1.2))
    mc, mr = min_enclosing_circle(np.array([A, B, C])); ax.add_patch(Circle(mc, mr, fill=False, edgecolor=GREEN, linewidth=1.4))
    ax.scatter([A[0], B[0], C[0]], [A[1], B[1], C[1]], color=INK, s=12)
    ax.text(C[0] + 0.8, C[1] + 0.6, "C", fontsize=9)
    ax.text(-1.2, -2.0, "A", fontsize=9); ax.text(40.0, -2.0, "B", fontsize=9)
    ax.set_aspect("equal"); ax.grid(True, color="#EEEEEE"); ax.set_xlim(-8, 52); ax.set_ylim(-8, 42)
    ax.set_xlabel("x / m"); ax.set_ylabel("y / m")
    ax.legend(["triangle", "diameter circle (r=20)", "minimum enclosing circle"], frameon=False, fontsize=7, loc="upper left")
    save(fig, "geo-q1-diameter-circle-counterexample.png")


def fig_q1_meetup():
    data_path = ROOT / "examples" / "q1" / "triangle.json"
    meas = json.loads(data_path.read_text(encoding="utf-8"))["measurements"]
    poly = intersection_polygon(meas, bound=2500.0)
    if len(poly) < 3:
        poly = np.array([[0.0, 0.0], [40.0, 0.0], [40.0, 30.0], [0.0, 30.0]])
    fig, axes = plt.subplots(1, 2, figsize=(8.2, 4.0))
    ax = axes[0]
    for m in meas:
        x, y = m["x"], m["y"]; th = math.radians(m["bearing_deg"]); d = math.radians(m["half_angle_deg"])
        for sgn in (-1.0, 1.0):
            ang = th + sgn * d
            ray = np.array([x, y]) + 3600.0 * np.array([math.cos(ang), math.sin(ang)])
            ax.plot([x, ray[0]], [y, ray[1]], color=LIGHT, linewidth=0.8)
        ax.scatter([x], [y], color=GRAY, s=14)
        ax.text(x + 40, y + 40, "S", fontsize=8)
    ax.add_patch(Polygon(poly, closed=True, facecolor=PALE, edgecolor=INK, linewidth=1.4))
    ax.set_aspect("equal"); ax.grid(True, color="#EEEEEE")
    ax.set_xlim(-1200, 1000); ax.set_ylim(-1000, 1100)
    ax.set_xlabel("x / m"); ax.set_ylabel("y / m")
    ax.set_title("Wedge intersection", fontsize=9)
    ax2 = axes[1]
    ax2.add_patch(Polygon(poly, closed=True, facecolor=PALE, edgecolor=INK, linewidth=1.4))
    c, r = min_enclosing_circle(poly)
    d = max(np.linalg.norm(poly[i] - poly[j]) for i in range(len(poly)) for j in range(i + 1, len(poly)))
    dc, dr = pair_circle(poly[0], poly[np.argmax(np.linalg.norm(poly - poly[0], axis=1))])
    ax2.add_patch(Circle(dc, dr, fill=False, edgecolor=GOLD, linestyle="--", linewidth=1.1))
    ax2.add_patch(Circle(c, r, fill=False, edgecolor=GREEN, linewidth=1.3))
    ax2.scatter(poly[:, 0], poly[:, 1], color=INK, s=10)
    lo = poly.min(axis=0) - 4; hi = poly.max(axis=0) + 4
    ax2.set_xlim(lo[0], hi[0]); ax2.set_ylim(lo[1], hi[1])
    ax2.set_aspect("equal"); ax2.grid(True, color="#EEEEEE")
    ax2.set_title("Zoom: diameter vs. enclosing circle", fontsize=9)
    ax2.set_xlabel("x / m"); ax2.set_ylabel("y / m")
    save(fig, "geo-q1-convex-intersection-zoom.png")

def fig_q2_safe_disk():
    delta = math.radians(1.0); rp = 1500.0; a = rp / (2.0 * math.cos(delta)); safe = 1000.0 - a
    fig, axes = plt.subplots(1, 2, figsize=(8.0, 3.6))
    ax = axes[0]
    ax.plot([0, rp * math.cos(delta)], [0, rp * math.sin(delta)], color=GRAY, linewidth=0.8)
    ax.plot([0, rp * math.cos(delta)], [0, -rp * math.sin(delta)], color=GRAY, linewidth=0.8)
    ax.add_patch(Circle((a, 0), a, fill=False, edgecolor=LIGHT, linestyle="--", linewidth=1.2))
    ax.add_patch(Circle((a, 0), safe, facecolor=PALE, edgecolor=INK, linewidth=1.2, alpha=0.8))
    ax.scatter([0], [0], color=GRAY, s=12); ax.scatter([a], [0], color=INK, s=12)
    ax.text(a + 20, 20, "c", fontsize=9); ax.text(10, 15, "s", fontsize=9)
    ax.set_aspect("equal"); ax.grid(True, color="#EEEEEE")
    ax.set_xlim(-80, rp + 80); ax.set_ylim(-320, 320)
    ax.set_xlabel("x / m"); ax.set_ylabel("y / m"); ax.set_title("Truncated wedge and safe disk", fontsize=9)
    ax = axes[1]
    ax.plot([0, 0], [0, 1], color=GRAY, linewidth=2)
    ax.plot([0, 1], [0, 0], color=INK, linewidth=4)
    ax.text(0.5, -0.08, "a = 1500/(2cos delta)", ha="center", fontsize=8)
    ax.text(0.5, 0.55, "safe radius = 1000 - a", ha="center", fontsize=8)
    ax.axis("off")
    save(fig, "geo-q2-guaranteed-reception-disk.png")


def fig_q3_coverage():
    a = 1159.7; R = 1800.0
    fig, ax = plt.subplots(figsize=(5.2, 5.0))
    ax.add_patch(Circle((0, 0), R, fill=False, edgecolor=INK, linewidth=1.4))
    stations = [(0.0, 0.0)] + [(a * math.cos(k * math.pi / 3.0), a * math.sin(k * math.pi / 3.0)) for k in range(6)]
    for x, y in stations:
        ax.add_patch(Circle((x, y), 1000.0, facecolor=PALE, edgecolor=LIGHT, linewidth=0.7, alpha=0.35))
        ax.scatter([x], [y], color=INK, s=16)
    ax.add_patch(Circle((0, 0), a, fill=False, edgecolor=GRAY, linestyle="--", linewidth=1.0))
    ax.text(0, 0, "O", fontsize=9, ha="right", va="top"); ax.text(a, 0, "a", fontsize=9, ha="center", va="bottom")
    ax.set_aspect("equal"); ax.grid(True, color="#EEEEEE"); ax.set_xlim(-2000, 2000); ax.set_ylim(-2000, 2000)
    ax.set_xlabel("x / m"); ax.set_ylabel("y / m"); ax.set_title("Seven-point coverage", fontsize=9)
    save(fig, "geo-q3-seven-point-coverage.png")


def fig_q4_plan():
    data = json.loads((ROOT / "src" / "mcm26b" / "guarantee" / "plans" / "q4_mincover_v2.json").read_text(encoding="utf-8"))
    pts = np.array(data["points"], dtype=float)
    fig, ax = plt.subplots(figsize=(5.4, 5.0))
    ax.add_patch(Circle((0, 0), 1800.0, facecolor=PALE, edgecolor=INK, linewidth=1.1, alpha=0.45))
    ax.scatter(pts[:, 0], pts[:, 1], color=INK, s=18, label="31 stations")
    ax.scatter([0], [0], color=GOLD, marker="*", s=50, label="origin")
    ax.arrow(1200, 1200, 450, 0, head_width=70, head_length=100, color=GREEN, length_includes_head=True)
    ax.add_patch(Wedge((1200, 1200), 650, -90, 90, facecolor=GREEN, alpha=0.10, edgecolor=GREEN, linewidth=0.8))
    ax.text(1250, 1250, "forward half-disk", fontsize=8)
    ax.set_aspect("equal"); ax.grid(True, color="#EEEEEE"); ax.set_xlim(-2800, 2800); ax.set_ylim(-2800, 2800)
    ax.set_xlabel("x / m"); ax.set_ylabel("y / m"); ax.legend(frameon=False, fontsize=8, loc="lower right")
    save(fig, "geo-q4-minimum-cover-31-points.png")


def main():
    fig_formal_total()
    fig_formal_time_bars()
    fig_flow_overall()
    fig_framework()
    fig_state_machine()
    fig_timing_attachment()
    fig_q1_counterexample()
    fig_q1_meetup()
    fig_q2_safe_disk()
    fig_q3_coverage()
    fig_q4_plan()


if __name__ == "__main__":
    main()