"""Earth-gravity note: oblateness, source/field geometry, Legendre generating function, multipole lobes, ring."""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse, Circle, Wedge
from numpy.polynomial import legendre as L
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine

def arrow(ax, p, q, color=INK, lw=1.8, ms=16, **kw):
    ax.annotate("", xy=q, xytext=p, arrowprops=dict(arrowstyle="-|>", color=color, lw=lw,
                shrinkA=0, shrinkB=0, mutation_scale=ms, **kw))

# ------------------------------------------------------------------ oblate Earth
fig, (a1, a2) = plt.subplots(1, 2, figsize=(11.5, 5.2))
f = 0.12
th = np.linspace(0, 2*np.pi, 400)
a1.add_patch(Circle((0, 0), 1, fill=False, ls="--", color=INK, lw=1))
a1.add_patch(Ellipse((0, 0), 2, 2*(1-f), color="#cfdcec", zorder=1))
a1.plot(np.cos(th), (1-f)*np.sin(th), color=BLUE, lw=2)
a1.plot([0, 0], [-1.35, 1.35], color=INK, lw=.8); a1.text(.03, 1.36, "spin axis", fontsize=10)
a1.plot([-1.35, 1.35], [0, 0], color=INK, lw=.8, ls=":")
arrow(a1, (0, 0), (1, 0), MAROON, 1.3, 12); a1.text(.35, -.10, r"$a=R$", color=MAROON, fontsize=13)
arrow(a1, (0, 0), (0, 1-f), MAROON, 1.3, 12); a1.text(.04, .40, r"$c$", color=MAROON, fontsize=13)
phi = np.radians(40); P = 1.32*np.array([np.cos(phi), np.sin(phi)])
a1.plot(*P, "o", color=AMBER, ms=9, zorder=6); a1.text(P[0]+.04, P[1]+.05, "satellite", fontsize=10.5)
d = -P/np.linalg.norm(P)
arrow(a1, P, P+0.55*d, INK, 2)
eq = np.array([-np.sin(phi), np.cos(phi)])*(-1)          # toward the equator
arrow(a1, P+0.55*d, P+0.55*d+0.22*eq, PURPLE, 2.4)
a1.text(.72, 1.04, "point-mass pull", fontsize=9.5, ha="right")
a1.text(.92, .43, "extra pull\ntoward the\nequatorial bulge", fontsize=9.5, color=PURPLE)
a1.set_xlim(-1.5, 1.85); a1.set_ylim(-1.5, 1.6); a1.set_aspect("equal"); a1.axis("off")
a1.set_title("An oblate Earth (flattening drawn ×36)", fontsize=12)

a2.add_patch(Circle((0, 0), 1, color="#e9e5da", zorder=1))
for sgn, col in ((1, "#e8a25d"), (-1, "#8fb3e0")):
    pass
a2.add_patch(Wedge((0, 0), 1, -25, 25, color="#e8a25d", zorder=2))
a2.add_patch(Wedge((0, 0), 1, 155, 205, color="#e8a25d", zorder=2))
a2.add_patch(Wedge((0, 0), 1, 65, 115, color="#8fb3e0", zorder=2))
a2.add_patch(Wedge((0, 0), 1, 245, 295, color="#8fb3e0", zorder=2))
a2.add_patch(Circle((0, 0), 1, fill=False, color=INK, lw=1.2, zorder=3))
a2.text(0, 0, "sphere\nof mass $M$", ha="center", va="center", fontsize=12, zorder=4)
a2.text(1.08, 0, r"$+$ excess" "\n(equatorial belt)", fontsize=10.5, color=AMBER, va="center")
a2.text(-1.08, 0, r"$+$ excess" "\n(equatorial belt)", fontsize=10.5, color=AMBER, va="center", ha="right")
a2.text(0, 1.1, "$-$ deficit (polar cap)", fontsize=10.5, color=BLUE, ha="center")
a2.text(0, -1.17, "$-$ deficit (polar cap)", fontsize=10.5, color=BLUE, ha="center")
a2.set_xlim(-1.9, 1.9); a2.set_ylim(-1.4, 1.35); a2.set_aspect("equal"); a2.axis("off")
a2.set_title("Same mass, redistributed:\nsphere + belt $-$ caps", fontsize=12)
fig.tight_layout(); fig.savefig("img/fig-grav-oblate.png", dpi=160, bbox_inches="tight"); plt.close(fig)

