"""Library that implements an effective density model for composition 
profiles (as described by Tolan) and also a home constructed model for
magnetic moment profiles (element specific). 
"""

from scipy.special import erf
from numpy import r_, cumsum, where, sqrt, exp, pi, array, zeros, sin, arange, bitwise_not, bitwise_or, append, newaxis

def erf_profile(z, z0, d, sigma0, sigma1):
    eta = (sigma1 * z0 + sigma0 * (z0 + d)) / (sigma0 + sigma1)
    p = 0.5 * (1 + where(z <= eta, erf((z - z0) / sqrt(2.0) / sigma0), -erf((z - z0 - d) / sqrt(2.0) / sigma1)))
    return p


def erf_interf(z, sigma):
    return 0.5 + 0.5 * erf(z / sqrt(2.0) / sigma)


def _lin(z, sigma):
    p = zeros(z.shape)
    p = (z > sqrt(3) * sigma) * 1.0 + (abs(z) <= sqrt(3) * sigma) * (0.5 + z / 2 / sqrt(3) / sigma)
    return p


def lin_profile(z, z0, d, sigma0, sigma1):
    eta = (sigma1 * z0 + sigma0 * (z0 + d)) / (sigma0 + sigma1)
    p = where(z < eta, _lin(z - z0, sigma0), _lin(-(z - z0 - d), sigma1))
    return p


def _exp(z, sigma):
    p = (z <= 0) * (0.5 * exp(sqrt(2) * z / sigma)) + (z > 0) * (1 - 0.5 * exp(-sqrt(2) * z / sigma))
    return p


def exp_profile(z, z0, d, sigma0, sigma1):
    eta = (sigma1 * z0 + sigma0 * (z0 + d)) / (sigma0 + sigma1)
    p = where(z < eta, _exp(z - z0, sigma0), _exp(-(z - z0 - d), sigma1))
    return p


def _sin(z, sigma):
    a = pi / sqrt(pi**2 - 8) * sigma
    p = (z > a) * 1.0 + (abs(z) <= a) * (0.5 + 0.5 * sin(pi * z / 2 / a))
    return p


def sin_profile(z, z0, d, sigma0, sigma1):
    eta = (sigma1 * z0 + sigma0 * (z0 + d)) / (sigma0 + sigma1)
    p = where(z < eta, _sin(z - z0, sigma0), _sin(-(z - z0 - d), sigma1))
    return p


def compress_profile_old(z, p, delta_max):
    pnew = p.copy()
    znew = z.copy()
    i = 0
    index = array([False])
    while any(bitwise_not(index)):
        # print i
        i += 1
        index = abs(pnew[:-1] - pnew[1:]) > delta_max
        index[::2] = True
        index = r_[index, True]
        # print pnew.shape, index.shape
        pnew = pnew[index]
        znew = znew[index]

    return znew, pnew


def compress_profile(z, p, delta_max):
    """Compresses a profile by merging the neighbouring layers that
    has an density difference of maximum delta_max"""
    pnew = p.copy()
    znew = z.copy()
    inew = arange(len(p))
    i = 0
    index = array([False])
    while any(bitwise_not(index)):
        # print i
        i += 1
        index = array([True] * len(pnew))
        index[1:-1:2] = abs(pnew[:-2:2] - pnew[2::2]) > delta_max
        # print pnew.shape, index.shape
        pnew = pnew[index]
        znew = znew[index]
        inew = inew[index]
    # print inew
    inew = append(inew, inew[-1])
    pret = array([p[inew[i] : inew[i + 1] + 1].mean() for i in range(len(inew) - 1)])
    return znew, pret


