__author__ = 'wack'

# part of the magwire package

# calculate magnetic fields arising from electrical current through wires of arbitrary shape
# with the law of Biot-Savart

# written by Michael Wack
# wack@geophysik.uni-muenchen.de

# tested with python 3.4.3


import numpy as np
import time
#import multiprocessing as mp
from matplotlib import pyplot as plt
import mpl_toolkits.mplot3d.axes3d as ax3d

class BiotSavart:
    '''
    calculates the magnetic field generated by currents flowing through wires
    '''

    def __init__(self, wire=None):
        self.wires = []
        if wire is not None:
            self.wires.append(wire)

    def AddWire(self, wire):
        self.wires.append(wire)

    def CalculateB(self, points):
        """
        calculate magnetic field B at given points
        :param points: numpy array of n points (xyz)
        :return: numpy array of n vectors representing the B field at given points
        """

        print("found {} wire(s).".format(len(self.wires)))
        c = 0
        # generate list of IdL and r1 vectors from all wires
        for w in self.wires:
            c += 1
            _IdL, _r1 = w.IdL_r1
            print("wire {} has {} segments".format(c, len(_IdL)))
            if c == 1:
                IdL = _IdL
                r1 = _r1
            else:
                IdL = np.vstack((IdL, _IdL))
                r1 = np.vstack((r1, _r1))
        print("total number of segments: {}".format(len(IdL)))
        print("number of field points: {}".format(len(points)))
        print("total number of calculations: {}".format(len(points)*len(IdL)))

        # now we have
        # all segment vectors multiplied by the flowing current in IdL
        # and all vectors to the central points of the segments in r1

        # calculate vector B*1e7 for each point in space
        t1 = time.process_time()
        # simple list comprehension to calculate B at each point r
        B = np.array([BiotSavart._CalculateB1(r, IdL, r1) * 1e-7 for r in points])

        # multi processing
        # slower than single processing?
        #pool = mp.Pool(processes=16)
        #B = np.array([pool.apply(_CalculateB1, args=(r, IdL, r1)) for r in points])

        t2 = time.process_time()
        print("time needed for calculation: {} s".format(t2-t1))

        return B

    def vv_PlotWires(self):
        for w in self.wires:
            w.vv_plot_path()

    def mpl3d_PlotWires(self, ax):
        for w in self.wires:
            w.mpl3d_plot_path(show=False, ax=ax)




    @staticmethod
    def _CalculateB1(r, IdL, r1):
        '''
        calculate magnetic field B for one point r in space
        :param r: 3 component numpy array representing the location where B will be calculated
        :param IdL:  all segment vectors multiplied by the flowing current
        :param r1: all vectors to the central points of the segments
        :return: numpy array of 3 component vector of B multiplied by 1e7
        '''

        # calculate law of biot savart for all current elements at given point r
        r2 = r - r1
        r25 = np.linalg.norm(r2, axis=1)**3
        r3 = r2 / r25[:, np.newaxis]

        cr = np.cross(IdL, r3)

        # claculate sum of contributions from all current elements
        s = np.sum(cr, axis=0)

        return s