#!/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 numpy as np
from math import *


def Bfield_sphere(xy, center, phi, R, m):
    '''
    xy: position where the field is probed
    center: position of the sphere
    phi: rotation angle from the vertical countercolockwise
    R: radius of the sphere
    m: magnetic moment
    '''
    r = np.array(xy) - np.array(center)
    rabs = np.linalg.norm(r)
    mvec = m * np.array([-sin(phi), cos(phi)])
    
    if rabs >= R:
        B = (3 * r * mvec.dot(r) - mvec * rabs**2) / rabs**5 / (4*pi)
    else:
        B = 2 * mvec / R**3 / (4*pi)
    return B


name = 'Magnetsphere_compasses'
size = 600, 600
spheres = [{'c':[0, 0], 'R':85., 'phi':0., 'm':1.}]
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 = np.sum([Bfield_sphere([x, y],
            s['c'], s['phi'], s['R'], s['m']) for s in spheres], axis=0)
        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 the sphere magnets
for isp, s in enumerate(spheres):
    R = s['R']
    magnet = g.add(doc.g(id='magnet' + str(isp),
        transform='translate({:.3f},{:.3f}) rotate({:.2f})'.format(
        s['c'][0], s['c'][1], degrees(s['phi']))))
    mgrad = doc.defs.add(doc.radialGradient(id='magnetGrad' + str(isp), r=1.4*R,
        center=(0,.2*R), focal=(-.4*R,.6*R), gradientUnits='userSpaceOnUse'))
    for of, c, op in ((0, '#ffffff', 0.7), (0.04, '#ffffff', 0.6),
            (0.11, '#ffffff', 0.4), (0.22, '#ffffff', 0.2),
            (0.7, '#666666', 0.3), (1, '#000000', 0.6)):
        mgrad.add_stop_color(of, c, op)
    
    magnet.add(doc.circle(center=(0, 0), r=R, fill='#00cc00', stroke='none'))
    magnet.add(doc.path(d='M -{0},0 A {0},{0} 0 0 0 {0},0 L -{0},0 Z'.format(R),
        fill='#ff0000', stroke='none'))
    magnet.add(doc.circle(center=(0, 0), r=R, stroke_width=4.,
        stroke='#000000', fill='url(#magnetGrad' + str(isp) + ')',
        transform='rotate({})'.format(degrees(-s['phi']))))
    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.22 * R, 0.005*R),
            y=[1.4 * s * R], text_anchor='middle', font_family='Bitstream Vera Sans'))

doc.save(pretty=True)