"""Original vector figures generated from the book's stated models."""
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle, Circle, Polygon, Arc
OUT=Path(__file__).resolve().parents[1]/'dist/figures'
OUT.mkdir(parents=True,exist_ok=True)
plt.rcParams.update({'font.family':'DejaVu Sans','font.size':10,'axes.spines.top':False,'axes.spines.right':False,'svg.fonttype':'none','pdf.fonttype':42,'axes.labelcolor':'#1b314c','text.color':'#1b314c','axes.edgecolor':'#8293a8'})
BLUE='#2364cc'; INK='#1b314c'; TEAL='#087e88'; LIGHT='#eaf1fb'; GRAY='#a4b4c7'
def save(fig,name):
 for ext in ('svg','pdf'):fig.savefig(OUT/f'{name}.{ext}',bbox_inches='tight',pad_inches=.13,facecolor='white')
 plt.close(fig)
def circuit(labels,width=8):
 n=len(labels);fig,ax=plt.subplots(figsize=(7.0,.55*n+.5));ax.set_xlim(-1,width+.55);ax.set_ylim(n-.4,-.65);ax.axis('off')
 for j,l in enumerate(labels):ax.plot([0,width],[j,j],color=GRAY,lw=1.25,zorder=0);ax.text(-.15,j,l,ha='right',va='center')
 return fig,ax

def gate(ax,x,q,label,w=.6,h=.58):
 ax.add_patch(Rectangle((x-w/2,q-h/2),w,h,facecolor=LIGHT,edgecolor=BLUE,lw=1.4,zorder=3));ax.text(x,q,label,ha='center',va='center',fontsize=10,zorder=4)
def cx(ax,x,c,t):
 ax.plot([x,x],[c,t],color=INK,lw=1.4);ax.add_patch(Circle((x,c),.07,color=INK,zorder=4));ax.add_patch(Circle((x,t),.19,facecolor='white',edgecolor=INK,lw=1.4,zorder=4));ax.plot([x-.19,x+.19],[t,t],color=INK,lw=1.4,zorder=5);ax.plot([x,x],[t-.19,t+.19],color=INK,lw=1.4,zorder=5)
def control(ax,x,c,t,label):
 ax.plot([x,x],[c,t],color=INK,lw=1.4);ax.add_patch(Circle((x,c),.07,color=INK,zorder=4));gate(ax,x,t,label)
def swap(ax,x,a,b):
 ax.plot([x,x],[a,b],color=INK,lw=1.4)
 for q in (a,b):
  ax.plot([x-.13,x+.13],[q-.13,q+.13],color=INK,lw=1.6);ax.plot([x-.13,x+.13],[q+.13,q-.13],color=INK,lw=1.6)

fig,ax=circuit([r'$q_0:|0\rangle$',r'$q_1:|0\rangle$'],5)
gate(ax,1,0,'H');cx(ax,2.6,0,1);ax.text(5.2,.5,r'$|\Phi^+\rangle$',va='center');save(fig,'bell')

fig,ax=circuit([r'$q_0:|\psi\rangle$',r'$q_1:|0\rangle$',r'$q_2:|0\rangle$'],9)
gate(ax,.6,1,'H');cx(ax,1.7,1,2);cx(ax,3.0,0,1);gate(ax,4.1,0,'H');gate(ax,5.2,0,r'$M_Z$');gate(ax,5.2,1,r'$M_Z$')
for y,l,x in [(1,'b',6.5),(0,'a',7.9)]:
 ax.plot([5.55,x,x],[y,y,2-.33],color=TEAL,ls='--',lw=1.4);ax.text(5.85,y-.17,l,color=TEAL)
gate(ax,6.5,2,r'$X^b$');gate(ax,7.9,2,r'$Z^a$');ax.text(9.15,2,r'$|\psi\rangle$',va='center');save(fig,'teleport')

fig,ax=circuit([r'$|0\rangle^{\otimes n}$',r'$|1\rangle$'],6.8)
gate(ax,1,0,r'$H^{\otimes n}$',1.0);gate(ax,1,1,'H')
ax.add_patch(Rectangle((2.6,-.38),1,1.76,facecolor=LIGHT,edgecolor=BLUE,lw=1.4,zorder=3));ax.text(3.1,.5,r'$U_f$',ha='center',va='center',fontsize=13)
gate(ax,4.8,0,r'$H^{\otimes n}$',1.0);gate(ax,6.1,0,r'$M_Z$');ax.text(6.9,1,r'$|-\rangle$',va='center');save(fig,'oracles')

