#!/usr/bin/env python3
"""
Fitting example: fit along slices
"""

from matplotlib import pyplot as plt
import bornagain as ba
from bornagain import deg, nm, ba_plot
import model1_cylinders as model

phi_slice_value = 0.0  # position of vertical slice in deg
alpha_slice_value = 0.2  # position of horizontal slice in deg

def get_masked_simulation(P):
    """
    Create and return GISAXS simulation with beam and detector defined
    """
    simulation = model.get_simulation(P)
    """
        At this point we mask all the detector and then unmask two areas
        corresponding to the vertical and horizontal lines. This will make
        simulation/fitting to be performed along slices only.
        """
    simulation.detector().maskAll()
    simulation.detector().addMask(ba.HorizontalLine(alpha_slice_value*deg), False)
    simulation.detector().addMask(ba.VerticalLine(phi_slice_value*deg), False)

    return simulation


def fake_data():
    """
    Generating "real" data by adding noise to the simulated data.
    """
    # initial values which we will have to find later during the fit
    P = {'radius': 5*nm, 'height': 10*nm}

    simulation = model.get_simulation(P)
    simulation.setBackground(ba.PoissonBackground())
    result = simulation.simulate()

    return result.noisy(0.1, 0.1)


class PlotObserver:
    """
    Draws fit progress every nth iteration. Here we plot slices along real
    and simulated images to see fit progress.
    """

    def __init__(self):
        self.fig = plt.figure(figsize=(10.25, 7.69))
        self.fig.canvas.draw()

    def __call__(self, fit_objective):
        self.update(fit_objective)

    @staticmethod
    def plot_data(data):
        """
        Plot experimental data as heatmap with horizontal/vertical lines
        representing slices on top.
        """
        plt.subplots_adjust(wspace=0.2, hspace=0.2)
        data.setTitle("Experimental data")
        ba_plot.plot_simulation_result(data)
        # line representing vertical slice
        plt.plot([phi_slice_value, phi_slice_value],
                 [data.yAxis().min(), data.yAxis().max()],
                 color='gray',
                 linestyle='-',
                 linewidth=1)
        # line representing horizontal slice
        plt.plot([data.xAxis().min(), data.xAxis().max()],
                 [alpha_slice_value, alpha_slice_value],
                 color='gray',
                 linestyle='-',
                 linewidth=1)

    @staticmethod
    def plot_slices(slices, title):
        """
        Plots vertical and horizontal projections.
        """
        plt.subplots_adjust(wspace=0.2, hspace=0.3)
        plt.yscale('log')
        for label, slice in slices:
            plt.plot(slice.xAxis().binCenters(), slice.flatVector(), label=label)
            plt.xlim(slice.xAxis().min(), slice.xAxis().max())
            plt.ylim(1, slice.maxVal()*10)
        plt.legend(loc='upper right')
        plt.title(title)

    @staticmethod
    def display_fit_parameters(fit_objective):
        """
        Displays fit parameters, chi and iteration number.
        """
        plt.title('Parameters')
        plt.axis('off')

        iteration_info = fit_objective.iterationInfo()

        plt.text(
            0.01, 0.85, "Iterations  " +
            '{:d}'.format(iteration_info.iterationCount()))
        plt.text(0.01, 0.75,
                 "Chi2       " + '{:8.4f}'.format(iteration_info.chi2()))
        for index, P in enumerate(iteration_info.parameters()):
            plt.text(
                0.01, 0.55 - index*0.1,
                '{:30.30s}: {:6.3f}'.format(P.name(), P.value))

        plt.tight_layout()
        plt.draw()
        plt.pause(0.01)

    def update(self, fit_objective):
        """
        Callback to access fit_objective on every n'th iteration.
        """
        self.fig.clf()

        data = fit_objective.experimentalData()
        simul_data = fit_objective.simulationResult()

        # plot real data
        plt.subplot(2, 2, 1)
        self.plot_data(data)

        # horizontal slices
        slices = [("real", data.xProjection(alpha_slice_value)),
                  ("simul", simul_data.xProjection(alpha_slice_value))]
        title = ("Horizontal slice at alpha =" +
                 '{:3.1f}'.format(alpha_slice_value))
        plt.subplot(2, 2, 2)
        self.plot_slices(slices, title)

        # vertical slices
        slices = [("real", data.yProjection(phi_slice_value)),
                  ("simul", simul_data.yProjection(phi_slice_value))]
        title = "Vertical slice at phi =" + '{:3.1f}'.format(
            phi_slice_value)
        plt.subplot(2, 2, 3)
        self.plot_slices(slices, title)

        # display fit parameters
        plt.subplot(2, 2, 4)
        self.display_fit_parameters(fit_objective)


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

    data = fake_data()

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

    # creating custom observer which will draw fit progress
    plotter = PlotObserver()
    fit_objective.initPlot(10, plotter)

    P = ba.Parameters()
    P.add("radius", 6.*nm, min=4, max=8)
    P.add("height", 9.*nm, min=8, max=12)

    minimizer = ba.Minimizer()
    result = minimizer.minimize(fit_objective.evaluate, P)
    fit_objective.finalize(result)


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