#!/usr/bin/env python3
"""
Fit example with data by M. Fitzsimmons et al,
https://doi.org/10.5281/zenodo.4072376.
Sample is a ~50 nm Pt film on a Si substrate.
Single event data from Spallation Neutron Source
Beamline-4A (MagRef) with 60 Hz pulses and a wavelength
band of roughly 4-7 Å in 100 steps of 2theta.
"""

import bornagain as ba, numpy as np, os, matplotlib.pyplot as plt
from bornagain import angstrom

datadir = os.getenv('BA_DATA_DIR', '')

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

# Use fixed values for the SLD of the substrate and Pt layer
sldPt = (6.3568e-06, 1.8967e-09)
sldSi = (2.0728e-06, 2.3747e-11)

def get_sample(P):

    vacuum = ba.MaterialBySLD("Vacuum", 0, 0)
    material_layer = ba.MaterialBySLD("Pt", *sldPt)
    material_substrate = ba.MaterialBySLD("Si", *sldSi)

    ambient_layer = ba.Layer(vacuum)
    layer = ba.Layer(material_layer, P["t_pt/nm"])
    substrate_layer = ba.Layer(material_substrate)

    r_si = ba.LayerRoughness(P["r_si/nm"])
    r_pt = ba.LayerRoughness(P["r_pt/nm"])

    sample = ba.MultiLayer()
    sample.addLayer(ambient_layer)
    sample.addLayerWithTopRoughness(layer, r_pt)

    sample.addLayerWithTopRoughness(substrate_layer, r_si)

    return sample


def get_simulation(q_axis, P):
    sample = get_sample(P)

    scan = ba.QzScan(q_axis)
    scan.setIntensity(P["intensity"])
    scan.setOffset(P["q_offset"])

    distr = ba.DistributionGaussian(0., 1., 25, 4.)
    scan.setAbsoluteQResolution(distr, P["q_res/q"])

    simulation = ba.SpecularSimulation(scan, sample)

    return simulation

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

def plot(q, r, data, P):
    fig = plt.figure()
    ax = fig.add_subplot(111)

    ax.errorbar(data.npXcenters(),
                data.npArray(),
                # xerr=data.xxx, TODO restore
                yerr=data.npErrors(),
                label="R",
                fmt='.',
                markersize=1.,
                linewidth=0.6,
                color='r')

    ax.plot(q, r, label="Simulation", color='C0', linewidth=0.5)

    ax.set_yscale('log')

    ax.set_xlabel("$q\;$(nm$^{-1}$)")
    ax.set_ylabel("$R$")

    y = 0.5
    if P is not None:
        for n, v in P.items():
            plt.text(0.7, y, f"{n} = {v:.3g}", transform=ax.transAxes)
            y += 0.05

    plt.tight_layout()
    

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

if __name__ == '__main__':

    # Parameters and bounds:

    fixedPnB = {
        # to keep some parameters fixed, move lines here from startPnB
    }

    startPnB = {
        "intensity": (1., 0.8, 1.2),
        "q_offset": (0.01, -0.02, 0.02),
        "q_res/q": (0.01, 0, 0.02),
        "t_pt/nm": (50, 45, 55),
        "r_si/nm": (1.22, 0, 5),
        "r_pt/nm": (0.25, 0, 5),
    }

    fixedP = {d: v[0] for d, v in fixedPnB.items()}
    initialP = {d: v[0] for d, v in startPnB.items()}

    # Set q axis, load data:

    qmin = 0.18
    qmax = 2.4
    qzs = np.linspace(qmin, qmax, 1500)

    fpath = os.path.join(datadir, "specular/RvsQ_36563_36662.dat.gz")
    flags = ba.ImportSettings1D("q (1/angstrom)", "#", "", 1, 2, 3, 4)
    data = ba.readData1D(fpath, ba.csv1D, flags)

    # Initial plot

    r = get_simulation(qzs, initialP | fixedP).simulate().npArray()
    plot(qzs, r, data, initialP)

    # Restrict data to given q range

    data = data.crop(qmin, qmax)

    # Fit:

    fit_objective = ba.FitObjective()
    fit_objective.setObjectiveMetric("chi2")
    fit_objective.initPrint(10)
    fit_objective.addFitPair(
        lambda P: get_simulation(data.npXcenters(), P | fixedP), data, 1)

    P = ba.Parameters()
    for name, p in startPnB.items():
        P.add(name, p[0], min=p[1], max=p[2])

    minimizer = ba.Minimizer()
    result = minimizer.minimize(fit_objective.evaluate, P)
    fit_objective.finalize(result)

    finalP = {r.name(): r.value for r in result.parameters()}

    # Print and plot fit outcome:

    print("Fit Result:")
    print(finalP)

    r = get_simulation(qzs, finalP | fixedP).simulate().npArray()
    plot(qzs, r, data, finalP)

    plt.show()
