'''
Author: acse-xy721 xy721@ic.ac.uk
Date: 2022-07-11 23:43:23
LastEditors: acse-xy721 xy721@ic.ac.uk
LastEditTime: 2022-08-30 13:58:50
FilePath: /YXYIPR/meanfield/meanField.py
'''
from .Exchange import Exchange
from .DMI import DMI
from .Zeeman import Zeeman
from .Anisotropy import Anisotropy
from .Math import Math
# from Exchange import Exchange
# from DMI import DMI
# from Zeeman import Zeeman
# from Anisotropy import Anisotropy
# from Math import Math
import numpy as np
import discretisedfield as df
import micromagneticmodel as mm
import oommfc as oc


def Langevin(x):
    """Langevin function
    It's the Langevin function, which is is a stochastic differential equation
    describing the time evolution of a subset of degrees of freedom.
    .. math::
    L(x)=\\operatorname{coth}(x)-\\frac{1}{x}

    Parameters
    ----------
    x: ndarray
    Magnetisation field

    Returns
    -------
    Magnetisation field after calculation using Langevin equation

    Examples
    --------
    1. Generate a Maagnetisation field of all 1's
    and then use the Langevin formula to calculate

    >>> Magnetisation_field = np.ones((5, 5, 5, 3))
    >>> result = Langevin(Magnetisation_field)
    """
    return (1/np.tanh(x)) - 1/x


def end_condition(iter_num, m, Hef, tol, maxiter, Ms):
    """
    If the maximum of the cross product of the new magnetisation field
    and the effective field is less than the tolerance, or if the
    number of iterations is greater than the maximum number of
    iterations, then the function returns True.
    Otherwise, it returns False

    Parameters
    ----------
    iter_num: int
        the number of iterations
    m: ndarray
        magnetisation field
    Hef: ndarray
        Unit effective field
    tol: float
        tolerance for the end condition
    maxiter: int
        the maximum number of iterations to run the algorithm.
    Ms: float, int
        Saturation magnetisation

    Returns
    -------
    True or False

    Examples
    --------
    1. the cross product of the new magnetisation field
    and the effective field is less than the tolerance
    >>> iter_num = 1
    >>> tol = 1e-4
    >>> maxiter = 20000
    >>> Ms = 8e5

    Assuming the values of the magnetisation field and the effective field,
    the latitude is a vector field of 3 * 3 * 3 * 3

    >>> m = np.ones((3,3,3,3))
    >>> Hef = np.ones((3,3,3,3))

    Apply the end condition

    >>> end_condition(iter_num, m, Hef, tol, maxiter, Ms)
    True

    2. the number of iterations is greater than
    the maximum number of iterations
    >>> iter_num = 20000000
    >>> tol = 1e-4
    >>> maxiter = 20000
    >>> Ms = 8e5
    >>> m = np.ones((3,3,3,3))
    >>> Hef = np.ones((3,3,3,3))
    >>> end_condition(iter_num, m, Hef, tol, maxiter, Ms)
    True
    """
    cross_result_max = np.max(np.linalg.norm(np.cross(m/Ms, Hef), axis=3))
    return cross_result_max < tol or iter_num >= maxiter


