import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.patches import Ellipse
from matplotlib import patches
from src import coil_class as BC
# %%
my_colors = {'light_green': '#97e144',
             'orange': '#FF914D',
             'light_grey': '#545454',
             'pastel_blue': '#1b64d1',
             'light_blue': '#71C8F4',
             'purple': '#7c588c'}

# %%
def myarrow(ax,posX,posY,direction=1):
    h=0.3
    w=0.35 * direction
    ax.arrow(posX,posY,w,h,linewidth=3,head_width=0,head_length=0)
    ax.arrow(posX,posY,w,-h,linewidth=3,head_width=0,head_length=0)

# %%
AHH = BC.BCoil(HH=-1, distance=3*np.sqrt(3), radius=3, layers=1,windings=1,wire_width=0.1
               ,wire_height=0.1,is_round=True)
ylim = 4
xlim = 3
x = np.arange(-xlim, xlim+0.1, 0.5)

z = np.arange(-ylim, ylim+0.1, 0.5)

x_m_AHH, z_m_AHH = np.meshgrid(x, z)

I_current = 1
B = AHH.B_multiple_3d(I_current, x, z, raster=1)

B_tot = BC.BCoil.B_tot_3d(B)
for xx in range(0, len(x)):
    for zz in range(0, len(z)):
        if B_tot[zz, xx] > 1.5:
            B[zz, xx, :] /= B_tot[zz, xx] / 1.5
# %%

HH = BC.BCoil(HH=1, distance=3, radius=3, layers=1,windings=1,wire_width=0.1
               ,wire_height=0.1,is_round=True)
ylim = 3
xlim = 3
x = np.arange(-xlim, xlim+0.1, 0.5)

z = np.arange(-ylim, ylim+0.1, 0.5)

x_m, z_m = np.meshgrid(x, z)

I_current = 1
B_HH = HH.B_multiple_3d(I_current, x, z, raster=1)

B_tot_HH = BC.BCoil.B_tot_3d(B_HH)
for xx in range(0, len(x)):
    for zz in range(0, len(z)):
        if B_tot_HH[zz, xx] > 1.5:
            B_HH[zz, xx, :] /= B_tot_HH[zz, xx] / 1.5

# %%
mpl.rcParams.update({'font.size':11})

fig = plt.figure(figsize=(5,2.5),dpi=600)


ax = fig.add_subplot(121,aspect=1)

ypos = 1.5
width = 6
height = 1.3
lwidth = 4

#axes
ylim = 4
xlim = 3.5


e1 = Ellipse((0,ypos), width, height, linewidth = lwidth, fill = False, capstyle='projecting', zorder=2)
e2 = Ellipse((0,-ypos), width, height, linewidth = lwidth, fill = False, capstyle='projecting', zorder=2)


ax.add_patch(e1)
ax.add_patch(e2)

#add arrows to ellipse
myarrow(ax,0.5,0.86)
myarrow(ax,0.5,-2.14, direction=1)

# make axis
x_min = 3.3
ax.arrow(-x_min,0,2*x_min,0,head_width=0.2, length_includes_head=True, facecolor='black',zorder=0)
ax.text(x_min + 0.16 ,-0.15,'y',alpha=1)

y_min = 3.2
ax.arrow(0,-y_min,0, 2*y_min,head_width=0.2, length_includes_head=True, facecolor='black', zorder=0)
ax.text(-0.2, y_min +0.25, 'z')


ax.set_ylim(-ylim,ylim)
ax.set_xlim(-xlim,xlim)

# plot field
x_m, z_m = np.meshgrid(x, z)

arr_color = my_colors['orange']
zord=3

lim1=0
lim2=3
ax.quiver(x_m[lim1:lim2, :], z_m[lim1:lim2, :], B_HH[lim1:lim2, :, 0], B_HH[lim1:lim2, :, 1], color=arr_color,zorder=1)
lim1=3
lim2=6
ax.quiver(x_m[lim1:lim2, :], z_m[lim1:lim2, :], B_HH[lim1:lim2, :, 0], B_HH[lim1:lim2, :, 1], color=arr_color,zorder=3)
lim1=6
lim2=9
ax.quiver(x_m[lim1:lim2, :], z_m[lim1:lim2, :], B_HH[lim1:lim2, :, 0], B_HH[lim1:lim2, :, 1], color=arr_color,zorder=1)
lim1=9
lim2=12
ax.quiver(x_m[lim1:lim2, :], z_m[lim1:lim2, :], B_HH[lim1:lim2, :, 0], B_HH[lim1:lim2, :, 1], color=arr_color,zorder=3)

# add distance + radius annotation

d = patches.FancyArrowPatch((-width/2 - 0.5, -ypos-0.2), (-width/2 - 0.5, ypos+0.2), arrowstyle='<->',mutation_scale=10,
                           linewidth=1, color=my_colors['pastel_blue'])

ax.text(-width/2-1, -0.18, 'd', color=my_colors['pastel_blue'])

