"""Animate the relaxation of a dipole cluster towards a local minimum."""
# Ingo Rehberg, Mar 23, 2022
# This program illustrates a comment on:
# 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 matplotlib.pyplot as plt
from matplotlib.widgets import Slider, Button, RadioButtons
import numpy as np
from matplotlib import animation


class Cluster:
    """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


class mg:
    """My globals, just to avoid unnamed globals."""

    alpha_max = 3  # the ratio of dipole strength as defined by the authors
    alpha = alpha_max/2  # for the slider
    sequential = True
    cl = Cluster(7)
    cl.set_data(phi2=0., alpha=alpha, model='circular')
    length = np.sqrt(pow(cl.px, 2)+pow(cl.py, 2))
    steps = 100  # calculation steps between plots
    xpl, ypl, yplm, yplu = [], [], [], []


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(ax0, cl):
    """Plot the orientation of the 7 dipoles."""
    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_ylabel("y")
    ax0.set_xlabel("x")
    ax0.quiver(cl.rx, cl.ry, cl.px, cl.py, pivot='middle',
               scale=10, color='r', 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 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.steps):
        if mg.sequential:
            such_seq()  # seems ok, but slow for small alpha
        else:
            such_par()  # seems less robust
    plot_dipoles(ax[0], mg.cl)
    plt_alfa_phi()


def update(val):
    """React to change of alpha."""
    mg.alpha = 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(ax[0], mg.cl)
    fig.canvas.draw_idle()
    mg.length = np.sqrt(pow(mg.cl.px, 2)+pow(mg.cl.py, 2))


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


def clear_dat(event):
    """Take a data pair."""
    mg.xpl.clear()
    mg.ypl.clear()
    mg.yplu.clear()
    mg.yplm.clear()
    plt_alfa_phi()


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())
    ax[1].cla()
    ax[1].set_xlabel(r'central dipole moment $\alpha$')
    ax[1].set_ylabel(r'$\phi_0$ (°)')
    ax[1].plot(mg.xpl, mg.ypl, 'b.')

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

    ax[3].cla()
    ax[3].set_ylabel(r'$U$')
    ax[3].plot(mg.xpl, mg.yplu, 'c.')
    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
    fig.canvas.draw()


def layout():
    """Arrange subplots and location of widgets."""
    fig = plt.figure(num=1, figsize=(12, 9))
    ax = [None, None, None, None]
    ax[0] = fig.add_axes([0.0, 0.4, 0.5, 0.5])
    ax[3] = fig.add_axes([0.55, 0.65, 0.43, 0.25], sharex=ax[1])
    ax[2] = fig.add_axes([0.55, 0.40, 0.43, 0.25], sharex=ax[1])
    ax[1] = fig.add_axes([0.55, 0.15, 0.43, 0.25], sharex=ax[1])
    alpha_slider = Slider(ax=plt.axes([0.05, 0.02, 0.85, 0.03]),
                          label=r'$\alpha$', valmin=0.0, valmax=mg.alpha_max,
                          valinit=mg.alpha)
    c = 'lightgoldenrodyellow'
    rotate = Button(plt.axes([0.05, 0.10, 0.06, 0.06]),
                    'rotate', color=c, hovercolor='0.975')
    clr = Button(plt.axes([0.15, 0.10, 0.06, 0.06]),
                 'clr', color=c, hovercolor='0.975')
    seq_par = RadioButtons(plt.axes([0.25, 0.1, 0.06, 0.06], fc=c),
                           ('seq', 'par'), active=1*(not(mg.sequential)))
    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)
    return fig, ax, alpha_slider, rotate, clr, seq_par


fig, ax, alpha_slider, rotate, clr, seq_par = layout()
alpha_slider.on_changed(update)
rotate.on_clicked(rota)
clr.on_clicked(clear_dat)
seq_par.on_clicked(toggle_seq_par)
plt_alfa_phi()
s = animation.FuncAnimation(fig, step, frames=1, interval=1)
plt.show()
