#!/usr/bin/python
# -*- coding: utf8 -*-

import numpy as np
import scipy.special as sp
import matplotlib.pyplot as plt
import matplotlib as mpl
from math import *

mpl.style.use("classic")


# fix elliptic integrals for negative argument in case of old scipy version
if sp.ellipe(-1) > 0:
    E = sp.ellipe
    K = sp.ellipk
else:
    def E(m):
        if m >= 0.:
            return sp.ellipe(m)
        else:
            return sp.ellipe(-m / (1. - m)) * sqrt(1. - m)
    
    def K(m):
        if m >= 0.:
            return sp.ellipk(m)
        else:
            return sp.ellipk(-m / (1. - m)) / sqrt(1. - m)


def force_between_disks(z):
    '''
    Exact formula for the force between two homogeneously charged round disks
    aligned on their axis of symmetry.
    z is the distance relative to the disk radius.
    The force is returned in units of Q^2 / (4pi epsilon_0 R^2)
    in case of an electric charge Q on each disk.
    The solution requires elliptical integrals
    '''
    if z == 0.:
        return 2.
    elif z > 0.:
        s = 1.0
    else:
        s = -1.0
        z = -z
    m = 4 / (4. + z**2)
    return s * (2 + 4/pi * z / sqrt(m) * (E(m) - K(m)))


def force_between_magnets(z, R, L):
    '''
    Exact formula for the force between two axially aligned identical
    cylindrical magnets, as long as they are homogeneously magnetized.
    '''
    zR = z / R
    F = force_between_disks(zR)
    F += force_between_disks(zR + 2*L / R)
    F -= 2 * force_between_disks(zR + L / R)
    return F


mpl.rcParams['font.sans-serif'] = 'DejaVu Sans'
mpl.rc('mathtext', default='regular')
mpl.rc('lines', linewidth=2.4)

colors = ['#0000ff', '#00aa00', '#ff0000', '#ee9900', '#cccc00']
L = [(r'$\infty$', float('inf')), ('4R', 4.), ('2R', 2.), ('R', 1.), ('R/2', 0.5)]


plt.figure()
zmax = 4
zspace = np.linspace(0., zmax**0.5, 5001)**2
for i in range(len(L)):
    if L[i][1] == float('inf'):
        f = lambda z: force_between_disks(z)
    else:
        f = lambda z: force_between_magnets(z, 1., L[i][1])
    plt.plot(zspace, [f(z) for z in zspace], '-',
             color=colors[i], label=r'L = ' + L[i][0], zorder=-i-len(L))
    plt.plot(0, f(0), 'o', color=colors[i], mew=1.2, zorder=-i)

plt.xlabel('z / R')
plt.ylabel(r'$F\ [\pi/4\;\mu_0M^2R^4]$')
plt.title('Force between two cylindrical magnets with magnetization M,\nlength L, radius R and axial end-to-end distance z')
plt.legend(loc='upper right')
plt.xlim(-0.05, zmax)
plt.ylim(0, 2.1)
plt.grid(True)
plt.tight_layout()
plt.savefig('Cylindrical-magnet-force-diagram.svg')


plt.figure()
zmax = 20
zspace = np.linspace(0., zmax**0.5, 5001)**2
for i in range(len(L)):
    if L[i][1] == float('inf'):
        f = lambda z: force_between_disks(z)
    else:
        f = lambda z: force_between_magnets(z, 1., L[i][1])
    plt.plot(zspace, [f(z) for z in zspace], '-',
             color=colors[i], label=r'L = ' + L[i][0], zorder=-i-len(L))
    plt.plot(0, f(0), 'o', color=colors[i], mew=1.2, zorder=-i)

plt.xlabel('z / R')
plt.ylabel(r'$F\ [\pi/4\;\mu_0M^2R^2]$')
plt.title('Force between two cylindrical magnets with\nmagnetization M, length L, radius R and axial distance z')
plt.gca().set_yscale('log')
plt.legend(loc='upper right')
plt.xlim(-0.5, zmax)
plt.ylim(1e-5, 2.5)
plt.grid(True)
plt.tight_layout()
plt.savefig('Cylindrical-magnet-force-diagram_logscale.svg')