# paste this code at the end of VectorFieldPlot 1.4
import scipy.optimize as op

M = 1.0 # magnetic moment
r = 0.65; l = 1.2
n = 10
fieldB = Field({'coils':[[0, 0, 0, r, l, M/(r**2*pi)]]})
fieldH = Field({'charged_discs':[[-l, -r, -l, r, -0.5*M/l],
                                 [ l, -r,  l, r,  0.5*M/l]]})
inside = lambda p: -min(1-fabs(p[0]/l), 1-fabs(p[1]/r))
outside = lambda p: min(1-fabs(p[0]/l), 1-fabs(p[1]/r))

for field in ['B', 'H']:
    fname = 'VFPt_magnet_' + field
    doc = FieldplotDocument(fname, width=400, height=300, commons=True)
    
    for i in range(n):
        r_rel = 2 * (0.5 + i) / n - 1
        y0 = r * (r_rel + 0.16 * (r_rel - r_rel**3))
        p0 = sc.array([l, y0])
        line1 = FieldLine(fieldB, p0*(1+1e-9), directions='forward',
                          maxr=1e4, maxn=1e4, stop_funcs=2*[outside])
        doc.draw_line(line1, arrows_style={})
        
        # B-field inside magnet
        line2 = FieldLine(fieldB, p0*(1-1e-9), directions='backward',
                          stop_funcs=2*[inside])
        if field == 'B':
            doc.draw_line(line2, arrows_style={})
        
        # H-field inside magnet
        line3 = FieldLine(fieldH, p0*(1-1e-9), directions='forward',
            stop_funcs=2*[inside])
        if field == 'H':
            doc.draw_line(line3, arrows_style={
                          'min_arrows':0, 'dist':0.6, 'max_arrows':2})
        
        p1 = line3.nodes[-1]['p'] # where lines leave the magnet on the side
        if fabs(p1[1]) >= r * (1-1e-8):
            p1[1] = copysign(max(r, fabs(p1[1])), p1[1])
            line4 = FieldLine(fieldH, p1, directions='forward',
                              stop_funcs=2*[outside])
            doc.draw_line(line4, arrows_style={'max_arrows':2})
            
            p2 = line4.nodes[-1]['p'] # where lines re-enter the magnet
            line5 = FieldLine(fieldH, p2, directions='forward',
                              stop_funcs=2*[inside])
            if field == 'H':
                doc.draw_line(line5, arrows_style={
                              'min_arrows':0, 'dist':0.6, 'max_arrows':2})
            
            line6 = FieldLine(fieldB, p1*(1-1e-9), directions='backward',
                              stop_funcs=2*[inside])
            if field == 'B':
                doc.draw_line(line6, arrows_style={})
    
    if field == 'H':
        doc.draw_object('path', {'d':'M {},{} V {} M {},{} V {}'.format(
        -l, -r, r, l, -r, r), 'style':'stroke:#999999; stroke-width:0.02'})
    
    doc.draw_magnets(fieldB)
    
    # postprocessing: strip text and add opacity
    for el in doc.svg.iter():
        if el.tag.endswith('g'):
            if el.get('id') != None and 'magnet' in el.get('id'):
                el.set('opacity', '0.4')
        if el.tag.endswith('text'):
            el.getparent().remove(el)
    
    doc.write()