#!/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 os
import numpy as np
from matplotlib import pyplot as plt
import bornagain as ba
from bornagain import angstrom, ba_fitmonitor


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

    expdata = np.loadtxt(data_fname, usecols=(0, 1), skiprows=3)
    # convert double incident angle (degs) to incident angle (radians)
    expdata[:, 0] *= np.pi/360

    return (expdata[:, 0], expdata[:, 1])


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(ba.ListScan("", exp_x))
    scan.setWavelength(1.54*angstrom)
    sample = get_sample(P)

    return ba.SpecularSimulation(scan, sample)


if __name__ == '__main__':
    global exp_x
    exp_x, exp_y = load_data()

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

    fit_objective = ba.FitObjective()
    fit_objective.addSimulationAndData(get_simulation, exp_y, 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()
