"""Original teaching diagrams. Requires matplotlib; no application code changes."""
import math
from pathlib import Path

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib import font_manager
from matplotlib.patches import Circle, FancyArrowPatch, Rectangle
import numpy as np
from PIL import Image, ImageChops

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / 'docs/teaching/rtcp-five-axis/assets'
OUT.mkdir(parents=True, exist_ok=True)
font_manager.fontManager.addfont('/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc')
plt.rcParams.update({'font.family': 'Noto Sans CJK JP', 'font.size': 15, 'axes.unicode_minus': False})
BLUE, TEAL, RED, GRAY, INK = '#2563b8', '#00877b', '#d45549', '#8a969f', '#23303a'

def save(fig, name):
    path = OUT / f'{name}.png'
    fig.savefig(path, dpi=190, facecolor='white', bbox_inches='tight', pad_inches=.18)
    plt.close(fig)
    if name == 'bc-frame':
        image = Image.open(path).convert('RGB')
        ink = ImageChops.difference(image, Image.new('RGB', image.size, 'white')).convert('L').point(lambda v: 255 if v > 8 else 0)
        left, top, right, bottom = ink.getbbox()
        image.crop((max(0,left-25),max(0,top-25),min(image.width,right+25),min(image.height,bottom+25))).save(path)

def axis2(ax, limit=125):
    ax.set(xlim=(-30, limit), ylim=(-30, limit), aspect='equal')
    ax.axis('off')
    ax.annotate('', (limit-5, 0), (-22, 0), arrowprops={'arrowstyle': '-|>', 'color': BLUE, 'lw': 2})
    ax.annotate('', (0, limit-5), (0, -22), arrowprops={'arrowstyle': '-|>', 'color': BLUE, 'lw': 2})
    ax.text(limit-8, -12, '$X_M$', color=BLUE, ha='right')
    ax.text(-12, limit-9, '$Y_M$', color=BLUE, va='top')
    ax.text(-10, -13, 'O', color=INK)

def point(ax, x, y, label, color, offset=(7, 7), marker='o'):
    ax.scatter([x], [y], s=150, c=color, marker=marker, zorder=5)
    ax.annotate(label, (x, y), xytext=offset, textcoords='offset points', color=color, fontsize=14)

def arc(ax, radius, end=90):
    angles = np.linspace(0, np.deg2rad(end), 101)
    ax.plot(radius*np.cos(angles), radius*np.sin(angles), color=TEAL, lw=2, ls='--')

# The same marked workpiece point, before and after turning the C table.
for follow in [False, True]:
    fig, axes = plt.subplots(1, 2, figsize=(10, 4.3))
    for ax in axes: axis2(ax); ax.add_patch(Circle((0, 0), 100, fill=False, ec='#d8e1e6', lw=2))
    axes[0].set_title('起点：C = 0°', pad=10, color=INK)
    point(axes[0], 100, 0, '工件标记点 + 刀尖', RED, (-138, 17), '*')
    axes[0].plot([0, 100], [0, 0], color=TEAL, lw=4)
    axes[1].set_title('终点：C = 90°', pad=10, color=INK)
    arc(axes[1], 100)
    axes[1].plot([0, 0], [0, 100], color=TEAL, lw=4)
    point(axes[1], 0, 100, '同一个工件点', TEAL, (12, -3), '*')
    if follow:
        axes[1].scatter([0], [100], s=380, facecolors='none', edgecolors=RED, linewidths=3)
        axes[1].text(40, 45, '刀尖由 XYZ 移到这里\n相对工件仍是同一点', color=RED, ha='center')
    else:
        point(axes[1], 100, 0, '刀尖留在原位置', RED, (-125, -22))
        axes[1].plot([0, 100], [100, 0], color=RED, lw=1.6, ls=':')
        axes[1].text(43, 52, '相对位置改变', color=RED, rotation=-45, ha='center')
    fig.subplots_adjust(wspace=.22)
    save(fig, 'tcp-follow' if follow else 'identity-no-follow')

# Rotated workpiece axes while the machine axes stay fixed.
fig, ax = plt.subplots(figsize=(6.2, 5.5)); axis2(ax, 150)
origin = np.array([60., 45.]); theta = math.pi/4
u, v = np.array([math.cos(theta), math.sin(theta)]), np.array([-math.sin(theta), math.cos(theta)])
corners = [origin+u*a+v*b for a,b in [(-22,-22),(32,-22),(32,32),(-22,32),(-22,-22)]]
ax.fill(*np.array(corners).T, color='#e3f3ef', zorder=0)
for direction, label, offset in [(u, '$X_W$', (3,0)), (v, '$Y_W$', (-12,0))]:
    end = origin+direction*70
    ax.annotate('', end, origin, arrowprops={'arrowstyle':'-|>', 'color':TEAL,'lw':3})
    ax.annotate(label,end,xytext=offset,textcoords='offset points',color=TEAL)