def compress_profile2(z, p1, p2, delta_max):
    """Compresses a profile by merging the neighbouring layers that
    has an density difference of maximum delta_max"""
    p1new = p1.copy()
    p2new = p2.copy()
    znew = z.copy()
    inew = arange(len(p1))
    i = 0
    index = array([False])
    while any(bitwise_not(index)):
        # print i
        i += 1
        index = array([True] * len(p1new))
        index[1:-1:2] = bitwise_or(
            abs(p1new[:-2:2] - p1new[2::2]) > delta_max, abs(p2new[:-2:2] - p2new[2::2]) > delta_max
        )
        # print pnew.shape, index.shape
        p1new = p1new[index]
        p2new = p2new[index]
        znew = znew[index]
        inew = inew[index]
    # print inew
    inew = append(inew, inew[-1])
    p1ret = array([p1[inew[i] : inew[i + 1] + 1].mean() for i in range(len(inew) - 1)])
    p2ret = array([p2[inew[i] : inew[i + 1] + 1].mean() for i in range(len(inew) - 1)])
    return znew, p1ret, p2ret


def compress_profile_n(z, ps, delta_maxs):
    """Compresses multiple profiles, each profile: ps[i],
    by merging the neighbouring layers that
    has an density difference of maximum delta_maxs[i]"""
    inew, znew = compress_profile_index_n(z, ps, delta_maxs)
    psnew = create_compressed_profile(ps, inew)
    return znew, psnew


def compare_layers(ps, delta_max):
    """Compare consecutive layers and see if ps is smaller than delta_max"""
    if len(ps.shape) > 1:
        return (abs(ps[:-2:2] - ps[2::2]) > delta_max).any(axis=1)
    else:
        return abs(ps[:-2:2] - ps[2::2]) > delta_max


def compress_profile_index_n(z, ps, delta_maxs):
    znew = z.copy()
    inew = arange(len(ps[0]))
    psnew = [ps_i.copy() for ps_i in ps]
    i = 0
    index = array([False])
    while any(bitwise_not(index)):
        i += 1
        index = array([True] * len(psnew[0]))
        test = array([compare_layers(ps_i, delta_maxs_i) for ps_i, delta_maxs_i in zip(psnew, delta_maxs)])
        index[1:-1:2] = test.any(axis=0)

        psnew = [ps_i[index] for ps_i in psnew]
        znew = znew[index]
        inew = inew[index]
    inew = append(inew, inew[-1])
    return inew, znew


def create_compressed_profile(ps, inew):
    psret = [array([ps_i[inew[i] : inew[i + 1] + 1].mean(0) for i in range(len(inew) - 1)]) for ps_i in ps]
    return psret


def create_profile(d, sigma, dens, prof_funcs, dz=0.01, mult=3, buffer=0, delta=20):
    zlay = r_[0, cumsum(d)]
    z = arange(-sigma[0] * mult - buffer, zlay[-1] + sigma[-1] * mult + buffer, dz)
    ptot = z * 0
    pdens = z * 0
    ps = []
    d = r_[delta + sigma[0] * mult + buffer, d, sigma[-1] * mult + delta + buffer]
    sigma = r_[0, sigma, 0]
    zlay = cumsum(d)
    # print zlay
    z0 = buffer + sigma[1] * mult
    zlay = r_[0, zlay] - z0
    # print zlay
    # for i in range(len(d)):
    #    print z0
    #    p = make_profile(z, z0, d[i], sigma[i], sigma[i+1])
    #    z0 += d[i]
    #    ptot += p
    #    pdens += p*dens[i]
    #    ps.append(p)
    ps = array(
        [prof(z, zp, di, s_low, s_up) for prof, zp, di, s_low, s_up in zip(prof_funcs, zlay, d, sigma[:-1], sigma[1:])]
    )
    ptot = ps.sum(0)
    pdens_indiv = (ps * dens[:, newaxis]) / ptot
    pdens = pdens_indiv.sum(0)

    return z, pdens, pdens_indiv


