#!/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 / (4 epsilon_0 R^2)
    in case of an electric charge Q on each disk.
    The solution requires elliptical integrals
    '''
    if z == 0.:
        return pi/2
    return pi/2 + 0.5 * (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
    F *= pi / 4
    return F


def dipole_force(z, m1, m2):
    '''
    Axial force between axially aligned dipoles with magnetic moments m1,m2
    z: axial distance
    Assume mu0=1
    '''
    F = 3. * m1 * m2 / (2. * pi * z**4)
    return F


mpl.style.use('classic')
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 = [('8R', 8.), ('4R', 4.), ('2R', 2.), ('R', 1.), ('R/2', 0.5)]
dash = [6.8, 2.4]
dot = [2.4, 5.8]
plt.figure()
z0, z1 = 0.4, 100
for i in range(len(L)):
    f = lambda z: force_between_magnets(z-L[i][1], 1., L[i][1])
    zspace = np.logspace(log10(max(z0, L[i][1])), log10(z1), 5001)
    plt.plot(zspace, [f(z) for z in zspace], '-',
         color=colors[i], label=r'L = ' + L[i][0], zorder=-i-len(L))
    plt.plot(L[i][1], f(L[i][1]), 'o', color=colors[i], mew=1.2, zorder=-i)

plt.xlabel('z / R')
plt.ylabel(r'$F\ [\mu_0M^2R^2]$')
plt.title('Force between two cylindrical magnets with magnetization M,\nlength L, radius R and axial center-of-mass distance z')
plt.gca().set_xscale('log')
plt.gca().set_yscale('log')
plt.legend(loc='upper right')
plt.xlim(z0, z1)
plt.ylim(1e-6, 1e1)
plt.grid(True)
plt.tight_layout()
plt.savefig('Cylindrical-magnet-force-diagram_loglog.svg')