#!/usr/bin/python3
# -*- coding: utf8 -*-
try:
    import svgwrite
except ImportError:
    print('requires svgwrite library: https://pypi.org/project/svgwrite/')
    # documentation at https://svgwrite.readthedocs.io/
    exit(1)


import mpmath
import numpy as np
from scipy.special import ellipk, ellipe
from math import *

ellippi = mpmath.ellippi # pretty slow. If not available, try sympy.


def Bfield_barmagnet(xy, R, L, M):
    '''
    xy: position where the field is probed
    R: radius of the magnet
    L: length of the magnet
    M: magnetisation
    '''
    rho = xy[0]
    z = xy[1]
    Brho, Bz = 0., 0.
    for s in -1., 1.:
        zeta = z - s * L / 2.
        r = sqrt((rho + R)**2 + zeta**2)
        n = 4. * R * rho / (rho + R)**2
        m = 4. * R * rho / r**2
        K = ellipk(m)
        E = ellipe(m)
        Pi = float(ellippi(n, m))
        if rho != 0.:
            Brho += s * r / rho * ((2. - m) * K - 2. * E)
        Bz += s * zeta / r * ((rho - R) / (rho + R) * Pi - K)
    Brho *= M / (4. * pi)
    Bz *= M / (2. * pi)
    return np.array([Brho, Bz])


name = 'Magnet_compasses'
size = 600, 600
R = 60.
L = 200.
needles_d = 40.
needle_w = 6.
needle_l = 16.
needle_c = 2.5

doc = svgwrite.Drawing(name + '.svg', profile='full', size=size)
doc.set_desc(name, 'https://commons.wikimedia.org/wiki/File:' + name +
    '.svg\nrights: Creative Commons Attribution ShareAlike license')
doc.add(doc.rect(id='background', insert=(0, 0), size=size, fill='#ffffff', stroke='none'))
g = doc.add(doc.g(id='image',
    transform='translate({:.0f}, {:.0f}) scale(1,-1)'.format(size[0]/2., size[1]/2.)))

# draw some compass needles
needle = doc.defs.add(doc.g(id='needle'))
needle.add(doc.path(d='M {:.3f},{:.3f} L {:.3f},{:.3f} L {:.3f},{:.3f} L {:.3f},{:.3f} Z'.format(
    -needle_w, 0, 0, needle_l, needle_w, 0, 0, -needle_l),
    fill='#00cc00', stroke='none'))
needle.add(doc.path(d='M {:.3f},{:.3f} L {:.3f},{:.3f} L {:.3f},{:.3f} Z'.format(
    -needle_w, 0, 0, needle_l, needle_w, 0),
    fill='#ff0000', stroke='none'))
needle.add(doc.path(d='M {:.3f},{:.3f} L {:.3f},{:.3f} L {:.3f},{:.3f} L {:.3f},{:.3f} Z'.format(
    -needle_w, 0, 0, needle_l, needle_w, 0, 0, -needle_l),
    fill='none', stroke='#000000', stroke_width=2,
    stroke_linejoin='miter', stroke_miterlimit=10))
needle.add(doc.circle(center=(0, 0), r='{:.3f}'.format(needle_c),
    fill='#ffffff', stroke='#000000', stroke_width=2))

needles_nx = round(size[0] / needles_d)
needles_ny = round(size[1] / needles_d)
needles_x = (np.arange(needles_nx) + 0.5) * needles_d - size[0] / 2.
needles_y = (np.arange(needles_ny) + 0.5) * needles_d - size[1] / 2.

needles = g.add(doc.g(id='needles'))

for y in needles_y:
    for x in needles_x:
        B = Bfield_barmagnet([x, y], R, L, 1.)
        direction = atan2(B[1], B[0])
        needles.add(doc.use(href='#needle', insert=(0, 0),
            transform='translate({:.3f},{:.3f}) rotate({:.2f})'.format(
            x, y, degrees(direction-pi/2))))

# draw a bar magnet
magnet = g.add(doc.g(id='magnet'))
mgrad = doc.defs.add(doc.linearGradient(id="magnetGrad",
    start=(0,0), end=(1,0), gradientUnits="objectBoundingBox"))
for of, c, op in ((0, '#000000', 0.125), (0.07, '#ffffff', 0.125),
        (0.25, '#ffffff', 0.5), (0.6, '#ffffff', 0.2), (1, '#000000', 0.33)):
    mgrad.add_stop_color(of, c, op)

magnet.add(doc.rect(insert=(-R, -L/2), size=(2*R, L), fill='#00cc00', stroke='none'))
magnet.add(doc.rect(insert=(-R, 0), size=(2*R, L/2.), fill='#ff0000', stroke='none'))
magnet.add(doc.rect(insert=(-R, -L/2), size=(2*R, L), fill='url(#magnetGrad)',
    stroke='#000000', stroke_width=4.))
for s, txt in ((1, 'S'), (-1, 'N')):
    magnet.add(doc.text(txt, font_size='120px', stroke='none', fill='#000000',
        transform='translate(0, {0}) scale({1},-{1})'.format(-0.11 * L, 0.0025*L),
        y=[0.62 * s * L], text_anchor='middle', font_family='Bitstream Vera Sans'))

doc.save(pretty=True)