# ------------------------------------------------------------------ geometry: source, field point
fig, ax = plt.subplots(figsize=(6.6, 4.6))
O = np.array([0., 0.]); Pp = 1.6*np.array([np.cos(.35), np.sin(.35)])
Q = 3.9*np.array([np.cos(.95), np.sin(.95)])*1.0
Q = np.array([2.2, 3.5]) ; Pp = np.array([1.5, .5])
ax.add_patch(Circle((0, 0), 2.1, fill=False, ls="--", color=INK, lw=1))
ax.text(-2.05, -1.0, "body of mass $M$\n(all mass inside\n$r'\\leq R$)", fontsize=10, ha="center", color=INK)
arrow(ax, O, Pp, GREEN); arrow(ax, O, Q, BLUE); arrow(ax, Pp, Q, MAROON, 2.2)
ax.plot(*O, "o", color=INK, ms=7); ax.plot(*Pp, "o", color=GREEN, ms=8); ax.plot(*Q, "o", color=BLUE, ms=8)
ax.text(-.25, -.3, "O (centre of mass)", fontsize=10)
ax.text(*(Pp+[.05, -.4]), r"$P'$: mass element $dm'$", color=GREEN, fontsize=11)
ax.text(*(Q+[.1, .1]), r"field point $P$", color=BLUE, fontsize=11)
ax.text(*(Pp*.5+[.05,-.32]), r"$r'$", color=GREEN, fontsize=15)
ax.text(*(Q*.5+[-.5, .1]), r"$r$", color=BLUE, fontsize=15)
ax.text(*((Pp+Q)/2+[.12, -.05]), r"$s$", color=MAROON, fontsize=15)
a0, a1_ = np.arctan2(*Pp[::-1]), np.arctan2(*Q[::-1])
t = np.linspace(a0, a1_, 40); ax.plot(.85*np.cos(t), .85*np.sin(t), color=INK, lw=1.2)
ax.text(.95, .78, r"$\gamma$", fontsize=15)
ax.set_xlim(-3.0, 3.4); ax.set_ylim(-2.3, 4.1); ax.set_aspect("equal"); ax.axis("off")
fig.tight_layout(); fig.savefig("img/fig-grav-geometry.png", dpi=160, bbox_inches="tight"); plt.close(fig)

# ------------------------------------------------------------------ generating function convergence
fig, (a1, a2) = plt.subplots(1, 2, figsize=(11.5, 4.3))
g = np.linspace(0, np.pi, 400); x = np.cos(g)
eps = 0.5
exact = (1 - 2*eps*x + eps**2)**-.5
a1.plot(np.degrees(g), exact, color=INK, lw=3, label=r"exact  $r/s$")
for N, c in zip((0, 1, 2, 4), (BLUE, GREEN, AMBER, MAROON)):
    part = sum(eps**n * L.legval(x, [0]*n + [1]) for n in range(N+1))
    a1.plot(np.degrees(g), part, color=c, lw=1.6, ls="--", label=rf"sum up to $n={N}$")
a1.set_xlabel(r"angle $\gamma$ between $\mathbf{r}$ and $\mathbf{r}'$ (deg)"); a1.set_ylabel(r"$r/s$")
a1.set_title(r"$\varepsilon=r'/r=0.5$: each new term hugs the exact curve", fontsize=11.5)
a1.legend(frameon=False, fontsize=9.5); despine(a1)
Ns = np.arange(0, 21)
for e, c in ((0.3, BLUE), (0.6, AMBER), (0.9, MAROON)):
    ex = 1/(1-e); part = np.cumsum(e**Ns)
    a2.semilogy(Ns, np.abs(ex-part), "o-", color=c, ms=4, label=rf"$\varepsilon={e}$")
a2.set_xlabel("highest degree $N$ kept"); a2.set_ylabel(r"error in $r/s$ (worst case $\gamma=0$)")
a2.set_title("Error shrinks like $\\varepsilon^{N+1}$: fast when $r'\\ll r$", fontsize=11.5)
a2.legend(frameon=False); despine(a2)
fig.tight_layout(); fig.savefig("img/fig-grav-legendre-gen.png", dpi=160, bbox_inches="tight"); plt.close(fig)