def mean_field(math, D, Ms, miu0, A, H, beta, lamda, maxiter,
               tol, dim, dx, dy, dz, K, u):
    """Mean-field model
    Mean-field model can be used to investigate the qualitative behaviour of
    phase transitions in the classical spin model,
    and this equation also introduces the parameter of temperature,
    which supports the exploration of the physical phenomena
    generated by the spin model at different temperatures.

    The function takes in the initial magnetisation field,
    the energy constant parameters, the external field, the temperature,
    the convergence factor, the maximum number of iterations, the tolerance,
    the dimensions of the system, and the discretization in each direction.
    It then iterates until the magnetization converges to a stable state

    Parameters
    ----------
    math: class object
        The initial magnetisation field put in Math class
    D : float, int
        DMI energy constant
    Ms: float, int
        Saturation magnetisation
    miu0: float, int
        the permeability of free space
    A : float, int
        The exchange energy constant
    H: tuple
        External magnetic field(A/m)
    beta: float
        .. math:: \beta=1 / kb T
        kB being the Boltzmann constant
        T the temperature
    lamda: float
        convergence factor
    maxiter: int
        maximum number of iterations
    tol: float
        tolerance for the end condition
    dim: int
        Dimension of the field's value
        if dim = 1, this is a scalar field
        if dim = 3, this is a vector field
    dx: int
        Cell size of this discrete field in the x-axis direction
    dy: int
        Cell size of this discrete field in the y-axis direction
    dz: int
        Cell size of this discrete field in the z-axis direction
    K : float, int
        Anisotropy energy constant
    u : tuple, list
        Axis unit-vector

    Returns
    -------
    Magnetisation field at steady state

    Examples
    --------
    1. Initially, a 5*5*5*3 magnetisation field is generated, and then
    the steady state magnetisation field is calculated using
    the mean-field model

    Generated a 5*5*5*3 magnetisation field

    >>> import discretisedfield as df
    >>> region = df.Region(p1=(0, 0, 0), p2=(50e-9, 50e-9, 50e-9))
    >>> mesh = df.Mesh(region=region, cell=(10e-9, 10e-9, 10e-9))
    >>> def value_fun(point):
    ...     vec = np.random.randn(3)
    ...     unit_vec = vec / np.linalg.norm(vec)
    ...     return  (unit_vec[0], unit_vec[1], unit_vec[2])
    >>> Ms = 3.84e5
    >>> m = df.Field(mesh, dim=3, value=value_fun, norm=Ms)

    Define the parameters needed for mean-field

    >>> import micromagneticmodel as mm
    >>> dim = 3
    >>> dx = 10e-9
    >>> dy = 10e-9
    >>> dz = 10e-9
    >>> math = Math(m=m.array, dim=dim, dx=dx, dy=dy, dz=dz)
    >>> D = 1.58e-3
    >>> miu0 = mm.consts.mu0
    >>> A = 8.78e-12
    >>> H = (0, 0, 1)
    >>> T = 100
    >>> Kb = 1.380649e-23
    >>> beta = 1/(Kb*T)
    >>> lamda = 0.005
    >>> maxiter = 100000
    >>> tol = 1e-4
    >>> K = 1e8
    >>> u = [0, 0, 1]

    Apply the mean-field model

    >>> steady_m = mean_field(math, D, Ms, miu0, A, H, beta, lamda, maxiter,\
    tol, dim, dx, dy, dz, K, u)
    """
    iter_num = 0
    m_old = Math(math.m/Ms, dim, dx, dy, dz)
    while (True):
        Heff_all = DMI(D, miu0, Ms).effective_field(m_old) +  \
            Exchange(A, miu0, Ms).effective_field(m_old) + \
            Zeeman(H).effective_field(m_old) + \
            Anisotropy(K, u, miu0, Ms).effective_field(m_old)
        shape = (m_old.m.shape[0], m_old.m.shape[1], m_old.m.shape[2], 1)
        Heff_norm = np.linalg.norm(Heff_all, axis=3).reshape(shape)
        if np.isinf(beta):
            langen_output = 1
        else:
            langen_output = Langevin(beta*Heff_norm*miu0)
        M_new = Ms*langen_output*(Heff_all / Heff_norm)
        M_lamda = m_old.m*Ms + lamda*(M_new - m_old.m*Ms)
        M_new = (M_lamda / np.linalg.norm(M_lamda, axis=3).reshape(shape)) * \
            np.linalg.norm(M_new, axis=3).reshape(shape)
        if end_condition(iter_num, M_new, (Heff_all / Heff_norm),
                         tol, maxiter, Ms):
            return M_new
        iter_num += 1
        m_old = Math(M_new/Ms, dim, dx, dy, dz)