fig,ax=circuit([r'$q_0$',r'$q_1$'],6.3)
gate(ax,.9,0,'H');control(ax,2.5,1,0,r'$P_{\pi/2}$');gate(ax,4.0,1,'H');swap(ax,5.5,0,1);save(fig,'qft')

fig,ax=circuit([r'$q_0:|0\rangle$',r'$q_1:|0\rangle$',r'$q_2:|0\rangle$',r'$|u\rangle$'],9)
for q in range(3):gate(ax,.7,q,'H')
for q,x,p in [(0,2.1,4),(1,3.5,2),(2,4.9,1)]:control(ax,x,q,3,r'$U^{'+str(p)+'}$')
ax.add_patch(Rectangle((5.9,-.35),1,2.7,facecolor=LIGHT,edgecolor=BLUE,lw=1.4,zorder=3));ax.text(6.4,1,r'$F_8^\dagger$',ha='center',va='center',fontsize=12)
for q in range(3):gate(ax,8,q,r'$M_Z$')
ax.text(9.13,1,'y / 8',va='center');save(fig,'qpe')

fig,ax=circuit([r'$|\psi\rangle$',r'$|0\rangle$',r'$|0\rangle$',r'$a_0:|0\rangle$',r'$a_1:|0\rangle$'],9)
cx(ax,.65,0,1);cx(ax,1.65,0,2)
for x,c,t in [(3,0,3),(4,1,3),(5.2,1,4),(6.2,2,4)]:cx(ax,x,c,t)
gate(ax,7.5,3,r'$M_Z$');gate(ax,7.5,4,r'$M_Z$');ax.text(8.2,3,r'$s_0$',va='center');ax.text(8.2,4,r'$s_1$',va='center');save(fig,'repetition')

fig=plt.figure(figsize=(6,4.4));ax=fig.add_subplot(111,projection='3d');ax.view_init(elev=23,azim=36)
t=np.linspace(0,2*np.pi,160)
for x,y,z in [(np.cos(t),np.sin(t),0*t),(np.cos(t),0*t,np.sin(t)),(0*t,np.cos(t),np.sin(t))]:ax.plot(x,y,z,color=GRAY,lw=.8)
for v,l in [(np.array([1.,0,0]),r'$X:\ |+\rangle$'),(np.array([0,1.,0]),r'$Y:\ |+i\rangle$'),(np.array([0,0,1.]),r'$Z:\ |0\rangle$')]:
 ax.plot(*np.array([-v,v]).T,color=GRAY,lw=1);ax.text(*(1.18*v),l)
r=np.array([0,12/13,-5/13]);ax.quiver(0,0,0,*r,color=BLUE,linewidth=2.5,arrow_length_ratio=.1);ax.scatter(*r,color=BLUE,s=32)
ax.text(.35,.95,-.65,r'$(0,12/13,-5/13)$',fontsize=10);ax.text(0,0,-1.25,r'$|1\rangle$');ax.set(xlim=(-1.3,1.3),ylim=(-1.3,1.3),zlim=(-1.3,1.3));ax.set_box_aspect((1,1,1));ax.set_axis_off();save(fig,'bloch')

fig,axes=plt.subplots(1,3,figsize=(7,2.6),sharey=True)
for ax,v,title in zip(axes,[[.5]*4,[.5,.5,-.5,.5],[0,0,1,0]],['Uniform preparation','Mark item 10','Invert about mean']):
 ax.axhline(0,color=GRAY,lw=.8);ax.bar(range(4),v,color=[BLUE if j==2 else '#78a0db' for j in range(4)],width=.6);ax.set_xticks(range(4),['00','01','10','11']);ax.set_title(title,fontsize=10);ax.set_ylim(-.65,1.2);ax.set_xlabel('Basis state')
axes[0].set_ylabel('Real amplitude');fig.tight_layout();save(fig,'grover')

