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


name = 'Ringmagnet_compasses'
size = 600, 600
R1 = 100
R2 = 200 - 1e-6
L = 100
needles_d = 40.
needle_w = 6.
needle_l = 16.
needle_c = 2.5


def cel(kc, p, a, b):
    """
    Bulirsch complete elliptic integral
    """
    
    if kc == 0.:
        return nan
    
    tol = 1e-9 # actual relative error will be tol**2
    k = kc = fabs(kc)
    m = 1.
    
    if p > 0.:
        p = sqrt(p)
        b /= p
    else:
        f = kc * kc
        g = 1. - p
        q = (1. - f) * (b - a * p)
        f -= p
        p = sqrt(f / g)
        a = (a - b) / g
        b = a * p - q / (g * g * p)
    
    for i in range(11):
        f = a
        a += b / p
        g = k / p
        b = 2. * (b + f * g)
        p += g
        g = m
        m += kc
        
        if fabs(g - kc) <= g * tol:
            break
        
        kc = 2. * sqrt(k)
        k = kc * m
    
    return pi * .5 * (a * m + b) / (m * (m + p))


def Bfield_barmagnet(xy, R, L, M):
    # www.doi.org/10.1119/1.3256157
    rho, z = xy
    rho0 = 1.
    if rho < 0.:
        rho = -rho
        rho0 = -1.
    Rm, Rp = R - rho, R + rho
    zm, zp = z - L / 2, z + L / 2
    Rmzm = hypot(Rm, zm)
    Rmzp = hypot(Rm, zp)
    Rpzm = hypot(Rp, zm)
    Rpzp = hypot(Rp, zp)
    g = Rm / Rp
    km = Rmzm / Rpzm
    kp = Rmzp / Rpzp
    Frhom = cel(km, 1., 1., -1.) / Rpzm
    Frhop = cel(kp, 1., 1., -1.) / Rpzp
    Fzm = cel(km, g * g, 1., g) * zm / Rpzm
    Fzp = cel(kp, g * g, 1., g) * zp / Rpzp
    return M * R / pi * np.array((rho0 * (Frhop - Frhom), (Fzp - Fzm) / Rp))


def Bfield(xy):
    return Bfield_barmagnet(xy, R2, L, 1.) - Bfield_barmagnet(xy, R1, L, 1.)


# draw the magnet
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.)))
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))
magnet_back = g.add(doc.g(id='magnet_back'))
needles = g.add(doc.g(id='needles'))
magnet_front = g.add(doc.g(id='magnet_front'))
mgrad = doc.defs.add(doc.linearGradient(id="magnetGrad",
    start=(0,0), end=(1,0), gradientUnits="objectBoundingBox"))
for c, of, op in [['#000000', '0', '0.33'], ['#ffffff', '0.4', '0.2'],
        ['#ffffff', '0.75', '0.5'], ['#ffffff', '0.93', '0.125'],
        ['#000000', '1', '0.125']]:
    mgrad.add_stop_color(of, c, op)
for x0, x1, isbg in [[-R2, -R1, False], [-R1, R1, True], [R1, R2, False]]:
    if isbg:
        magnet = magnet_back
        fill = 'url(#magnetGrad)'
        colors = ['#00cc00', '#ff0000']
    else:
        magnet = magnet_front
        fill = 'none'
        colors = ['#49da49', '#ff4949']
    for i in [0, 1]:
        magnet.add(doc.rect(insert=(x0, [-L/2, 0][i]), size=(x1-x0, [L, L/2][i]),
            fill=colors[i], stroke='none'))
    magnet.add(doc.rect(insert=(x0, -L/2), size=(x1 - x0, L), fill=fill,
            stroke='#000000', stroke_width=4, stroke_linejoin='miter'))


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.

for y in needles_y:
    for x in needles_x:
        B = Bfield([x, y])
        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))))


doc.save(pretty=True)