#!/usr/bin/env python3
"""
This example demonstrates how to fit a complex experimental setup using BornAgain.
It is based on real data published in  https://doi.org/10.1002/advs.201700856
by A. Glavic et al.
In this example we utilize the scalar reflectometry engine to fit polarized
data without spin-flip for performance reasons.
"""

import bornagain as ba, numpy as np, os, matplotlib.pyplot as plt, scipy
from bornagain import nm, sample_tools as st


datadir = os.getenv('BA_DATA_DIR', '')
if not datadir:
    raise Exception("Environment variable BA_DATA_DIR not set")

####################################################################
#  Sample and simulation model
####################################################################

def get_sample(P, sign, T):

    if T < 200:
        ms150 = P["ms150"]
    else:
        ms150 = 1

    air_color = (0.90, 0.93, 0.97)
    air_mat = ba.SLDMaterial("Air", air_color, 0, 0)
    pyox_color = (0.62, 0.68, 0.72)
    pyox_mat = ba.SLDMaterial("PyOx", pyox_color,
                               (P["sld_PyOx_real"] + \
                                 sign * ms150 * P["msld_PyOx"] )* 1e-6,
                               P["sld_PyOx_imag"] * 1e-6)
    py2_color = (0.58, 0.40, 0.74)
    py2_mat = ba.SLDMaterial("Py2", py2_color,
                               ( P["sld_Py2_real"] + \
                                 sign * ms150 * P["msld_Py2"] ) * 1e-6,
                               P["sld_Py2_imag"] * 1e-6)
    py1_color = (0.05, 0.62, 0.55)
    py1_mat = ba.SLDMaterial("Py1", py1_color,
                               ( P["sld_Py1_real"] + \
                                 sign * ms150 * P["msld_Py1"] ) * 1e-6,
                               P["sld_Py1_imag"] * 1e-6)
    sio2_color = (0.25, 0.74, 0.42)
    sio2_mat = ba.SLDMaterial("SiO2", sio2_color, P["sld_SiO2_real"]*1e-6,
                                     P["sld_SiO2_imag"]*1e-6)
    si_color = (0.28, 0.57, 0.82)
    si_mat = ba.SLDMaterial("Substrate", si_color, P["sld_Si_real"]*1e-6,
                                   P["sld_Si_imag"]*1e-6)

    transient_model = ba.ErfTransient()

    rPyOx_autocorr = ba.SelfAffineFractalModel(P["rPyOx"]*nm, 0.7, 25*nm)
    rPy2_autocorr = ba.SelfAffineFractalModel(P["rPy2"]*nm, 0.7, 25*nm)
    rPy1_autocorr = ba.SelfAffineFractalModel(P["rPy1"]*nm, 0.7, 25*nm)
    rSiO2_autocorr = ba.SelfAffineFractalModel(P["rSiO2"]*nm, 0.7, 25*nm)
    rSi_autocorr = ba.SelfAffineFractalModel(P["rSi"]*nm, 0.7, 25*nm)

    rPyOx = ba.Roughness(rPyOx_autocorr, transient_model)
    rPy2 = ba.Roughness(rPy2_autocorr, transient_model)
    rPy1 = ba.Roughness(rPy1_autocorr, transient_model)
    rSiO2 = ba.Roughness(rSiO2_autocorr, transient_model)
    rSi = ba.Roughness(rSi_autocorr, transient_model)

    l_Air = ba.Layer(air_mat)
    l_PyOx = ba.Layer(pyox_mat, P["t_PyOx"]*nm, rPyOx)
    l_Py2 = ba.Layer(py2_mat, P["t_Py2"]*nm, rPy2)
    l_Py1 = ba.Layer(py1_mat, P["t_Py1"]*nm, rPy1)
    l_SiO2 = ba.Layer(sio2_mat, P["t_SiO2"]*nm, rSiO2)
    l_Si = ba.Layer(si_mat, rSi)

    sample = ba.Sample()

    sample.addLayer(l_Air)
    sample.addLayer(l_PyOx)
    sample.addLayer(l_Py2)
    sample.addLayer(l_Py1)
    sample.addLayer(l_SiO2)
    sample.addLayer(l_Si)

    return sample


