import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from src import coil_class as BC
from src import physical_constants as cs
import logging as log


def main():
    AHH_Coil = BC.BCoil(HH=-1, distance=69.313, radius=47, layers=8, windings=16,
                        wire_height=0.5, wire_width=0.5, insulation_thickness=0.068 / 2,
                        is_round=True, winding_scheme=2)
    
    step = 0.03
    xlim_mm = 30
    xlim = int(xlim_mm/step) + 1
    ylim_mm = 30
    ylim = int(ylim_mm/step) + 1
    zlim_mm = 20
    zlim = int(zlim_mm/step) + 1


    # x = np.linspace(0,xlim_mm,xlim) * 1e-3
    # y = np.linspace(0,ylim_mm,ylim) * 1e-3
    # z = np.linspace(0,zlim_mm,zlim) * 1e-3

    x = np.linspace(-xlim_mm, xlim_mm, 2 *xlim - 1) * 1e-3
    y = np.linspace(-ylim_mm, ylim_mm, 2 *ylim - 1) * 1e-3
    z = np.linspace(-zlim_mm, zlim_mm, 2 *zlim - 1) * 1e-3



    b_plot = np.load('output/B_quarter.npy')
    #b_raw = np.load('output/b_cart_1_step_5.npy')
    #b_raw = np.load('output/b_full_test.npy')

    b_plot = np.swapaxes(b_raw,0,1) # Meshgrid uses different order of axis (y,x,z)
    # x_m, z_m = np.meshgrid(x, z)
    #
    # plt.figure(34)
    #
    # plt.quiver(x_m,z_m,np.transpose(b_plot[0, :, :, 0]), np.transpose(b_plot[0, :, :, 2]))
    # plt.xlabel("x-axis [mm]")
    # plt.ylabel("z-axis [mm]")
    # plt.title("x-z plane")
    # plt.show()
    #
    # x_m, y_m = np.meshgrid(x, y)
    #
    # plt.figure(35)
    #
    # plt.quiver(x_m, y_m, b_plot[:, :, 0, 0], b_plot[:, :, 0, 1])
    # plt.xlabel("x-axis [mm]")
    # plt.ylabel("y-axis [mm]")
    # plt.title("x-y plane")
    # plt.show()
    #
    # y_m, z_m = np.meshgrid(y, z)
    #
    # plt.figure(36)
    #
    # plt.quiver(y_m, z_m, np.transpose(b_plot[0, :, :, 1]), np.transpose(b_plot[0, :, :, 2]))
    # plt.xlabel("y-axis [mm]")
    # plt.ylabel("z-axis [mm]")
    # plt.title("y-z plane")
    # plt.show()

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

    print(np.shape(x_m))
    print(np.shape(b_plot))
    fig = plt.figure(37)
    ax = fig.gca(projection = '3d')
    #print(b_plot[:, :, :, 0])
    #print(b_plot[:, 0, :, 0])
    scaling = 0.0001
    ax.quiver(x_m, y_m, z_m, scaling * b_plot[:, :, :, 0], scaling * b_plot[:, :, :, 1], scaling * b_plot[:, :, :, 2])
    plt.xlabel("x-axis [mm]")
    plt.ylabel("y-axis [mm]")
    #plt.zlabel("z-axis [mm]")
    #plt.title("full 3d")
    plt.show()


    # AHH_Coil.B_quick_plot(5,abs = False)
    #AHH_Coil.plot_3d(5,30,20)





if __name__ == '__main__':
    main()