#!/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 *


name = 'Ironfilings_loopcurrent'
size = 600, 600
loops = [{'p_in':[120, 0], 'p_out':[-120, 0], 'I':1.}]
n_filings = 2500
l_filings = 24

D = 24 # ring thickness
D2 = 4
tilt = 0.25


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


def rot(xy, phi):
    s, c = sin(phi), cos(phi)
    return np.array([c * xy[0] - s * xy[1], c * xy[1] + s * xy[0] ])


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


def Bfield_ringcurrent(xy, p_in, p_out, I):
    '''
    p_in: position where the loop current goes into the plane
    p_out: position where the loop current comes out of the plane
    I: current
    '''
    p_c = (np.array(p_in) + np.array(p_out)) / 2
    Rvec = np.array(p_in) - p_c
    R = hypot(*Rvec)
    phi = atan2(Rvec[1], Rvec[0])
    xyrel = xy - p_c
    # change into a new coordinate system with z and r aligned to the ring
    r, z = rot(xyrel, -phi)
    rplus = (r + R)**2 + z*z
    rminus = (r - R)**2 + z*z
    dplus = R*R + r*r + z*z
    dminus =  R*R - r*r - z*z
    
    if (fabs(rplus) < 1e-14 or fabs(rminus) < 1e-14):
        # constrain maximum field near conductor
        Fz = 1e7
        Fr = 0.
    else:
        # Integrate the Biot-Savart law using elliptic integrals:
        kk = 4 * r * R / rplus
        elle = ellipe(kk)
        ellk = ellipk(kk)
        
        Fz = (dminus / rminus * elle + ellk) / sqrt(rplus)
        
        if fabs(r) < 1e-4 * (R + z*z / R):
            # elliptic integral formula is unstable near r=0, use Taylor series
            Fr = 1.5 * pi * r * z * R*R / sqrt(R*R + z*z)**5
        else:
            Fr = z / r * (dplus / rminus * elle - ellk) / sqrt(rplus)

    # backtransform
    return rot([Fr, Fz], phi) * I / (2*pi)


def Bfield(xy):
    return np.sum([Bfield_ringcurrent(xy,
        l['p_in'], l['p_out'], l['I']) for l in loops], 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'))
img = 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 = doc.g(id='iron-filings', fill='none', stroke='black',
    stroke_width=2, stroke_linecap='round')
back = img.add(doc.g(id='back'))
img.add(filings)
front = img.add(doc.g(id='front'))

Bmax = 2 * hypot(*Bfield([2,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)
    
    B = Bfield([x, y])
    
    Brel = hypot(*B) / Bmax
    line_density = Brel**(2/3)
    line_width = 2.7 * 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)-l['p_in'])) < D/2 or hypot(*(np.array(p)-
        l['p_out'])) < D/2 for l in loops]) 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
    print(i_filings, end=' ', flush=True)


# draw the loop