start = origin+u*13; end = start+u*40
ax.annotate('',end,start,arrowprops={'arrowstyle':'-|>','color':RED,'lw':4})
ax.plot([start[0],end[0],end[0]],[start[1],start[1],end[1]],ls=':',color=RED,lw=2)
ax.text(78,49,'+7.071',color=RED,fontsize=13)
ax.text(107,75,'+7.071',color=RED,fontsize=13)
ax.text(22,126,'C = 45°：工件 +X 方向倾斜',color=INK,fontsize=14)
ax.text(42,21,'W 原点',color=TEAL,fontsize=13)
save(fig,'rotated-direction')

# A point on a tilting BC trunnion; explicit right-hand matrix convention.
fig = plt.figure(figsize=(7.8,5.9)); ax=fig.add_subplot(111,projection='3d')
b,c=np.deg2rad([30,45]); ry=np.array([[np.cos(b),0,np.sin(b)],[0,1,0],[-np.sin(b),0,np.cos(b)]])
rz=np.array([[np.cos(c),-np.sin(c),0],[np.sin(c),np.cos(c),0],[0,0,1]])
r=ry@rz; p=np.array([100,0,20]); q=r@p
for idx,label in enumerate(['$X_M$','$Y_M$','$Z_M$']):
    e=np.eye(3)[idx]*110;ax.quiver(0,0,0,*e,color=BLUE,arrow_length_ratio=.07,lw=1.5);ax.text(*e,label,color=BLUE,fontsize=13)
for idx,label in enumerate(['$X_W$','$Y_W$','$Z_W$']):
    e=r[:,idx]*75;ax.quiver(0,0,0,*e,color=TEAL,arrow_length_ratio=.09,lw=2.5);ax.text(*(e*1.13),label,color=TEAL,fontsize=15)
phi=np.linspace(0,2*np.pi,150);ring=r@np.array([100*np.cos(phi),100*np.sin(phi),np.zeros(len(phi))]);ax.plot(*ring,color=TEAL,alpha=.45)
ax.scatter(*q,color=RED,s=65);ax.plot([0,q[0]],[0,q[1]],[0,q[2]],color=RED,lw=2)
ax.text(*(q+np.array([0,0,9])),'刀尖目标',color=RED,fontsize=14)
ax.set(xlim=(-110,130),ylim=(-110,130),zlim=(-65,120));ax.set_box_aspect((1,1,.8));ax.view_init(elev=23,azim=-62);ax.set_axis_off()
save(fig,'bc-frame')

# Parent-child hierarchy, two separate mechanical branches.
fig,ax=plt.subplots(figsize=(11.2,4.5));ax.set(xlim=(0,11.2),ylim=(0,4.5));ax.axis('off')
def node(x,y,label,color=BLUE,w=1.6):
    ax.add_patch(Rectangle((x,y),w,.65,fc='#f1f6f8',ec=color,lw=1.8))
    ax.text(x+w/2,y+.325,label,ha='center',va='center',color=color,fontsize=15)
def link(a,b,color=GRAY):
    style = 'arc3,rad=0' if abs(a[1]-b[1]) < 1e-9 else 'angle,angleA=0,angleB=90,rad=8'
    ax.annotate('',b,a,arrowprops={'arrowstyle':'-|>','color':color,'lw':2,'connectionstyle':style})
node(.3,1.9,'机床底座',INK,1.8)
for x,label in [(3,'X 轴'),(5.5,'Y 轴'),(8,'Z 轴 + 刀具')]:node(x,3,label,BLUE,2.4 if x==8 else 1.6)
for x,label in [(3,'A 或 B 轴'),(6,'C 轴 + 工件')]:node(x,.8,label,TEAL,2.2)
link((2.1,2.2),(3,3.325),BLUE);link((4.6,3.325),(5.5,3.325),BLUE);link((7.1,3.325),(8,3.325),BLUE)
link((2.1,2.2),(3,1.125),TEAL);link((5.2,1.125),(6,1.125),TEAL)
ax.text(5.8,4.08,'刀具侧：直线运动链',ha='center',color=BLUE,fontsize=17)
ax.text(6,.22,'工件侧：倾转轴带着 C 转台一起运动',ha='center',color=TEAL,fontsize=16)
save(fig,'hierarchy')

