#!/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.
    return 2 + 2/pi * (z**2 * E(-4./z**2) - (4+z**2) * K(-4./z**2))


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 -= 2 * force_between_disks(zR + L / R)
    F += force_between_disks(zR + 2*L / R)
    return F


def force_between_magnets_approx(z, L):
    '''
    Asymptotic formula for the force between two axially aligned identical
    cylindrical magnets for the case z >> R, assuming magnetic point charges
    '''
    F = 1. / z**2
    F -= 2. / (z + L)**2
    F += 1. / (z + 2*L)**2
    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)
        f2 = lambda z: 1. / z**2
    else:
        f = lambda z: force_between_magnets(z, 1., L[i][1])
        f2 = lambda z: force_between_magnets_approx(z, 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.plot(zspace[1:], [f2(z) for z in zspace[1:]], '--', dashes=[2.4, 4.8],
             color=colors[i], zorder=-i-2*len(L))
plt.plot([], [], '--', dashes=[2.4, 4.8], color='gray', label='two-pole\napproximation')

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-approx.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)
        f2 = lambda z: 1. / z**2
    else:
        f = lambda z: force_between_magnets(z, 1., L[i][1])
        f2 = lambda z: force_between_magnets_approx(z, 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.plot(zspace[1:], [f2(z) for z in zspace[1:]], '--', dashes=[2.4, 4.8],
             color=colors[i], zorder=-i-2*len(L))
plt.plot([], [], '--', dashes=[2.4, 4.8], color='gray', label='two-pole\napproximation')

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-approx_logscale.svg')