#!/usr/bin/env python3
"""
Fitting example: two-stage global+local fitting.

Stage 1 uses scipy differential_evolution for a global search over large
parameter space. Stage 2 uses lmfit (Levenberg-Marquardt) to refine
the result to a precise minimum.
"""

from matplotlib import pyplot as plt
import bornagain as ba
from bornagain import ba_fitmonitor, deg, angstrom, nm, nm2
import lmfit
import scipy.optimize


def get_sample(P):
    """
    A sample with uncorrelated cylinders and pyramids on a substrate.
    """
    radius = P["radius"]
    height = P["height"]

    vacuum = ba.Vacuum()
    substrate_color = (0.28, 0.57, 0.82)
    substrate_mat = ba.RefractiveMaterial("Substrate", substrate_color, 6e-6, 2e-8)
    particle_color = (0.86, 0.24, 0.18)
    particle_mat = ba.RefractiveMaterial("Particle", particle_color, 6e-4, 2e-8)

    ff = ba.Cylinder(radius, height)
    particle = ba.Particle(particle_mat, ff)

    vacuum_layer = ba.Layer(vacuum)
    vacuum_layer.deposit2D(ba.Dilute2D(0.01/nm2, particle))

    substrate_layer = ba.Layer(substrate_mat)
    sample = ba.Sample()
    sample.addLayer(vacuum_layer)
    sample.addLayer(substrate_layer)
    return sample


def get_simulation(P):
    """
    A GISAXS simulation with beam and detector defined.
    """
    beam = ba.Beam(1e8, 1*angstrom, 0.2*deg)
    n = 100
    detector = ba.SphericalDetector(n, 0., 2*deg, n, 0., 2*deg)
    simulation = ba.ScatteringSimulation(beam, get_sample(P), detector)
    simulation.options().setUseAvgMaterials(False)
    return simulation


def fake_data():
    """
    Generating "real" data by adding noise to the simulated data.
    """
    P = {'radius': 5*nm, 'height': 5*nm}

    simulation = get_simulation(P)
    result = simulation.simulate()

    return result.noisy(0.3, 0.5)


def run_fitting():
    data = fake_data()

    fit_objective = ba.FitObjective()
    fit_objective.addFitPair(get_simulation, data, 1)
    fit_objective.initPrint(30)
    observer = ba_fitmonitor.PlotterGISAS()

    plt.close() # hide plot

    P_names = ["height", "radius"]
    bounds = [(0.01, 30), (0.01, 30)]

    # 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=3,
        popsize=15,
        seed=42)

    # Stage 2: local refinement seeded from global search result
    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))


if __name__ == '__main__':
    run_fitting()
    plt.show()
