import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.path import Path
from matplotlib.patches import PathPatch
import os, sys

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "figs")
os.makedirs(OUT, exist_ok=True)

import argparse

p = argparse.ArgumentParser(description="Render smooth enrobing-circle hero shapes.")
p.add_argument("--rho", type=float, default=75.0, help="phantom (bridge-carving) circle radius")
p.add_argument("--R", type=float, default=55.0, help="visible source-circle radius")
p.add_argument("--m", type=float, default=8.0, help="coat margin added to R")
p.add_argument("--s", type=float, default=13.0, help="smooth-max fillet width (cusp smoothing)")
p.add_argument("--D", type=float, default=None,
               help="center-to-center distance between the two duo circles "
                    "(default: scales with R as in the original layout)")
p.add_argument("--D_tri", type=float, default=None,
               help="edge length of the equilateral triangle for the tri layout "
                    "(default: scales with R as in the original layout)")
p.add_argument("--rho_duo1", type=float, default=None,
               help="phantom radius for the duo's first (above-centerline) bridge "
                    "(default: falls back to --rho)")
p.add_argument("--rho_duo2", type=float, default=None,
               help="phantom radius for the duo's second (below-centerline) bridge "
                    "(default: falls back to --rho)")
p.add_argument("--rho_tri_a", type=float, default=None,
               help="phantom radius for triangle edge 0-1 (default: falls back to --rho)")
p.add_argument("--rho_tri_b", type=float, default=None,
               help="phantom radius for triangle edge 1-2 (default: falls back to --rho)")
p.add_argument("--rho_tri_c", type=float, default=None,
               help="phantom radius for triangle edge 2-0 (default: falls back to --rho)")
p.add_argument("--which", choices=["duo", "tri", "both"], default="both",
               help="which shape(s) to render")
p.add_argument("--rot_duo", type=float, default=0.0,
               help="rotation angle in degrees (CCW) for the duo output figure (default: 0)")
p.add_argument("--rot_tri", type=float, default=0.0,
               help="rotation angle in degrees (CCW) for the tri output figure (default: 0)")
args = p.parse_args()

RHO, R_PARAM, M_PARAM, S_PARAM, WHICH = args.rho, args.R, args.m, args.s, args.which
D_PARAM = args.D
D_TRI_PARAM = args.D_tri
RHO_DUO1, RHO_DUO2 = args.rho_duo1, args.rho_duo2
RHO_TRI_A, RHO_TRI_B, RHO_TRI_C = args.rho_tri_a, args.rho_tri_b, args.rho_tri_c
ROT_DUO, ROT_TRI = args.rot_duo, args.rot_tri

# ---------------------------------------------------------------- math kernels
def smin(a, b, k):
    if k <= 0:
        return np.minimum(a, b)
    h = np.maximum(1 - np.abs(a - b) / k, 0.0)
    return np.minimum(a, b) - 0.25 * k * h * h

def smax(a, b, k):
    if k <= 0:
        return np.maximum(a, b)
    h = np.maximum(1 - np.abs(a - b) / k, 0.0)
    return np.maximum(a, b) + 0.25 * k * h * h

def dist_seg(X, Y, a, b):
    ax, ay = a; bx, by = b
    dx, dy = bx - ax, by - ay
    l2 = dx * dx + dy * dy
    t = ((X - ax) * dx + (Y - ay) * dy) / l2 if l2 > 0 else np.zeros_like(X)
    t = np.clip(t, 0, 1)
    return np.hypot(X - (ax + t * dx), Y - (ay + t * dy))

def _sgn(X, Y, a, b):
    return (X - b[0]) * (a[1] - b[1]) - (a[0] - b[0]) * (Y - b[1])

def in_tri(X, Y, A, B, C):
    d1 = _sgn(X, Y, A, B); d2 = _sgn(X, Y, B, C); d3 = _sgn(X, Y, C, A)
    neg = (d1 < 0) | (d2 < 0) | (d3 < 0)
    pos = (d1 > 0) | (d2 > 0) | (d3 > 0)
    return ~(neg & pos)

def dist_hull(X, Y, C):
    if len(C) == 2:
        return dist_seg(X, Y, C[0], C[1])
    A, B, D = C
    d = np.minimum(np.minimum(dist_seg(X, Y, A, B), dist_seg(X, Y, B, D)), dist_seg(X, Y, D, A))
    return np.where(in_tri(X, Y, A, B, D), 0.0, d)

