#!/usr/bin/env python3
"""
Basic example how to fit specular data.
The sample consists of twenty alternating Ti and Ni layers.
Reference data was generated with GenX.
We fit just one parameter, the thickness of the Ti layers,
which has an original value of 30 angstroms.
"""

import bornagain as ba, numpy as np, os
from bornagain import angstrom, ba_fitmonitor


def load_data():
    datadir = os.getenv('BA_DATA_DIR', '')
    fname = os.path.join(datadir, "specular/genx_alternating_layers.dat.gz")

    flags = ba.ImportSettings1D("2alpha (deg)", "#", "", 1, 2)
    return ba.readData1D(fname, ba.csv1D, flags)


def get_sample(P):
    # Materials
    vacuum = ba.MaterialBySLD()
    material_Ti = ba.MaterialBySLD("Ti", -1.9493e-06, 0)
    material_Ni = ba.MaterialBySLD("Ni", 9.4245e-06, 0)
    material_Si = ba.MaterialBySLD("Si", 2.0704e-06, 0)

    # Multilayer
    layer_Ti = ba.Layer(material_Ti, P["thickness_Ti"])
    layer_Ni = ba.Layer(material_Ni, 70*angstrom)

    sample = ba.MultiLayer()
    sample.addLayer(ba.Layer(vacuum))
    for _ in range(10):
        sample.addLayer(layer_Ti)
        sample.addLayer(layer_Ni)
    sample.addLayer(ba.Layer(material_Si))

    return sample


def get_simulation(P):
    scan = ba.AlphaScan(data.xAxis())
    scan.setWavelength(1.54*angstrom)
    sample = get_sample(P)

    return ba.SpecularSimulation(scan, sample)


if __name__ == '__main__':
    data = load_data()

    P = ba.Parameters()
    P.add("thickness_Ti", 50*angstrom, min=10*angstrom, max=60*angstrom)

    fit_objective = ba.FitObjective()
    fit_objective.addFitPair(get_simulation, data, 1)

    fit_objective.initPrint(10)
    plot_observer = ba_fitmonitor.PlotterSpecular(pause=0.5)
    fit_objective.initPlot(10, plot_observer)

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

    fit_objective.finalize(result)

    plot_observer.show()
