#!/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 scipy.special import ellipk, ellipe
from math import *

import mpmath
ellippi = mpmath.fp.ellippi


name = 'Ironfilings_cylindermagnet'
size = 600, 600
R = 60.
L = 200.
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_barmagnet(xy, R, L, M):
    '''
    xy: position where the field is probed
    R: radius of the magnet
    L: length of the magnet
    M: magnetization
    '''
    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 fabs(rho) > 5e-8 * R:
            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])


def Bfield(xy):
    return Bfield_barmagnet(xy, R, L, 1.)


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\nrights: 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 = 1.2 * 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)
    
    if inside([x, y], R, L/2. - l/2):
        continue
    
    B = Bfield([x, y])
    
    Brel = hypot(*B) / Bmax
    line_density = Brel**(2/3)
    line_width = 3.2 * 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([inside(p, R, L/2.) for p in [[x, y]] + points]):
        continue
    if all([not inside(p, size[0]/2., size[1]/2.) for p in [[x, y]] + points]):
        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
    print(i_filings, end=' ', flush=True)


# 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)