def phantoms(C, R, m, rho_list):
    """rho_list: per-phantom radii, in bridge order.
       duo -> [rho_above, rho_below] (length 2)
       tri -> [rho_edge01, rho_edge12, rho_edge20] (length 3)"""
    Rc = R + m
    out = []
    if len(C) == 2:
        A, B = C
        D = np.hypot(B[0] - A[0], B[1] - A[1]); Mx, My = (A[0] + B[0]) / 2, (A[1] + B[1]) / 2
        ux, uy = (B[0] - A[0]) / D, (B[1] - A[1]) / D; nx, ny = -uy, ux
        rho_above, rho_below = rho_list[0], rho_list[1]
        rad_a = Rc + rho_above
        hh_a = rad_a * rad_a - (D / 2) ** 2
        if hh_a > 0:
            h = np.sqrt(hh_a)
            out.append((Mx + nx * h, My + ny * h, rho_above))
        rad_b = Rc + rho_below
        hh_b = rad_b * rad_b - (D / 2) ** 2
        if hh_b > 0:
            h = np.sqrt(hh_b)
            out.append((Mx - nx * h, My - ny * h, rho_below))
    else:
        G = (sum(c[0] for c in C) / 3, sum(c[1] for c in C) / 3)
        for idx, (i, j) in enumerate([(0, 1), (1, 2), (2, 0)]):
            A, B = C[i], C[j]
            D = np.hypot(B[0] - A[0], B[1] - A[1]); Mx, My = (A[0] + B[0]) / 2, (A[1] + B[1]) / 2
            ux, uy = (B[0] - A[0]) / D, (B[1] - A[1]) / D; nx, ny = -uy, ux
            if (Mx - G[0]) * nx + (My - G[1]) * ny < 0:
                nx, ny = -nx, -ny
            rho_e = rho_list[idx]
            rad = Rc + rho_e
            hh = rad * rad - (D / 2) ** 2
            if hh > 0:
                h = np.sqrt(hh)
                out.append((Mx + nx * h, My + ny * h, rho_e))
    return out

def carve_field(X, Y, C, R, m, rho_list, s):
    v = dist_hull(X, Y, C) - (R + m)
    for px, py, pr in phantoms(C, R, m, rho_list):
        v = smax(v, pr - np.hypot(X - px, Y - py), s)
    return v

def rotate_points(pts, deg):
    """Rotate an array of (x, y) points about the origin by deg degrees, CCW.
       Always returns an (N,2) numpy array, including the deg==0 identity case."""
    pts = np.asarray(pts, dtype=float)
    if deg == 0:
        return pts
    th = np.radians(deg)
    c, s = np.cos(th), np.sin(th)
    pts = np.asarray(pts, dtype=float)
    rot = np.array([[c, -s], [s, c]])
    return pts @ rot.T

def grid(xmin, xmax, ymin, ymax, n=720):
    xs = np.linspace(xmin, xmax, n)
    ys = np.linspace(ymin, ymax, int(n * (ymax - ymin) / (xmax - xmin)))
    return np.meshgrid(xs, ys)

def contour_segs(X, Y, Z, level):
    f0 = plt.figure(); a0 = f0.add_subplot(111)
    cs = a0.contour(X, Y, Z, levels=[level])
    segs = [np.asarray(s) for s in cs.allsegs[0] if len(s) > 2]
    plt.close(f0)
    return segs

# ---------------------------------------------------------------- Lightfall-styled hero
def hero_bg(extent, w=1100):
    xmin, xmax, ymin, ymax = extent
    h = int(w * (ymax - ymin) / (xmax - xmin))
    yy = np.linspace(0, 1, h)[:, None, None]           # 0 = top
    xx = np.linspace(0, 1, w)[None, :, None]
    top = np.array([0.082, 0.060, 0.155]); bot = np.array([0.215, 0.105, 0.262])
    base = top * (1 - yy) + bot * yy
    base = np.broadcast_to(base, (h, w, 3)).copy()
    cx, cy = 0.5, 1.06
    r = np.sqrt((xx - cx) ** 2 + (yy - cy) ** 2)
    bloom = np.clip(1 - r / 0.8, 0, 1) ** 1.4 * np.array([0.62, 0.20, 0.40]) * 0.95
    r2 = np.sqrt((xx - 0.85) ** 2 + (yy - 0.0) ** 2)
    bloom2 = np.clip(1 - r2 / 0.55, 0, 1) ** 1.6 * np.array([0.16, 0.13, 0.34]) * 0.8
    return np.clip(base + bloom + bloom2, 0, 1)

def glow_line(ax, x, y, base="#e9e5ff", col=(0.78, 0.74, 0.98), lw=2.4):
    for w, al in [(lw * 7, 0.05), (lw * 4.2, 0.09), (lw * 2.4, 0.16)]:
        ax.plot(x, y, color=col, lw=w, alpha=al, solid_capstyle="round", solid_joinstyle="round")
    ax.plot(x, y, color=base, lw=lw, solid_capstyle="round", solid_joinstyle="round")

