"""Animate the relaxation of a dipole cluster towards a local minimum."""
# Version 1.2, Ingo.Rehberg, University of Bayreuth, May 9, 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_3d:
    """Define an ensemble of magnetic dipoles, together with some functions."""

    def __init__(self, n, model='test'):
        self.n = n
        self.exclude = n+1  # calculate field and potential without this one
        rx = np.zeros(self.n, dtype='double')
        self.R = np.array([rx, rx, rx]).T
        px = np.zeros(self.n, dtype='double')
        self.P = np.array([px, px, px]).T
        dxet = np.zeros(self.n*self.n, dtype='double')
        self.DET = np.array([dxet, dxet, dxet]).T
        self.dp3t = np.zeros(self.n*self.n, dtype='double')
        self.set_data(model=model)

    def set_data(self, phi2=0., alpha=2.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.
        if model == 'circular':
            self.P[0] = (alpha, 0, 0)
            self.R[0] = (0, 0, 0.5)
            for i in range(1, 7):
                self.R[i] = (np.cos((i-1)*np.pi/3.)*r,
                             np.sin((i-1)*np.pi/3.)*r, 0)
                self.P[i] = (-self.R[i, 1]/r, self.R[i, 0]/r, 0)
        elif model == 'test':
            for i in range(self.n):
                self.R[i] = ((i % 2)*2 - 1,
                             ((i//2) % 2)*2 - 1,
                             ((i//4) % 2)*2 - 1)
                self.P[i] = self.R[i]/np.sqrt(3)
        else:
            print('no such model')
        self.refresh_tables()

    def refresh_tables(self):
        """Call after position or length change."""
        self.length = np.sqrt(np.sum(self.P**2, axis=1))  # 3d
        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
                D = self.R[i] - self.R[p]
                dist = np.sqrt(np.dot(D, D))
                self.DET[pnt] = D * np.sqrt(3.) / dist
                self.dp3t[pnt] = 1. / pow(dist, 3)

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

    def calc_b_loc(self, i):
        """Get local field vector."""
        if i == 0:  # because 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(3)  # 3-dimensional version
            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 (np.dot(self.P[i], self.P[p])
                - np.dot(self.P[i], self.DET[pnt])
                * np.dot(self.P[p], self.DET[pnt])
                ) * self.dp3t[pnt]

    def w_total(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 x, y, z."""
        pot = 0.
        for i in range(self.n):
            if i == self.exclude:
                continue
            rx = x-self.R[i, 0]
            ry = y-self.R[i, 1]
            rz = z-self.R[i, 2]
            r2 = pow(rx, 2)+pow(ry, 2)+pow(rz, 2)
            nen = np.clip(pow(r2, 1.5), 1.e-10, None)
            scpr = self.P[i, 0]*rx + self.P[i, 1]*ry + self.P[i, 2]*rz
            pot += scpr/nen
        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):
            if i == self.exclude:
                continue
            rx = x-self.R[i, 0]
            ry = y-self.R[i, 1]
            rz = z-self.R[i, 2]
            r2 = pow(rx, 2)+pow(ry, 2)+pow(rz, 2)
            scpr = self.P[i, 0]*rx + self.P[i, 1]*ry + self.P[i, 2]*rz
            nen = np.clip(pow(r2, 2.5), 1.e-10, None)  # avoid zero
            lbx += (3.*scpr*rx - r2*self.P[i, 0])/nen
            lby += (3.*scpr*ry - r2*self.P[i, 1])/nen
            lbz += (3.*scpr*rz - r2*self.P[i, 2])/nen
        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."""

    cl = Cluster_3d(7, model='circular')
    alpha_max = 3  # the ratio of dipole strength as defined in the paper
    alpha = cl.length[0]  # for the slider setting this bifurcation parameter
    z0_max = 3  # hight of the central dipole suggested in the comment

    z0 = cl.R[0, 2]  # the height is the second interactive control parameter
    two_d = False  # determines if the plot is 2d or 3d
    free = False  # determines weather the central dipole rotates freely or not
    dsp_alf = True  # the x-axis for the 3 plots as set by alf_z0
    d_p = 0.5  # height of the plane to display potential or energy density
    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)  # class for the plane
    xpl, ypl, yplm, yplu = [], [], [], []  # parameters taken in a plot_interv
    intervals = {'$10^0$': 1, '$10^1$': 10, '$10^2$': 100,
                 '$10^3$': 1000, '$10^4$': 10000, '$10^5$': 100000}
    plot_interv = intervals.get('$10^0$')  # no. of iterations between plots
    fig, ax = None, []  # the figure with its subplots
    calc_stop = None  # Slider to stop the relaxtion code
    alpha_slider, rotate, clr, free_fx, d2_d3 = None, None, None, None, None
    back_off, rs_decr, rs_incr, trp_slider = None, None, None, None
    z0_slider, alf_z0 = None, None
    s = None  # the animation
    calc = True  # stop or continue the relaxation code
    field = False  # 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
    show_pot = False  # determines weather to show potential or energy density


def relaxation():
    """Turn all dipoles into the local B-field."""
    start = 0 if mg.free else 1
    for i in range(start, mg.cl.n):
        xy = mg.cl.calc_b_loc(i)
        xy *= mg.cl.length[i]/np.sqrt(np.dot(xy, xy))  # unity vector
        mg.cl.P[i] = xy


def plot_dipoles_2d(ax0, cl):
    """Plot the orientation of the n dipoles."""
    ax0.cla()
    s1 = r'$\phi_3=$' + str(int(np.arctan2(cl.P[3, 1],
                                           cl.P[3, 0])*180/np.pi) % 360)
    s2 = r'$\phi_4=$' + str(int(np.arctan2(cl.P[4, 1],
                                           cl.P[4, 0])*180/np.pi) % 360)
    s3 = r'$\phi_5=$' + str(int(np.arctan2(cl.P[5, 1],
                                           cl.P[5, 0])*180/np.pi) % 360)
    s4 = '$U=$' + f"{cl.w_total():.2f}"
    s5 = r'$m_\mathrm{tor,z}=$' + f"{tor_moment(mg.cl)[2]:.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)  # 2d-projection of the B-field
        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.R[:, 0], cl.R[:, 1], cl.P[:, 0], cl.P[:, 1], pivot='middle',
               scale=10, color='c', linewidth=1)
    circle1 = plt.Circle((0, 0), 1., color='lightgrey', fill=False)
    ax0.add_patch(circle1)
    for n in range(cl.n):
        ax0.add_patch(plt.Circle((cl.R[n, 0], cl.R[n, 1]),
                                 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, R, P):
    """Plot tube arrows."""
    x, y, z = R[:, 0], R[:, 1], R[:, 2]
    u, v, w = np.copy(P[:, 0]), np.copy(P[:, 1]), np.copy(P[:, 2])  # not deep!
    le = np.power(u*u+v*v+w*w, 1/6)  # 3rd root of dipole moment
    uni_sc = 1/np.power(u*u+v*v+w*w, 1/2)  # unity_scale
    u0 = u * uni_sc
    v0 = v * uni_sc
    w0 = w * uni_sc
    color = ['c', 'r', 'r', 'r', 'r', 'r', 'r', 'r']
    t_x = w/mg.cl.length
    t_x = np.arccos(np.clip(t_x, -1, 1))
    for i in range(len(x)):
        arrow3d(a, offset=(x[i]-u0[i]*le[i]/2, y[i]-v0[i]*le[i]/2,
                           z[i]-w0[i]*le[i]/2),
                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*t_x[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.P[3, 1],
                                           cl.P[3, 0])*180/np.pi) % 360)
    s2 = r'$\phi_4=$' + str(int(np.arctan2(cl.P[4, 1],
                                           cl.P[4, 0])*180/np.pi) % 360)
    s3 = r'$\phi_5=$' + str(int(np.arctan2(cl.P[5, 1],
                                           cl.P[5, 0])*180/np.pi) % 360)
    s4 = '$U=$' + f"{cl.w_total():.2f}"
    s5 = r'$m_\mathrm{tor,z}=$' + f"{tor_moment(mg.cl)[2]:.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_zlabel("z")
    quiver_3d(ax0, cl.R, cl.P)
    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 filled hex dipoles such that phi_0 is within +- 15°."""
    pxn = np.copy(cl.P[:, 0])
    pyn = np.copy(cl.P[:, 1])
    pzn = np.copy(cl.P[:, 2])
    pr = np.sqrt(pow(pxn, 2) + pow(pyn, 2))
    phiall = (np.arctan2(pyn, pxn))  # 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, cl.n):  # 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]
        pzn[npos] = cl.P[p, 2]
    pxn[0] = np.cos(phiall[0])*pr[0]
    pyn[0] = np.sin(phiall[0])*pr[0]
    cl.P = np.array([pxn, pyn, pzn]).T


def tor_moment(cp):
    """Get the toroidial moment of the ring."""
    return np.sum(np.cross(cp.R[1:], cp.P[1:]), axis=0)


def step(self):
    """Animate the cluster by showing the dipoles after a relaxation step."""
    for i in range(mg.plot_interv):
        relaxation()  # most primitive version
    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
    mg.cl.P[0] *= mg.alpha / mg.cl.length[0]
    mg.cl.length[0] = mg.alpha
    plot_dipoles(mg.ax[0], mg.cl)
    mg.fig.canvas.draw_idle()


def update_z0(val):
    """React to change of z0."""
    mg.z0 = mg.z0_slider.val
    mg.cl.R[0, 2] = mg.z0
    mg.cl.refresh_tables()
    plot_dipoles(mg.ax[0], mg.cl)
    mg.fig.canvas.draw_idle()


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


def update_cli(val):
    """Change the clip level of the color map."""
    mg.cli = mg.cli_slider.val


def update_d_p(val):
    """Change the heigt of the display plane."""
    mg.g.d_p = mg.d_p_slider.val
    mg.g.new_geo()


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


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


def clear_dat(event):
    """React to the clear button."""
    clear_old_data()


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.ax[1].cla()
    if mg.dsp_alf:
        mg.xpl.append(mg.alpha)
        mg.ypl.append(np.arctan2(mg.cl.P[0, 1], mg.cl.P[0, 0])*180/np.pi)
        mg.ax[1].set_xlabel(r'central dipole moment $\alpha$')
        mg.ax[1].set_ylabel(r'$\phi_0$ (°)')
    else:
        mg.xpl.append(mg.cl.R[0, 2])
        mg.ax[1].set_xlabel(r'height of the central dipole $z_0$')
        if mg.free:
            hhh = np.clip(mg.cl.P[0, 2]/mg.cl.length[0], -1, 1)
            mg.ypl.append(np.arccos(hhh)*180/np.pi)
            mg.ax[1].set_ylabel(r'$\theta_0$ (°)')
        else:
            mg.ypl.append(tor_moment(mg.cl)[1])
            mg.ax[1].set_ylabel(r'$m_\mathrm{tor,y}$')
    mg.ax[1].plot(mg.xpl, mg.ypl, 'b.')

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

    mg.yplu.append(mg.cl.w_total())
    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_free_fx(label):
    """Toggle between sequential and parallel updating."""
    mg.free = (label == 'free')
    clear_old_data()
    layout_and_start()


def toggle_alf_z0(label):
    """Toggle between sequential and parallel updating."""
    mg.dsp_alf = (label == r'$\alpha$')
    clear_old_data()
    layout_and_start()


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.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
        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()
    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.01  # 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.z0_slider = Slider(ax=plt.axes([o+0.02, l0+0.03, 0.85, 0.03]),
                          label=r'$z_0$', valmin=0.0,
                          valmax=mg.z0_max, valinit=mg.cl.R[0, 2])
    mg.z0_slider.on_changed(update_z0)

    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.free_fx = RadioButtons(plt.axes([o+3*d, l1, s, s], fc=c),
                              ('free', 'fixed'), active=1*(not(mg.free)))
    mg.free_fx.on_clicked(toggle_free_fx)

    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.alf_z0 = RadioButtons(plt.axes([o+5*d, l1, s, s], fc=c),
                             (r'$\alpha$', r'$z_0$'),
                             active=1*(not(mg.dsp_alf)))
    mg.alf_z0.on_clicked(toggle_alf_z0)

    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."""
    if mg.fig is not None:
        mg.fig.clf()
    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()