# Offset pivot: the same 90-degree rotation produces different positions.
fig,axes=plt.subplots(1,2,figsize=(10.2,4.3))
for ax in axes: axis2(ax,145)
for ax, o, title in [(axes[0],np.array([0,0]),'旋转中心 O = (0, 0)'),(axes[1],np.array([30,20]),'旋转中心 O = (30, 20)')]:
    point(ax,*o,'O',TEAL,(-12,-20));start=np.array([100,20]);end=o+np.array([-(start-o)[1],(start-o)[0]])
    ax.plot([o[0],start[0]],[o[1],start[1]],color=GRAY,ls='--');ax.plot([o[0],end[0]],[o[1],end[1]],color=TEAL,lw=2)
    point(ax,*start,'起点 (100, 20)',GRAY,(-94,-22));point(ax,*end,f'终点 ({end[0]}, {end[1]})',RED,(8,4))
    ax.set_title(title,color=INK,pad=14,fontsize=16)
save(fig,'pivot-offset')

# Correct TCP arc versus independently interpolated joint chord.
fig,ax=plt.subplots(figsize=(6.5,5.2));axis2(ax,125)
t=np.linspace(0,np.pi/2,121);ax.plot(100*np.cos(t),100*np.sin(t),color=TEAL,lw=3)
ax.plot([100,0],[0,100],color=RED,lw=2.5)
q=100/np.sqrt(2);ax.scatter([q,50],[q,50],color=[TEAL,RED],s=95)
ax.plot([50,q],[50,q],color=INK,ls=':',lw=2)
ax.text(67,84,'RTCP：跟随工件点的圆弧',color=TEAL,ha='center',fontsize=14)
ax.text(35,33,'非 RTCP：独立轴线性插补',color=RED,ha='center',fontsize=14)
ax.annotate('中途偏离 29.289 mm',(61,61),xytext=(4,113),arrowprops={'arrowstyle':'->','color':INK},fontsize=14)
save(fig,'arc-chord')

# Tilted table and fixed vertical spindle: tool vector relative to workpiece changes.
fig,axes=plt.subplots(1,2,figsize=(10.2,4.1))
for ax,angle in zip(axes,[0,30]):
    ax.set(xlim=(-70,70),ylim=(-40,105),aspect='equal');ax.axis('off')
    phi=np.deg2rad(angle);u=np.array([np.cos(phi),-np.sin(phi)]);v=np.array([np.sin(phi),np.cos(phi)])
    ends=np.array([-60*u,60*u]);ax.plot(*ends.T,color=TEAL,lw=7)
    ax.add_patch(Rectangle((-9,58),18,35,fc='#b5c1c9',ec=INK,lw=1.5))
    ax.plot([0,0],[6,58],color=RED,lw=5);ax.scatter([0],[6],color=RED,s=45)
    ax.annotate('',(0,32),(0,78),arrowprops={'arrowstyle':'-|>','lw':2,'color':BLUE})
    ax.annotate('',tuple(v*57),(0,0),arrowprops={'arrowstyle':'-|>','lw':2,'color':TEAL})
    ax.text(18,80,'主轴方向固定',color=BLUE,fontsize=14);ax.text(*(v*62),'工件法向',color=TEAL,fontsize=14)
    ax.set_title(f'B = {angle}°',color=INK,pad=10)
save(fig,'orientation')

# Tool length: preserving spindle reference alone does not preserve the tip.
fig,ax=plt.subplots(figsize=(8.2,4));ax.set(xlim=(0,10),ylim=(0,5));ax.axis('off')
for x,length,label in [(1.7,1.1,'短刀'),(5,2.0,'长刀，主轴未补偿'),(8.3,2.0,'长刀，主轴上移')]:
    gauge=3.7 if x<7 else 4.6
    ax.add_patch(Rectangle((x-.45,gauge),.9,.3,fc='#bac8ce',ec=INK))
    ax.plot([x,x],[gauge-length,gauge],color=RED,lw=8)
    ax.scatter([x],[gauge-length],color=RED,s=70)
    ax.text(x,.28,label,ha='center',fontsize=14,color=INK)
    ax.text(x+.24,gauge-length/2,'L',color=RED,fontsize=17)
ax.axhline(2.6,color=TEAL,ls='--',lw=2);ax.text(.1,2.77,'目标刀尖高度',color=TEAL,fontsize=13)
ax.annotate('',(8.9,4.6),(8.9,3.7),arrowprops={'arrowstyle':'-|>','color':BLUE,'lw':2});ax.text(9.05,4.1,'ΔL',color=BLUE)
save(fig,'tool-length')

print('Created 9 original diagrams in',OUT)