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
from matplotlib.ticker import AutoMinorLocator
# %%
my_colors = {'light_green': '#97e144',
             'orange': '#FF914D',
             'light_grey': '#545454',
             'pastel_blue': '#1b64d1',
             'light_blue': '#71C8F4',
             'purple': '#7c588c'}


mpl.rcParams.update({'font.size': 10, 'axes.linewidth': 1, 'lines.linewidth': 2, 'text.usetex': False, 'font.family': 'arial'})
mpl.rcParams['xtick.direction'] = 'in'
mpl.rcParams['ytick.direction'] = 'in'

mpl.rcParams['xtick.top'] = True
mpl.rcParams['ytick.right'] = True

def a_scat(B, a_bg = 1, B_0 = 0, DeltaB = 1):
    return a_bg * (1 - (DeltaB/(B - B_0)))

mm = 1/25.4
# %%

HH_Coil = BC.BCoil(HH = 1, distance = 54, radius = 48, layers = 8, windings = 8, wire_height = 0.5,
                   wire_width = 0.5, insulation_thickness = (0.546-0.5)/2, is_round = True,
                   winding_scheme= 2)

HH_Coil.set_R_inner(45.6)
HH_Coil.set_d_min(47.4)
HH_Coil.print_info()

AHH_Coil = BC.BCoil(HH = -1, distance = 54, radius = 48, layers = HH_Coil.get_layers, windings=2 * HH_Coil.get_windings,
                    wire_height = 0.5, wire_width=0.5, insulation_thickness=(0.546-0.5)/2,
                    is_round = True, winding_scheme= 2, red_N=1)

AHH_Coil.set_R_inner(45.6)
AHH_Coil.set_d_max(77.6)
AHH_Coil.print_info()

# %%
x = np.linspace(-80, 80, 1000)
z = np.linspace(-80, 80, 1000)

I = 1
Bz, Bx = HH_Coil.B_tot_along_axis(I, x, z, raster = 7)
Bz_grad = np.abs(BC.BCoil.grad(Bz, z))
Bx_grad = np.abs(BC.BCoil.grad(Bx, x))


# %%
#HH plots
fig, axes = plt.subplots(2, 2, figsize=(5.8, 4), gridspec_kw={'height_ratios':[12,5]},dpi=600)
#fig.tight_layout()

lim1 = 55
lim1 =(-lim1, lim1)
lim2= 5
lim2 = (-lim2, lim2)

axs = axes[0, 0]

axs.plot(x, Bx, color=my_colors['pastel_blue'], label=r'$x/y-$axis')
axs.plot(z, Bz, color=my_colors['orange'], label=r'$z-$axis')

axs.set_ylabel(r'$B_{tot}$ (G/cm)')
axs.set_xlim(lim1)
axs.set_xticks(np.arange(-40,42,20))
axs.legend(loc=8)



axs = axes[1, 0]

axs.plot(x, Bx, color=my_colors['pastel_blue'])
axs.plot(z, Bz, color=my_colors['orange'])
axs.set_ylabel(r'$B_{tot}$ (G)')
axs.set_xlim(lim2)
axs.set_ylim(10.725,10.775)
axs.set_xlabel('distance to center (mm)')
axs.set_yticks([10.73,10.75,10.77])
axs.set_xticks(np.arange(-4,5,2))

axs = axes[0, 1]

axs.plot(x, Bx_grad, color=my_colors['pastel_blue'])
axs.plot(z, Bz_grad, color=my_colors['orange'])
axs.set_xlim(lim1)
axs.set_ylabel(r'$| ∇ B_{tot}|$ (G/cm)')
axs.set_xticks(np.arange(-40,42,20))


axs = axes[1, 1]

axs.plot(x, Bx_grad, color=my_colors['pastel_blue'])
axs.plot(z, Bz_grad, color=my_colors['orange'])
axs.set_ylabel(r'$| ∇ B_{tot}|$ (G/cm)')
axs.set_xlim(lim2)
axs.set_ylim(-0.01,0.09)
axs.set_yticks([0.0,0.04, 0.08])
axs.set_xticks(np.arange(-4,5,2))

