"""Animate the relaxation of a dipole cluster towards a local minimum."""
# Version 1.1, Ingo Rehberg, Apr 1, 2022
# This program illustrates a comment available at
#  https://doi.org/10.48550/arXiv.2203.13670
# on the paper: Andrew D.P.Smith, Peter T.Haugen, Boyd F.Edwards:
# Hysteretic transition between states
# of a filled hexagonal magnetic dipole cluster,
# Journal of Magnetism and Magnetic Materials 549 (2022): 168991


import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, Button, RadioButtons
from matplotlib import animation
from matplotlib import cm


class Cluster_2d:
    """Define an ensemble of magnetic dipoles, together with some functions."""

    def __init__(self, n):
        self.n = n
        self.rx = np.zeros(self.n, dtype='double')
        self.ry = np.zeros(self.n, dtype='double')
        self.px = np.zeros(self.n, dtype='double')
        self.py = np.zeros(self.n, dtype='double')
        self.dxet = np.zeros(self.n*self.n, dtype='double')
        self.dyet = np.zeros(self.n*self.n, dtype='double')
        self.dp3t = np.zeros(self.n*self.n, dtype='double')
        self.set_data()

    def set_data(self, phi2=0., alpha=0.5, model='circular'):
        """Define the filled hexagon from angle in rad, alpha as dip-ratio."""
        r = 1.  # the radius of the ring, double that of a sphere.
        ang = np.zeros(self.n, dtype='double')
        if model == 'dipolar':
            self.px[0] = alpha
            self.py[0] = 0.
            ang[1] = 0.
            ang[2] = phi2
            ang[3] = np.arctan2(-np.sin(phi2), np.cos(phi2))
            ang[4] = 0.
            ang[5] = ang[2]
            ang[6] = ang[3]
            for i in range(1, 7):
                self.rx[i] = np.cos((i-1)*np.pi/3.)*r
                self.ry[i] = np.sin((i-1)*np.pi/3.)*r
                self.px[i] = np.cos(ang[i])
                self.py[i] = np.sin(ang[i])
        elif model == 'circular':
            self.px[0] = alpha
            self.py[0] = 0.
            for i in range(1, 7):
                self.rx[i] = np.cos((i-1)*np.pi/3.)*r
                self.ry[i] = np.sin((i-1)*np.pi/3.)*r
                self.px[i] = -self.ry[i]/r
                self.py[i] = self.rx[i]/r
        elif model == 'test':
            for i in range(self.n):
                self.rx[i] = np.cos((i)*np.pi/3.)*r
                self.ry[i] = np.sin((i)*np.pi/3.)*r
                self.px[i] = -self.ry[i]/r
                self.py[i] = self.rx[i]/r
        else:
            print('no such model')
        for i in range(0, self.n):  # calcuate geometrical cluster constants
            for p in range(0, self.n):  # redundancies are ok here
                if p == i:
                    continue
                pnt = p*self.n + i
                dx = self.rx[i] - self.rx[p]
                dy = self.ry[i] - self.ry[p]
                dist = np.sqrt(pow(dx, 2)+pow(dy, 2))
                self.dxet[pnt] = dx * np.sqrt(3.) / dist
                self.dyet[pnt] = dy * np.sqrt(3.) / dist
                self.dp3t[pnt] = 1. / pow(dist, 3)

    def reset_circular_angles(self, phi3, alpha):
        """Set the angles according to the symmetries of the circular state."""
        self.px[0] = alpha
        self.py[0] = 0
        for i in range(3, 6):  # angles 3,4,5 are stacked in a numpy array
            self.px[i] = np.cos(phi3[i-3])
            self.py[i] = np.sin(phi3[i-3])
        self.px[6] = self.px[5]
        self.py[6] = -self.py[5]
        self.px[1] = self.px[4]
        self.py[1] = -self.py[4]
        self.px[2] = self.px[3]
        self.py[2] = -self.py[3]

    def reset_dipolar_angles(self, phi2):
        """Set the angles according to the symmetries of the dipolar state."""
        self.px[2] = np.cos(phi2)
        self.py[2] = np.sin(phi2)
        self.px[3] = self.px[2]
        self.py[3] = -self.py[2]
        self.px[5] = self.px[2]
        self.py[5] = self.py[2]
        self.px[6] = self.px[2]
        self.py[6] = -self.py[2]

    def b(self, i, p):
        """B-field at i caused by p."""
        pnt = p*self.n + i
        dot = self.px[p]*self.dxet[pnt]+self.py[p]*self.dyet[pnt]
        return np.array([(self.dxet[pnt]*dot-self.px[p]) * self.dp3t[pnt],
                         (self.dyet[pnt]*dot-self.py[p]) * self.dp3t[pnt]])

    def calc_b_loc(self, i):
        """Get local field vector."""
        if i == 0:  # the field can be tiny here
            b_acc = self.b(0, 1) + self.b(0, 4)
            b_acc += self.b(0, 2) + self.b(0, 5)
            b_acc += self.b(0, 3) + self.b(0, 6)
        else:
            b_acc = np.zeros(2)
            for p in range(self.n):
                if p == i:
                    continue
                b_acc += self.b(i, p)
        return b_acc

    def w(self, i, p):
        """Energy between the dipoles with number i and p."""
        pnt = p*self.n + i
        return (self.px[i]*self.px[p]+self.py[i]*self.py[p]
                - (self.px[i]*self.dxet[pnt] + self.py[i]*self.dyet[pnt])
                * (self.px[p]*self.dxet[pnt] + self.py[p]*self.dyet[pnt])
                ) * self.dp3t[pnt]

    def w_gesamt(self):
        """Get the total energy."""
        w_acc = 0.
        for ldip in range(1, self.n):  # no self interaction
            for partner in range(0, ldip):  # partner < ldip
                w_acc += self.w(ldip, partner)
        return w_acc

    def potential(self, x, y, z):
        """Calculate the potential at position px, py, pz."""
        pot = 0.
        for i in range(self.n):
            rx = x-self.rx[i]
            ry = y-self.ry[i]
            rz = z
            r2 = pow(rx, 2)+pow(ry, 2)+pow(rz, 2) + 0.000001
            scpr = self.px[i]*rx+self.py[i]*ry
            pot += scpr/pow(r2, 1.5)
        return pot

    def b_flux(self, x, y, z):
        """Calculate the 2-d field at position x, y, z."""
        lbx, lby, lbz = 0., 0., 0.
        for i in range(self.n):
            rx = x-self.rx[i]
            ry = y-self.ry[i]
            rz = z
            r2 = pow(rx, 2)+pow(ry, 2)+pow(rz, 2) + 0.000001  # avoid zero
            scpr = self.px[i]*rx+self.py[i]*ry
            lbx += (3.*scpr*rx - r2*self.px[i])/pow(r2, 2.5)
            lby += (3.*scpr*ry - r2*self.py[i])/pow(r2, 2.5)
            lbz += (3.*scpr*rz)/pow(r2, 2.5)
        return lbx, lby, lbz

    def b_flux_2(self, x, y, z):
        """Calculate the field energy density at position x, y, z."""
        lbx, lby, lbz = self.b_flux(x, y, z)
        return lbx*lbx + lby*lby + lbz*lbz

    def radial_flux(self, x, y, z):
        """Calculate the r-component of the field at position x, y, z."""
        lbx, lby, lbz = self.b_flux(x, y, z)
        return (lbx*x+lby*y+lbz*z)/np.sqrt(x*x+y*y+z*z)


