#!/usr/bin/env python3
"""
Fitting example: simultaneous fit of two datasets
"""
import numpy as np
import matplotlib
import bornagain as ba
from bornagain import ba_plot as bp, deg, nm, nm2
import lmfit


def get_sample(P):
    """
    A sample with uncorrelated cylinders and pyramids.
    """
    radius_a = P["radius_a"]
    radius_b = P["radius_b"]
    radius_c = P["radius_c"]

    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.EllipsoidalSegment(radius_a, radius_b, radius_c, 0, radius_c)
    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.
    """
    incident_angle = P["incident_angle"]

    beam = ba.Beam(1e8, 0.1*nm, incident_angle)
    n = 100
    detector = ba.SphericalDetector(n, -1.5*deg, 1.5*deg, n, 0, 2*deg)
    return ba.ScatteringSimulation(beam, get_sample(P), detector)


def simulation1(P):
    P["incident_angle"] = 0.1*deg
    return get_simulation(P)


def simulation2(P):
    P["incident_angle"] = 0.4*deg
    return get_simulation(P)


def fake_data(incident_alpha):
    """
    Generating "real" data by adding noise to the simulated data.
    """
    P = {
        'radius_a': 5*nm,
        'radius_b': 6*nm,
        'radius_c': 8*nm,
        "incident_angle": incident_alpha
    }

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

    return result.noisy(0.1, 0.1)


class PlotObserver():
    """
    Draws fit progress every nth iteration. Real data, simulated data
    and chi2 map will be shown for both datasets.
    """

    def __init__(self):
        self.fig = bp.plt.figure(figsize=(12.8, 10.24))
        self.fig.canvas.draw()

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

    @staticmethod
    def plot_dataset(fit_objective, canvas):
        for j in range(0, fit_objective.nPairs()):
            data = fit_objective.experimentalData(j)
            simul_data = fit_objective.simulationResult(j)
            chi2_map = fit_objective.relativeDifference(j)

            zmax = data.maxVal()

            bp.plt.subplot(canvas[j*3])
            bp.plot_heatmap(data,
                            title="\"Real\" data - #" +
                            str(j + 1),
                            intensity_min=1,
                            intensity_max=zmax,
                            zlabel="")
            bp.plt.subplot(canvas[1 + j*3])
            bp.plot_heatmap(simul_data,
                            title="Simulated data - #" +
                            str(j + 1),
                            intensity_min=1,
                            intensity_max=zmax,
                            zlabel="")
            bp.plt.subplot(canvas[2 + j*3])
            bp.plot_heatmap(chi2_map,
                            title="Chi2 map - #" + str(j + 1),
                            intensity_min=0.001,
                            intensity_max=10,
                            zlabel="")

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

        iteration_info = fit_objective.iterationInfo()

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

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

        iteration_info = fit_objective.iterationInfo()

        bp.plt.text(
            0.01, 0.95, "Iterations  " +
            '{:d}'.format(iteration_info.iterationCount()))
        bp.plt.text(0.01, 0.70,
                 "Chi2       " + '{:8.4f}'.format(iteration_info.chi2()))
        for index, (name, val) in enumerate(
                iteration_info.parameterMap().items()):
            bp.plt.text(
                0.01, 0.30 - index*0.3,
                '{:30.30s}: {:6.3f}'.format(name, val))

    def update(self, fit_objective):
        self.fig.clf()

        # we divide figure to have 3x3 subplots, with two first rows occupying
        # most of the space
        canvas = matplotlib.gridspec.GridSpec(3,
                                              3,
                                              width_ratios=[1, 1, 1],
                                              height_ratios=[4, 4, 1])
        canvas.update(left=0.05, right=0.95, hspace=0.5, wspace=0.2)

        self.plot_dataset(fit_objective, canvas)
        bp.plt.subplot(canvas[6:])
        self.plot_fit_parameters(fit_objective)

        bp.plt.draw()
        bp.plt.pause(0.01)


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

    data1 = fake_data(0.1*deg)
    data2 = fake_data(0.4*deg)

    fit_objective = ba.FitObjective()
    fit_objective.addFitPair(simulation1, data1, 1)
    fit_objective.addFitPair(simulation2, data2, 1)
    fit_objective.initPrint(10)

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

    P = lmfit.Parameters()
    P.add("radius_a", 4.*nm, min=2, max=10)
    P.add("radius_b", 6.*nm, vary=False)
    P.add("radius_c", 4.*nm, min=2, max=10)

    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)


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