#!/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 angstrom, sample_tools as st
from bornagain.numpyutil import Arrayf64Converter as dac

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

    material_Air = ba.MaterialBySLD("Air", 0, 0)
    material_PyOx = ba.MaterialBySLD("PyOx",
                               (P["sld_PyOx_real"] + \
                                 sign * ms150 * P["msld_PyOx"] )* 1e-6,
                               P["sld_PyOx_imag"] * 1e-6)
    material_Py2 = ba.MaterialBySLD("Py2",
                               ( P["sld_Py2_real"] + \
                                 sign * ms150 * P["msld_Py2"] ) * 1e-6,
                               P["sld_Py2_imag"] * 1e-6)
    material_Py1 = ba.MaterialBySLD("Py1",
                               ( P["sld_Py1_real"] + \
                                 sign * ms150 * P["msld_Py1"] ) * 1e-6,
                               P["sld_Py1_imag"] * 1e-6)
    material_SiO2 = ba.MaterialBySLD("SiO2", P["sld_SiO2_real"]*1e-6,
                                     P["sld_SiO2_imag"]*1e-6)
    material_Si = ba.MaterialBySLD("Substrate", P["sld_Si_real"]*1e-6,
                                   P["sld_Si_imag"]*1e-6)

    interlayer_model = ba.ErfInterlayer()

    rPyOx_autocorr = ba.K_CorrelationModel(P["rPyOx"]*angstrom)
    rPy2_autocorr = ba.K_CorrelationModel(P["rPy2"]*angstrom)
    rPy1_autocorr = ba.K_CorrelationModel(P["rPy1"]*angstrom)
    rSiO2_autocorr = ba.K_CorrelationModel(P["rSiO2"]*angstrom)
    rSi_autocorr = ba.K_CorrelationModel(P["rSi"]*angstrom)

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

    l_Air = ba.Layer(material_Air)
    l_PyOx = ba.Layer(material_PyOx, P["t_PyOx"]*angstrom, rPyOx)
    l_Py2 = ba.Layer(material_Py2, P["t_Py2"]*angstrom, rPy2)
    l_Py1 = ba.Layer(material_Py1, P["t_Py1"]*angstrom, rPy1)
    l_SiO2 = ba.Layer(material_SiO2, P["t_SiO2"]*angstrom, rSiO2)
    l_Si = ba.Layer(material_Si, 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 dac.npArray(simulation.simulate().dataArray())

####################################################################
#  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(dac.npArray(exp.xCenters()),
                    dac.asNpArray(exp.dataArray()) / shift,
                    yerr=dac.asNpArray(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$")


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()

####################################################################
#  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": (77, 60, 100),
        "t_Py2": (56, 46, 66),
        "t_Py1": (56, 46, 66),
        "t_SiO2": (22, 15, 29),
    }

    # 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": (27, 15, 35),
        "rPy2": (12, 2, 20),
        "rPy1": (12, 2, 20),
        "rSiO2": (15, 5, 25),
        "rSi": (15, 5, 25),
        "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 = [dac.npArray(d.xCenters()) for d in data]
    rdata = [dac.asNpArray(d.dataArray()) 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,
        maxiter=2,
        popsize=2,
        tol=1e-2,
        mutation=(0.5, 1.5),
        seed=0,
        disp=True,
        polish=False
    )

    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)

    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)

