#!/usr/bin/env python3
"""
This simulation example demonstrates how to replicate the
fitting example "Magnetically Dead Layers in Spinel Films"
given at the Nist website:
https://www.nist.gov/ncnr/magnetically-dead-layers-spinel-films

For simplicity, here we only reproduce the first part of that
demonstration without the magnetically dead layer.
"""

import os
import numpy
import bornagain as ba
from bornagain import nm, ba_plot as bp, deg, R3

# q-range on which the simulation and fitting are to be performed
qmin = 0.05997
qmax = 1.96

# number of points on which the computed result is plotted
scan_size = 1500

# The SLD of the substrate is kept constant
sldMao = (5.377e-06, 0)

# constant to convert between B and magnetic SLD
RhoMconst = 2.910429812376859e-12

####################################################################
#                  Create Sample and Simulation                    #
####################################################################

def get_sample(P):
    """
    construct the sample with the given parameters
    """
    BMagnitude = P["rhoM_Mafo"]*1e-6/RhoMconst
    angle = 0
    B = R3(
        BMagnitude*numpy.sin(angle*deg),
        BMagnitude*numpy.cos(angle*deg), 0)

    vacuum = ba.Vacuum()
    layer_color = (0.45, 0.32, 0.80)
    layer_mat = ba.SLDMaterial(
        "(Mg,Al,Fe)3O4", layer_color, P["rho_Mafo"]*1e-6, 0, B)
    substrate_color = (0.28, 0.57, 0.82)
    substrate_mat = ba.SLDMaterial("MgAl2O4", substrate_color, *sldMao)

    r_Mafo_autocorr = ba.SelfAffineFractalModel(P["r_Mafo"]*nm, 0.7, 25*nm)
    r_substrate_autocorr = ba.SelfAffineFractalModel(P["r_Mao"]*nm, 0.7, 25*nm)

    transient = ba.TanhTransient()

    r_Mafo = ba.Roughness(r_Mafo_autocorr, transient)
    r_substrate = ba.Roughness(r_substrate_autocorr, transient)

    ambient_layer = ba.Layer(vacuum)
    layer = ba.Layer(layer_mat, P["t_Mafo"]*nm, r_Mafo)
    substrate_layer = ba.Layer(substrate_mat, r_substrate)

    sample = ba.Sample()
    sample.addLayer(ambient_layer)
    sample.addLayer(layer)
    sample.addLayer(substrate_layer)

    return sample


def get_simulation(sample, q_axis, parameters, polarizer_vec,
                   analyzer_vec):
    """
    A simulation object.
    Polarization, analyzer and resolution are set
    from given parameters
    """
    q_axis = q_axis + parameters["q_offset"]
    distr = ba.DistributionGaussian(0., 1., 25, 4.)

    scan = ba.QzScan(q_axis)
    scan.setAbsoluteQResolution(distr, parameters["q_res"])
                                       # TODO CHECK not parameters["q_res"]*q_axis ??

    scan.setPolarization(polarizer_vec)
    scan.setAnalyzer(analyzer_vec)

    return ba.SpecularSimulation(scan, sample)


def run_simulation(q_axis, fitP, *, polarizer_vec, analyzer_vec):
    """
    Run a simulation on the given q-axis, where the sample is
    constructed with the given parameters.
    Vectors for polarization and analyzer need to be provided
    """
    parameters = dict(fitP, **fixedP)

    sample = get_sample(parameters)
    simulation = get_simulation(sample, q_axis, parameters, polarizer_vec,
                                analyzer_vec)

    return simulation.simulate()


def qr(result):
    """
    Returns two arrays that hold the q-values as well as the
    reflectivity from a given simulation result
    """
    q = result.xCenters()
    r = result.intensities()

    return q, r

####################################################################
#                         Plot Handling                            #
####################################################################

