#!/usr/bin/env python3
"""
Basic real-life example of fitting specular data.
The sample consists of single Ni film on SiO2 substrate.
"""

import math
import bornagain as ba, os
from bornagain import nm, deg, ba_fitmonitor


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/MLZ-TREFF-Ni58.txt")

    flags = ba.ImportSettings1D("qz (1/angstrom)", "#", "", 1, 2)
    return ba.readData1D(fname, ba.csv1D, flags)


def get_sample(P):
    # Materials
    vacuum = ba.MaterialBySLD()
    material_Ni_58 = ba.MaterialBySLD("Ni", 9.408e-06, 0)
    material_SiO2 = ba.MaterialBySLD("SiO2", 2.0704e-06, 0)

    # Layers and interfaces
    interlayer = ba.TanhInterlayer()

    Ni_autocorr = ba.K_CorrelationModel(P["sigma_Ni"])
    roughness_Ni = ba.LayerRoughness(Ni_autocorr, interlayer)

    sub_autocorr = ba.K_CorrelationModel(P["sigma_Substrate"])
    roughness_Substrate = ba.LayerRoughness(sub_autocorr, interlayer)

    layer_Ni = ba.Layer(material_Ni_58, P["thickness"], roughness_Ni)
    substrate = ba.Layer(material_SiO2, roughness_Substrate)

    sample = ba.Sample()
    sample.addLayer(ba.Layer(vacuum))
    sample.addLayer(layer_Ni)
    sample.addLayer(substrate)

    return sample


def get_simulation(P):
    scan = ba.QzScan(data.xAxis())

    # Finite resolution due to beam divergence
    n_samples = 5
    rel_sampling_width = 2.0
    res_distr = ba.DistributionGaussian(0, 1, n_samples, rel_sampling_width)

    wavelength = 0.473*nm
    res_alpha = 0.006*deg
    res_q = 4*math.pi*res_alpha/wavelength
    scan.setAbsoluteQResolution(res_distr, res_q)

    sample = get_sample(P)
    simulation = ba.SpecularSimulation(scan, sample)
    simulation.setBackground(ba.ConstantBackground(1e-4))

    return simulation


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

    P = ba.Parameters()
    P.add("thickness", 59*nm, min=50*nm, max=100*nm)
    P.add("sigma_Ni", 2*nm, min=0.01*nm, max=3*nm)
    P.add("sigma_Substrate", 2.5*nm, min=0.01*nm, max=3*nm)

    fit_objective = ba.FitObjective()
    fit_objective.addFitPair(get_simulation, data, 1)
    fit_objective.setObjectiveMetric("log", "l2")

    fit_objective.initPrint(50)

    plot_observer = ba_fitmonitor.PlotterSpecular(pause=0.001)
    fit_objective.initPlot(50, plot_observer)

    minimizer = ba.Minimizer()

    # Main minimization
    minimizer.setMinimizer("Genetic", "",
                           "MaxIterations=5;PopSize=300"
                           )
    result = minimizer.minimize(fit_objective.evaluate, P)
    fit_objective.finalize(result)

    # Polish results with another minimizer
    best_P_so_far = result.parameters()
    minimizer.setMinimizer("Minuit2", "Migrad",
                            "MaxFunctionCalls=0")
    result = minimizer.minimize(fit_objective.evaluate, best_P_so_far)
    fit_objective.finalize(result)

    plot_observer.show()