def create_profile_cm(
    d,
    sigma_c,
    sigma_m,
    prof_funcs,
    prof_funcs_mag,
    dmag_dens_l,
    dmag_dens_u,
    mag_dens,
    dd_m,
    dz=0.01,
    mult=3,
    buffer=0,
    delta=20,
):
    """Create a scattering length profile for magnetic x-ray scattering for the charge and
    magnetic part of the scattering lengths. sigma is roughnesses.
    Note that prof_funcs relate to layer profiles and prof_funcs_mag relates to
    magnetic interface profiles.
    """
    zlay = r_[0, cumsum(d)]
    s_max_bot = max(sigma_c[0], sigma_m[0])
    s_max_up = max(sigma_c[-1], sigma_m[-1])
    z = arange(-s_max_bot * mult - buffer, zlay[-1] + s_max_up * mult + buffer + dz, dz)
    ptot = z * 0
    pdens = z * 0
    ps = []
    d = r_[delta + s_max_bot * mult + buffer, d, s_max_up * mult + delta + buffer]
    sigma_c = r_[0, sigma_c, 0] + 1e-20
    sigma_m = r_[0, sigma_m, 0] + 1e-20
    dd_m = r_[0, dd_m, 0] + 1e-20
    prof_funcs_mag = [prof_funcs_mag[0]] + prof_funcs_mag + [prof_funcs_mag[-1]]
    zlay = cumsum(d)
    z0 = delta + buffer + s_max_bot * mult
    zlay = r_[0, zlay] - z0
    ps = array(
        [
            r_[
                prof(z, zp, di, s_c_low, s_c_up),
                # prof(z, zp, di, s_m_low, s_m_up)]
                prof_m_l(-(z - zp - dd_low), s_m_low) * dm_l + prof_m_u(z - zp - di - dd_up, s_m_up) * dm_u + m,
            ]
            for prof, zp, di, s_c_low, s_c_up, s_m_low, s_m_up, prof_m_l, prof_m_u, dm_l, dm_u, m, dd_low, dd_up in zip(
                prof_funcs,
                zlay,
                d,
                sigma_c[:-1],
                sigma_c[1:],
                sigma_m[:-1],
                sigma_m[1:],
                prof_funcs_mag[:-1],
                prof_funcs_mag[1:],
                dmag_dens_l,
                dmag_dens_u,
                mag_dens,
                dd_m[:-1],
                dd_m[1:],
            )
        ]
    )
    ptot = ps.sum(0)
    # print ps.shape, ptot.shape
    pdens_indiv = (ps[:, : len(z)]) / ptot[: len(z)]
    pdens_indiv_c = pdens_indiv  # *dens_c[:,newaxis]
    pdens_indiv_m = ps[:, len(z) :]  # *dens_m[:,newaxis]
    # pdens_c = pdens_indiv_c.sum(0)
    # pdens_m = (pdens_indiv_m*pdens_indiv).sum(0)

    return z, pdens_indiv_c, pdens_indiv_m


def create_profile_cm2(
    d,
    sigma_c,
    sigma_ml,
    sigma_mu,
    prof_funcs,
    prof_funcs_mag,
    dmag_dens_l,
    dmag_dens_u,
    mag_dens,
    dd_l,
    dd_u,
    dz=0.01,
    mult=3,
    buffer=0,
    delta=20,
):
    """Create a scattering length profile for magnetic x-ray scattering for the charge and
    magnetic part of the scattering lengths. sigma is roughnesses.
    Note that prof_funcs relate to layer profiles and prof_funcs_mag relates to
    magnetic interface profiles.
    """
    zlay = r_[0, cumsum(d)]
    s_max_bot = max(sigma_c[0], sigma_ml[0])
    s_max_up = max(sigma_c[-1], sigma_mu[-1])
    z = arange(-s_max_bot * mult - buffer, zlay[-1] + s_max_up * mult + buffer + dz, dz)
    ptot = z * 0
    pdens = z * 0
    ps = []
    d = r_[delta + s_max_bot * mult + buffer, d, s_max_up * mult + delta + buffer]
    sigma_c = r_[0, sigma_c, 0] + 1e-20
    sigma_mu = r_[sigma_mu, 0] + 1e-20
    sigma_ml = r_[0, sigma_ml] + 1e-20

    dd_l = r_[dd_l, 0] + 1e-20
    dd_u = r_[dd_u, 0] + 1e-20

    prof_funcs_mag = [prof_funcs_mag[0]] + prof_funcs_mag + [prof_funcs_mag[-1]]
    zlay = cumsum(d)
    z0 = delta + buffer + s_max_bot * mult
    zlay = r_[0, zlay] - z0
    ps = array(
        [
            r_[
                prof(z, zp, di, s_c_low, s_c_up),
                # prof(z, zp, di, s_m_low, s_m_up)]
                prof_m_l(-(z - zp - dd_low), s_m_low) * dm_l * m
                + prof_m_u((z - zp - di + dd_up), s_m_up) * dm_u * m
                + m,
            ]
            for prof, zp, di, s_c_low, s_c_up, s_m_low, s_m_up, prof_m_l, prof_m_u, dm_l, dm_u, m, dd_low, dd_up in zip(
                prof_funcs,
                zlay,
                d,
                sigma_c[:-1],
                sigma_c[1:],
                sigma_ml,
                sigma_mu,
                prof_funcs_mag[:-1],
                prof_funcs_mag[1:],
                dmag_dens_l,
                dmag_dens_u,
                mag_dens,
                dd_l,
                dd_u,
            )
        ]
    )
    ptot = ps.sum(0)
    # print ps.shape, ptot.shape
    pdens_indiv = (ps[:, : len(z)]) / ptot[: len(z)]
    pdens_indiv_c = pdens_indiv  # *dens_c[:,newaxis]
    pdens_indiv_m = ps[:, len(z) :]  # *dens_m[:,newaxis]
    # pdens_c = pdens_indiv_c.sum(0)
    # pdens_m = (pdens_indiv_m*pdens_indiv).sum(0)

    return z, pdens_indiv_c, pdens_indiv_m


