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


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

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

    formfactor = ba.HemiEllipsoid(radius_a, radius_b, height)
    particle = ba.Particle(material_particle, formfactor)

    layout = ba.ParticleLayout()
    layout.addParticle(particle)

    vacuum_layer = ba.Layer(vacuum)
    vacuum_layer.addLayout(layout)

    substrate_layer = ba.Layer(material_substrate)
    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.
    """
    incident_angle = P["incident_angle"]

    beam = ba.Beam(1e8, 0.1*nm, incident_angle)
    n = 11
    detector = ba.Detector2D(3*deg, 2*deg, n, n, 0, 1*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,
        'height': 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 = 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 i_dataset in range(0, fit_objective.fitObjectCount()):
            data = fit_objective.experimentalData(i_dataset)
            simul_data = fit_objective.simulationResult(i_dataset)
            chi2_map = fit_objective.relativeDifference(i_dataset)

            zmax = data.maxVal()

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

    @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))

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

        iteration_info = fit_objective.iterationInfo()

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

    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)
        plt.subplot(canvas[6:])
        self.plot_fit_parameters(fit_objective)

        plt.draw()
        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()
    plt.close() # (hide plot) fit_objective.initPlot(10, plotter.update)

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

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


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