#!/usr/bin/env python3
"""
Fitting example: running same fit using various minimizer and their settings.
"""
import bornagain as ba
from bornagain import deg, angstrom, nm


def get_sample(P):
    """
    A sample with uncorrelated cylinders and prisms on a substrate.
    """
    cylinder_height = P["cylinder_height"]
    cylinder_radius = P["cylinder_radius"]
    prism_height = P["prism_height"]
    prism_base_edge = P["prism_base_edge"]

    # defining materials
    vacuum = ba.RefractiveMaterial("Vacuum", 0, 0)
    material_substrate = ba.RefractiveMaterial("Substrate", 6e-6, 2e-8)
    material_particle = ba.RefractiveMaterial("Particle", 6e-4, 2e-8)

    # collection of particles
    cylinder_ff = ba.Cylinder(cylinder_radius, cylinder_height)
    cylinder = ba.Particle(material_particle, cylinder_ff)
    prism_ff = ba.Prism3(prism_base_edge, prism_height)
    prism = ba.Particle(material_particle, prism_ff)
    layout = ba.ParticleLayout()
    layout.addParticle(cylinder, 0.5)
    layout.addParticle(prism, 0.5)
    interference = ba.InterferenceNone()
    layout.setInterference(interference)

    # vacuum layer with particles and substrate form multilayer
    vacuum_layer = ba.Layer(vacuum)
    vacuum_layer.addLayout(layout)
    substrate_layer = ba.Layer(material_substrate, 0)
    sample = ba.MultiLayer()
    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 # bp.simargs['n']
    detector = ba.Detector2D(2*deg, 2*deg, n, n, 0, 1*deg)
    return ba.ScatteringSimulation(beam, get_sample(P), detector)


def fake_data():
    """
    Generating "real" data from simulated image with default parameters.
    """

    P = {
        'cylinder_height': 5*nm,
        'cylinder_radius': 5*nm,
        'prism_height': 5*nm,
        'prism_base_edge': 5*nm
    }

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

    return result


def run_fitting():
    """
    main function to run fitting
    """

    data = fake_data()

    # prints info about available minimizers
    print(ba.MinimizerFactory().catalogToString())

    # prints detailed info about available minimizers and their options
    print(ba.MinimizerFactory().catalogDetailsToString())

    fit_objective = ba.FitObjective()
    fit_objective.addFitPair(get_simulation, data, 1)
    fit_objective.initPrint(10)

    P = ba.Parameters()
    P.add("cylinder_height", 4.*nm, min=0.01)
    P.add("cylinder_radius", 6.*nm, min=0.01)
    P.add("prism_height", 4.*nm, min=0.01)
    P.add("prism_base_edge", 12.*nm, min=0.01)

    minimizer = ba.Minimizer()
    minimizer.setMinimizer("Minuit2", "Migrad", "MaxFunctionCalls=5") # (fast test)
    # Uncomment one of the line below to adjust minimizer settings
    """
    Setting Minuit2 minimizer with Migrad algorithm, limiting number of iterations.
    Minimization will try to respect MaxFunctionCalls value.
    """
    # minimizer.setMinimizer("Minuit2", "Migrad", "MaxFunctionCalls=50")
    """
    Setting two options at once.
    Strategy=2 promises more accurate fit.
    """
    # minimizer.setMinimizer("Minuit2", "Migrad", "MaxFunctionCalls=500;Strategy=2")
    """
    Setting Minuit2 minimizer with Fumili algorithm.
    """
    # minimizer.setMinimizer("Minuit2", "Fumili")
    """
    Setting Levenberg-Marquardt algorithm.
    """
    # minimizer.setMinimizer("GSLLMA")

    result = minimizer.minimize(fit_objective.evaluate_residuals, P)

    fit_objective.finalize(result)

    print("Fitting completed.")
    print("chi2:", result.minValue())
    for fitPar in result.parameters():
        print(fitPar.name(), fitPar.value, fitPar.error)


if __name__ == '__main__':
    run_fitting()