class Geo:
    """Class describing the geometry to be examined."""

    def __init__(self, richtung=np.array([1., 1., 0.]),
                 origin=np.array([0, 0, 0]), length=20., d_p=2.,
                 width=2, st=0.1):  # line and perpendicular plane
        self.px = self.py = self.pz = np.arange(0, length, length/100.)
        # position vectors x,y,z
        self.iii = int(2.*width/st+1.1)  # resolution of the plane surfaces
        self.X = self.Y = self.Z = self.XPL = self.YPL = np.zeros(
            (self.iii, self.iii))  # coordinates for plane surfaces, 3d and 2d
        self.richtung = richtung  # the direction of the straight line
        self.origin = origin  # the start position of the line
        self.length = length  # the length of the line
        self.d_p = d_p  # the distance of the plane from ...
        self.width = width  # the quadratic plane extends from -width to width
        self.st = st  # the step size
        self.new_geo()

    def new_geo(self):
        """Set a new geometry to be examined."""
        self.direction = self.richtung/np.sqrt(np.dot(self.richtung,
                                                      self.richtung))
        self.px = np.arange(0, self.length,
                            self.length/100.)*self.direction[0]+self.origin[0]
        self.py = np.arange(0, self.length,
                            self.length/100.)*self.direction[1]+self.origin[1]
        self.pz = np.arange(0, self.length,
                            self.length/100.)*self.direction[2]+self.origin[2]
        self.ps = np.sqrt(pow(self.px-self.origin[0], 2) +
                          pow(self.py-self.origin[1], 2) +
                          pow(self.pz-self.origin[2], 2))  # distance from orig
        self.cdist = np.sqrt(pow(self.px, 2) +
                             pow(self.py, 2) +
                             pow(self.pz, 2))  # distance from center
        # generate a plane perpendicular to "direction", within the d_p:
        self.d_cross_y = (self.direction[2], 0, -self.direction[0])/np.sqrt(
            self.direction[2]*self.direction[2] +
            self.direction[0]*self.direction[0])
        self.perp = np.cross(self.direction, self.d_cross_y)
        self.X = np.zeros((self.iii, self.iii))  # set to zero
        self.Y = np.zeros((self.iii, self.iii))  # before the accumulation
        self.Z = np.zeros((self.iii, self.iii))  # in the double loop below
        self.XPL, self.YPL = np.meshgrid(
            np.arange(-self.width, self.width+self.st/2., self.st),
            np.arange(-self.width, self.width+self.st/2., self.st))
        for i in range(self.iii):
            fi = -self.width+i*self.st
            for j in range(self.iii):
                fj = -self.width+j*self.st
                self.X[j, i] += (fi*self.d_cross_y[0] +
                                 fj*self.perp[0] +
                                 self.direction[0]*self.d_p)
                self.Y[j, i] += (fi*self.d_cross_y[1] +
                                 fj*self.perp[1] +
                                 self.direction[1]*self.d_p)
                self.Z[j, i] += (fi*self.d_cross_y[2] +
                                 fj*self.perp[2] +
                                 self.direction[2]*self.d_p)