# ------------------------------------------------------------------ Legendre polynomials
fig, ax = plt.subplots(figsize=(5.4, 4.2))
x = np.linspace(-1, 1, 400)
for n, c in zip(range(6), (INK, BLUE, GREEN, AMBER, MAROON, PURPLE)):
    ax.plot(x, L.legval(x, [0]*n + [1]), color=c, lw=1.8, label=rf"$P_{n}$")
ax.set_xlabel(r"$x=\cos\gamma$ (or $\sin\varphi$)"); ax.set_ylim(-1.1, 1.1)
ax.axhline(0, color=INK, lw=.6); ax.legend(frameon=False, ncol=3, fontsize=9, loc="lower center"); despine(ax)
fig.tight_layout(); fig.savefig("img/fig-grav-legendre-poly.png", dpi=160, bbox_inches="tight"); plt.close(fig)

# ------------------------------------------------------------------ 3-D lobes (surfaces of revolution)
from matplotlib.colors import to_rgba
fig = plt.figure(figsize=(11.5, 3.3))
names = ["monopole", "dipole", "quadrupole", "octupole", "hexadecapole"]
th, ph = np.meshgrid(np.linspace(0, np.pi, 90), np.linspace(0, 1.5*np.pi, 90))   # 90-deg wedge cut away
for n in range(5):
    ax = fig.add_subplot(1, 5, n+1, projection="3d")
    p = L.legval(np.cos(th), [0]*n + [1]); rho = np.abs(p)
    X, Y, Z = rho*np.sin(th)*np.cos(ph), rho*np.sin(th)*np.sin(ph), rho*np.cos(th)
    fc = np.where((p >= 0)[..., None], np.array(to_rgba(BLUE, .85)), np.array(to_rgba(MAROON, .85)))
    ax.plot_surface(X, Y, Z, facecolors=fc, rstride=1, cstride=1, linewidth=0, antialiased=False, shade=False)
    ax.plot([0, 0], [0, 0], [-1.25, 1.25], color=INK, lw=.8)
    ax.set_box_aspect((1, 1, 1)); ax.set_xlim(-1, 1); ax.set_ylim(-1, 1); ax.set_zlim(-1, 1)
    ax.view_init(elev=22, azim=35); ax.set_axis_off()
    ax.set_title(rf"$n={n}$: {names[n]}", fontsize=10.5, pad=-2)
fig.subplots_adjust(left=0, right=1, wspace=0, top=.92, bottom=0)
fig.savefig("img/fig-grav-lobes.png", dpi=160, bbox_inches="tight"); plt.close(fig)

# ------------------------------------------------------------------ ring on the axis
fig, ax = plt.subplots(figsize=(5.2, 4.4))
a = 1.0; t = np.linspace(0, 2*np.pi, 300)
ax.plot(a*np.cos(t), .28*a*np.sin(t), color=AMBER, lw=3)
z = 1.9; ax.plot([0, 0], [-.3, z+.3], color=INK, lw=.8)
ax.plot(0, z, "o", color=BLUE, ms=9); ax.text(.08, z, "$P$ on the axis", color=BLUE, fontsize=11)
R0 = np.array([np.cos(.5), .28*np.sin(.5)])
ax.plot(*R0, "o", color=AMBER, ms=6)
arrow(ax, R0, (0, z), MAROON, 1.6, 13); ax.text(.62, 1.0, r"$s=\sqrt{a^2+z^2}$", color=MAROON, fontsize=11)
ax.plot([0, 0], [0, z], color=BLUE, lw=2); ax.text(-.22, .8, r"$z$", color=BLUE, fontsize=14)
ax.plot([0, R0[0]], [0, R0[1]], color=AMBER, lw=1.2); ax.text(.3, -.13, r"$a$", color=AMBER, fontsize=13)
ax.text(-1.05, -.55, "ring of mass $m$", color=AMBER, fontsize=11)
ax.set_xlim(-1.4, 1.7); ax.set_ylim(-.8, 2.3); ax.set_aspect("equal"); ax.axis("off")
fig.tight_layout(); fig.savefig("img/fig-grav-ring.png", dpi=160, bbox_inches="tight"); plt.close(fig)
print("basics done")