y_off = -3.4
R = patches.FancyArrowPatch((-width/2 , y_off), (0.13, y_off), arrowstyle='<->',mutation_scale=10,
                           linewidth=1, color=my_colors['pastel_blue'])

ax.text(-width/4-0.1,y_off-0.5, 'r', color=my_colors['pastel_blue'])

ax.add_patch(d)
ax.add_patch(R)

# general plot settings
ax.set_ylim(-ylim-0.5,ylim+0.5)
ax.set_xlim(-xlim-0.5,xlim+0.5)

ax.axis('off')
#ax.title.set_text('(a)',pos='left')
ax.set_title('(a)',loc='left')


# plt.savefig("Out/magnetic_configuration.png")
##########################################################################################################

ax = fig.add_subplot(122, aspect=1)


width = 6
scale = width/4 *np.sqrt(3) - ypos
ypos = ypos + scale
height = 1.3
print(ypos)
#axes


e1 = Ellipse((0,ypos), width, height, linewidth = lwidth, fill = False, capstyle='projecting', zorder=2)
e2 = Ellipse((0,-ypos), width, height, linewidth = lwidth, fill = False, capstyle='projecting', zorder=2)


ax.add_patch(e1)
ax.add_patch(e2)

#add arrows to ellipse
myarrow(ax,0.5,0.86+scale)
myarrow(ax,-0.5,-2.14-scale, direction=-1)

# make axis
x_min = 3.2
ax.arrow(-x_min,0,2*x_min,0,head_width=0.2, length_includes_head=True, facecolor='black',zorder=0)
ax.text(x_min + 0.16 ,-0.15,'y',alpha=1)

y_min = 4.1
ax.arrow(0,-y_min,0, 2*y_min,head_width=0.2, length_includes_head=True, facecolor='black', zorder=0)
ax.text(-0.2, y_min +0.25, 'z')


ax.set_ylim(-ylim,ylim)
ax.set_xlim(-xlim,xlim)

ax.axis('off')

#ax = fig.add_subplot(222,aspect=1)

x_m = x_m_AHH
z_m = z_m_AHH
print(np.shape(B))
arr_color = my_colors['orange']
zord=3

lim1=1
lim2=3
ax.quiver(x_m[lim1:lim2, :], z_m[lim1:lim2, :], B[lim1:lim2, :, 0], B[lim1:lim2, :, 1], color=arr_color,zorder=1)

ax.quiver(x_m[3, 1:-1], z_m[3, 1:-1], B[3, 1:-1, 0], B[3, 1:-1, 1], color=arr_color,zorder=1)
ax.quiver(x_m[3, 0], z_m[3, 0], B[3, 0, 0], B[3, 0, 1], color=arr_color,zorder=3)
ax.quiver(x_m[3, -1], z_m[3, -1], B[3, -1, 0], B[3, -1, 1], color=arr_color,zorder=3)
lim1=4
lim2=6
ax.quiver(x_m[lim1:lim2, :], z_m[lim1:lim2, :], B[lim1:lim2, :, 0], B[lim1:lim2, :, 1], color=arr_color,zorder=3)
lim1=6
lim2=13
ax.quiver(x_m[lim1:lim2, :], z_m[lim1:lim2, :], B[lim1:lim2, :, 0], B[lim1:lim2, :, 1], color=arr_color,zorder=1)
lim1=14
lim2=len(x_m)-1
ax.quiver(x_m[lim1:lim2, :], z_m[lim1:lim2, :], B[lim1:lim2, :, 0], B[lim1:lim2, :, 1], color=arr_color,zorder=3)


ax.quiver(x_m[13, 1:-1], z_m[13, 1:-1], B[13, 1:-1, 0], B[13, 1:-1, 1], color=arr_color,zorder=3)
ax.quiver(x_m[13, 0], z_m[13, 0], B[13, 0, 0], B[13, 0, 1], color=arr_color,zorder=1)
ax.quiver(x_m[13, -1], z_m[13, -1], B[13, -1, 0], B[13, -1, 1], color=arr_color,zorder=1)

# Add d + R annotations
d = patches.FancyArrowPatch((-width/2 - 0.5, -ypos-0.2), (-width/2 - 0.5, ypos+0.2), arrowstyle='<->',mutation_scale=10,
                           linewidth=1, color=my_colors['pastel_blue'])

ax.text(-width/2-1, -0.18, 'd', color=my_colors['pastel_blue'])

y_off = -4.2
R = patches.FancyArrowPatch((-width/2 , y_off), (0.13, y_off), arrowstyle='<->',mutation_scale=10,
                           linewidth=1, color=my_colors['pastel_blue'])

ax.text(-width/4-0.1,y_off-0.5, 'r', color=my_colors['pastel_blue'])

ax.add_patch(d)
ax.add_patch(R)

ax.set_ylim(-ylim-0.5,ylim+0.5)
ax.set_xlim(-xlim-0.5,xlim+0.5)

ax.axis('off')
ax.set_title('(b)',loc='left')

plt.savefig("C:/Users/Joschka/Desktop/Magnetic_field_coils_project/Thesis_Plots/Coil_Design/Out/magnetic_configuration.png")

plt.show()