class mg:
    """My globals, avoiding unnamed globals and simplifying replots."""

    alpha_max = 3  # the ratio of dipole strength as defined by the authors
    alpha = 0.5  # for the interactive slider of the bifurcation parameter
    two_d = True  # determines if the plot is 2d or 3d
    sequential = True
    cl = Cluster_2d(7)
    cl.set_data(phi2=0., alpha=alpha, model='circular')
    d_p = 0.5
    g = Geo(richtung=np.array([0., 0., 1.]),
            origin=np.array([0, 0., 0]),
            length=10., d_p=d_p, width=1.6, st=0.1)
    length = np.sqrt(pow(cl.px, 2)+pow(cl.py, 2))
    xpl, ypl, yplm, yplu = [], [], [], []
    fig, ax = None, []
    calc_stop = None
    alpha_slider, rotate, clr, seq_par, d2_d3 = None, None, None, None, None
    back_off, rs_decr, rs_incr, trp_slider = None, None, None, None
    s = None
    calc = True
    field = True  # sphere in 3d, plane in 2d
    sphere_resol = 64  # only for 3d, interactive
    trp = 0.5  # transparency, interactive
    cli = 5  # factor for the clip level, interactive
    intervals = {'$10^0$': 1, '$10^1$': 10, '$10^2$': 100,
                 '$10^3$': 1000, '$10^4$': 10000, '$10^5$': 100000}
    plot_interv = intervals.get('$10^3$')  # no. of iterations between plots
    show_pot = False


