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


def load_data():
    datadir = os.getenv('BA_DATA_DIR', '')
    if not datadir:
        raise Exception("Environment variable BA_DATA_DIR not set")
    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.Vacuum()
    ti_color = (0.05, 0.62, 0.55)
    ti_mat = ba.SLDMaterial("Ti", ti_color, -1.9493e-06, 0)
    ni_color = (0.93, 0.48, 0.14)
    ni_mat = ba.SLDMaterial("Ni", ni_color, 9.4245e-06, 0)
    si_color = (0.30, 0.62, 0.86)
    si_mat = ba.SLDMaterial("Si", si_color, 2.0704e-06, 0)

    # Layers
    layer_Ti = ba.Layer(ti_mat, P["thickness_Ti"])
    layer_Ni = ba.Layer(ni_mat, 70*angstrom)

    # Periodic stack
    n_repetitions = 10
    stack = ba.LayerStack(n_repetitions)
    stack.addLayer(layer_Ti)
    stack.addLayer(layer_Ni)

    # Sample
    sample = ba.Sample()
    sample.addLayer(ba.Layer(vacuum))
    sample.addStack(stack)
    sample.addLayer(ba.Layer(si_mat))

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

    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)

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

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

    plot_observer.show()
