#!/usr/bin/env python3
"""
An example of using the Bayesian sampling library emcee with BornAgain.

Author: Andrew McCluskey (andrew.mccluskey@ess.eu)
"""

import bornagain as ba, corner, emcee, numpy as np, os, scipy
import matplotlib.pyplot as plt
from bornagain import nm


np.random.seed(1)

datadir = os.getenv('BA_DATA_DIR', '')
if not datadir:
    raise Exception("Environment variable BA_DATA_DIR not set")


def get_sample(ni_thickness, ti_thickness):
    # pure real scattering-length densities (in angstrom^-2)
    si_sld_real = 2.0704e-06  # Si (substrate)
    ni_sld_real = 9.4245e-06  # Ni
    ti_sld_real = -1.9493e-06  # Ti

    # materials
    vacuum = ba.Vacuum()
    ni_color = (0.93, 0.48, 0.14)
    ni_mat = ba.SLDMaterial("Ni", ni_color, ni_sld_real, 0)
    ti_color = (0.05, 0.62, 0.55)
    ti_mat = ba.SLDMaterial("Ti", ti_color, ti_sld_real, 0)
    substrate_color = (0.28, 0.57, 0.82)
    substrate_mat = ba.SLDMaterial("SiSubstrate", substrate_color, si_sld_real, 0)

    # layers
    vacuum_layer = ba.Layer(vacuum)
    ni_layer = ba.Layer(ni_mat, ni_thickness)
    ti_layer = ba.Layer(ti_mat, ti_thickness)
    substrate_layer = ba.Layer(substrate_mat)

    # periodic stack
    n_repetitions = 10
    stack = ba.LayerStack(n_repetitions)
    stack.addLayer(ti_layer)
    stack.addLayer(ni_layer)

    # sample
    sample = ba.Sample()
    sample.addLayer(vacuum_layer)
    sample.addStack(stack)
    sample.addLayer(substrate_layer)

    return sample


def get_simulation(sample, points):
    scan = ba.AlphaScan(ba.ListScan("alpha_i (rad)", points))
    scan.setWavelength(0.154 * nm);
    return ba.SpecularSimulation(scan, sample)


def run_simulation(points, ni_thickness, ti_thickness):
    sample = get_sample(ni_thickness, ti_thickness)
    simulation = get_simulation(sample, points)

    result = simulation.simulate()
    return result.intensities()



if __name__ == '__main__':
    filepath = os.path.join(datadir, "specular/genx_alternating_layers.dat.gz")
    flags = ba.ImportSettings1D("2alpha (deg)", "#", "", 1, 2)
    data = ba.readData1D(filepath, ba.csv1D, flags)

    q = data.xCenters()
    y = data.intensities()
    dy = y * 0.1 # arbitrary uncertainties

    def log_likelihood(P):
        """
        Calculates the log-likelihood for the normal uncertainties
        :tuple sim_var: the variable parameters
        :array x: the abscissa data (q-values)
        :array y: the ordinate data (R-values)
        :array yerr: the ordinate uncertainty (dR-values)
        :return: log-likelihood
        """
        y_sim = run_simulation(q, *P)
        sigma2 = dy**2 + y_sim**2
        return -0.5*np.sum((y - y_sim)**2/sigma2 + np.log(sigma2))

    # Use differential evolution to find the maximum likelihood estimate
    objective_function = lambda *args: -log_likelihood(*args) # just change sign
    initial = np.array([9.0, 1.0]) + 0.1*np.random.randn(2)
    solution = scipy.optimize.differential_evolution(
        objective_function,
        bounds=((5.0, 9.0), (1.0, 10.0)))

    ni_thickness_ml, ti_thickness_ml = solution.x
    print('MLE Ni Thickness', ni_thickness_ml, 'nm')
    print('MLE Ti Thickness', ti_thickness_ml, 'nm')

    # Perform the likelihood sampling
    nwalkers = 32
    ndim = 2
    pos = solution.x + 1e-4*np.random.randn(nwalkers, ndim)
    sampler = emcee.EnsembleSampler(nwalkers,
                                    ndim,
                                    log_likelihood)
    sampler.run_mcmc(pos, 1000, progress=True)

    # Plot and show corner plot of samples
    flat_samples = sampler.get_chain(flat=True)
    corner.corner(flat_samples,
                  labels=['Ni-thickness/nm', 'Ti-thickness/nm'])


    # Plot and show MLE and data of reflectivity
    plt.errorbar(q, y, dy, marker='.', ls='')
    plt.plot(
        q,
        run_simulation(q, *flat_samples.mean(axis=0)),
        '-')
    plt.xlabel('$\\alpha$/rad')
    plt.ylabel('$R$')
    plt.yscale('log')
    plt.close()
