#!/usr/bin/env python3
"""
Basic GISAS 2D fitting example.
The model is in gisas_model1.
Fake experimental data are generated by gisas_fake1.
"""

import os
import numpy as np
import gisas_model1 as model
import bornagain as ba
from bornagain import ba_fitmonitor, nm
import lmfit

from matplotlib import pyplot as plt


if __name__ == '__main__':
    datadir = os.getenv('BA_DATA_DIR', '')
    if not datadir:
        raise Exception("Environment variable BA_DATA_DIR not set")
    fname = os.path.join(datadir, "scatter2d/faked_gisas1.int")
    data = ba.readData2D(fname)

    fit_objective = ba.FitObjective()
    fit_objective.addFitPair(model.get_simulation, data)

    fit_objective.initPrint(10)  # Print on every 10th iteration.
    observer = ba_fitmonitor.PlotterGISAS()
    <%- if test_mode or figure_mode -%>
    plt.close() # hide plot
    <%- else -%>
    fit_objective.initPlot(10, observer) # Plot every 10th, slow!
    <%- end -%>

    P = model.start_parameters_1()
    result = lmfit.minimize(fit_objective.evaluate_residuals, P)

    print("Fit completed.")
    print(lmfit.fit_report(result))
    <%- unless test_mode || figure_mode -%>
    ba.showSample3D(model.get_sample(result.params), sample_size=120*nm, seed=0)
    <%- end -%>

    # Save simulation image corresponding to the best fit parameters
    np.savetxt("fit.txt",
               fit_objective.simulationResult().intensities())
    plt.show()