# if __name__ == "__main__":
#     from mpl_toolkits.axes_grid import Grid
#     from mpl_toolkits.axes_grid.anchored_artists import AnchoredText
#     from pylab import *
#
#     # sigma0 = 5.
#     # sigma1 = 1.
#     # mult = 3
#     # buffer = 0
#     # delta = 20
#     # z0 = sigma0*mult
#     # d = 9.#sigma0*mult + sigma1*mult
#     # z = arange(0, z0 + d + sigma1*mult, 0.1)
#     # p = make_profile(z, z0, d, sigma0, sigma1)
#     # plot(z,p)
#     sigma_c = array([4, 4, 2]) * 1.0  # sigma_c = array([2,8,2])*1.0
#     sigma_m = array([1, 1, 1]) * 1.0  # sigma_m = array([2,5,4])*1.0
#     dmag_l = array([0.0, 0.0, 0.0])
#     dmag_u = array([0.0, 0.0, 0.0])
#     mag_dens = array([0.0, 2.25, 0.0])
#     # ddm = array([10.0, -10.0, 0.0])
#     ddm = array([10.0, -10.0, 10.0])
#     d = array([0.0, 50, 0.0]) * 1.0  # d = array([5, 60])*1.0
#     dens_c = array([-1.79e-3 + 0.0029j, -4e-23])  # dens_c = array([5,50,25,0])
#     dens_m = array([1.4458e-3 - 0.00144j, 2.12e-23])  # dens_m = array([0,6,10,0])
#     prof_funcs = [erf_profile] * 3  # prof_funcs = [erf_profile]*4
#     mag_prof_funcs = [erf_interf] * 3
#     # z, pdens, pdens_indiv = create_profile(d, sigma_c, dens_cprof_funcs, dz = 0.1,
#     #                                       mult = 5.0,
#     #                                       buffer = 0, delta = 5.)
#     z, c_den, m_den = create_profile_cm(
#         d[1:-1],
#         sigma_c[:-1],
#         sigma_m[:-1],
#         prof_funcs,
#         mag_prof_funcs,
#         dmag_l,
#         dmag_u,
#         mag_dens,
#         ddm,
#         dz=1.0,
#         mult=4.0,
#         buffer=20,
#         delta=5.0,
#     )
#
#     def plot_fig(filename, fig_title=""):
#         fig = figure(1)
#         grid = Grid(fig, 111, nrows_ncols=(3, 1), axes_pad=0.15)
#         # subplot(311)
#         grid[0].set_title(fig_title)
#         grid[0].plot(z, c_den[0], lw=2.0, c="b")
#         grid[0].plot(z, c_den[1], lw=2.0, c="r")
#         grid[0].plot(z, c_den[2], lw=2.0, c="g")
#         grid[0].axvline(0, linestyle="--", c="k")
#         grid[0].axvline(d[1], linestyle="--", c="k")
#         grid[0].set_ylabel("Comp.")
#         grid[0].legend(("Substrate", "Fe layer", "Ambient"))
#         # subplot(312)
#         grid[1].plot(z, m_den[1], lw=2.0, c="r")
#         grid[1].axvline(0, linestyle="--", c="k")
#         grid[1].axvline(d[1], linestyle="--", c="k")
#         grid[1].axvline(0 + ddm[0], linestyle=":", c="k")
#         grid[1].axvline(d[1] + ddm[1], linestyle=":", c="k")
#         # grid[1].annotate('Sub.dd_m = %.1f'%ddm[0], xy=(ddm[0], mag_dens[1]), xycoords = 'data',
#         #                 xytext = (-50, 0), textcoords = 'offset points',
#         #                 arrowprops=dict(arrowstyle="->", connectionstyle = 'arc3'))
#         # grid[1].annotate('Fe.dd_m = %.1f'%ddm[1], xy=(d[1]+ddm[1], mag_dens[1]), xycoords = 'data',
#         #                 xytext = (50, 0), textcoords = 'offset points',
#         #                 arrowprops=dict(arrowstyle="<-", connectionstyle = 'arc3'))
#         at = AnchoredText(
#             "dd_ml = %.1f" % (ddm[0]),
#             prop=dict(size=10),
#             frameon=True,
#             loc=6,
#         )
#         at.patch.set_boxstyle("round,pad=0.,rounding_size=0.2")
#         grid[1].add_artist(at)
#         at = AnchoredText(
#             "dd_mu = %.1f" % (ddm[1]),
#             prop=dict(size=10),
#             frameon=True,
#             loc=7,
#         )
#         at.patch.set_boxstyle("round,pad=0.,rounding_size=0.2")
#         grid[1].add_artist(at)
#         grid[1].set_ylabel("Mag. mom.\n per res. atom.")
#         grid[1].set_ylim(0, 3.1)
#         # subplot(313)
#         grid[2].plot(z, m_den[1] * c_den[1], "r", z, 2.25 * c_den[1], "b", lw=2.0)
#         grid[2].axvline(0, linestyle="--", c="k")
#         grid[2].axvline(d[1], linestyle="--", c="k")
#         grid[2].axvline(0 + ddm[0], linestyle=":", c="k")
#         grid[2].axvline(d[1] + ddm[1], linestyle=":", c="k")
#         grid[2].set_ylabel("Tot. Mag mom.")
#         grid[2].legend(("Profile", "Bulk Fe"))
#         grid[2].set_xlabel("Depth [AA]")
#         savefig(filename)
#         clf()
#
#     plot_fig("Bulk_Fe_sigma", "Bulk Fe with no change at the interfaces")
#     dmag_l = array([0.0, 0.75, 0.0])
#     dmag_u = array([0.0, 0.75, 0.0])
#     z, c_den, m_den = create_profile_cm(
#         d[1:-1],
#         sigma_c[:-1],
#         sigma_m[:-1],
#         prof_funcs,
#         mag_prof_funcs,
#         dmag_l,
#         dmag_u,
#         mag_dens,
#         ddm,
#         dz=1.0,
#         mult=4.0,
#         buffer=20,
#         delta=5.0,
#     )
#     plot_fig("Enhanced_Fe_sigma", "Bulk Fe with enhanced interfaces")
#     dmag_l = array([0.0, -1.0, 0.0])
#     dmag_u = array([0.0, -1.0, 0.0])
#     z, c_den, m_den = create_profile_cm(
#         d[1:-1],
#         sigma_c[:-1],
#         sigma_m[:-1],
#         prof_funcs,
#         mag_prof_funcs,
#         dmag_l,
#         dmag_u,
#         mag_dens,
#         ddm,
#         dz=1.0,
#         mult=4.0,
#         buffer=20,
#         delta=5.0,
#     )
#     plot_fig("Depleted_Fe_sigma", "Bulk Fe with depleted interfaces")