def such_seq():
    """Turn all dipoles sequentially into the local B-field."""
    for i in range(mg.cl.n):
        xy = mg.cl.calc_b_loc(i)
        xy *= mg.length[i]/np.sqrt(np.dot(xy, xy))  # unity vector
        mg.cl.px[i] = xy[0]
        mg.cl.py[i] = xy[1]


def such_par():
    """Turn all dipoles in parallel into the local B-field."""
    pxn = np.copy(mg.cl.px)
    pyn = np.copy(mg.cl.py)
    for i in range(mg.cl.n):
        xy = mg.cl.calc_b_loc(i)
        xy *= mg.length[i]/np.sqrt(np.dot(xy, xy))
        pxn[i] = xy[0]
        pyn[i] = xy[1]
    wa = np.arctan2(mg.cl.py, mg.cl.px)+2*np.pi
    wn = np.arctan2(pyn,     pxn)+2*np.pi
    a = 0.5*wa+0.5*wn
    mg.cl.px = np.cos(a)*mg.length
    mg.cl.py = np.sin(a)*mg.length


def plot_dipoles_2d(ax0, cl):
    """Plot the orientation of the 7 dipoles."""
    ax0.cla()
    s1 = r'$\phi_3=$' + str(int(np.arctan2(cl.py[3],
                                           cl.px[3])*180/np.pi) % 360)
    s2 = r'$\phi_4=$' + str(int(np.arctan2(cl.py[4],
                                           cl.px[4])*180/np.pi) % 360)
    s3 = r'$\phi_5=$' + str(int(np.arctan2(cl.py[5],
                                           cl.px[5])*180/np.pi) % 360)
    s4 = '$U=$' + f"{cl.w_gesamt():.2f}"
    s5 = r'$m_\mathrm{tor}=$' + f"{tor_moment(mg.cl):.2f}"
    ax0.set_title(s1+', '+s2+', '+s3 + '; ' + s4 + '; ' + s5, size=14)
    d = mg.g.width
    ax0.set_xlim(-d, d)
    ax0.set_ylim(-d, d)
    ax0.set_axis_off()
    if mg.field:
        if mg.show_pot:
            H4 = mg.cl.potential(mg.g.X, mg.g.Y, mg.g.Z)
        else:
            H4 = mg.cl.b_flux_2(mg.g.X, mg.g.Y, mg.g.Z)
            H4 = np.log(H4)
        med = np.mean(H4)
        std = np.std(H4)
        cli_up = med + mg.cli*std
        cli_dn = med - mg.cli*std
        H4 = np.clip(H4, cli_dn, cli_up)
        H4 -= H4.min()
        H4 /= H4.max()
        U, V, _ = mg.cl.b_flux(mg.g.X, mg.g.Y, mg.g.Z)
        ax0.quiver(mg.g.X, mg.g.Y, U, V)
        ax0.contour(mg.g.XPL, mg.g.YPL, H4, cmap=cm.gray)
        ax0.imshow(H4, interpolation='bilinear',
                   origin='lower', cmap=cm.bwr,
                   extent=(-d, d, -d, d))
    ax0.quiver(cl.rx, cl.ry, cl.px, cl.py, pivot='middle',
               scale=10, color='c', linewidth=1)
    circle1 = plt.Circle((0, 0), 1., color='grey', fill=False)
    ax0.add_patch(circle1)
    for n in range(len(cl.rx)):
        ax0.add_patch(plt.Circle((cl.rx[n], cl.ry[n]),
                                 0.5, color='g', fill=False))