def hero(C, fname, extent, deco, rho_list, R=55.0, m=8.0, s=13.0, rot=0.0):
    C_rot = [tuple(p) for p in rotate_points(C, rot)] if rot != 0 else C
    fig, ax = plt.subplots(figsize=(7.8, 7.8 * (extent[3] - extent[2]) / (extent[1] - extent[0])))
    ax.set_aspect("equal"); ax.axis("off")
    ax.set_xlim(extent[0], extent[1]); ax.set_ylim(extent[2], extent[3])
    rho_max = max(rho_list)
    pad = max(60.0, (R + m + rho_max) - 55.0 + 60.0)   # keep phantoms/coat inside the grid
    X, Y = grid(extent[0] - pad, extent[1] + pad, extent[2] - pad, extent[3] + pad, 760)
    Z = carve_field(X, Y, C_rot, R, m, rho_list, s)
    segs = contour_segs(X, Y, Z, 0.0)
    # faint fill
    if segs:
        big = max(segs, key=len)
        ax.add_patch(PathPatch(Path(big), facecolor=(0.74, 0.69, 0.96),
                               alpha=0.05, lw=0, zorder=2))
    # source circles, faint
    for c in C_rot:
        ax.add_patch(plt.Circle(c, R, fill=False, ec=(0.70, 0.66, 0.90),
                                lw=1.0, alpha=0.5, zorder=2))
    for seg in segs:
        glow_line(ax, seg[:, 0], seg[:, 1], lw=2.6)
    fig.patch.set_alpha(0.0)
    ax.patch.set_alpha(0.0)
    fig.savefig(os.path.join(OUT, fname), bbox_inches="tight",
                transparent=True, pad_inches=0.05)
    plt.close(fig)
    print(f"  {fname}: rho={rho_list}, phantoms={len(phantoms(C_rot, R, m, rho_list))}, rot={rot}")

def fig_heroes(which="both", rho=75.0, R=55.0, m=8.0, s=13.0, D=None, D_tri=None,
                rho_duo=(None, None), rho_tri=(None, None, None), rot_duo=0.0, rot_tri=0.0):
    k = R / 55.0   # scale factor relative to the original R=55 layout (still used for deco)
    if which in ("duo", "both"):
        Dd = D if D is not None else 190.0 * k
        half = Dd / 2.0
        coast = R + m  # coated radius, for bounding-box purposes
        rho_list = [r if r is not None else rho for r in rho_duo]
        bulge = max(R, max(rho_list))
        C_base = [(-half, 0), (half, 0)]
        C_rot = rotate_points(C_base, rot_duo)
        pad = coast + bulge + 60 * k
        xs = C_rot[:, 0]; ys = C_rot[:, 1]
        ext = (xs.min() - pad, xs.max() + pad, ys.min() - pad, ys.max() + pad)
        hero(C_base, "hero_duo.svg", ext,
             deco=[(-180 * k, 120 * k, 150 * k), (190 * k, -90 * k, 175 * k), (40 * k, -190 * k, 150 * k)],
             rho_list=rho_list, R=R, m=m, s=s, rot=rot_duo)
    if which in ("tri", "both"):
        Dt = D_tri if D_tri is not None else 208.0 * k   # original edge length ~208 at k=1
        Rc = Dt / np.sqrt(3.0)   # circumradius of an equilateral triangle with edge Dt
        angles = [90.0, 210.0, 330.0]   # apex up, matching the original orientation
        tri = [(Rc * np.cos(np.radians(a)), Rc * np.sin(np.radians(a))) for a in angles]
        coast = R + m
        rho_list = [r if r is not None else rho for r in rho_tri]
        bulge = max(R, max(rho_list))
        tri_rot = rotate_points(tri, rot_tri)
        pad = coast + bulge + 40
        xs = tri_rot[:, 0]; ys = tri_rot[:, 1]
        ext = (xs.min() - pad, xs.max() + pad, ys.min() - pad, ys.max() + pad)
        hero(tri, "hero_tri.svg", ext,
             deco=[(-185 * k, 150 * k, 150 * k), (200 * k, 120 * k, 170 * k), (0, -240 * k, 180 * k), (210 * k, -120 * k, 120 * k)],
             rho_list=rho_list, R=R, m=m, s=s, rot=rot_tri)

D_display = D_PARAM if D_PARAM is not None else f"{190.0 * (R_PARAM/55.0):.2f} (auto, from R)"
D_tri_display = D_TRI_PARAM if D_TRI_PARAM is not None else f"{208.0 * (R_PARAM/55.0):.2f} (auto, from R)"
rho_duo_args = (RHO_DUO1, RHO_DUO2)
rho_tri_args = (RHO_TRI_A, RHO_TRI_B, RHO_TRI_C)
print(f"Rendering heroes with rho={RHO}, R={R_PARAM}, m={M_PARAM}, s={S_PARAM}, "
      f"D={D_display}, D_tri={D_tri_display}, which={WHICH}, "
      f"rho_duo={rho_duo_args}, rho_tri={rho_tri_args}, "
      f"rot_duo={ROT_DUO}, rot_tri={ROT_TRI}")
fig_heroes(WHICH, rho=RHO, R=R_PARAM, m=M_PARAM, s=S_PARAM, D=D_PARAM, D_tri=D_TRI_PARAM,
           rho_duo=rho_duo_args, rho_tri=rho_tri_args, rot_duo=ROT_DUO, rot_tri=ROT_TRI)
print("done:", sorted(os.listdir(OUT)))