fig,axes=plt.subplots(1,2,figsize=(6.8,3.5));pos={j:(j%3,2-j//3) for j in range(9)}
xsets=[[0,1],[1,2,4,5],[3,4,6,7],[7,8]];zsets=[[3,6],[0,1,3,4],[4,5,7,8],[2,5]]
for ax,sets,col,typ in zip(axes,[xsets,zsets],[BLUE,TEAL],['X','Z']):
 for k,ss in enumerate(sets):
  pts=np.array([pos[j] for j in ss]);xmin,ymin=pts.min(0);xmax,ymax=pts.max(0)
  if len(ss)==4:
   ax.add_patch(Rectangle((xmin,ymin),xmax-xmin,ymax-ymin,facecolor=col,alpha=.13,edgecolor=col,lw=2));label=((xmin+xmax)/2,(ymin+ymax)/2)
  else:
   ax.plot(pts[:,0],pts[:,1],color=col,lw=7,alpha=.28);label=((xmin+xmax)/2,(ymin+ymax)/2)
   label=(label[0]+(.32 if label[0]==2 else -.32 if label[0]==0 else 0),label[1]+(.3 if label[1]==2 else -.3 if label[1]==0 else 0))
  ax.text(*label,f'{typ}{k}',ha='center',va='center',color=col,fontsize=11,fontweight='bold')
 for j,(x,y) in pos.items():ax.add_patch(Circle((x,y),.14,facecolor='white',edgecolor=INK,lw=1.3,zorder=3));ax.text(x,y,str(j),ha='center',va='center',fontsize=9,zorder=4)
 ax.set_title(f'{typ}-type check supports');ax.set_aspect('equal');ax.set_xlim(-.55,2.55);ax.set_ylim(-.5,2.5);ax.axis('off')
fig.tight_layout();save(fig,'surface')

fig,axes=plt.subplots(1,2,figsize=(6.8,2.9))
for ax,phase in zip(axes,[0,np.pi]):
 alpha=1/np.sqrt(2);beta=np.exp(1j*phase)/np.sqrt(2)
 ax.axhline(0,color=GRAY,lw=.8);ax.axvline(0,color=GRAY,lw=.8)
 for z,c,l,dy in [(alpha,BLUE,r'$\alpha$',.09),(beta,TEAL,r'$\beta$',-.15)]:
  ax.annotate('',xy=(np.real(z),np.imag(z)),xytext=(0,0),arrowprops={'arrowstyle':'->','color':c,'lw':2});ax.text(np.real(z)*.65,np.imag(z)+dy,l,color=c)
 ax.set(xlim=(-1,1),ylim=(-.7,.7),xlabel='Real part',ylabel='Imaginary part',title='Relative phase '+('0' if phase==0 else 'π'));ax.set_aspect('equal');ax.set_xticks([-1,0,1]);ax.set_yticks([-.5,0,.5])
fig.tight_layout();save(fig,'phase')

fig,ax=plt.subplots(figsize=(6.4,3.1));p=np.linspace(0,.5,151)
from math import comb
for d,col in [(1,INK),(3,BLUE),(5,TEAL),(7,'#8e56a5')]:
 y=sum(comb(d,k)*p**k*(1-p)**(d-k) for k in range(d//2+1,d+1));ax.plot(p,y,label=f'd = {d}',color=col,lw=1.8)
ax.set(xlabel='Independent X-error probability p',ylabel='Logical failure probability',xlim=(0,.5),ylim=(0,.5));ax.legend(frameon=False);ax.grid(alpha=.14);fig.tight_layout();save(fig,'repetition-scaling')

fig,ax=plt.subplots(figsize=(6.4,3.1));n=np.arange(1,51);ax.semilogy(n,16.*2**n/2**30,color=BLUE,lw=2,label='State vector');ax.semilogy(n,16.*4**n/2**30,color=TEAL,lw=2,label='Dense density matrix');ax.axhline(32,color=INK,ls='--',lw=1);ax.text(2,65,'32 GiB reference',fontsize=9);ax.set(xlabel='Qubits n',ylabel='Memory (GiB)',xlim=(1,50),ylim=(1e-8,1e23));ax.legend(frameon=False);ax.grid(alpha=.14);fig.tight_layout();save(fig,'memory')
print('Created',len(list(OUT.glob('*.svg'))),'original figures in SVG and PDF.')
