#!/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
import lmfit
import scipy.optimize


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("q_z (1/angstrom)", "#", "", 1, 2)
    return ba.readData1D(fname, ba.csv1D, flags)


def get_sample(P):
    # Materials
    vacuum = ba.Vacuum()
    material_Ni_58_color = (0.93, 0.48, 0.14)
    material_Ni_58 = ba.SLDMaterial("Ni", material_Ni_58_color, 9.408e-06, 0)
    sio2_color = (0.20, 0.55, 0.72)
    sio2_mat = ba.SLDMaterial("SiO2", sio2_color, 2.0704e-06, 0)

    # Layers and interfaces
    transient = ba.TanhTransient()

    Ni_autocorr = ba.SelfAffineFractalModel(P["sigma_Ni"], 0.7, 25*nm)
    roughness_Ni = ba.Roughness(Ni_autocorr, transient)

    sub_autocorr = ba.SelfAffineFractalModel(P["sigma_Substrate"], 0.7, 25*nm)
    roughness_Substrate = ba.Roughness(sub_autocorr, transient)

    layer_Ni = ba.Layer(material_Ni_58, P["thickness"], roughness_Ni)
    substrate = ba.Layer(sio2_mat, 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_names = ["thickness", "sigma_Ni", "sigma_Substrate"]
    bounds = [(50*nm, 100*nm), (0.01*nm, 3*nm), (0.01*nm, 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)

    # Stage 1: global search with differential evolution
    def de_objective(x):
        P = lmfit.Parameters()
        for name, val in zip(P_names, x):
            P.add(name, val)
        return fit_objective.evaluate(P)

    de_result = scipy.optimize.differential_evolution(
        de_objective, bounds,
        maxiter=5,
        popsize=15,
        seed=42)

    # Stage 2: local refinement with lmfit
    P = lmfit.Parameters()
    for name, val, (lo, hi) in zip(P_names, de_result.x, bounds):
        P.add(name, val, min=lo, max=hi)

    result = lmfit.minimize(fit_objective.evaluate_residuals, P)
    print(lmfit.fit_report(result))
    ba.showSample3D(get_sample(result.params), sample_size=120*nm, seed=0)

    plot_observer.show()