def arrow3d(ax, length=1, width=0.05, head=0.2, headwidth=2,
            theta_x=0, theta_z=0, offset=(0, 0, 0), **kw):
    """Plot a tube arrow."""
    w = width
    h = head
    hw = headwidth
    theta_x = np.deg2rad(theta_x)
    theta_z = np.deg2rad(theta_z)

    a = [[0, 0], [w, 0], [w, (1-h)*length], [hw*w, (1-h)*length], [0, length]]
    a = np.array(a)

    r, theta = np.meshgrid(a[:, 0], np.linspace(0, 2*np.pi, 30))
    z = np.tile(a[:, 1], r.shape[0]).reshape(r.shape)
    x = r*np.sin(theta)
    y = r*np.cos(theta)

    rot_x = np.array([[1, 0, 0], [0, np.cos(theta_x), -np.sin(theta_x)],
                      [0, np.sin(theta_x), np.cos(theta_x)]])
    rot_z = np.array([[np.cos(theta_z), -np.sin(theta_z), 0],
                      [np.sin(theta_z),  np.cos(theta_z), 0], [0, 0, 1]])

    b1 = np.dot(rot_x, np.c_[x.flatten(), y.flatten(), z.flatten()].T)
    b2 = np.dot(rot_z, b1)
    b2 = b2.T + np.array(offset)
    x = b2[:, 0].reshape(r.shape)
    y = b2[:, 1].reshape(r.shape)
    z = b2[:, 2].reshape(r.shape)
    ax.plot_surface(x, y, z, **kw)


def quiver_3d(a, x, y, u, v):
    """Plot tube arrows with 2d input."""
    w = np.zeros(len(u))
    le = np.power(u*u+v*v, 1/4)
    color = ['c', 'r', 'r', 'r', 'r', 'r', 'r']
    for i in range(len(x)):
        arrow3d(a, offset=(x[i]-u[i]/le[i]/2, y[i]-v[i]/le[i]/2, 0),
                length=le[i],
                width=0.05*le[i],
                head=4*0.05*le[i],
                headwidth=2*le[i],
                theta_z=180./np.pi*np.arctan2(v[i], u[i])+90,
                theta_x=180./np.pi*np.arccos(w[i]),
                color=color[i])


def plot_dipoles(ax0, cl):
    """Plot the dipoles in 2d or 3d."""
    if mg.two_d:
        plot_dipoles_2d(ax0, cl)
    else:
        plot_dipoles_3d(ax0, cl)


