#!/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
import random
from scipy.integrate import solve_ivp
from scipy.optimize import minimize
from math import *


name = 'Ironfilings_magnetsphere'
size = 600, 600
spheres = [{'c':[0, 0], 'R':85., 'phi':0., 'm':1.}]
n_filings = 2500
l_filings = 24


def vnorm(x):
    return x / hypot(*x)


def inside(p, w2, h2):
    return fabs(p[0]) <= w2 and fabs(p[1]) <= h2


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 counterclockwise
    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


def Bfield(xy):
    return np.sum([Bfield_sphere(xy,
        s['c'], s['phi'], s['R'], s['m']) for s in spheres], axis=0)


def bezier(field, p, l):
    # returns control points of a Bezier curve approximating the Bfield at p
    f = lambda t, xy: vnorm(field(xy))
    p2, p0 = solve_ivp(f, (0, -l/2), p, t_eval=[-l/4, -l/2]).y.T
    p3, p1 = solve_ivp(f, (0, l/2), p, t_eval=[l/4, l/2]).y.T
    v0 = vnorm(field(p0))
    v1 = vnorm(field(p1))
    
    if hypot(*(v1 - v0)) < 0.01:
        l0, l1 = l/3, l/3
    else:
        def err(x):
            l0, l1, t2, t4, t3 = x
            dist = 0.
            for t, pref in ((t2, p2), (t4, p), (t3, p3)):
                pc = (1-t)**3 * p0 + 3*(1-t)**2*t * (p0+l0*v0) + 3*(1-t)*t**2 * (p1-l1*v1) + t**3 * p1
                dist += hypot(*(pc - pref))**2
            return dist
        
        l0, l1 = minimize(err, [l/3., l/3., 0.25, 0.5, 0.75],
            bounds=((0, l), (0, l), (0, 1), (0, 1), (0, 1))).x[:2]
    
    p0c = p0 + v0 * l0
    p1c = p1 - v1 * l1
    
    return [p0, p0c, p1c, p1]


doc = svgwrite.Drawing(name + '.svg', profile='full', size=size)
doc.set_desc(name, 'https://commons.wikimedia.org/wiki/File:' + name + '.svg\n'
    'rights: Creative Commons Attribution ShareAlike license')
clip = doc.defs.add(doc.clipPath(id='image_clip'))
clip.add(doc.rect(insert=(-size[0]/2., -size[1]/2.), size=size))
doc.add(doc.rect(id='background', insert=(0, 0), size=size, fill='#ffffff', stroke='none'))
g = doc.add(doc.g(id='image', clip_path='url(#image_clip)',
    transform='translate({:.0f}, {:.0f}) scale(1,-1)'.format(size[0]/2., size[1]/2.)))
filings = g.add(doc.g(id='iron-filings', fill='none', stroke='black',
    stroke_width=2, stroke_linecap='round'))


Bmax = hypot(*Bfield([0,0]))

i_filings = 0
while i_filings < n_filings:
    x = random.uniform(-size[0]/2. - l_filings/2., size[0]/2. + l_filings/2.)
    y = random.uniform(-size[1]/2. - l_filings/2., size[1]/2. + l_filings/2.)
    l = random.uniform(l_filings*0.5, l_filings*1.0)
    
    if any([hypot(*(np.array([x, y]) - s['c'])) < s['R'] - l/2 for s in spheres]):
        continue
    
    B = Bfield([x, y])
    
    Brel = hypot(*B) / Bmax
    line_density = Brel**(2/3)
    line_width = 2.4 * Brel**(1/3)
    
    # use rejection sampling to reproduce field line density
    if random.random() >= line_density:
        continue
    
    points = bezier(Bfield, [x, y], l)
    
    if all([any([hypot(*(np.array(p)-s['c'])) < s['R'] for s in spheres]) for p in [ [x, y], points[0], points[-1] ] ]):
        continue
    if all([not inside(p, size[0]/2., size[1]/2.) for p in [ [x, y], points[0], points[-1] ] ]):
        continue
    
    filings.add(doc.path(
        d='M {:.1f},{:.1f} C {:.1f},{:.1f} {:.1f},{:.1f} {:.1f},{:.1f}'.format(
        *points[0], *points[1], *points[2], *points[3]),
        stroke_width='{:.1f}'.format(line_width)))
    i_filings += 1


# draw the sphere magnet
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)