def plotData(qs, rs, exps, labels, colors):
    """
    Plot the simulated result together with the experimental data
    """
    fig = bp.plt.figure()
    ax = fig.add_subplot(111)

    for q, r, exp, l, c in zip(qs, rs, exps, labels, colors):
        ax.errorbar(exp.xCenters(),
                    exp.intensities(),
                    # xerr=TODO i742,
                    yerr=exp.errorSigmas(),
                    fmt='.',
                    markersize=0.75,
                    linewidth=0.5,
                    color=c[1])
        ax.plot(q, r, label=l, color=c[0])

    ax.set_yscale('log')
    bp.plt.legend()

    bp.plt.xlabel(r"$q$ (nm$^{-1}$)")
    bp.plt.ylabel("$R$")

    bp.plt.tight_layout()

def plotSpinAsymmetry(data_pp, data_mm, q, r_pp, r_mm):
    """
    Plot the simulated spin asymmetry as well its
    experimental counterpart with errorbars
    """
    Yp = data_pp.intensities()
    Ym = data_mm.intensities()
    Ep = data_pp.errors()
    Em = data_mm.errors()
    # compute the errorbars of the spin asymmetry
    delta = numpy.sqrt(4 * (Yp**2 * Em**2 + Ym**2 * Ep**2 ) / ( Yp + Ym )**4 )

    fig = bp.plt.figure()
    ax = fig.add_subplot(111)

    ax.errorbar(data_pp.xCenters(),
                (Yp - Ym) / (Yp + Ym),
                # xerr=TODO i742,
                yerr=delta,
                fmt='.',
                markersize=0.75,
                linewidth=0.5)

    ax.plot(q, (r_pp - r_mm)/(r_pp + r_mm))

    bp.plt.gca().set_ylim((-0.3, 0.5))

    bp.plt.xlabel(r"$q$ (nm$^{-1}$)")
    bp.plt.ylabel("Spin asymmetry")

    bp.plt.tight_layout()

####################################################################
#                          Data Handling                           #
####################################################################

def load_data(fname):
    flags = ba.ImportSettings1D("q_z (1/nm)", "", "", 1, 2, 3, 4)
    return ba.readData1D(fname, ba.csv1D, flags)

####################################################################
#                          Main Function                           #
####################################################################

if __name__ == '__main__':
    datadir = os.getenv('BA_DATA_DIR', '')
    if not datadir:
        raise Exception("Environment variable BA_DATA_DIR not set")
    fname_stem = os.path.join(datadir, "specular/MAFO_Saturated_")

    expdata_pp = load_data(fname_stem + "pp.tab")
    expdata_mm = load_data(fname_stem + "mm.tab")

    fixedP = {
        # parameters from our own fit run
        'q_res': 0.010542945012551425,
        'q_offset': 7.971243487467318e-05,
        'rho_Mafo': 6.370140108715461,
        'rhoM_Mafo': 0.27399566816062926,
        't_Mafo': 13.746913056084736,
        'r_Mao': 0.860487712674644,
        'r_Mafo': 0.37844265311293483
    }

    def run_Simulation_pp(qzs, P):
        return run_simulation(qzs,
                              P,
                              polarizer_vec=R3(0, 1, 0),
                              analyzer_vec=R3(0, 1, 0))

    def run_Simulation_mm(qzs, P):
        return run_simulation(qzs,
                              P,
                              polarizer_vec=R3(0, -1, 0),
                              analyzer_vec=R3(0, -1, 0))

    qzs = numpy.linspace(qmin, qmax, scan_size)
    ba.showSample3D(get_sample(fixedP), sample_size=120*nm, seed=0)
    q_pp, r_pp = qr(run_Simulation_pp(qzs, fixedP))
    q_mm, r_mm = qr(run_Simulation_mm(qzs, fixedP))

    plotData([q_pp, q_mm], [r_pp, r_mm], [expdata_pp, expdata_mm],
             ["$++$", "$--$"], [['orange','red'], ['green','blue']])

    plotSpinAsymmetry(expdata_pp, expdata_mm, qzs, r_pp, r_mm)
    bp.plt.show()