def run_simulation(qaxis, P, *, sign, T):

    qdistr = ba.DistributionGaussian(0., 1., 25, 3.)

    dq = P["dq"]*qaxis
    scan = ba.QzScan(qaxis)
    scan.setVectorResolution(qdistr, dq)
    scan.setIntensity(P["intensity"])

    sample = get_sample(P, sign, T)

    simulation = ba.SpecularSimulation(scan, sample)
    simulation.setBackground(ba.ConstantBackground(5e-7))

    return simulation.simulate().intensities()

####################################################################
#  Experimental data
####################################################################

def load_data(fname, qmin, qmax):
    fpath = os.path.join(datadir, fname)
    flags = ba.ImportSettings1D("q (1/angstrom)", "#", "", 1, 3, 4, 5)
    data = ba.readData1D(fpath, ba.csv1D, flags)
    data = data.normalizedToMax()
    return data.crop(qmin, qmax)

####################################################################
#  Plotting
####################################################################

def plot(q, rs, data, shifts, labels):
    """
    Plot the simulated result together with the experimental data.
    """
    fig = plt.figure()
    ax = fig.add_subplot(111)

    for r, exp, shift, l in zip(rs, data, shifts, labels):

        ax.errorbar(exp.xCenters(),
                    exp.intensities() / shift,
                    yerr=exp.errors() / shift,
                    fmt='.',
                    markersize=0.75,
                    linewidth=0.5)

        ax.plot(q, r/shift, label=l)

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

    plt.xlabel(r"$q\; $(nm$^{-1}$)")
    plt.ylabel("$R$")
    <%- if test_mode or figure_mode -%>
    <%- else -%>
    plt.tight_layout()
    <%- end -%>


def plot_sld_profile(P):

    z_300p, sld_300p = st.materialProfile(get_sample(P, +1, 300))
    z_300m, sld_300m = st.materialProfile(get_sample(P, -1, 300))
    z_150p, sld_150p = st.materialProfile(get_sample(P, +1, 150))
    z_150m, sld_150m = st.materialProfile(get_sample(P, -1, 150))

    plt.figure()
    plt.plot(z_300p, np.real(sld_300p)*1e6, label=r"300K $+$")
    plt.plot(z_300m, np.real(sld_300m)*1e6, label=r"300K $-$")
    plt.plot(z_150p, np.real(sld_150p)*1e6, label=r"150K $+$")
    plt.plot(z_150m, np.real(sld_150m)*1e6, label=r"150K $-$")

    plt.xlabel(r"$z\;$(Å)")
    plt.ylabel(r"$\delta(z) \cdot 10^6$")

    plt.legend()
    <%- if test_mode or figure_mode -%>
    <%- else -%>
    plt.tight_layout()
    <%- end -%>

####################################################################
#  Main
####################################################################