axs.set_xlabel('distance to center (mm)')

for i in [0,1]:
    for j in [0,1]:
        axes[1,j].xaxis.set_minor_locator(AutoMinorLocator(2))
        axes[0, j].xaxis.set_minor_locator(AutoMinorLocator(2))
        axes[i,j].yaxis.set_minor_locator(AutoMinorLocator(2))

        axes[i, 0].yaxis.set_label_coords(-0.19, 0.5)
        axes[i, 1].yaxis.set_label_coords(-0.17, 0.5)
fig.tight_layout()
fig.savefig('Thesis_Plots/Coil_Design/Out/HH_field.png')
plt.show()

# %%
x = np.linspace(-80, 80, 1000)
z = np.linspace(-80, 80, 1000)

I = 1
Bz, Bx = AHH_Coil.B_field(I, x, z, raster = 7)
Bz_grad = np.abs(BC.BCoil.grad(Bz, z))
Bx_grad = np.abs(BC.BCoil.grad(Bx, x))

Bz = np.abs(Bz)
Bx = np.abs(Bx)


# %%
#AHH plots
fig, axes = plt.subplots(2, 2, figsize=(5.8, 4), gridspec_kw={'height_ratios':[12,5]},dpi=600)
#fig.tight_layout()

lim1 = 55
lim1 =(-lim1, lim1)
lim2= 5
lim2 = (-lim2, lim2)

axs = axes[0, 0]

axs.plot(x, Bx, color=my_colors['pastel_blue'], label=r'$x/y-$axis')
axs.plot(z, Bz, color=my_colors['orange'], label=r'$z-$axis')

axs.set_ylabel(r'$B_{tot}$ (G/cm)')
axs.set_xlim(lim1)
axs.set_xticks(np.arange(-40,42,20))




axs = axes[1, 0]

axs.plot(x, Bx, color=my_colors['pastel_blue'])
axs.plot(z, Bz, color=my_colors['orange'])
axs.set_ylabel(r'$B_{tot}$ (G)')
axs.set_xlim(lim2)
axs.set_ylim(-0.2, 3.5)
axs.set_xlabel('distance to center (mm)')
axs.set_yticks([0,1,2,3])
axs.set_xticks(np.arange(-4,5,2))

axs = axes[0, 1]

axs.plot(x, Bx_grad, color=my_colors['pastel_blue'], label=r'$x/y-$axis')
axs.plot(z, Bz_grad, color=my_colors['orange'], label=r'$z-$axis')
axs.set_xlim(lim1)
axs.set_ylabel(r'$| ∇ B_{tot}|$ (G/cm)')
axs.set_xticks(np.arange(-40,42,20))
axs.legend()


axs = axes[1, 1]

axs.plot(x, Bx_grad, color=my_colors['pastel_blue'])
axs.plot(z, Bz_grad, color=my_colors['orange'])
axs.set_ylabel(r'$| ∇ B_{tot}|$ (G/cm)')
axs.set_xlim(lim2)
axs.set_ylim(1.5,6.4)
#axs.set_yticks([0.0,0.04, 0.08])
axs.set_xticks(np.arange(-4,5,2))

axs.set_xlabel('distance to center (mm)')

for i in [0,1]:
    for j in [0,1]:
        axes[1,j].xaxis.set_minor_locator(AutoMinorLocator(2))
        axes[0, j].xaxis.set_minor_locator(AutoMinorLocator(2))
        axes[i,j].yaxis.set_minor_locator(AutoMinorLocator(2))

        axes[i,0].yaxis.set_label_coords(-0.15,0.5)
        axes[i, 1].yaxis.set_label_coords(-0.08, 0.5)
fig.tight_layout()
fig.savefig('Thesis_Plots/Coil_Design/Out/AHH_field.png')
plt.show()