def plot_dipoles_3d(ax0, cl):
    """Plot the orientation of the 7 dipoles in three dimensions."""
    # ax0.set_aspect('equal', adjustable='box')
    ax0.cla()
    s1 = r'$\phi_3=$' + str(int(np.arctan2(cl.py[3],
                                           cl.px[3])*180/np.pi) % 360)
    s2 = r'$\phi_4=$' + str(int(np.arctan2(cl.py[4],
                                           cl.px[4])*180/np.pi) % 360)
    s3 = r'$\phi_5=$' + str(int(np.arctan2(cl.py[5],
                                           cl.px[5])*180/np.pi) % 360)
    s4 = '$U=$' + f"{cl.w_gesamt():.2f}"
    s5 = r'$m_\mathrm{tor}=$' + f"{tor_moment(mg.cl):.2f}"
    ax0.set_title(s1+', '+s2+', '+s3 + '; ' + s4 + '; ' + s5, size=14)
    d = 1.6
    ax0.set_xlim(-d, d)
    ax0.set_ylim(-d, d)
    ax0.set_zlim(-d, d)
    ax0.set_ylabel("y")
    ax0.set_xlabel("x")
    ax0.set_xlabel("z")
    quiver_3d(ax0, cl.rx, cl.ry, cl.px, cl.py)
    if mg.field:  # show the field on sherical surface around cuboid
        radius = d  # radius of the sphere
        us = np.linspace(0, 2 * np.pi, mg.sphere_resol)
        vs = np.linspace(0, np.pi, mg.sphere_resol//2)
        xs = radius * np.outer(np.cos(us), np.sin(vs))
        ys = radius * np.outer(np.sin(us), np.sin(vs))
        zs = radius * np.outer(np.ones(np.size(us)), np.cos(vs))
        if mg.show_pot:
            ens = mg.cl.potential(xs, ys, zs)  # potential
        else:
            ens = mg.cl.radial_flux(xs, ys, zs)  # flux comp. radial direction
        ens -= ens.min()
        ens /= ens.max()
        mycolors = np.empty(xs.shape)  # the size of the color map
        mycolors = plt.cm.bwr(ens)  # the type of color coding
        mycolors[:, :, 3] = mg.trp
        ax0.plot_surface(xs, ys, zs, facecolors=mycolors,
                         rcount=mg.sphere_resol, ccount=mg.sphere_resol//2)


def rotate_cluster(cl):
    """Rotate all dipoles such that phi_0 is within +- 15°."""
    rxn = np.copy(cl.rx)
    ryn = np.copy(cl.ry)
    pxn = np.copy(cl.px)
    pyn = np.copy(cl.py)
    pr = np.sqrt(pow(cl.px, 2) + pow(cl.py, 2))
    phiall = (np.arctan2(cl.py, cl.px))  # original angles
    r30 = np.pi*30./180.
    r60 = np.pi/3.
    phi0 = np.copy(phiall[0])  # original angle center, copy is important !!
    ipos = np.round(phi0/r60)
    phid = ((phi0+r30) % r60) - r30  # the desired angle
    phidiff = phi0-phid  # turn all angles by subracting phidiff:
    phiall -= phidiff
    for p in range(1, 7):  # rotate the positions, if larger than 30°
        npos = p-int(ipos)  # minus seems ok here?
        npos = ((npos-1) % 6) + 1  # exclude the adress zero here
        pxn[npos] = np.cos(phiall[p])*pr[p]
        pyn[npos] = np.sin(phiall[p])*pr[p]
    pxn[0] = np.cos(phiall[0])*pr[0]
    pyn[0] = np.sin(phiall[0])*pr[0]
    cl.rx, cl.ry, cl.px, cl.py = rxn, ryn, pxn, pyn


def tor_moment(cp):
    """Get toroidial moment."""
    return np.dot(cp.rx[1:], cp.py[1:]) - np.dot(cp.ry[1:], cp.px[1:])


def step(self):
    """Animate the cluster by showing the dipoles after a relaxation step."""
    for i in range(mg.plot_interv):
        if mg.sequential:
            such_seq()  # seems ok, but slow for small alpha
        else:
            such_par()  # seems less robust
    plot_dipoles(mg.ax[0], mg.cl)
    plt_alfa_phi()


def update_alpha(val):
    """React to change of alpha."""
    mg.alpha = mg.alpha_slider.val
    f = mg.alpha / np.sqrt(pow(mg.cl.px[0], 2)+pow(mg.cl.py[0], 2))
    mg.cl.px[0] *= f
    mg.cl.py[0] *= f
    plot_dipoles(mg.ax[0], mg.cl)
    mg.fig.canvas.draw_idle()
    mg.length = np.sqrt(pow(mg.cl.px, 2)+pow(mg.cl.py, 2))


def update_trp(val):
    """Change the transparency of the background or sphere."""
    mg.trp = mg.trp_slider.val


def update_cli(val):
    """Change the transparency of the background or sphere."""
    mg.cli = mg.cli_slider.val


def update_d_p(val):
    """Change the transparency of the background or sphere."""
    mg.g.d_p = mg.d_p_slider.val
    mg.g.new_geo()


def rota(event):
    """Rotate."""
    rotate_cluster(mg.cl)


def clear_dat(event):
    """Clear the plots by deleting the older data."""
    mg.xpl.clear()
    mg.ypl.clear()
    mg.yplu.clear()
    mg.yplm.clear()
    plt_alfa_phi()


def resolution_decrease(event):
    """Clear the plots by deleting the older data."""
    mg.sphere_resol = max(mg.sphere_resol//2, 8)


def resolution_increase(event):
    """Clear the plots by deleting the older data."""
    mg.sphere_resol = min(mg.sphere_resol*2, 256)


def plt_alfa_phi():
    """Plot a data point in response to button."""
    mg.xpl.append(mg.alpha)
    mg.ypl.append(np.arctan2(mg.cl.py[0], mg.cl.px[0])*180/np.pi)
    mg.yplm.append(tor_moment(mg.cl))
    mg.yplu.append(mg.cl.w_gesamt())
    mg.ax[1].cla()
    mg.ax[1].set_xlabel(r'central dipole moment $\alpha$')
    mg.ax[1].set_ylabel(r'$\phi_0$ (°)')
    mg.ax[1].plot(mg.xpl, mg.ypl, 'b.')

    mg.ax[2].cla()
    mg.ax[2].set_ylabel(r'$m_\mathrm{tor}$')
    mg.ax[2].plot(mg.xpl, mg.yplm, 'm.')

    mg.ax[3].cla()
    mg.ax[3].set_ylabel(r'$U$')
    mg.ax[3].plot(mg.xpl, mg.yplu, 'c.')
    mg.fig.canvas.draw_idle()


def toggle_seq_par(label):
    """Toggle between sequential and parallel updating."""
    if (label == 'seq'):
        mg.sequential = True
    else:
        mg.sequential = False
    mg.fig.canvas.draw_idle()


def set_my_rc():
    """Adjust to your needs."""
    plt.rc('axes', facecolor='white', edgecolor='black', titlesize=20,
           labelsize=15, grid=False)  # set the format of the plotaxes
    plt.rc('legend', fontsize=12, frameon=True, numpoints=1, labelspacing=0.3,
           columnspacing=0)    # set the format of the legend
    plt.rc('xtick', labelsize=12, top=True, direction='in')
    plt.rc('ytick', labelsize=12, right=True, direction='in')
    plt.rc('font', family="sans-serif", size=12)
    plt.rc('text', usetex=False)


def toggle_2d_3d(label):
    """Toggle between a 2d-display and a 3d-display of the dipoles."""
    mg.s.event_source.stop()    # stop the animation
    mg.fig.clf()                # clear the figure
    mg.two_d = (label == '2d')  # set two_d True or False
    layout_and_start()          # replot everything with appropriate widgets


def calculate_stop(label):
    """Stop and continue the calculation."""
    if (label == 'stop'):
        mg.calc = False
        mg.s.event_source.stop()
        mg.fig.canvas.draw_idle()

    elif (label == 'calc'):
        mg.calc = True
        mg.fig.clf()
        layout_and_start()


def set_plot_interval(label):
    """Set the plot interval with discrete steps according to dictionary."""
    mg.plot_interv = mg.intervals.get(label)


def toggle_bgr(label):
    """Turn field display on or off - for quicker relaxation."""
    mg.field = (label == 'field')
    mg.s.event_source.stop()
    mg.fig.clf()
    layout_and_start()


def select_background(label):
    """Select between potential and field energy."""
    mg.show_pot = (label == 'pot')


def set_sliders_and_buttons():
    """Set the graphical elements for the interactive part."""
    c = 'lightgoldenrodyellow'
    l0 = 0.02  # bottom line for slider
    l1 = 0.08  # first line for buttons
    l2 = 0.16  # second line for buttons
    l3 = 0.28  # third line for slider
    l4 = 0.34  # third line for slider
    s = 0.06   # size
    o = 0.02   # offset
    d = 0.08   # distance

    mg.alpha_slider = Slider(ax=plt.axes([o+0.02, l0, 0.85, 0.03]),
                             label=r'$\alpha$', valmin=0.0,
                             valmax=mg.alpha_max, valinit=mg.alpha)
    mg.alpha_slider.on_changed(update_alpha)

    mg.calc_stop = RadioButtons(plt.axes([o+0*d, l1, s, s], fc=c),
                                ('calc', 'stop'), active=1*(not(mg.calc)))
    mg.calc_stop.on_clicked(calculate_stop)

    mg.rotate = Button(plt.axes([o+1*d, l1, s, s]),
                       'rotate', color=c, hovercolor='0.975')
    mg.rotate.on_clicked(rota)

    mg.clr = Button(plt.axes([o+2*d, l1, s, s]),
                    'clr', color=c, hovercolor='0.975')
    mg.clr.on_clicked(clear_dat)

    mg.seq_par = RadioButtons(plt.axes([o+3*d, l1, s, s], fc=c),
                              ('seq', 'par'), active=1*(not(mg.sequential)))
    mg.seq_par.on_clicked(toggle_seq_par)

    mg.d2_d3 = RadioButtons(plt.axes([o+4*d, l1, s, s], fc=c),
                            ('2d', '3d'), active=1*(not(mg.two_d)))
    mg.d2_d3.on_clicked(toggle_2d_3d)

    mg.interv_set = RadioButtons(plt.axes([o+0*d, l2, s, s*1.8], fc=c),
                                 mg.intervals.keys(),
                                 active=list(mg.intervals.values()).index(
                                     mg.plot_interv))
    mg.interv_set.on_clicked(set_plot_interval)

    mg.back_off = RadioButtons(plt.axes([o+1*d, l2, s, s], fc=c),
                               ('field', 'off'), active=1*(not(mg.field)))
    mg.back_off.on_clicked(toggle_bgr)

    if (mg.field and mg.two_d):
        mg.backg_sel = RadioButtons(plt.axes([o+2*d, l2, s, s], fc=c),
                                    ('pot', '$B^2$'),
                                    active=1*(not(mg.show_pot)))
        mg.backg_sel.on_clicked(select_background)

    if (mg.field and not mg.two_d):
        mg.backg_sel = RadioButtons(plt.axes([o+2*d, l2, s, s], fc=c),
                                    ('pot', r'$B\perp$'),
                                    active=1*(not(mg.show_pot)))
        mg.backg_sel.on_clicked(select_background)

        mg.rs_decr = Button(plt.axes([o+3*d, l2, s, s]),
                            'rs<', color=c, hovercolor='0.975')
        mg.rs_decr.on_clicked(resolution_decrease)
        mg.rs_incr = Button(plt.axes([o+4*d, l2, s, s]),
                            'rs>', color=c, hovercolor='0.975')
        mg.rs_incr.on_clicked(resolution_increase)

        mg.trp_slider = Slider(ax=plt.axes([o+0.02, l3, 0.2, 0.03]),
                               label='trp', valmin=0.0,
                               valmax=1, valinit=mg.trp)
        mg.trp_slider.on_changed(update_trp)
    if (mg.field and mg.two_d):
        mg.cli_slider = Slider(ax=plt.axes([o+0.08, l3, 0.2, 0.03]),
                               label='clip level', valmin=0.1,
                               valmax=10, valinit=mg.cli)
        mg.cli_slider.on_changed(update_cli)
        mg.d_p_slider = Slider(ax=plt.axes([o+0.08, l4, 0.2, 0.03]),
                               label='plane dist.', valmin=0.1,
                               valmax=10, valinit=mg.d_p)
        mg.d_p_slider.on_changed(update_d_p)


def layout_and_start():
    """Arrange subplots and location of widgets."""
    mg.fig = plt.figure(num=1, figsize=(12, 9))
    set_my_rc()
    mg.ax = [None, None, None, None]
    if mg.two_d:
        mg.ax[0] = mg.fig.add_axes([0.0, 0.4, 0.5, 0.5])
        mg.ax[0].set_aspect('equal', adjustable='box')
        mg.ax[0].set_axis_off()
    else:
        mg.ax[0] = mg.fig.add_axes([0.0, 0.4, 0.5, 0.5], projection='3d')

    mg.ax[3] = mg.fig.add_axes([0.55, 0.65, 0.43, 0.25], sharex=mg.ax[1])
    mg.ax[2] = mg.fig.add_axes([0.55, 0.40, 0.43, 0.25], sharex=mg.ax[1])
    mg.ax[1] = mg.fig.add_axes([0.55, 0.15, 0.43, 0.25], sharex=mg.ax[1])
    set_sliders_and_buttons()
    mg.s = animation.FuncAnimation(mg.fig, step, frames=1, interval=1)


layout_and_start()
plt.show()
