#!/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 nm
import lmfit

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

####################################################################
#  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.Vacuum()
    layer_color = (0.93, 0.72, 0.25)
    layer_mat = ba.SLDMaterial("Pt", layer_color, *sldPt)
    substrate_color = (0.30, 0.62, 0.86)
    substrate_mat = ba.SLDMaterial("Si", substrate_color, *sldSi)

    transient = ba.TanhTransient()

    si_autocorr = ba.SelfAffineFractalModel(P["r_si_nm"]*nm, 0.7, 25*nm)
    pt_autocorr = ba.SelfAffineFractalModel(P["r_pt_nm"]*nm, 0.7, 25*nm)

    r_si = ba.Roughness(si_autocorr, transient)
    r_pt = ba.Roughness(pt_autocorr, transient)

    ambient_layer = ba.Layer(vacuum)
    layer = ba.Layer(layer_mat, P["t_pt_nm"]*nm, r_pt)
    substrate_layer = ba.Layer(substrate_mat, r_si)

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

    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_over_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.xCenters(),
                data.intensities(),
                # xerr=data.xxx, TODO restore
                yerr=data.errors(),
                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()
    plt.close() # hide plot

####################################################################
#  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_over_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

    res = get_simulation(qzs, initialP | fixedP).simulate()
    r = res.intensities()
    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.xCenters(), P | fixedP), data, 1)

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

    result = lmfit.minimize(fit_objective.evaluate_residuals, P)
    print(lmfit.fit_report(result))

    finalP = {name: result.params[name].value for name in result.params}

    # Print and plot fit outcome:

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

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

    plt.show()