if __name__ == '__main__':

    # Parameters and bounds.

    # We start with rather good values so that the example takes not too much time
    startPnB = {
        "intensity": (0.5, 0.4, 0.6),
        "t_PyOx": (7.7, 6.0, 10.0),
        "t_Py2": (5.6, 4.6, 6.6),
        "t_Py1": (5.6, 4.6, 6.6),
        "t_SiO2": (2.2, 1.5, 2.9),
    }

    # For fixed parameters, bounds are ignored. We leave them here just
    # to facilitate moving entries between startPnB and fixedPnB.
    fixedPnB = {
        "sld_PyOx_imag": (0, 0, 0),
        "sld_Py2_imag": (0, 0, 0),
        "sld_Py1_imag": (0, 0, 0),
        "sld_SiO2_imag": (0, 0, 0),
        "sld_Si_imag": (0, 0, 0),
        "sld_SiO2_real": (3.47, 3, 4),
        "sld_Si_real": (2.0704, 2, 3),
        "dq": (0.018, 0, 0.1),
    # Start by moving the following back to startPnB:
        "sld_PyOx_real": (1.995, 1.92, 2.07),
        "sld_Py2_real": (5, 4.7, 5.3),
        "sld_Py1_real": (4.62, 4.32, 4.92),
        "rPyOx": (2.7, 1.5, 3.5),
        "rPy2": (1.2, .2, 2.0),
        "rPy1": (1.2, .2, 2.0),
        "rSiO2": (1.5, .5, 2.5),
        "rSi": (1.5, .5, 2.5),
        "msld_PyOx": (0.25, 0, 1),
        "msld_Py2": (0.63, 0, 1),
        "msld_Py1": (0.64, 0, 1),
        "ms150": (1.05, 1.0, 1.1),
    }

    fixedP = {d: v[0] for d, v in fixedPnB.items()}
    P = {d: v[0] for d, v in startPnB.items()} | fixedP
    bounds = [(par[1], par[2]) for par in startPnB.values()]
    freeParNames = [name for name in startPnB.keys()]

    # Restrict the q range for fitting and plotting
    qmin = 0.08
    qmax = 1.4

    data = [
        load_data("specular/honeycomb300p.dat", qmin, qmax),
        load_data("specular/honeycomb300m.dat", qmin, qmax),
        load_data("specular/honeycomb150p.dat", qmin, qmax),
        load_data("specular/honeycomb150m.dat", qmin, qmax)]

    simFunctions = [
        lambda q, P: run_simulation(q, P, sign=+1, T=300),
        lambda q, P: run_simulation(q, P, sign=-1, T=300),
        lambda q, P: run_simulation(q, P, sign=+1, T=150),
        lambda q, P: run_simulation(q, P, sign=-1, T=150)]

    qzs = np.linspace(qmin, qmax, 1500) # x-axis for plot R vs q

    # Plot data with initial model

    simResults = [ f(qzs, P) for f in simFunctions ]
    plot(qzs, simResults, data, [1, 1, 10, 10],
         ["300K $+$", "300K $-$", "150K $+$", "150K $-$"])
    plot_sld_profile(P)

    # Fit

    qaxes = [d.xCenters() for d in data]
    rdata = [d.intensities() for d in data]

    def par_dict(*args):
        return {name: value for name, value in zip(freeParNames, *args)} | fixedP

    def objective_function(*args):
        """
        Returns fit objective, i.e. sum of weighted squared relative differences.
        """
        fullP = par_dict(*args)
        result = 0
        for q, r, sim_fct in zip(qaxes, rdata, simFunctions):
            t = sim_fct(q, fullP)
            reldiff = (r - t) / (r + t)
            result += np.sum(reldiff**2/len(t))
        return result

    result = scipy.optimize.differential_evolution(
        objective_function,
        bounds,
        <%- if test_mode or figure_mode -%>
        maxiter=2,
        popsize=2,
        <%- else -%>
        maxiter=5, # for a serious DE fit, choose 500
        popsize=3, # for a serious DE fit, choose 10
        <%- end -%>
        tol=1e-2,
        mutation=(0.5, 1.5),
        seed=0,
        disp=True,
        <%- if test_mode or figure_mode -%>
        polish=False
        <%- else -%>
        polish=True
        <%- end -%>
    )

    print(f"Final chi2: {result.fun}")
    print("Fit Result:")
    for name, value in zip(freeParNames, result.x):
        print(f'   {name} = {value}')

    # Plot data with fit result

    P = par_dict(result.x)
    <%- unless test_mode || figure_mode -%>
    ba.showSample3D(get_sample(P, 1, 300), sample_size=120*nm, seed=0)
    <%- end -%>

    sim_results = [ f(qzs, P) for f in simFunctions ]
    plot(qzs, simResults, data, [1, 1, 10, 10],
         ["300K $+$", "300K $-$", "150K $+$", "150K $-$"])
    plot_sld_profile(P)

    <%- if test_mode or figure_mode -%>
    <%- else -%>
    plt.show()
    <%- end -%>