for il, l in enumerate(loops):
    m = (np.array(l['p_in']) + np.array(l['p_out'])) / 2.
    r = hypot(*(np.array(l['p_in']) - np.array(l['p_out']))) / 2.
    a = degrees(atan2(l['p_in'][1] - l['p_out'][1], l['p_in'][0] - l['p_out'][0]))
    mask_back = back.add(doc.mask(id="ring-back-mask"+str(il)))
    ring_back = back.add(doc.g(transform="translate({},{}) rotate({})".format(*m, a)))
    mask_front = front.add(doc.mask(id="ring-front-mask"+str(il)))
    mask_contour = front.add(doc.mask(id="ring-contour-mask"+str(il)))
    ring_front = front.add(doc.g(transform="translate({},{}) rotate({})".format(*m, a)))
    
    d = 'M {},0 A {},{} 0 1 0 {},0'.format(-r, r, tilt*r, r)
    ring_back.add(doc.path(d=d, fill='none', stroke='#000000', stroke_width=D+2*D2,
        stroke_linecap='round'))
    ring_back.add(doc.path(d=d, fill='none', stroke='#999999', stroke_width=D))
    mask_back.add(doc.path(d=d, fill='none', stroke='white', stroke_width=D))
    
    mask_contour.add(doc.circle(center=(0,0), r=r+D+D2, fill='white',stroke='none'))
    mask_contour.add(doc.path(d=d, fill='none', stroke='black', stroke_width=D))
    
    d = 'M {},0 A {},{} 0 1 1 {},0'.format(-r, r, tilt*r, r)
    ring_front.add(doc.path(d=d, fill='none', stroke='#000000', stroke_width=D+2*D2,
        mask="url(#ring-contour-mask"+str(il)+")"))
    ring_front.add(doc.path(d=d, fill='none', stroke='#999999', stroke_width=D,
        stroke_linecap='round'))
    mask_front.add(doc.path(d=d, fill='none', stroke='white', stroke_width=D,
        stroke_linecap='round'))
    
    r2 = r*(tilt*r+D/2)/(tilt*r)
    r3 = r*(tilt*r+D/2+D/2*r)/(tilt*r)
    grad1 = doc.defs.add(doc.radialGradient(id='grad1_'+str(il), r=r2,
        center=(0, 0), gradientUnits='userSpaceOnUse'))
    for of, c, op in (((tilt*r-D/2)/(tilt*r+D/2), '#666666', 1),
            ((tilt*r)/(tilt*r+D/2), '#aaaaaa', 1), (1, '#666666', 1)):
        grad1.add_stop_color(of, c, op)
    grad2 = doc.defs.add(doc.radialGradient(id='grad2_'+str(il), r=r,
        center=(0, 0), gradientUnits='userSpaceOnUse'))
    for of, c, op in ((0., '#ffffff', 0.8), (0.3, '#ffffff', 0.4), (1, '#ffffff', 0)):
        grad2.add_stop_color(of, c, op)
    
    g_back = ring_back.add(doc.g(mask="url(#ring-back-mask"+str(il)+")"))
    g_back.add(doc.rect(insert=(-r3, -r3), size=(2*r3, 2*r3), 
        fill='url(#grad1_'+str(il)+')', transform="translate(0,{}) scale(1, {})".format(0.25*D, tilt)))
    g_back.add(doc.circle(center=(0, 0), r=r, fill='url(#grad2_'+str(il)+')',
        transform="translate({},{}) rotate(-6) scale(1, 0.12)".format(0.4*r, tilt*sqrt(.84)*r+.25*D)))
    g_front = ring_front.add(doc.g(mask="url(#ring-front-mask"+str(il)+")"))
    g_front.add(doc.rect(insert=(-r3, -r3), size=(2*r3, 2*r3),
        fill='url(#grad1_'+str(il)+')', transform="translate(0,{}) scale(1, {})".format(0.25*D, tilt)))
    g_front.add(doc.circle(center=(0, 0), r=r, fill='url(#grad2_'+str(il)+')',
        transform="translate({},{}) rotate(-6) scale(1, 0.12)".format(-0.4*r, -tilt*sqrt(.84)*r+.25*D)))
    
    marker = doc.defs.add(doc.marker(id="arrow", insert=(1,1.5), size=(5,3), orient="auto"))
    marker.add(doc.path(id="arrow", stroke="none", fill="#000000",
        d="M 4,1.5 L 0,3 L 0,0 L 4,1.5 Z"))
    a2 = copysign(0.22, l['I'])
    d = 'M {},{} A {},{} 0 0 {} {},{}'.format(
        -r*sin(a2), -tilt*r*cos(a2), r, tilt*r,
        {True:1, False:0}[l['I']>0], r*sin(a2), -tilt*r*cos(a2))
    line = ring_front.add(doc.path(d=d, fill='none', stroke='#000000', stroke_width=5))
    line.set_markers((None, False, marker))

doc.save(pretty=True)